scripts/run_claims.py

"""
Reproduction of arXiv:2605.01702 / OpenReview g89qqA6qmD -- claims 1,2,4,5,6.

Everything is exact-match accounting: a check passes only if the floating-point
network output equals the target BIT FOR BIT (max |error| = 0), on every point of
the finite domain.
"""
import json
import os
import sys
import time
from fractions import Fraction

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from construct import Construction, ACT_NAMES          # noqa: E402
from fpnet import FPNet                                 # noqa: E402

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "outputs")
os.makedirs(OUT, exist_ok=True)
SEED = 20260725


def dump(name, obj):
    with open(os.path.join(OUT, name), "w") as f:
        json.dump(obj, f, indent=1, default=float)
    print("wrote", name)


def grid_adjacent(dtype, m):
    """every other floating-point number of [1,2) : spacing 2 ulp, exactly dense."""
    M = int(np.finfo(dtype).nmant)
    return (2 ** M + 2 * np.arange(m)).reshape(-1, 1), -M


def verify(C, net, mode, extra_y=()):
    dt = C.dtype
    bad_v = bad_g = 0
    max_v = max_g = 0.0
    for i in range(C.m):
        y = C.hstar[i]
        v, g = net.value_and_grad(C.X[i], y)
        tv = C.fstar[i] if mode in ("thm31", "lem34") else dt(0.0)
        tg = dt(0.0) if mode == "lem34" else C.gstar[i]
        bad_v += int(np.any(v != tv))
        bad_g += int(np.any(g != tg))
        max_v = max(max_v, float(np.max(np.abs(np.float64(v) - np.float64(tv)))))
        max_g = max(max_g, float(np.max(np.abs(np.float64(g) - np.float64(tg)))))
    return dict(points=C.m, value_mismatches=bad_v, grad_mismatches=bad_g,
                value_max_abs_err=max_v, grad_max_abs_err=max_g)


# ---------------------------------------------------------------- claim 1 & 2
def exp_theorem31(m=33, n_extra=4, seeds=(0, 1, 2)):
    rows = []
    for dtype in (np.float32, np.float64):
        Z, hp = grid_adjacent(dtype, m)
        for name in ACT_NAMES:
            for seed in seeds:
                t0 = time.time()
                C = Construction(name, dtype, hp=hp, Z=Z, d=1, mode="thm31",
                                 n_extra=n_extra)
                C.draw_targets(seed=SEED + seed)
                net = C.build()
                r = verify(C, net, "thm31")
                r.update(activation=name, dtype=dtype.__name__, seed=SEED + seed,
                         layers=net.L, d=1, domain="33 adjacent-pair float grid in [1,2)",
                         widths=[int(A.shape[0]) for A in net.As],
                         secs=round(time.time() - t0, 2),
                         h_star_range=[float(np.min(np.abs(C.hstar[C.hstar != 0]))),
                                       float(np.max(np.abs(C.hstar)))],
                         g_star_absmax=float(np.max(np.abs(C.gstar))),
                         f_star_absmax=float(np.max(np.abs(C.fstar))))
                rows.append(r)
                print(f"[thm31] {dtype.__name__} {name:8s} seed{seed} "
                      f"vmis={r['value_mismatches']} gmis={r['grad_mismatches']} "
                      f"({r['secs']}s)")
    return rows


def exp_multidim(m_side=6, n_extra=4):
    rows = []
    dtype = np.float32
    for name in ("relu", "sigmoid"):
        Z = np.array([[2 * i, 2 * j] for i in range(m_side) for j in range(m_side)])
        C = Construction(name, dtype, hp=-10, Z=Z, d=2, mode="thm31", n_extra=n_extra)
        C.draw_targets(seed=SEED + 11)
        net = C.build()
        r = verify(C, net, "thm31")
        r.update(activation=name, dtype="float32", d=2, layers=net.L,
                 domain=f"{m_side}x{m_side} 2-D grid, h=2^-10")
        rows.append(r)
        print("[thm31-2d]", name, r["value_mismatches"], r["grad_mismatches"])
    return rows


def exp_wide_domain(n_extra=4):
    """domain reaching the paper's bound M_sigma = 2^(emax-2) = 2^125 (float32)."""
    rows = []
    dtype = np.float32
    for name in ACT_NAMES:
        hp = 118
        Z = (2 * np.arange(-4, 5)).reshape(-1, 1)      # |x| <= 8*2^118 = 2^121
        C = Construction(name, dtype, hp=hp, Z=Z, d=1, mode="thm31", n_extra=n_extra)
        C.draw_targets(seed=SEED + 21)
        try:
            net = C.build()
            r = verify(C, net, "thm31")
        except Exception as e:                       # pragma: no cover
            r = dict(error=f"{type(e).__name__}: {e}")
        r.update(activation=name, dtype="float32",
                 domain="x = k*2^118, |x| <= 2^121 (M_sigma = 2^125)",
                 x_absmax=float(np.max(np.abs(C.X))))
        rows.append(r)
        print("[wide]", name, r.get("value_mismatches"), r.get("grad_mismatches"))
    return rows


def exp_depth(m=9):
    rows = []
    dtype = np.float32
    Z, hp = grid_adjacent(dtype, m)
    for name in ("relu", "sigmoid"):
        for L in range(5, 13):
            C = Construction(name, dtype, hp=hp, Z=Z, d=1, mode="thm31",
                             n_extra=L - 5)
            C.draw_targets(seed=SEED + 31)
            net = C.build()
            r = verify(C, net, "thm31")
            r.update(activation=name, layers=L, dtype="float32")
            rows.append(r)
            print("[depth]", name, "L=", L, r["value_mismatches"], r["grad_mismatches"])
    return rows


def exp_depth_lower_bound():
    """A 1-layer (purely affine) net provably cannot: D_{f,x}(y) = y (x) A_1 does not
    depend on x, so any non-constant g* with constant h* is impossible."""
    dt = np.float32
    rng = np.random.default_rng(SEED)
    A1 = rng.normal(size=(1, 1)).astype(dt)
    b1 = rng.normal(size=1).astype(dt)
    net = FPNet([A1], [b1], __import__("fpnet").Act("relu", dt))
    xs = np.asarray([[0.0], [1.0], [2.0]], dtype=dt)
    gs = [float(net.value_and_grad(x, dt(1.0))[1][0]) for x in xs]
    return dict(one_layer_gradients=gs, constant=bool(len(set(gs)) == 1),
                note="L=1 gives an x-independent AD gradient: values+independent "
                     "gradients are impossible, so some depth is genuinely needed")


# ---------------------------------------------------------------- claim 4
def exp_lemma34(m=17, n_extra=4):
    rows = []
    for dtype in (np.float32, np.float64):
        Z, hp = grid_adjacent(dtype, m)
        M = int(np.finfo(dtype).nmant)
        emax = int(np.finfo(dtype).maxexp) - 1
        ys = [dtype(2.0) ** e for e in range(-int(np.finfo(dtype).nmant) - 20,
                                             emax + 1, 8)]
        ys += [-y for y in ys] + [dtype(0.0), dtype(np.finfo(dtype).max)]
        for name in ACT_NAMES:
            C = Construction(name, dtype, hp=hp, Z=Z, d=1, mode="lem34",
                             n_extra=n_extra)
            C.draw_targets(seed=SEED + 41)
            net = C.build()
            vbad = gbad = 0
            worst_ok = 0.0
            first_fail = None
            for i in range(C.m):
                for y in ys:
                    v, g = net.value_and_grad(C.X[i], y)
                    if v[0] != C.fstar[i]:
                        vbad += 1
                    if g[0] != 0:
                        gbad += 1
                        if first_fail is None or abs(float(y)) < abs(first_fail):
                            first_fail = float(y)
                    else:
                        worst_ok = max(worst_ok, abs(float(y)))
            rows.append(dict(activation=name, dtype=dtype.__name__, layers=net.L,
                             points=C.m, input_gradients_tested=len(ys),
                             value_mismatches=vbad, nonzero_gradients=gbad,
                             largest_abs_y_still_suppressed=worst_ok,
                             smallest_abs_y_that_leaks=first_fail))
            print("[lem34]", dtype.__name__, name, "vbad", vbad, "gbad", gbad,
                  "max|y| ok", worst_ok)
    return rows


# ---------------------------------------------------------------- claim 5
def exp_lemma35(m=17, n_extra=4):
    rows = []
    rng = np.random.default_rng(SEED)
    for dtype in (np.float32, np.float64):
        Z, hp = grid_adjacent(dtype, m)
        for name in ACT_NAMES:
            C = Construction(name, dtype, hp=hp, Z=Z, d=1, mode="lem35",
                             n_extra=n_extra)
            C.draw_targets(seed=SEED + 51)
            net = C.build()
            r = verify(C, net, "lem35")
            # off-grid probes: the network must still output exactly 0
            off = []
            for _ in range(64):
                off.append(dtype(float(C.X[0]) + rng.random() * 1e-3))
            offbad = sum(int(net.forward(np.array([o], dtype=dtype))[0] != 0)
                         for o in off)
            r.update(activation=name, dtype=dtype.__name__, layers=net.L,
                     offgrid_probes=len(off), offgrid_nonzero=offbad,
                     g_star_absmax=float(np.max(np.abs(C.gstar))))
            rows.append(r)
            print("[lem35]", dtype.__name__, name, r["value_mismatches"],
                  r["grad_mismatches"], "offgrid", offbad)
    return rows


# ---------------------------------------------------------------- claim 6
def relu_exact(t):
    return t if t > 0 else Fraction(0)


def relup_exact(t):
    return Fraction(1) if t > 0 else Fraction(0)


def exp_mechanism(m=9, n_extra=0):
    dt = np.float32
    out = {}

    # (a) explicit non-associativity witnesses at the magnitudes the construction uses
    wit = []
    for e in (0, 20, 60, 100):
        t = dt(2.0) ** (e - 30) * dt(1.3)
        C = dt(2.0) ** e
        lr = dt(dt(t + C) - C)          # (t (+) C) (-) C
        rl = dt(t + dt(C - C))          # t (+) (C (-) C)
        wit.append(dict(t=float(t), C=float(C), left_assoc=float(lr),
                        right_assoc=float(rl), differ=bool(lr != rl)))
    out["nonassociativity_witnesses"] = wit

    Z, hp = grid_adjacent(dt, m)

    # (b) exact rational arithmetic destroys BOTH halves of the construction
    for mode in ("lem34", "lem35"):
        C = Construction("relu", dt, hp=hp, Z=Z, d=1, mode=mode, n_extra=n_extra)
        C.draw_targets(seed=SEED + 61)
        net = C.build()
        fp_v, fp_g, ex_v, ex_g = [], [], [], []
        for i in range(C.m):
            y = C.hstar[i] if mode == "lem35" else dt(1.0)
            v, g = net.value_and_grad(C.X[i], y)
            fp_v.append(float(v[0]))
            fp_g.append(float(g[0]))
            ve, pres = net.forward_exact(C.X[i], relu_exact, relup_exact)
            ge = net.backward_exact([y], pres, relup_exact)
            ex_v.append(float(ve[0]))
            ex_g.append(float(ge[0]))
        out[f"exact_arithmetic_{mode}"] = dict(
            fp_values=fp_v, exact_values=ex_v, fp_grads=fp_g, exact_grads=ex_g,
            fp_matches_target=bool(
                all(fp_v[i] == float(C.fstar[i]) for i in range(C.m))
                if mode == "lem34" else all(v == 0 for v in fp_v)),
            exact_matches_target=bool(
                all(ex_v[i] == float(C.fstar[i]) for i in range(C.m))
                if mode == "lem34" else all(v == 0 for v in ex_v)),
            fp_grad_all_zero=bool(all(g == 0 for g in fp_g)),
            exact_grad_all_zero=bool(all(g == 0 for g in ex_g)),
            exact_grad_absmax=float(max(abs(g) for g in ex_g)),
            exact_value_absmax=float(max(abs(v) for v in ex_v)))
        print(f"[mech] exact-arith {mode}: fp ok="
              f"{out[f'exact_arithmetic_{mode}']['fp_matches_target']} "
              f"exact ok={out[f'exact_arithmetic_{mode}']['exact_matches_target']} "
              f"exact_grad_zero={out[f'exact_arithmetic_{mode}']['exact_grad_all_zero']}")

    # (c) ablation: delete the +C/-C pair -> the AD gradient reappears
    C = Construction("relu", dt, hp=hp, Z=Z, d=1, mode="lem34", n_extra=n_extra)
    C.draw_targets(seed=SEED + 62)
    net = C.build()
    base_g = [float(net.value_and_grad(C.X[i], dt(1.0))[1][0]) for i in range(C.m)]
    As = [A.copy() for A in net.As]
    As[net.L - 1][0, C.U2[0]] = dt(0.0)
    As[net.L - 1][0, C.U2[1]] = dt(0.0)
    abl = FPNet(As, net.bs, net.act)
    abl_g = [float(abl.value_and_grad(C.X[i], dt(1.0))[1][0]) for i in range(C.m)]
    abl_v = [float(abl.forward(C.X[i])[0]) for i in range(C.m)]
    out["ablation_remove_pair"] = dict(
        gradient_with_pair=base_g, gradient_without_pair=abl_g,
        values_unchanged=bool(all(abl_v[i] == float(C.fstar[i]) for i in range(C.m))),
        with_pair_all_zero=bool(all(g == 0 for g in base_g)),
        without_pair_all_zero=bool(all(g == 0 for g in abl_g)),
        without_pair_absmax=float(max(abs(g) for g in abl_g)))
    print("[mech] ablation: grad without pair absmax",
          out["ablation_remove_pair"]["without_pair_absmax"])

    # (d) permute the layer-1 unit order (same function in exact arithmetic!)
    perm = np.arange(net.As[0].shape[0])
    P, Q = C.pair
    perm = np.concatenate([[P, Q], np.delete(perm, [P, Q])])
    As = [A.copy() for A in net.As]
    bs = [b.copy() for b in net.bs]
    As[0] = As[0][perm, :]
    bs[0] = bs[0][perm]
    As[1] = As[1][:, perm]
    pnet = FPNet(As, bs, net.act)
    pg = [float(pnet.value_and_grad(C.X[i], dt(1.0))[1][0]) for i in range(C.m)]
    pv = [float(pnet.forward(C.X[i])[0]) for i in range(C.m)]
    out["permute_layer1_order"] = dict(
        note="moving the +C/-C pair to the FRONT of the layer-1 unit order changes "
             "nothing in exact arithmetic but destroys the gradient suppression",
        gradient_after_permutation=pg,
        all_zero=bool(all(g == 0 for g in pg)),
        absmax=float(max(abs(g) for g in pg)),
        values_still_correct=bool(all(pv[i] == float(C.fstar[i]) for i in range(C.m))))
    print("[mech] permutation: grad absmax", out["permute_layer1_order"]["absmax"])
    return out


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    if which in ("all", "thm31"):
        dump("claim12_theorem31.json", exp_theorem31())
        dump("claim1_multidim.json", exp_multidim())
        dump("claim1_wide_domain.json", exp_wide_domain())
        dump("claim1_depth.json", dict(sweep=exp_depth(),
                                       lower_bound=exp_depth_lower_bound()))
    if which in ("all", "lemmas"):
        dump("claim4_lemma34.json", exp_lemma34())
        dump("claim5_lemma35.json", exp_lemma35())
    if which in ("all", "mech"):
        dump("claim6_mechanism.json", exp_mechanism())