scripts/construct_old.py

"""
Independent re-derivation of the constructions behind arXiv:2605.01702 (OpenReview g89qqA6qmD).

We do not transcribe the paper's proof.  We build our own floating-point sigma-networks
that satisfy the *statements*:

  Lemma 3.4  :  f = f*   and   D_{f,x}(y) = 0             (values kept, AD gradient killed)
  Lemma 3.5  :  f = 0    and   D_{f,x}(h*(x)) = g*(x)     (values killed, AD gradient arbitrary)
  Theorem 3.1:  f = f*   and   D_{f,x}(h*(x)) = g*(x)     (f = f2 # f1)

for every sigma in {ReLU, ELU, GELU, Swish, Sigmoid, tanh}, in genuine IEEE-754 arithmetic.

Both halves rest on the non-associativity of floating-point summation:

  wipe identity :  (t (+) C) (-) C = 0   whenever  |t| < ulp(C)/2 ,
                   while the same pair leaves an accumulator that is *reset* to 0 alone.

  * put the +C/-C pair inside the AD accumulation grad(x_k) = (+)_r s_r (x) A_1[r,k]
    -> the value path's gradient is annihilated, the forward pass never sees it.
  * put it inside a forward accumulation -> the value is annihilated, AD never sees it.
  * for the saturating activations sigma'=round(sigma_hat') underflows to exactly 0
    where sigma is exactly +-1: values survive, gradients do not.

Layer-1 unit order (this order IS the summation order of the AD accumulation):

    [f1 detector features][f2 detector features][P Q][f2 carriers][K1 K2]
                                                 ^^^ wipes everything to its left
"""

import numpy as np
from fpnet import Act, Builder, FPNet

SAT = {"sigmoid", "tanh"}
ACT_NAMES = ["relu", "elu", "gelu", "swish", "sigmoid", "tanh"]


def acfg(name, dtype, act=None):
    """Operating points of the construction; every entry is verified numerically.

    Saturating sigma: BP is the smallest power of two at which sigma is exactly
    +-1 AND the rounded derivative sigma' has underflowed to exactly 0 -- this is
    the paper's Condition-1 regime |sigma'(gamma)| << |sigma(gamma)|.
    Non-saturating sigma: we work in the exact-identity region (sigma(t)=t, sigma'(t)=1).
    """
    a = act if act is not None else Act(name, dtype)
    c = dict(name=name, sat=(name in SAT), act=a, dtype=dtype)

    def S(t):
        return a.s(np.array([t], dtype=dtype))[0]

    def Sp(t):
        return a.sp(np.array([t], dtype=dtype))[0]

    if name in SAT:
        BP = None
        for e in range(2, 200):
            P = dtype(2.0) ** e
            if (
                S(P) == 1
                and abs(S(-P)) == (0 if name == "sigmoid" else 1)
                and Sp(P) == 0
                and Sp(-P) == 0
            ):
                BP = P
                break
        GP = None
        for e in range(1, 200):
            P = dtype(2.0) ** e
            if S(P) == 1 and Sp(P) != 0:
                GP = P
                break
        assert BP is not None and GP is not None, name
        c["BP"] = BP
        c["ONA"] = dtype(1.0)
        c["gate_on_pre"] = GP
        c["const_pre"] = GP
    else:
        c["BP"] = dtype(2.0) ** 59
        c["ONA"] = dtype(2.0) ** 100
        c["gate_on_pre"] = dtype(2.0) ** 80
        c["const_pre"] = dtype(2.0) ** 60
        for p in (c["ONA"], c["gate_on_pre"], c["const_pre"], c["BP"]):
            assert S(p) == p and Sp(p) == 1, (name, p)
    c["gate_off_pre"] = -c["BP"]
    c["D_on"] = S(c["BP"])
    c["D_off"] = S(-c["BP"])
    assert Sp(-c["BP"]) == 0, name
    c["const_val"] = S(c["const_pre"])
    return c


def _relay(v_on, v_off, p_on, p_off, dtype):
    v_on, v_off = np.float64(v_on), np.float64(v_off)
    w = (np.float64(p_on) - np.float64(p_off)) / (v_on - v_off)
    b = np.float64(p_on) - w * v_on
    return dtype(w), dtype(b)


def _pow2_ceil(v, dtype):
    if v == 0:
        return dtype(1.0)
    e = int(np.ceil(np.log2(float(abs(v)))))
    return dtype(2.0) ** e


class Construction:
    """f = f2 # f1 for one activation / dtype / finite grid domain."""

    def __init__(
        self,
        name,
        dtype,
        hp,
        Z,
        d=1,
        mode="thm31",
        n_carrier=4,
        alpha_exp=-30,
        margin=None,
    ):
        self.name, self.dtype, self.hp, self.d = name, dtype, hp, d
        self.cfg = acfg(name, dtype)
        self.Z = np.asarray(Z, dtype=np.int64).reshape(-1, d)
        self.m = self.Z.shape[0]
        self.mode = mode
        self.nc = n_carrier
        self.alpha = dtype(2.0) ** alpha_exp
        self.margin = int(np.finfo(dtype).nmant + 6) if margin is None else int(margin)
        self.h = dtype(2.0) ** hp
        self.X = np.asarray(self.Z, dtype=np.float64) * float(self.h)
        self.X = self.X.astype(dtype)
        self.use_f1 = mode in ("thm31", "lem34")
        self.use_f2 = mode in ("thm31", "lem35")

    # ---------------- targets ------------------------------------------
    def draw_targets(
        self, seed=0, fexp=(-8, 8), gexp=(-8, 8), hexp=(-4, 4), zero_frac=0.15
    ):
        rng = np.random.default_rng(seed)
        dt, m, d = self.dtype, self.m, self.d

        def rf(lo, hi, size):
            e = rng.integers(lo, hi + 1, size=size)
            man = 1.0 + rng.integers(0, 2**20, size=size) / 2.0**20
            s = rng.choice([-1.0, 1.0], size=size)
            return (s * man * 2.0**e).astype(dt)

        self.fstar = rf(*fexp, m)
        self.hstar = rf(*hexp, m)
        zero = rng.random(m) < zero_frac
        self.hstar[zero] = dt(0.0)
        self.gstar = rf(*gexp, (m, d)).astype(dt)
        self.gstar[zero, :] = dt(0.0)
        self.zero_mask = zero
        return self

    def set_targets(self, fstar=None, hstar=None, gstar=None):
        dt = self.dtype
        if fstar is not None:
            self.fstar = np.asarray(fstar, dtype=dt)
        if hstar is not None:
            self.hstar = np.asarray(hstar, dtype=dt)
        if gstar is not None:
            self.gstar = np.asarray(gstar, dtype=dt).reshape(self.m, self.d)
        self.zero_mask = self.hstar == 0
        return self

    # ---------------- skeleton -----------------------------------------
    def _skeleton(self):
        dt, cfg, d, m = self.dtype, self.cfg, self.d, self.m
        sat = cfg["sat"]
        S60 = dt(2.0) ** 60
        W1 = dt(2.0) ** (60 - self.hp)
        bld = Builder(d, 5)
        i1 = list(range(m)) if self.use_f1 else []
        i2 = list(range(m)) if self.use_f2 else []
        self.i1, self.i2 = i1, i2

        # ---- layer 1 ----
        self.hat1 = {i: self._features(bld, i, W1, S60) for i in i1}
        self.hat2 = {i: self._features(bld, i, W1, S60) for i in i2}
        self.pair = None
        if (not sat) and (self.use_f1 or self.use_f2):
            P = bld.add(0, {k: 1.0 for k in range(d)}, float(dt(2.0) ** 100))
            Q = bld.add(0, {k: 1.0 for k in range(d)}, float(dt(2.0) ** 100))
            self.pair = (P, Q)
        self.carr = {}
        for i in i2:
            self.carr[i] = [
                [bld.add(0, {k: 0.0}, float(cfg["const_pre"])) for _ in range(self.nc)]
                for k in range(d)
            ]
        self.k12 = None
        if self.use_f2:
            self.k12 = (
                bld.add(0, {}, float(cfg["const_pre"])),
                bld.add(0, {}, float(cfg["const_pre"])),
            )

        # ---- layer 2 ----
        self.sink = None
        if self.pair is not None:
            P, Q = self.pair
            self.sink = (
                bld.add(1, {P: 1.0, Q: -1.0}, float(cfg["const_pre"])),
                bld.add(1, {}, float(cfg["const_pre"])),
            )
        self.det1 = {i: self._detector(bld, self.hat1[i], S60) for i in i1}
        self.det2 = {i: self._detector(bld, self.hat2[i], S60) for i in i2}
        self.wipe = {}
        for i in i2:
            ins = {}
            for k in range(d):
                for c in self.carr[i][k]:
                    ins[c] = float(self.alpha)
            ins[self.k12[0]] = 0.0
            ins[self.k12[1]] = 0.0
            self.wipe[i] = bld.add(1, ins, float(cfg["const_pre"]))

        # ---- layer 3 ----
        self.T = None
        if self.sink is not None:
            s1, s2 = self.sink
            self.T = (
                bld.add(2, {s1: 1.0, s2: -1.0}, float(cfg["const_pre"])),
                bld.add(2, {}, float(cfg["const_pre"])),
            )
        self.V = {}
        for i in i1:
            w, b = self._value_relay()
            self.V[i] = bld.add(2, {self.det1[i]: float(w)}, float(b))
        self.gate = {}
        for i in i2:
            self.gate[i] = bld.add(2, {self.det2[i]: 0.0, self.wipe[i]: 0.0}, 0.0)

        # ---- layer 4 ----
        self.U = None
        if self.T is not None:
            t1, t2 = self.T
            self.U = (
                bld.add(3, {t1: 1.0, t2: -1.0}, float(cfg["const_pre"])),
                bld.add(3, {}, float(cfg["const_pre"])),
            )
        self.R = {i: bld.add(3, {self.gate[i]: 0.0}, 0.0) for i in i2}
        self.kout = None
        if self.use_f2:
            self.kout = (
                bld.add(3, {}, float(cfg["const_pre"])),
                bld.add(3, {}, float(cfg["const_pre"])),
            )
        self.A = {i: bld.add(3, {self.V[i]: 0.0}, 0.0) for i in i1}

        # ---- layer 5 (output) ----
        ins = {}
        if self.U is not None:
            ins[self.U[0]] = 1.0
            ins[self.U[1]] = -1.0
        for i in i2:
            ins[self.R[i]] = 0.0
        if self.kout is not None:
            ins[self.kout[0]] = 0.0
            ins[self.kout[1]] = 0.0
        for i in i1:
            ins[self.A[i]] = 0.0
        bld.add(4, ins, 0.0)
        self.bld = bld
        return bld

    def _features(self, bld, i, W1, S60):
        out = []
        for k in range(self.d):
            nz = int(self.Z[i, k])
            offs = [0.5, -0.5] if self.cfg["sat"] else [1.0, 0.0, -1.0]
            out.append(
                [
                    bld.add(0, {k: float(W1)}, float((np.float64(o) - nz) * float(S60)))
                    for o in offs
                ]
            )
        return out

    def _detector(self, bld, hats, S60):
        """detector pre-activation is +BP at x=z_i and <= -BP elsewhere."""
        ins = {}
        BP = float(self.cfg["BP"])
        if self.cfg["sat"]:
            # sigmoid: u1-u2 in {0,1};  tanh: u1-u2 in {0,2}
            sc = 2.0 * BP if self.name == "sigmoid" else BP
            for k in range(self.d):
                ins[hats[k][0]] = sc
                ins[hats[k][1]] = -sc
            theta = -(2.0 * self.d - 1.0) * BP
        else:
            for k in range(self.d):
                ins[hats[k][0]] = 1.0
                ins[hats[k][1]] = -2.0
                ins[hats[k][2]] = 1.0
            theta = -(float(self.d) - 0.5) * float(S60)
        return bld.add(1, ins, theta)

    def _value_relay(self):
        dt = self.dtype
        return {
            "sigmoid": (dt(2.0) * dt(self.cfg["BP"]), -dt(self.cfg["BP"])),
            "tanh": (dt(self.cfg["BP"]) / dt(2.0), dt(self.cfg["BP"]) / dt(2.0)),
            "elu": (dt(2.0) ** 41, dt(2.0) ** 41),
        }.get(self.name, (dt(2.0) ** 41, dt(0.0)))

    # ---------------- calibration ---------------------------------------
    def build(self, verbose=False):
        dt, cfg, d, m = self.dtype, self.cfg, self.d, self.m
        bld = self._skeleton()
        act = cfg["act"]
        ONA = cfg["ONA"]

        # --- f1 output weights: w5 = f*/ONA (exact: ONA is a power of two) ---
        for i in self.i1:
            aw, ab = {
                "sigmoid": (dt(2.0) * dt(cfg["BP"]), -dt(cfg["BP"])),
                "tanh": (dt(cfg["BP"]), dt(0.0)),
            }.get(self.name, (dt(1.0), dt(0.0)))
            bld.set_weight(3, self.A[i], self.V[i], float(aw))
            bld.set_bias(3, self.A[i], float(ab))
            bld.set_weight(4, 0, self.A[i], float(dt(self.fstar[i]) / ONA))

        # --- f2 gate + R relays ---
        for i in self.i2:
            self._set_gate(bld, i)

        net = bld.build(act)
        if not self.use_f2:
            self._calibrate_pair(bld, net)
            return bld.build(act)

        # --- output-weight scale so that the carrier multiplier ~ max|g*| ---
        for it in range(3):
            net = bld.build(act)
            if self._calibrate_scale(bld, net):
                break
        # --- carriers (iterative, exact) ---
        for it in range(8):
            net = bld.build(act)
            ok = self._calibrate_carriers(bld, net)
            if ok:
                break
        # --- forward wipe constants ---
        net = bld.build(act)
        self._calibrate_wipes(bld, net)
        net = bld.build(act)
        self._calibrate_pair(bld, net)
        # carriers again (pair changes nothing upstream, but be safe)
        for it in range(4):
            net = bld.build(act)
            if self._calibrate_carriers(bld, net):
                break
        net = bld.build(act)
        self._calibrate_wipes(bld, net)
        return bld.build(act)

    def _set_gate(self, bld, i):
        dt, cfg = self.dtype, self.cfg
        aW = dt(2.0) ** -20
        wg, bg = _relay(
            cfg["D_on"], cfg["D_off"], cfg["gate_on_pre"], cfg["gate_off_pre"], dt
        )
        bld.set_weight(2, self.gate[i], self.det2[i], float(wg))
        bld.set_weight(2, self.gate[i], self.wipe[i], float(aW))
        # the wipe unit is a constant sigma(const_pre); subtract its contribution
        bld.set_bias(2, self.gate[i], float(dt(bg) - dt(aW) * dt(cfg["const_val"])))
        # R : gate value -> {ident, 0}
        g_on = cfg["act"].s(np.array([cfg["gate_on_pre"]], dtype=dt))[0]
        g_off = cfg["act"].s(np.array([cfg["gate_off_pre"]], dtype=dt))[0]
        p_off = dt(0.0) if self.name != "sigmoid" else -(dt(2.0) ** 59)
        wr, br = _relay(g_on, g_off, cfg["const_pre"], p_off, dt)
        bld.set_weight(3, self.R[i], self.gate[i], float(wr))
        bld.set_bias(3, self.R[i], float(br))
        bld.set_weight(4, 0, self.R[i], 1.0)

    # gradient at the pre-activation of a layer-1 unit
    def _grad_l1(self, net, x, gin):
        net.forward(x)
        _, tr = net.backward(np.array([gin], dtype=self.dtype), debug=True)
        return tr[("post_act", 0)]

    def _calibrate_scale(self, bld, net):
        """choose the layer-5 weight of each R_i so the AD multiplier arriving at
        the carriers is ~ max_k |g*(z_i)_k| (keeps carrier weights O(1))."""
        dt = self.dtype
        stable = True
        for i in self.i2:
            hs = self.hstar[i]
            if hs == 0:
                continue
            gl1 = self._grad_l1(net, self.X[i], hs)
            mu = abs(dt(gl1[self.carr[i][0][0]]))
            desired = _pow2_ceil(np.max(np.abs(self.gstar[i])), dt)
            if mu == 0:
                stable = False
                continue
            e = int(round(np.log2(float(desired) / float(mu))))
            if e != 0:
                stable = False
                cur = dt(net.As[4][0, self.R[i]])
                bld.set_weight(4, 0, self.R[i], float(cur * dt(2.0) ** e))
        return stable

    def _calibrate_carriers(self, bld, net):
        dt, d = self.dtype, self.d
        done = True
        for i in self.i2:
            x, hs = self.X[i], self.hstar[i]
            if hs == 0:
                for k in range(d):
                    for c in self.carr[i][k]:
                        bld.set_weight(0, c, k, 0.0)
                continue
            gl1 = self._grad_l1(net, x, hs)
            for k in range(d):
                mult = [dt(gl1[c]) for c in self.carr[i][k]]
                target = dt(self.gstar[i, k])
                if all(mu == 0 for mu in mult):
                    done = False
                    continue
                ws, acc = [], dt(0.0)
                for j, mu in enumerate(mult):
                    if mu == 0:
                        ws.append(dt(0.0))
                        continue
                    r = dt(target - acc) if j else target
                    w = dt(np.float64(r) / np.float64(mu))
                    ws.append(w)
                    acc = dt(acc + dt(mu * w))
                if acc != target:
                    done = False
                for j, c in enumerate(self.carr[i][k]):
                    bld.set_weight(0, c, k, float(ws[j]))
        return done

    def _calibrate_wipes(self, bld, net):
        """set the +C/-C constants that annihilate the forward accumulations."""
        dt, cfg = self.dtype, self.cfg
        vconst = dt(cfg["const_val"])
        # (a) layer-2 wipe rows: kill the carrier contributions
        As = net.As
        worst = {i: dt(0.0) for i in self.i2}
        for xi in range(self.m):
            v1 = self._layer_values(net, self.X[xi], 0)
            for i in self.i2:
                s = dt(0.0)
                for k in range(self.d):
                    for c in self.carr[i][k]:
                        s = dt(s + dt(self.alpha * dt(v1[c])))
                worst[i] = max(worst[i], abs(s))
        for i in self.i2:
            C0 = _pow2_ceil(worst[i], dt) * (dt(2.0) ** self.margin)
            w = float(dt(C0) / vconst)
            bld.set_weight(1, self.wipe[i], self.k12[0], w)
            bld.set_weight(1, self.wipe[i], self.k12[1], -w)
        # (b) output row: kill the accumulated f2 contributions
        net2 = bld.build(cfg["act"])
        worst_o = dt(0.0)
        for xi in range(self.m):
            v4 = self._layer_values(net2, self.X[xi], 3)
            s = dt(0.0)
            for i in self.i2:
                s = dt(s + dt(dt(net2.As[4][0, self.R[i]]) * dt(v4[self.R[i]])))
            worst_o = max(worst_o, abs(s))
        C1 = _pow2_ceil(worst_o, dt) * (dt(2.0) ** self.margin)
        w = float(dt(C1) / vconst)
        bld.set_weight(4, 0, self.kout[0], w)
        bld.set_weight(4, 0, self.kout[1], -w)

    def _layer_values(self, net, x, l):
        net.forward(x)
        return net.act.s(net._pres[l])

    def _calibrate_pair(self, bld, net):
        """scale the P/Q pair so that its +C/-C annihilates the detector gradients."""
        if self.pair is None:
            return
        dt = self.dtype
        P, Q = self.pair
        # with the pair's output weight at 0 the pair contributes nothing
        bld.set_weight(4, 0, self.U[0], 0.0)
        bld.set_weight(4, 0, self.U[1], 0.0)
        base = bld.build(net.act)
        worst_t, worst_s = dt(0.0), None
        ratios = []
        for xi in range(self.m):
            for hs in set([dt(self.hstar[xi]), dt(1.0), dt(-1.0)]):
                if hs == 0:
                    continue
                base.forward(self.X[xi])
                g, tr = base.backward(np.array([hs], dtype=dt), debug=True)
                t = np.max(np.abs(g)) if g.size else dt(0.0)
                # gradient reaching P with unit output weight
                bld.set_weight(4, 0, self.U[0], 1.0)
                bld.set_weight(4, 0, self.U[1], -1.0)
                probe = bld.build(net.act)
                probe.forward(self.X[xi])
                _, tr2 = probe.backward(np.array([hs], dtype=dt), debug=True)
                sP = abs(dt(tr2[("post_act", 0)][P]))
                bld.set_weight(4, 0, self.U[0], 0.0)
                bld.set_weight(4, 0, self.U[1], 0.0)
                if sP > 0:
                    ratios.append(float(t) / float(sP) if sP else 0.0)
        need = max(ratios) if ratios else 0.0
        w5p = (
            _pow2_ceil(dt(need), dt) * (dt(2.0) ** self.margin)
            if need > 0
            else dt(2.0) ** self.margin
        )
        bld.set_weight(4, 0, self.U[0], float(w5p))
        bld.set_weight(4, 0, self.U[1], float(-w5p))