checkmlp.py

"""Check the neural model in the browser against the Python reference.

Bit-identical output is not available here and pretending otherwise would be
dishonest: training runs thousands of floating-point operations deep, and
tanh, exp and log differ in their last bits between engines. So this checks
three things instead, which together are a stronger claim than "the numbers
matched once":

  1. Each implementation's analytic gradients match its own finite differences.
     A wrong derivative — the actual risk in hand-written backprop — fails here
     immediately, in whichever language it was written in.
  2. The two agree to a tolerance on initialisation and the forward pass, so
     they are the same model and not merely two models that both train.
  3. Training reduces loss in both, from the same start, by a similar amount.
"""

import json
import math
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 mlp  # noqa: E402
from corpora import BERRIES, CODE, JAPANESE  # noqa: E402

CORPORA = [("berries", BERRIES), ("code", CODE), ("japanese", JAPANESE)]

INIT_TOL = 1e-9      # same PRNG, same order: only transcendentals can differ
FORWARD_TOL = 1e-9
GRAD_TOL = 1e-4      # finite differences are not exact; this is the usual bar


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, "learn", "mlp.js"), encoding="utf-8").read())
    return ctx


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

    for name, text in CORPORA:
        vocab = mlp.build_vocab(text)
        xs, ys = mlp.make_examples(text, vocab)
        model = mlp.Model(vocab, seed=1)

        ctx.set("_t", text)
        ctx.eval("var vocab = MLP.buildVocab(_t);"
                 "var ex = MLP.makeExamples(_t, vocab);"
                 "var m = new MLP.Model(vocab, 1);")

        js_vocab = json.loads(ctx.eval("JSON.stringify(vocab)"))
        if js_vocab != vocab:
            failures.append(f"{name}: vocabulary differs "
                            f"({len(js_vocab)} vs {len(vocab)})")
            continue
        js_xs = json.loads(ctx.eval("JSON.stringify(ex.xs)"))
        js_ys = json.loads(ctx.eval("JSON.stringify(ex.ys)"))
        if js_xs != xs or js_ys != ys:
            failures.append(f"{name}: training examples differ")
            continue

        # 2a. same initial weights
        js_C = json.loads(ctx.eval("JSON.stringify(m.C)"))
        worst_init = 0.0
        for a_row, b_row in zip(model.C, js_C):
            for a, b in zip(a_row, b_row):
                worst_init = max(worst_init, abs(a - b))
        js_W2 = json.loads(ctx.eval("JSON.stringify(m.W2)"))
        for a_row, b_row in zip(model.W2, js_W2):
            for a, b in zip(a_row, b_row):
                worst_init = max(worst_init, abs(a - b))
        if worst_init > INIT_TOL:
            failures.append(f"{name}: initial weights differ by {worst_init:.2e}")

        # 2b. same forward pass
        worst_fwd = 0.0
        for probe in xs[:25]:
            py_probs, _ = model.forward(probe)
            ctx.set("_c", json.dumps(probe))
            js_probs = json.loads(ctx.eval(
                "JSON.stringify(m.forward(JSON.parse(_c)).probs)"))
            for a, b in zip(py_probs, js_probs):
                worst_fwd = max(worst_fwd, abs(a - b))
        if worst_fwd > FORWARD_TOL:
            failures.append(f"{name}: forward pass differs by {worst_fwd:.2e}")

        # 1. each side checks its own gradients
        py_worst, py_fail = mlp.verify_gradients(model, xs[:10], ys[:10])
        ctx.eval("var sub = {xs: ex.xs.slice(0,10), ys: ex.ys.slice(0,10)};"
                 "var gc = m.gradcheck(sub.xs, sub.ys, {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})")

        # 3. both actually learn
        py_before = model.loss(xs[:120], ys[:120])
        model.train(xs, ys, 120, 0.5, seed=7)
        py_after = model.loss(xs[:120], ys[:120])

        ctx.eval("var before = m.loss(ex.xs.slice(0,120), ex.ys.slice(0,120));"
                 "var rng = new MLP.Rng(7);"
                 "m.train(ex.xs, ex.ys, 120, 0.5, 32, rng);"
                 "var after = m.loss(ex.xs.slice(0,120), ex.ys.slice(0,120));")
        js_before = float(ctx.eval("before"))
        js_after = float(ctx.eval("after"))

        if not (py_after < py_before and js_after < js_before):
            failures.append(f"{name}: loss did not fall "
                            f"(py {py_before:.3f}->{py_after:.3f}, "
                            f"js {js_before:.3f}->{js_after:.3f})")

        print(f"PASS  {name:<10} grad max rel err py {py_worst:.1e} / "
              f"js {js_worst:.1e}   init {worst_init:.1e}   "
              f"fwd {worst_fwd:.1e}   loss {py_before:.2f}->{py_after:.2f}")

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


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