scripts/run_thm32.py

"""
Claim 3 (Theorem 3.2):  a SINGLE floating-point network whose AD gradient varies
with the loss derivative y = phi'_x(f(x)), subject to the antisymmetry condition
g*(x,-y) = -g*(x,y).

Two independent halves:

 (A) NECESSITY of the antisymmetry condition.  IEEE-754 multiplication and addition
     are exactly sign-symmetric, so  D_{f,x}(-y) = -D_{f,x}(y)  holds identically for
     *every* floating-point network.  Any target g* violating g*(x,-y) = -g*(x,y) is
     therefore unrepresentable.  We verify the underlying identity EXHAUSTIVELY over
     all of binary16 (all 2^16 x 2^16 ordered pairs) and over random binary32 nets.

 (B) EXISTENCE of non-proportional dependence on y.  Under exact real arithmetic the
     chain rule forces  grad(phi o f) = phi'(f(x)) * grad f(x), i.e. exactly
     proportional to y.  We build a floating-point network whose AD gradient is a
     prescribed, arbitrary, antisymmetric, wildly non-proportional function of y on a
     finite set Y of input gradients, using the underflow / subnormal-collapse of the
     backward multiplications.
"""

import json
import os
import sys

import numpy as np

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

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


# ------------------------------------------------------------------ (A)
def exhaustive_sign_symmetry():
    """for all a,b in binary16 :  fl(-a*b) == -fl(a*b)  and  fl(-a-b) == -fl(a+b)."""
    allf = np.arange(2**16, dtype=np.uint16).view(np.float16)
    allf = allf[np.isfinite(allf)]
    n = allf.size
    bad_mul = bad_add = 0
    for i in range(0, n, 512):
        a = allf[i : i + 512][:, None]
        b = allf[None, :]
        m1, m2 = (-a) * b, -(a * b)
        bad_mul += int(np.sum((m1 != m2) & ~(np.isnan(m1) & np.isnan(m2))))
        s1, s2 = (-a) + (-b), -(a + b)
        bad_add += int(np.sum((s1 != s2) & ~(np.isnan(s1) & np.isnan(s2))))
    return dict(
        format="binary16",
        finite_values=int(n),
        ordered_pairs=int(n) ** 2,
        multiplication_violations=bad_mul,
        addition_violations=bad_add,
    )


def network_antisymmetry(trials=200, seed=7):
    """D_{f,x}(-y) == -D_{f,x}(y) on random float32 networks, exactly."""
    rng = np.random.default_rng(seed)
    bad = 0
    checked = 0
    for t in range(trials):
        name = ["relu", "elu", "gelu", "swish", "sigmoid", "tanh"][t % 6]
        act = Act(name, np.float32)
        dims = [2, 6, 5, 1]
        As = [
            rng.normal(scale=3, size=(dims[i + 1], dims[i])).astype(np.float32)
            for i in range(3)
        ]
        bs = [
            rng.normal(scale=3, size=dims[i + 1]).astype(np.float32) for i in range(3)
        ]
        net = FPNet(As, bs, act)
        x = rng.normal(size=2).astype(np.float32)
        for y in rng.normal(scale=10, size=8).astype(np.float32):
            net.forward(x)
            gp = net.backward(np.array([y], dtype=np.float32))
            net.forward(x)
            gm = net.backward(np.array([-y], dtype=np.float32))
            checked += 1
            if not np.array_equal(gp, -gm):
                bad += 1
    return dict(random_networks=trials, gradient_pairs=checked, violations=bad)


# ------------------------------------------------------------------ (B)
def build_y_dependent(dtype=np.float32, n_band=24, seed=11):
    """
    3-layer sigma-network f : F -> F  with  f(x) = 0 for every x (its forward
    accumulation is annihilated by a +C/-C constant pair) whose AD gradient

        D_{f,x}(y)  =  (+)_j  ( ( y (x) p_j ) (x) w ) (x) a_j

    is a prescribed function of y.  p_j w = omega * 2^-e_j, so band j contributes
    exactly 0 while |y| < 2^{e_j} (the backward product UNDERFLOWS) and contributes
    a_j * omega once |y| >= 2^{e_j}.  Under exact real arithmetic every term would be
    exactly proportional to y; under IEEE-754 they are not.
    """
    dt = dtype
    act = Act("relu", dt)
    fi = np.finfo(dt)
    P0 = int(np.log2(float(fi.smallest_subnormal))) + 79        # p_j = 2^(P0-e_j)
    w = dt(2.0) ** -79
    B = dt(2.0) ** 10
    exps = list(range(0, n_band))
    rep = 3                                 # units per band (residual refinement)
    ps, bld = [], Builder(1, 3)
    u, v = [], []
    for e in exps:
        for _ in range(rep):
            ps.append(dt(2.0) ** (P0 - e))
            u.append(bld.add(0, {0: 1.0}, float(B)))
    for j in range(len(u)):
        v.append(bld.add(1, {u[j]: float(w)}, float(B)))
    exps = [e for e in exps for _ in range(rep)]
    kp = bld.add(1, {}, float(B))
    km = bld.add(1, {}, float(B))
    ins = {v[j]: float(ps[j]) for j in range(len(exps))}
    ins[kp] = 0.0
    ins[km] = 0.0
    out = bld.add(2, ins, 0.0)
    return bld, exps, u, v, (kp, km), out, ps, act, dt


def _set_output_wipe(bld, exps, v, kpm, out, ps, act, dt, xs):
    kp, km = kpm
    net = bld.build(act)
    worst = dt(0.0)
    for x in xs:
        net.forward(np.array([x], dtype=dt))
        vv = act.s(net._pres[1])
        acc = dt(0.0)
        for j in range(len(exps)):
            acc = dt(acc + dt(dt(ps[j]) * dt(vv[v[j]])))
        worst = max(worst, abs(acc))
        vk = dt(vv[kp])
    e = int(np.ceil(np.log2(max(float(worst), 1e-300)))) + int(np.finfo(dt).nmant) + 6
    C = dt(2.0) ** e
    bld.set_weight(2, out, kp, float(dt(C) / vk))
    bld.set_weight(2, out, km, float(-dt(C) / vk))
    return bld.build(act)


def fit_y_dependent(targets, seed=11):
    """choose the layer-1 weights a_j so that D_{f,x}(y) hits `targets` exactly."""
    bld, exps, u, v, kpm, out, ps, act, dt = build_y_dependent(seed=seed)
    bands = sorted(set(exps))
    idx_of = {e: [j for j, ee in enumerate(exps) if ee == e] for e in bands}
    ys = [dt(2.0) ** e for e in bands]
    xs = [0.0, 0.25, 0.5, 1.0, -1.0]
    net = bld.build(act)
    x0 = np.array([1.0], dtype=dt)
    for _ in range(40):
        ok = True
        for bi, e in enumerate(bands):
          for idx in idx_of[e]:
            y = ys[bi]
            net.forward(x0)
            _, tr = net.backward(np.array([y], dtype=dt), debug=True)
            mult = dt(tr[("post_act", 0)][u[idx]])
            got = dt(net.backward(np.array([y], dtype=dt))[0])
            tgt = dt(targets[bi])
            if got == tgt:
                continue
            ok = False
            if mult == 0:
                continue
            cur = dt(net.As[0][u[idx], 0])
            new = dt(cur + dt(np.float64(dt(tgt - got)) / np.float64(mult)))
            if not np.isfinite(new):
                continue
            bld.set_weight(0, u[idx], 0, float(new))
            # keep the layer-1 pre-activation strictly positive (sigma' = 1) for every
            # probe x, whatever the weight: bias dominates.
            B_j = dt(2.0) ** int(np.ceil(np.log2(max(4.0 * abs(float(new)), 1024.0))))
            assert np.isfinite(B_j), "bias overflow: target too large for 3 layers"
            bld.set_bias(0, u[idx], float(B_j))
            net = bld.build(act)
        if ok:
            break
    net = _set_output_wipe(bld, exps, v, kpm, out, ps, act, dt, xs)
    rows = []
    for bi, e in enumerate(bands):
        for s_ in (1, -1):
            y = dt(s_) * ys[bi]
            net.forward(x0)
            g = dt(net.backward(np.array([y], dtype=dt))[0])
            tt = dt(s_) * dt(targets[bi])
            rows.append(dict(y=float(y), target=float(tt), got=float(g),
                             exact=bool(g == tt)))
    fwd = [float(net.forward(np.array([xx], dtype=dt))[0]) for xx in xs]
    ratios = [abs(r["got"] / r["y"]) for r in rows if r["y"] != 0 and r["got"] != 0]
    return net, dict(
        bands=len(bands), units=len(u), checks=rows,
        all_exact=bool(all(r["exact"] for r in rows)),
        n_exact=int(sum(r["exact"] for r in rows)),
        forward_values=fwd, forward_all_zero=bool(all(f == 0 for f in fwd)),
        ratio_min=min(ratios) if ratios else None,
        ratio_max=max(ratios) if ratios else None,
        ratio_spread_orders=float(np.log10(max(ratios) / min(ratios)))
        if ratios else None)


def exact_arithmetic_proportionality(net, targets):
    """under exact real arithmetic the same graph gives D(y) exactly proportional to y."""
    from fractions import Fraction

    dt = np.float32

    def relu(t):
        return t if t > 0 else Fraction(0)

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

    x0 = np.array([1.0], dtype=dt)
    _, pres = net.forward_exact(x0, relu, relup)
    out = []
    for e in (0, 5, 10, 15, 20):
        y = 2.0**e
        g = net.backward_exact([y], pres, relup)[0]
        out.append(
            dict(y=y, exact_grad=float(g), exact_grad_over_y=float(g / Fraction(y)))
        )
    ratios = [o["exact_grad_over_y"] for o in out]
    return dict(
        samples=out,
        exact_ratio_constant=bool(len(set(ratios)) == 1),
        exact_ratio=ratios[0] if ratios else None,
    )


if __name__ == "__main__":
    res = {}
    print("exhaustive binary16 sign symmetry ...")
    res["antisymmetry_necessity_exhaustive"] = exhaustive_sign_symmetry()
    print(res["antisymmetry_necessity_exhaustive"])
    res["antisymmetry_random_networks"] = network_antisymmetry()
    print(res["antisymmetry_random_networks"])

    rng = np.random.default_rng(3)
    n = 24
    targets = (
        rng.choice([-1.0, 1.0], size=n)
        * (1.0 + rng.integers(0, 2**20, size=n) / 2.0**20)
        * 2.0 ** rng.integers(-40, -28, size=n)
    ).astype(np.float32)
    net, r = fit_y_dependent(targets)
    res["y_dependent_gradient"] = r
    print(
        "bands",
        r["bands"],
        "exact",
        r["n_exact"],
        "/",
        2 * r["bands"],
        "fwd all zero",
        r["forward_all_zero"],
        "ratio spread (orders of magnitude)",
        r["ratio_spread_orders"],
    )
    res["exact_arithmetic_contrast"] = exact_arithmetic_proportionality(net, targets)
    print(res["exact_arithmetic_contrast"]["exact_ratio_constant"])

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