widgets/matmul.js

/* Matrix-multiply tester — walks C[i][j] = Σ A[i][k]·B[k][j] one multiply-accumulate TERM at a time,
 * showing every product and the running partial sum (the SAME steps the Manim `matmul` animation narrates).
 * Interconnected: if the paper wires an upstream matrix in (io.consumes, e.g. X), this becomes Q·Kᵀ = X·Xᵀ
 * and publishes the resulting scores downstream — so editing X upstream re-derives these scores live. */
(function () {
  const k = MANIMO.kit, bus = MANIMO.bus;
  function mount(root, concept) {
    const p = (concept && concept.params) || {};
    const io = (concept && concept.io) || {};
    let A0 = p.A || [[1, 2, 3], [4, 5, 6]];
    let B0 = p.B || [[7, 8], [9, 10], [11, 12]];
    const up = io.consumes ? bus.get(io.consumes) : null;        // Q·Kᵀ: A = X, B = Xᵀ
    if (up && up.length && (up[0] || []).length) { A0 = up.map((r) => r.slice()); B0 = k.transpose(up); }
    const m = A0.length, kk = (A0[0] || []).length, n = (B0[0] || []).length;
    const A = k.grid(m, kk, A0), B = k.grid(kk, n, B0, { readonly: !!io.consumes });  // B = Xᵀ is derived ⇒ locked
    const C = k.grid(m, n, zero(m, n), { readonly: true });
    const log = k.panel("steps");

    function plan() {                                            // one entry PER term (i,j,t)
      const a = A.get(), b = B.get(), seq = [];
      for (let i = 0; i < m; i++) for (let j = 0; j < n; j++) {
        let s = 0;
        for (let t = 0; t < kk; t++) { s += a[i][t] * b[t][j]; seq.push({ i, j, t, term: k.fmt(a[i][t]) + "·" + k.fmt(b[t][j]), prod: a[i][t] * b[t][j], partial: s, last: t === kk - 1 }); }
      }
      return seq;
    }
    let seq = plan();
    // When this is the attention Q·Kᵀ stage (A = X, B = Xᵀ), publish the SCALED scores S = Q·Kᵀ/√d so the
    // downstream softmax card softmaxes the SAME logits the attention card does (it scales by 1/√d too) —
    // keeping the "same numbers flow through every card" promise true. The card itself still shows raw A·B.
    const scale = io.consumes && kk > 0 ? 1 / Math.sqrt(kk) : 1;
    function publish() {
      if (!io.produces) return;
      const C0 = k.matmul(A.get(), B.get());
      bus.set(io.produces, scale === 1 ? C0 : C0.map((r) => r.map((v) => v * scale)));
    }

    function render(step) {
      A.highlight(all(m, kk), false); B.highlight(all(kk, n), false);
      k.clear(log); log.appendChild(k.el("div", { class: "steps-h" }, ["C = A · B   (" + step + " / " + seq.length + " terms)"]));
      for (let s = 0; s < step; s++) {
        const st = seq[s];
        if (st.last) C.set(setCell(C.get(), st.i, st.j, st.partial));
        const lead = "C[" + st.i + "][" + st.j + "] += " + st.term + " = " + k.fmt(st.prod);
        k.logLine(log, st.last ? lead + "  ⟹  total <b>" + k.fmt(st.partial) + "</b>" : lead + "  · running " + k.fmt(st.partial), s === step - 1);
      }
      if (step > 0 && step <= seq.length) {
        const st = seq[step - 1];
        A.highlight([[st.i, st.t]]); B.highlight([[st.t, st.j]]);
        C.cell(st.i, st.j).classList.add("hot");
      }
    }
    const ctrl = k.stepper(seq, render);

    function rebuild() { seq = plan(); C.set(zero(m, n)); ctrl.reset(); publish(); }
    [A.node, B.node].forEach((g) => g.addEventListener("input", rebuild));
    if (io.consumes) bus.on(io.consumes, (X) => { if (X && X.length === m && (X[0] || []).length === kk) { A.set(X); B.set(k.transpose(X)); rebuild(); } });
    publish();

    root.appendChild(k.row([col("A (" + m + "×" + kk + ")", A.node), times(), col("B (" + kk + "×" + n + ")", B.node), eq(), col("C (" + m + "×" + n + ")", C.node)], "mats"));
    root.appendChild(k.row([k.button("▸ Step", () => ctrl.next()), k.button("▶ Play", () => ctrl.play(640)), k.button("↺ Reset", () => ctrl.reset(), "ghost")], "ctrls"));
    root.appendChild(log);
    return k.player(ctrl, 460);                                  // many more (term-level) steps → quicker tick, but unhurried
  }
  function col(t, node) { return k.el("div", { class: "matcol" }, [k.el("div", { class: "matlbl" }, [t]), node]); }
  function times() { return k.el("div", { class: "op" }, ["×"]); }
  function eq() { return k.el("div", { class: "op" }, ["="]); }
  function all(r, c) { const o = []; for (let i = 0; i < r; i++) for (let j = 0; j < c; j++) o.push([i, j]); return o; }
  function zero(r, c) { return Array.from({ length: r }, () => Array(c).fill(0)); }
  function setCell(M, i, j, v) { M[i][j] = v; return M; }
  MANIMO.registry["matmul"] = { title: "Matrix Multiply", mount };
})();