checkattn.py

"""Check the attention head in the browser against the Python reference.

Same arrangement as checkmlp.py, and for the same honest reason: training is too
deep in floating point for bit-identical output across engines. So each side
checks its own analytic gradients against finite differences — the real risk in
hand-written backprop through a softmax — and the two are compared on
initialisation, forward pass and attention weights, which are shallow enough to
agree exactly.
"""

import json
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "pylib"))
ROOT = os.path.abspath(os.path.join(HERE, ".."))

import quickjs  # noqa: E402

import attn  # noqa: E402
import task  # noqa: E402
from corpora import BERRIES  # noqa: E402

TOL = 1e-9


def js_context():
    ctx = quickjs.Context()
    ctx.set_memory_limit(1 << 30)
    ctx.set_max_stack_size(1 << 22)
    ctx.eval("var globalThis = this;")
    ctx.eval(open(os.path.join(ROOT, "attention", "attn.js"),
                  encoding="utf-8").read())
    return ctx


def main():
    ctx = js_context()
    failures = []

    copy_text, _ = task.make_corpus(120, seed=1)
    cases = [("copy-task", copy_text, True), ("copy-task-plain", copy_text, False),
             ("berries", BERRIES, True)]

    for name, text, shift in cases:
        vocab = attn.build_vocab(text)
        xs, ys = attn.make_examples(text, vocab)
        model = attn.Model(vocab, seed=1, shift_values=shift)

        ctx.set("_t", text)
        ctx.set("_shift", shift)
        ctx.eval("var vocab = Attn.buildVocab(_t);"
                 "var ex = Attn.makeExamples(_t, vocab);"
                 "var m = new Attn.Model(vocab, 1, {shiftValues: _shift});")

        if json.loads(ctx.eval("JSON.stringify(vocab)")) != vocab:
            failures.append(f"{name}: vocabulary differs")
            continue
        if json.loads(ctx.eval("JSON.stringify(ex.xs)")) != xs:
            failures.append(f"{name}: examples differ")
            continue
        if int(ctx.eval("m.parameterCount()")) != model.parameter_count():
            failures.append(f"{name}: parameter count differs")

        # Initialisation
        worst_init = 0.0
        for tensor in ("C", "P", "Wq", "Wk", "Wv", "Wo"):
            js = json.loads(ctx.eval(f"JSON.stringify(m.{tensor})"))
            py = getattr(model, tensor)
            for a_row, b_row in zip(py, js):
                for a, b in zip(a_row, b_row):
                    worst_init = max(worst_init, abs(a - b))
        if worst_init > TOL:
            failures.append(f"{name}: initial weights differ by {worst_init:.2e}")

        # Forward pass and, importantly, the attention weights themselves —
        # they are what the page shows the reader.
        worst_fwd = worst_attn = 0.0
        for probe in xs[:20]:
            py_probs, py_cache = model.forward(probe)
            ctx.set("_c", json.dumps(probe))
            ctx.eval("var out = m.forward(JSON.parse(_c));")
            js_probs = json.loads(ctx.eval("JSON.stringify(out.probs)"))
            js_w = json.loads(ctx.eval("JSON.stringify(out.w)"))
            for a, b in zip(py_probs, js_probs):
                worst_fwd = max(worst_fwd, abs(a - b))
            for a, b in zip(py_cache["w"], js_w):
                worst_attn = max(worst_attn, abs(a - b))
        if worst_fwd > TOL:
            failures.append(f"{name}: forward differs by {worst_fwd:.2e}")
        if worst_attn > TOL:
            failures.append(f"{name}: attention weights differ by {worst_attn:.2e}")

        # Each side checks its own derivatives.
        py_worst, py_fail = attn.gradcheck(model, xs[:6], ys[:6])
        ctx.eval("var gc = m.gradcheck(ex.xs.slice(0,6), ex.ys.slice(0,6), "
                 "{checks: 40});")
        js_worst = float(ctx.eval("gc.worst"))
        js_fail = int(ctx.eval("gc.failures.length"))
        if py_fail:
            failures.append(f"{name}: python gradients wrong ({len(py_fail)})")
        if js_fail:
            failures.append(f"{name}: browser gradients wrong ({js_fail})")

        print(f"PASS  {name:<16} grad py {py_worst:.1e} / js {js_worst:.1e}   "
              f"init {worst_init:.1e}   fwd {worst_fwd:.1e}   "
              f"attn {worst_attn:.1e}")

    print()
    if failures:
        print(f"{len(failures)} problem(s):")
        for f in failures:
            print(f"  - {f}")
        return 1
    print("Browser attention head and Python reference agree; "
          "both gradient checks pass.")
    return 0


if __name__ == "__main__":
    sys.exit(main())