scripts/fpnet.py

"""
Floating-point neural network + automatic-differentiation simulator.

Faithful to the formalism of arXiv:2605.01702 (Park, Park, Hwang):

    f(x) = rho_L o sigma o ... o sigma o rho_1 (x),   rho_l(x) = A_l (x) (+) b_l

  * every scalar operation is an IEEE-754 operation in the working format
  * matrix-vector products are accumulated STRICTLY LEFT TO RIGHT
    ( (+)_{j} a_ij (x) x_j ) then (+) b_i
  * the AD ("backward") gradient is

        D_{f,x}(g) = g (x) A_L (x) sigma'(y_{L-1}) (x) ... (x) sigma'(y_1) (x) A_1

    again evaluated strictly left to right; every row-vector-times-matrix
    product is itself accumulated left to right.

Reproduction of ICML 2026 submission g89qqA6qmD.  Author: SabaPivot.
"""

import numpy as np
from fractions import Fraction


# --------------------------------------------------------------------------
# activation table:  sigma = round(sigma_hat), sigma' = round(sigma_hat')
# the real-valued sigma_hat / sigma_hat' are evaluated in float64 and then
# rounded to the working format (this is the paper's  ceil(sigma_hat)  ).
# --------------------------------------------------------------------------
def _sigmoid64(t):
    t = np.asarray(t, dtype=np.float64)
    out = np.empty_like(t)
    pos = t >= 0
    out[pos] = 1.0 / (1.0 + np.exp(-t[pos]))
    e = np.exp(t[~pos])
    out[~pos] = e / (1.0 + e)
    return out


def _phi64(t):  # standard normal cdf
    from scipy.special import ndtr

    return ndtr(np.asarray(t, dtype=np.float64))


def _npdf64(t):
    t = np.asarray(t, dtype=np.float64)
    return (
        np.exp(-0.5 * np.minimum(t * t, 1500.0)) / np.sqrt(2 * np.pi) * (np.abs(t) < 40)
    )


ACTS = {}


def _reg(name, s, sp):
    ACTS[name] = (s, sp)


_reg(
    "relu",
    lambda t: np.maximum(np.asarray(t, np.float64), 0.0),
    lambda t: (np.asarray(t, np.float64) > 0).astype(np.float64),
)
_reg(
    "elu",
    lambda t: np.where(
        np.asarray(t, np.float64) > 0,
        np.asarray(t, np.float64),
        np.expm1(np.minimum(np.asarray(t, np.float64), 0.0)),
    ),
    lambda t: np.where(
        np.asarray(t, np.float64) > 0,
        1.0,
        np.exp(np.minimum(np.asarray(t, np.float64), 0.0)),
    ),
)
_reg(
    "gelu",
    lambda t: np.asarray(t, np.float64) * _phi64(t),
    lambda t: _phi64(t) + np.asarray(t, np.float64) * _npdf64(t),
)
_reg(
    "swish",
    lambda t: np.asarray(t, np.float64) * _sigmoid64(t),
    lambda t: _sigmoid64(t) * (1.0 + np.asarray(t, np.float64) * (1.0 - _sigmoid64(t))),
)


def _dsigmoid64(t):
    """accurate sigmoid' = e^{-|t|} / (1+e^{-|t|})^2  (no cancellation)."""
    e = np.exp(-np.abs(np.asarray(t, np.float64)))
    return e / (1.0 + e) ** 2


def _dtanh64(t):
    """accurate tanh' = sech^2 = 4 e^{-2|t|} / (1+e^{-2|t|})^2."""
    e = np.exp(-2.0 * np.abs(np.asarray(t, np.float64)))
    return 4.0 * e / (1.0 + e) ** 2


_reg("sigmoid", lambda t: _sigmoid64(t), _dsigmoid64)
_reg("tanh", lambda t: np.tanh(np.asarray(t, np.float64)), _dtanh64)

# framework-style derivatives: PyTorch computes sigmoid/tanh backward from the
# *output*, so sigma' underflows to 0 even earlier than the rounded true derivative.
ACTS_FRAMEWORK = dict(ACTS)
ACTS_FRAMEWORK["sigmoid"] = (
    _sigmoid64,
    lambda t: _sigmoid64(t) * (1.0 - _sigmoid64(t)),
)
ACTS_FRAMEWORK["tanh"] = (
    lambda t: np.tanh(np.asarray(t, np.float64)),
    lambda t: 1.0 - np.tanh(np.asarray(t, np.float64)) ** 2,
)

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


class Act:
    def __init__(self, name, dtype):
        self.name = name
        self.dtype = dtype
        self._s, self._sp = ACTS[name]

    def s(self, t):
        return np.asarray(self._s(np.asarray(t, np.float64)), dtype=self.dtype)

    def sp(self, t):
        return np.asarray(self._sp(np.asarray(t, np.float64)), dtype=self.dtype)


# --------------------------------------------------------------------------
# network container
# --------------------------------------------------------------------------
class FPNet:
    """A floating-point sigma-network.  As[l] is (d_l x d_{l-1}); bs[l] is (d_l,)."""

    def __init__(self, As, bs, act):
        self.As = [np.asarray(A, dtype=act.dtype) for A in As]
        self.bs = [np.asarray(b, dtype=act.dtype) for b in bs]
        self.act = act
        self.dtype = act.dtype
        self.L = len(As)

    # ---- forward, strictly left-to-right accumulation -------------------
    def forward(self, x, keep=True):
        dt = self.dtype
        v = np.asarray(x, dtype=dt)
        pres = []
        for l in range(self.L):
            A, b = self.As[l], self.bs[l]
            acc = np.zeros(A.shape[0], dtype=dt)
            for k in range(A.shape[1]):
                col = A[:, k]
                if not col.any():
                    continue  # 0 (x) v = 0 and acc (+) 0 = acc, exactly
                acc = acc + col * v[k]
            acc = acc + b
            pres.append(acc if keep else None)
            v = self.act.s(acc) if l < self.L - 1 else acc
        self._pres = pres
        return v

    # ---- AD gradient, strictly left-to-right ----------------------------
    def backward(self, g_out, pres=None, debug=False):
        dt = self.dtype
        pres = self._pres if pres is None else pres
        g = np.asarray(np.atleast_1d(g_out), dtype=dt)
        trace = {}
        for l in range(self.L - 1, -1, -1):
            A = self.As[l]
            newg = np.zeros(A.shape[1], dtype=dt)
            for j in range(A.shape[0]):
                if g[j] == 0:
                    continue  # exact no-op
                row = A[j, :]
                if not row.any():
                    continue
                newg = newg + g[j] * row
            if debug:
                trace[("pre_act", l)] = newg.copy()  # grad wrt output of layer l-1
            if l > 0:
                newg = newg * self.act.sp(pres[l - 1])
                if debug:
                    trace[("post_act", l - 1)] = (
                        newg.copy()
                    )  # grad wrt pre-act of layer l-1
            g = newg
        return (g, trace) if debug else g

    def value_and_grad(self, x, gin):
        y = self.forward(x)
        g = self.backward(np.atleast_1d(np.asarray(gin, dtype=self.dtype)))
        return y, g

    # ---- exact-real-arithmetic reference (no rounding at all) -----------
    def forward_exact(self, x, sigma_exact, sigmap_exact):
        v = [Fraction(float(t)) for t in np.asarray(x, dtype=self.dtype)]
        pres = []
        for l in range(self.L):
            A, b = self.As[l], self.bs[l]
            acc = []
            for j in range(A.shape[0]):
                s = Fraction(0)
                for k in range(A.shape[1]):
                    if A[j, k] != 0:
                        s += Fraction(float(A[j, k])) * v[k]
                s += Fraction(float(b[j]))
                acc.append(s)
            pres.append(acc)
            v = [sigma_exact(t) for t in acc] if l < self.L - 1 else acc
        return v, pres

    def backward_exact(self, gin, pres, sigmap_exact):
        g = [Fraction(float(t)) for t in np.atleast_1d(gin)]
        for l in range(self.L - 1, -1, -1):
            A = self.As[l]
            newg = [Fraction(0)] * A.shape[1]
            for j in range(A.shape[0]):
                if g[j] == 0:
                    continue
                for k in range(A.shape[1]):
                    if A[j, k] != 0:
                        newg[k] += g[j] * Fraction(float(A[j, k]))
            if l > 0:
                newg = [
                    newg[k] * sigmap_exact(pres[l - 1][k]) for k in range(len(newg))
                ]
            g = newg
        return g


# --------------------------------------------------------------------------
# incremental builder:  units are appended to layers, order is preserved and
# IS SEMANTICALLY MEANINGFUL (it fixes the left-to-right summation order).
# --------------------------------------------------------------------------
class Builder:
    def __init__(self, d_in, L):
        self.d_in = d_in
        self.L = L
        self.units = [
            [] for _ in range(L)
        ]  # units[l] = list of (dict{in_idx: w}, bias)

    def add(self, l, ins, bias=0.0):
        self.units[l].append((dict(ins), float(bias)))
        return len(self.units[l]) - 1

    def set_weight(self, l, j, k, w):
        self.units[l][j][0][k] = float(w)

    def set_bias(self, l, j, b):
        ins, _ = self.units[l][j]
        self.units[l][j] = (ins, float(b))

    def width(self, l):
        return len(self.units[l])

    def build(self, act):
        As, bs = [], []
        prev = self.d_in
        for l in range(self.L):
            n = len(self.units[l])
            A = np.zeros((n, prev), dtype=act.dtype)
            b = np.zeros(n, dtype=act.dtype)
            for j, (ins, bias) in enumerate(self.units[l]):
                for k, w in ins.items():
                    A[j, k] = act.dtype(w)
                b[j] = act.dtype(bias)
            As.append(A)
            bs.append(b)
            prev = n
        return FPNet(As, bs, act)


# --------------------------------------------------------------------------
# format constants
# --------------------------------------------------------------------------
def fmt(dtype):
    fi = np.finfo(dtype)
    return dict(
        M=fi.nmant,
        emax=fi.maxexp - 1,
        emin=fi.minexp,
        Omega=dtype(fi.max),
        omega=dtype(fi.smallest_subnormal),
    )