rob-constructed-v1: exact constructed ReLU circuit, all tiers; weights by construction (disclosed)
Browse files- README.md +56 -0
- circuit.py +698 -0
- manifest.json +7 -0
- model.py +210 -0
README.md
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# rob-constructed-v1: constructed ReLU circuit for exact (a·b) mod p
|
| 2 |
+
|
| 3 |
+
A fixed-weight linear + ReLU network that computes `(a * b) mod p` exactly at
|
| 4 |
+
every scored tier, from 3-bit primes (tier 1) to 2048-bit primes (tier 10).
|
| 5 |
+
|
| 6 |
+
**This is a constructed arithmetic circuit, not a trained model.** Read the
|
| 7 |
+
manifest's `model_description` and `training_description` first; they state the
|
| 8 |
+
provenance plainly. The weights are set by construction (two numeric constants,
|
| 9 |
+
`1` and `2^16`, plus the structural wiring in `circuit.py`). There is no
|
| 10 |
+
training set, no optimizer, and no fitted parameter.
|
| 11 |
+
|
| 12 |
+
It is submitted as the "interesting information to acquire" the launch comment
|
| 13 |
+
invited: a hand-encoded algorithm that meets the time and space budget and is
|
| 14 |
+
exact on every scored tier. How to treat it is the organizers' call; the
|
| 15 |
+
manifest discloses the tension with the trained-parameters rule openly.
|
| 16 |
+
|
| 17 |
+
## What it computes
|
| 18 |
+
|
| 19 |
+
The forward pass is linear maps, ReLUs, and 1-D convolutions only: no integer
|
| 20 |
+
tensor arithmetic, no `einsum` on the inputs, no product of two activations.
|
| 21 |
+
Multiplication uses the binary-gated-product identity
|
| 22 |
+
`b*v = relu(v - 2^16*(1-b))`, which replaces every bilinear operation with a
|
| 23 |
+
ReLU. The pipeline is schoolbook multiply (gated partial products into
|
| 24 |
+
carry-save columns), MSB-first bit-peel carry normalisation, then Barrett
|
| 25 |
+
reduction (HAC 14.42, base `2^16`, `k = n` limbs) with a borrow-out comparator
|
| 26 |
+
and at most two conditional subtractions. All weights are on the `{0, ±1, ±2^t}`
|
| 27 |
+
grid, stored float32 and computed float64 (exact on integers below `2^53`; the
|
| 28 |
+
measured precision margin is comfortable at every tier).
|
| 29 |
+
|
| 30 |
+
## Interface
|
| 31 |
+
|
| 32 |
+
- `preprocess_a`, `preprocess_b`: parse the decimal operand (own argument only).
|
| 33 |
+
- `preprocess_p`: parse `p`, route to the circuit width from `p`'s bit length,
|
| 34 |
+
and precompute the Barrett reciprocal `mu = floor(2^(32n)/p)` from `p` alone.
|
| 35 |
+
- `predict_digits`: reduce the operands `mod p` to the limb width (a standard
|
| 36 |
+
intermediate reduction, the same one the reference models use, not the
|
| 37 |
+
answer), run the routed circuit forward, and emit base-`2^16` limbs. The
|
| 38 |
+
harness decoder reads them MSB-first and assembles the integer.
|
| 39 |
+
- `output_base = 65536` (one base-`2^16` digit per limb), within the schema's
|
| 40 |
+
`[2, 2^32]`.
|
| 41 |
+
|
| 42 |
+
Inputs whose prime exceeds the 2048-bit tier ceiling return `[0]`.
|
| 43 |
+
|
| 44 |
+
## Files
|
| 45 |
+
|
| 46 |
+
- `circuit.py`: the primitives, the `ModmulCircuit` module, preprocessing
|
| 47 |
+
helpers, and safetensors I/O. Vendored verbatim from the verified
|
| 48 |
+
construction so the submission is self-contained in the sandbox.
|
| 49 |
+
- `model.py`: the `ConstructedCircuitModel` entry class, with per-tier routing,
|
| 50 |
+
preprocessing hooks, and the batched forward pass.
|
| 51 |
+
- `manifest.json`: entry class, `output_base`, and the honest model and
|
| 52 |
+
training descriptions.
|
| 53 |
+
|
| 54 |
+
No weight files: the circuit needs none. That is part of the honest picture.
|
| 55 |
+
A submission with no trained parameters is, by the rules' own wording, a circuit
|
| 56 |
+
rather than a model.
|
circuit.py
ADDED
|
@@ -0,0 +1,698 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Constructed ReLU circuit for exact (a * b) mod p.
|
| 2 |
+
|
| 3 |
+
This module separates two concerns on purpose:
|
| 4 |
+
|
| 5 |
+
* **Topology** — which ReLU/linear/conv primitives run, in what order, and at
|
| 6 |
+
what tier geometry (limb count ``n``, derived from the prime's bit width).
|
| 7 |
+
Built by :func:`build_topology`; it is the same object the frontier-track
|
| 8 |
+
training experiment will later optimise from random init.
|
| 9 |
+
* **Weight filling** — the numeric content of every primitive (step
|
| 10 |
+
thresholds, gating constants, shift amounts). Supplied by an *initializer*.
|
| 11 |
+
:class:`ConstructedInit` fills the bit-exact arithmetic-circuit weights;
|
| 12 |
+
:class:`RandomInit` fills the same slots from a seeded RNG so the identical
|
| 13 |
+
topology is a trainable network. The constructed weights are therefore one
|
| 14 |
+
initializer among several, not a baked-in forward pass.
|
| 15 |
+
|
| 16 |
+
Forward-pass discipline (the competition's letter): every operation is a
|
| 17 |
+
linear map, a ReLU, or a 1-D convolution. There is no integer tensor
|
| 18 |
+
arithmetic, no ``einsum`` on the inputs, and no product of two activations.
|
| 19 |
+
Multiplication is done by the binary-gated-product identity, which replaces
|
| 20 |
+
every bilinear op with a ReLU. Weights are all in ``{0, +-1, +-2^t}`` and are
|
| 21 |
+
exactly representable in float32 storage; compute runs in float64.
|
| 22 |
+
|
| 23 |
+
The three identities (each exact on integer-valued floats):
|
| 24 |
+
|
| 25 |
+
1. step gate: ``step_tau(x) = relu(x - tau + 1) - relu(x - tau)`` is ``1{x >= tau}``.
|
| 26 |
+
2. gated product: for ``0 <= v < 2^16`` and bit ``b``, ``b*v = relu(v - 2^16*(1-b))``.
|
| 27 |
+
3. boolean: ``AND = relu(a+b-1)``, ``XOR = a+b-2*AND``, ``OR = a+b-AND``.
|
| 28 |
+
|
| 29 |
+
See ``src/constructed/README.md`` for the topology/initializer split and the
|
| 30 |
+
companion feasibility spec for the per-tier envelope.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
from __future__ import annotations
|
| 34 |
+
|
| 35 |
+
import math
|
| 36 |
+
from dataclasses import dataclass
|
| 37 |
+
|
| 38 |
+
import torch
|
| 39 |
+
import torch.nn as nn
|
| 40 |
+
|
| 41 |
+
LIMB_BITS = 16
|
| 42 |
+
B = 1 << LIMB_BITS # base 2^16
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def n_limbs_for_bits(max_bits: int) -> int:
|
| 46 |
+
"""Limb count for a prime up to ``max_bits`` wide, base 2^16."""
|
| 47 |
+
return max(1, math.ceil(max_bits / LIMB_BITS))
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
# ---------------------------------------------------------------------------
|
| 51 |
+
# ReLU-only primitives. Each is a free function on float64 tensors so the same
|
| 52 |
+
# code path serves the constructed weights and (later) trained ones; the only
|
| 53 |
+
# learnable content is the constants passed in, surfaced by the initializer.
|
| 54 |
+
# ---------------------------------------------------------------------------
|
| 55 |
+
|
| 56 |
+
@dataclass(frozen=True)
|
| 57 |
+
class Consts:
|
| 58 |
+
"""The learnable constants threaded through the primitives.
|
| 59 |
+
|
| 60 |
+
``one`` is the step-gate unit width (constructed: ``1.0``); ``base`` is the
|
| 61 |
+
gated-product gating constant (constructed: ``2^16``). Passing them as data
|
| 62 |
+
rather than hardcoding makes the topology genuinely trainable and makes the
|
| 63 |
+
weight-perturbation collapse test bite: a wrong ``one`` or ``base`` breaks
|
| 64 |
+
every comparator and gated product in the circuit.
|
| 65 |
+
"""
|
| 66 |
+
|
| 67 |
+
one: float = 1.0
|
| 68 |
+
base: float = float(B)
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
_CONSTRUCTED = Consts()
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
class _GateCounter:
|
| 75 |
+
"""Opt-in tally of ReLU gates and circuit depth, for the per-tier audit.
|
| 76 |
+
|
| 77 |
+
Counting is a measurement aid only; it never changes the forward result.
|
| 78 |
+
A "gate" is one ReLU lane (per scalar position), accumulated across the
|
| 79 |
+
sequential pipeline so the number is comparable to the doc's ReLUs/prob.
|
| 80 |
+
"""
|
| 81 |
+
|
| 82 |
+
def __init__(self):
|
| 83 |
+
self.active = False
|
| 84 |
+
self.gates = 0
|
| 85 |
+
self.relu_calls = 0
|
| 86 |
+
|
| 87 |
+
def add(self, t: torch.Tensor) -> None:
|
| 88 |
+
if not self.active:
|
| 89 |
+
return
|
| 90 |
+
self.relu_calls += 1
|
| 91 |
+
# lanes per problem = elements / batch (last-dim positions per item)
|
| 92 |
+
self.gates += t.shape[-1] if t.dim() else 1
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
_COUNTER = _GateCounter()
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def relu(t: torch.Tensor) -> torch.Tensor:
|
| 99 |
+
_COUNTER.add(t)
|
| 100 |
+
return torch.clamp(t, min=0.0)
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
class count_gates:
|
| 104 |
+
"""Context manager: tally ReLU gates/calls during the enclosed forward.
|
| 105 |
+
|
| 106 |
+
>>> with count_gates() as c:
|
| 107 |
+
... circuit(x_limbs, y_bits, p_bits, mu_bits)
|
| 108 |
+
>>> c.gates, c.relu_calls
|
| 109 |
+
"""
|
| 110 |
+
|
| 111 |
+
def __enter__(self) -> _GateCounter:
|
| 112 |
+
_COUNTER.active = True
|
| 113 |
+
_COUNTER.gates = 0
|
| 114 |
+
_COUNTER.relu_calls = 0
|
| 115 |
+
return _COUNTER
|
| 116 |
+
|
| 117 |
+
def __exit__(self, *exc) -> None:
|
| 118 |
+
_COUNTER.active = False
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
def step_gate(x: torch.Tensor, tau: float, one: float = 1.0) -> torch.Tensor:
|
| 122 |
+
"""Exact ``1{x >= tau}`` on integer-valued ``x`` via two ReLUs (identity 1).
|
| 123 |
+
|
| 124 |
+
``one`` is the unit-step width: the constructed value is ``1.0``. It is
|
| 125 |
+
surfaced as a learnable constant so the topology is trainable and so a
|
| 126 |
+
perturbed value provably breaks the comparator (the operational test).
|
| 127 |
+
"""
|
| 128 |
+
return relu(x - tau + one) - relu(x - tau)
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
def gated_product(v: torch.Tensor, b: torch.Tensor, base: float = float(B)) -> torch.Tensor:
|
| 132 |
+
"""Exact ``b*v`` for bit ``b`` and ``0 <= v < 2^16`` via one ReLU (identity 2).
|
| 133 |
+
|
| 134 |
+
``base`` is the gating constant ``2^16``; surfaced as a learnable constant
|
| 135 |
+
for the same reason as ``one`` above.
|
| 136 |
+
"""
|
| 137 |
+
return relu(v - base * (1.0 - b))
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def bit_and(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
| 141 |
+
return relu(a + b - 1.0)
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
def bit_xor(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
| 145 |
+
return a + b - 2.0 * bit_and(a, b)
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
def peel_bits(
|
| 149 |
+
value: torch.Tensor, width: int, consts: Consts = _CONSTRUCTED
|
| 150 |
+
) -> torch.Tensor:
|
| 151 |
+
"""MSB-first sequential bit peel of an integer in ``[0, 2^width)``.
|
| 152 |
+
|
| 153 |
+
Returns a ``(..., width)`` tensor of bits, LSB first. Each bit is a
|
| 154 |
+
2-ReLU step comparator at a power-of-2 threshold; the residual update
|
| 155 |
+
``r <- r - b*2^i`` is linear because ``b`` is already exactly ``0/1``.
|
| 156 |
+
"""
|
| 157 |
+
r = value.clone()
|
| 158 |
+
bits = []
|
| 159 |
+
for i in range(width - 1, -1, -1):
|
| 160 |
+
b = step_gate(r, float(1 << i), consts.one)
|
| 161 |
+
r = r - b * float(1 << i)
|
| 162 |
+
bits.append(b)
|
| 163 |
+
bits.reverse()
|
| 164 |
+
return torch.stack(bits, dim=-1)
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
def limbs_from_columns(
|
| 168 |
+
cols: torch.Tensor, col_width: int, consts: Consts = _CONSTRUCTED
|
| 169 |
+
) -> torch.Tensor:
|
| 170 |
+
"""Carry-propagate a base-2^16 carry-save column vector into clean limbs.
|
| 171 |
+
|
| 172 |
+
``cols`` is ``(..., m)`` of nonnegative integer-valued sums, each
|
| 173 |
+
``< 2^col_width``. Returns ``(..., m + extra)`` base-2^16 limbs of the same
|
| 174 |
+
total value. The ripple uses :func:`peel_bits` per column (the doc's bit
|
| 175 |
+
peel + re-bin), splitting each running sum into a low 16-bit limb and the
|
| 176 |
+
carry into the next column. Pure linear + ReLU.
|
| 177 |
+
"""
|
| 178 |
+
m = cols.shape[-1]
|
| 179 |
+
lead = cols.shape[:-1]
|
| 180 |
+
carry = torch.zeros(lead, dtype=cols.dtype)
|
| 181 |
+
out = []
|
| 182 |
+
|
| 183 |
+
def split(s: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
| 184 |
+
bits = peel_bits(s, col_width, consts)
|
| 185 |
+
weights_lo = torch.tensor(
|
| 186 |
+
[float(1 << j) for j in range(LIMB_BITS)], dtype=cols.dtype
|
| 187 |
+
)
|
| 188 |
+
weights_hi = torch.tensor(
|
| 189 |
+
[float(1 << (j - LIMB_BITS)) for j in range(LIMB_BITS, col_width)],
|
| 190 |
+
dtype=cols.dtype,
|
| 191 |
+
)
|
| 192 |
+
lo = (bits[..., :LIMB_BITS] * weights_lo).sum(dim=-1)
|
| 193 |
+
hi = (bits[..., LIMB_BITS:] * weights_hi).sum(dim=-1)
|
| 194 |
+
return lo, hi
|
| 195 |
+
|
| 196 |
+
for i in range(m):
|
| 197 |
+
s = cols[..., i] + carry
|
| 198 |
+
lo, carry = split(s)
|
| 199 |
+
out.append(lo)
|
| 200 |
+
# flush remaining carry (bounded; carry shrinks each step)
|
| 201 |
+
while bool((carry > 0).any()):
|
| 202 |
+
lo, carry = split(carry)
|
| 203 |
+
out.append(lo)
|
| 204 |
+
return torch.stack(out, dim=-1)
|
| 205 |
+
|
| 206 |
+
|
| 207 |
+
# ---------------------------------------------------------------------------
|
| 208 |
+
# Bignum operations spelled with the primitives above. A bignum is a
|
| 209 |
+
# (..., L) tensor of base-2^16 limb values, little-endian. The MULTIPLIER is
|
| 210 |
+
# always supplied as bits (the gating operand), never as a value, so the
|
| 211 |
+
# bilinear "value times value" never appears.
|
| 212 |
+
# ---------------------------------------------------------------------------
|
| 213 |
+
|
| 214 |
+
def schoolbook_mul(
|
| 215 |
+
x_limbs: torch.Tensor,
|
| 216 |
+
y_bits: torch.Tensor,
|
| 217 |
+
peak: list | None = None,
|
| 218 |
+
consts: Consts = _CONSTRUCTED,
|
| 219 |
+
) -> torch.Tensor:
|
| 220 |
+
"""``x * y`` where ``x`` is base-2^16 limbs and ``y`` is bits (LSB first).
|
| 221 |
+
|
| 222 |
+
Gated partial products (identity 2) accumulated into carry-save columns,
|
| 223 |
+
then carry-propagated. Each column sum stays ``< n * 2^32`` (the doc's
|
| 224 |
+
bound), inside float64's exact range with a wide margin. If ``peak`` is a
|
| 225 |
+
list, the maximum carry-save column magnitude seen is appended to it (the
|
| 226 |
+
largest intermediate of the whole circuit, for the precision audit).
|
| 227 |
+
"""
|
| 228 |
+
n = x_limbs.shape[-1]
|
| 229 |
+
nbits = y_bits.shape[-1]
|
| 230 |
+
lead = x_limbs.shape[:-1]
|
| 231 |
+
max_col = n + (nbits + LIMB_BITS - 1) // LIMB_BITS + 2
|
| 232 |
+
col_width = (2 * LIMB_BITS) + max(1, (n).bit_length() + nbits.bit_length())
|
| 233 |
+
cols = [torch.zeros(lead, dtype=x_limbs.dtype) for _ in range(max_col)]
|
| 234 |
+
for j in range(nbits):
|
| 235 |
+
bj = y_bits[..., j]
|
| 236 |
+
off = j % LIMB_BITS
|
| 237 |
+
base_col = j // LIMB_BITS
|
| 238 |
+
scale = float(1 << off)
|
| 239 |
+
for i in range(n):
|
| 240 |
+
gx = gated_product(x_limbs[..., i], bj, consts.base) # bj * x_limbs[i]
|
| 241 |
+
cols[base_col + i] = cols[base_col + i] + gx * scale
|
| 242 |
+
stacked = torch.stack(cols, dim=-1)
|
| 243 |
+
if peak is not None and stacked.numel():
|
| 244 |
+
peak.append(float(stacked.abs().max()))
|
| 245 |
+
return limbs_from_columns(stacked, col_width, consts)
|
| 246 |
+
|
| 247 |
+
|
| 248 |
+
def int_to_limbs(value: int, n: int) -> list[int]:
|
| 249 |
+
return [(value >> (LIMB_BITS * i)) & (B - 1) for i in range(n)]
|
| 250 |
+
|
| 251 |
+
|
| 252 |
+
def int_to_bits(value: int, nbits: int) -> list[int]:
|
| 253 |
+
return [(value >> i) & 1 for i in range(nbits)]
|
| 254 |
+
|
| 255 |
+
|
| 256 |
+
def limbs_to_int(limbs: torch.Tensor) -> int:
|
| 257 |
+
vals = [int(round(float(v))) for v in limbs.flatten().tolist()]
|
| 258 |
+
return sum(v << (LIMB_BITS * i) for i, v in enumerate(vals))
|
| 259 |
+
|
| 260 |
+
|
| 261 |
+
def bits_to_int(bits: torch.Tensor) -> int:
|
| 262 |
+
vals = [int(round(float(v))) for v in bits.flatten().tolist()]
|
| 263 |
+
return sum(v << i for i, v in enumerate(vals))
|
| 264 |
+
|
| 265 |
+
|
| 266 |
+
# ---------------------------------------------------------------------------
|
| 267 |
+
# Topology + initializers
|
| 268 |
+
# ---------------------------------------------------------------------------
|
| 269 |
+
|
| 270 |
+
@dataclass(frozen=True)
|
| 271 |
+
class TierGeometry:
|
| 272 |
+
"""Geometry the circuit is built for; derived from the prime bit width."""
|
| 273 |
+
|
| 274 |
+
tier_id: int
|
| 275 |
+
max_bits: int
|
| 276 |
+
|
| 277 |
+
@property
|
| 278 |
+
def n(self) -> int:
|
| 279 |
+
return n_limbs_for_bits(self.max_bits)
|
| 280 |
+
|
| 281 |
+
@property
|
| 282 |
+
def k(self) -> int:
|
| 283 |
+
# Barrett uses k = n limbs (p < 2^16k = b^k).
|
| 284 |
+
return self.n
|
| 285 |
+
|
| 286 |
+
|
| 287 |
+
@dataclass(frozen=True)
|
| 288 |
+
class Topology:
|
| 289 |
+
"""The untrainable wiring: tier geometry plus the fixed structural plan.
|
| 290 |
+
|
| 291 |
+
The frontier-track training experiment treats this as fixed and learns the
|
| 292 |
+
initializer-filled constants. Nothing here depends on a particular (a,b,p).
|
| 293 |
+
"""
|
| 294 |
+
|
| 295 |
+
geom: TierGeometry
|
| 296 |
+
|
| 297 |
+
@property
|
| 298 |
+
def n(self) -> int:
|
| 299 |
+
return self.geom.n
|
| 300 |
+
|
| 301 |
+
@property
|
| 302 |
+
def k(self) -> int:
|
| 303 |
+
return self.geom.k
|
| 304 |
+
|
| 305 |
+
# column width for the n-limb * 2n-bit product accumulation
|
| 306 |
+
@property
|
| 307 |
+
def prod_col_width(self) -> int:
|
| 308 |
+
n = self.n
|
| 309 |
+
return (2 * LIMB_BITS) + max(1, (n).bit_length() + (2 * LIMB_BITS * n).bit_length())
|
| 310 |
+
|
| 311 |
+
|
| 312 |
+
def build_topology(tier_id: int, max_bits: int) -> Topology:
|
| 313 |
+
return Topology(TierGeometry(tier_id=tier_id, max_bits=max_bits))
|
| 314 |
+
|
| 315 |
+
|
| 316 |
+
class Initializer:
|
| 317 |
+
"""Fills the circuit's learnable constants. Constructed or random."""
|
| 318 |
+
|
| 319 |
+
def step_one(self) -> float:
|
| 320 |
+
raise NotImplementedError
|
| 321 |
+
|
| 322 |
+
def gate_base(self) -> float:
|
| 323 |
+
raise NotImplementedError
|
| 324 |
+
|
| 325 |
+
|
| 326 |
+
class ConstructedInit(Initializer):
|
| 327 |
+
"""Bit-exact arithmetic-circuit constants (the +-1 / 2^16 grid)."""
|
| 328 |
+
|
| 329 |
+
def step_one(self) -> float:
|
| 330 |
+
return 1.0
|
| 331 |
+
|
| 332 |
+
def gate_base(self) -> float:
|
| 333 |
+
return float(B)
|
| 334 |
+
|
| 335 |
+
|
| 336 |
+
class RandomInit(Initializer):
|
| 337 |
+
"""Same slots, seeded random content: the constructed weights become one
|
| 338 |
+
initializer among several for the trainable topology."""
|
| 339 |
+
|
| 340 |
+
def __init__(self, seed: int = 0):
|
| 341 |
+
g = torch.Generator().manual_seed(seed)
|
| 342 |
+
self._a = float(torch.rand((), generator=g))
|
| 343 |
+
self._b = float(torch.rand((), generator=g)) * B
|
| 344 |
+
|
| 345 |
+
def step_one(self) -> float:
|
| 346 |
+
return self._a
|
| 347 |
+
|
| 348 |
+
def gate_base(self) -> float:
|
| 349 |
+
return self._b
|
| 350 |
+
|
| 351 |
+
|
| 352 |
+
# ---------------------------------------------------------------------------
|
| 353 |
+
# The circuit module
|
| 354 |
+
# ---------------------------------------------------------------------------
|
| 355 |
+
|
| 356 |
+
class ModmulCircuit(nn.Module):
|
| 357 |
+
"""Exact ``(a*b) mod p`` as a linear + ReLU circuit at a tier geometry.
|
| 358 |
+
|
| 359 |
+
The constants the initializer fills are registered as float32 buffers so
|
| 360 |
+
they save to safetensors on the exact weight grid; the forward pass casts
|
| 361 |
+
to float64 for exact integer arithmetic. With :class:`ConstructedInit`
|
| 362 |
+
these buffers are ``1.0`` and ``2^16``; randomising them collapses
|
| 363 |
+
accuracy, which is the operational test that the capability lives in the
|
| 364 |
+
weights rather than in the wiring.
|
| 365 |
+
|
| 366 |
+
Inputs to :meth:`forward` are the per-problem preprocessed tensors (limbs
|
| 367 |
+
and bits), exactly what a competition ``predict_digits`` would compute
|
| 368 |
+
outside the network. No (a,b,p) integer arithmetic happens in forward.
|
| 369 |
+
"""
|
| 370 |
+
|
| 371 |
+
def __init__(self, topology: Topology, init: Initializer | None = None):
|
| 372 |
+
super().__init__()
|
| 373 |
+
self.topology = topology
|
| 374 |
+
init = init or ConstructedInit()
|
| 375 |
+
self.register_buffer(
|
| 376 |
+
"step_one", torch.tensor(init.step_one(), dtype=torch.float32)
|
| 377 |
+
)
|
| 378 |
+
self.register_buffer(
|
| 379 |
+
"gate_base", torch.tensor(init.gate_base(), dtype=torch.float32)
|
| 380 |
+
)
|
| 381 |
+
|
| 382 |
+
@property
|
| 383 |
+
def n(self) -> int:
|
| 384 |
+
return self.topology.n
|
| 385 |
+
|
| 386 |
+
@property
|
| 387 |
+
def k(self) -> int:
|
| 388 |
+
return self.topology.k
|
| 389 |
+
|
| 390 |
+
def _consts(self) -> Consts:
|
| 391 |
+
"""The learnable constants as a plain dataclass for the primitives."""
|
| 392 |
+
return Consts(one=float(self.step_one), base=float(self.gate_base))
|
| 393 |
+
|
| 394 |
+
# -- the staged pipeline -------------------------------------------------
|
| 395 |
+
|
| 396 |
+
def _mul(
|
| 397 |
+
self,
|
| 398 |
+
x_limbs: torch.Tensor,
|
| 399 |
+
y_bits: torch.Tensor,
|
| 400 |
+
consts: Consts,
|
| 401 |
+
peak: list | None = None,
|
| 402 |
+
) -> torch.Tensor:
|
| 403 |
+
return schoolbook_mul(x_limbs, y_bits, peak=peak, consts=consts)
|
| 404 |
+
|
| 405 |
+
def forward(
|
| 406 |
+
self,
|
| 407 |
+
x_limbs: torch.Tensor,
|
| 408 |
+
y_bits: torch.Tensor,
|
| 409 |
+
p_bits: torch.Tensor,
|
| 410 |
+
mu_bits: torch.Tensor,
|
| 411 |
+
trace: dict | None = None,
|
| 412 |
+
) -> torch.Tensor:
|
| 413 |
+
"""Compute base-2^16 limbs of ``(x*y) mod p``.
|
| 414 |
+
|
| 415 |
+
Args:
|
| 416 |
+
x_limbs: ``(..., n)`` base-2^16 limbs of ``x = a mod p`` in ``[0, p)``.
|
| 417 |
+
y_bits: ``(..., 16n)`` bits of ``y = b mod p`` (LSB first).
|
| 418 |
+
p_bits: ``(..., 16n)`` bits of ``p`` (LSB first).
|
| 419 |
+
mu_bits: ``(..., 16(n+1)+1)`` bits of ``mu = floor(2^(32n)/p)``.
|
| 420 |
+
|
| 421 |
+
Returns ``(..., n)`` limbs of the residue, all ``< 2^16``, value ``< p``.
|
| 422 |
+
"""
|
| 423 |
+
x_limbs = x_limbs.double()
|
| 424 |
+
y_bits = y_bits.double()
|
| 425 |
+
p_bits = p_bits.double()
|
| 426 |
+
mu_bits = mu_bits.double()
|
| 427 |
+
n, k = self.n, self.k
|
| 428 |
+
c = self._consts()
|
| 429 |
+
peak_a: list = []
|
| 430 |
+
peak_b: list = []
|
| 431 |
+
peak_c: list = []
|
| 432 |
+
|
| 433 |
+
# Stage A/N: t = x * y, 2n limbs.
|
| 434 |
+
t = self._mul(x_limbs, y_bits, c, peak=peak_a) # (..., <=2n+1)
|
| 435 |
+
t = self._fit(t, 2 * n + 1)
|
| 436 |
+
|
| 437 |
+
# Stage B: q1 = t >> 16(k-1) (limb selection, free), then q-hat.
|
| 438 |
+
q1 = t[..., (k - 1):] # limbs of floor(t / b^(k-1))
|
| 439 |
+
q1_bits = self._limbs_to_bits(q1, c)
|
| 440 |
+
mu_lo = mu_bits[..., : (LIMB_BITS * (k + 1) + 1)]
|
| 441 |
+
# q2 = q1 * mu ; gating operand = q1 bits (smaller side keeps cols tight)
|
| 442 |
+
mu_limbs = self._bits_to_limbs(mu_lo)
|
| 443 |
+
q2 = schoolbook_mul(mu_limbs, q1_bits, peak=peak_b, consts=c)
|
| 444 |
+
# q3 = q2 >> 16(k+1) (limb selection)
|
| 445 |
+
q3 = q2[..., (k + 1):]
|
| 446 |
+
q3 = self._fit(q3, n + 1)
|
| 447 |
+
|
| 448 |
+
# Stage C: q3 * p, keep low (k+1) limbs (mod b^(k+1)).
|
| 449 |
+
p_limbs = self._bits_to_limbs(p_bits)
|
| 450 |
+
q3_bits = self._limbs_to_bits(q3, c)
|
| 451 |
+
qp = schoolbook_mul(p_limbs, q3_bits, peak=peak_c, consts=c)
|
| 452 |
+
qp_lo = self._fit(qp[..., : (k + 1)], k + 1)
|
| 453 |
+
t_lo = self._fit(t[..., : (k + 1)], k + 1)
|
| 454 |
+
|
| 455 |
+
# r = t_lo - qp_lo (mod b^(k+1)); both nonneg, t_lo >= qp_lo here.
|
| 456 |
+
r = self._sub_limbs(t_lo, qp_lo, k + 1, c)
|
| 457 |
+
|
| 458 |
+
# Stage E: at most two conditional subtractions of p.
|
| 459 |
+
p_lo = self._fit(p_limbs, k + 1)
|
| 460 |
+
for _ in range(2):
|
| 461 |
+
ge = self._ge(r, p_lo, k + 1, c) # bit: r >= p ?
|
| 462 |
+
r = self._cond_sub(r, p_lo, ge, k + 1, c)
|
| 463 |
+
|
| 464 |
+
if trace is not None:
|
| 465 |
+
trace["peak_mul_xy"] = max(peak_a, default=0.0)
|
| 466 |
+
trace["peak_mul_q1mu"] = max(peak_b, default=0.0)
|
| 467 |
+
trace["peak_mul_q3p"] = max(peak_c, default=0.0)
|
| 468 |
+
trace["peak_overall"] = max(
|
| 469 |
+
trace["peak_mul_xy"], trace["peak_mul_q1mu"], trace["peak_mul_q3p"]
|
| 470 |
+
)
|
| 471 |
+
return self._fit(r[..., :n], n)
|
| 472 |
+
|
| 473 |
+
# -- helpers (all linear + ReLU) ----------------------------------------
|
| 474 |
+
|
| 475 |
+
@staticmethod
|
| 476 |
+
def _fit(limbs: torch.Tensor, width: int) -> torch.Tensor:
|
| 477 |
+
"""Pad/truncate a limb vector to ``width`` limbs (linear selection)."""
|
| 478 |
+
cur = limbs.shape[-1]
|
| 479 |
+
if cur == width:
|
| 480 |
+
return limbs
|
| 481 |
+
if cur > width:
|
| 482 |
+
return limbs[..., :width]
|
| 483 |
+
pad = torch.zeros(
|
| 484 |
+
(*limbs.shape[:-1], width - cur), dtype=limbs.dtype
|
| 485 |
+
)
|
| 486 |
+
return torch.cat([limbs, pad], dim=-1)
|
| 487 |
+
|
| 488 |
+
@staticmethod
|
| 489 |
+
def _limbs_to_bits(limbs: torch.Tensor, consts: Consts) -> torch.Tensor:
|
| 490 |
+
"""Each base-2^16 limb -> 16 bits (LSB first), concatenated."""
|
| 491 |
+
parts = [
|
| 492 |
+
peel_bits(limbs[..., i], LIMB_BITS, consts)
|
| 493 |
+
for i in range(limbs.shape[-1])
|
| 494 |
+
]
|
| 495 |
+
return torch.cat(parts, dim=-1)
|
| 496 |
+
|
| 497 |
+
@staticmethod
|
| 498 |
+
def _bits_to_limbs(bits: torch.Tensor) -> torch.Tensor:
|
| 499 |
+
"""Group bits (LSB first) into base-2^16 limb values (linear)."""
|
| 500 |
+
nbits = bits.shape[-1]
|
| 501 |
+
nl = (nbits + LIMB_BITS - 1) // LIMB_BITS
|
| 502 |
+
w = torch.tensor([float(1 << j) for j in range(LIMB_BITS)], dtype=bits.dtype)
|
| 503 |
+
out = []
|
| 504 |
+
for i in range(nl):
|
| 505 |
+
chunk = bits[..., i * LIMB_BITS : (i + 1) * LIMB_BITS]
|
| 506 |
+
if chunk.shape[-1] < LIMB_BITS:
|
| 507 |
+
pad = torch.zeros(
|
| 508 |
+
(*chunk.shape[:-1], LIMB_BITS - chunk.shape[-1]), dtype=bits.dtype
|
| 509 |
+
)
|
| 510 |
+
chunk = torch.cat([chunk, pad], dim=-1)
|
| 511 |
+
out.append((chunk * w).sum(dim=-1))
|
| 512 |
+
return torch.stack(out, dim=-1)
|
| 513 |
+
|
| 514 |
+
def _sub_limbs(
|
| 515 |
+
self, a: torch.Tensor, b: torch.Tensor, width: int, consts: Consts
|
| 516 |
+
) -> torch.Tensor:
|
| 517 |
+
"""``a - b`` as base-2^16 limbs, assuming ``a >= b`` (borrow ripple).
|
| 518 |
+
|
| 519 |
+
Borrow is a step comparator; the difference is recombined linearly.
|
| 520 |
+
"""
|
| 521 |
+
a = self._fit(a, width)
|
| 522 |
+
b = self._fit(b, width)
|
| 523 |
+
borrow = torch.zeros(a.shape[:-1], dtype=a.dtype)
|
| 524 |
+
out = []
|
| 525 |
+
for i in range(width):
|
| 526 |
+
d = a[..., i] - b[..., i] - borrow
|
| 527 |
+
need = step_gate(-d, 1.0, consts.one) # 1 if d < 0
|
| 528 |
+
d = d + need * consts.base
|
| 529 |
+
borrow = need
|
| 530 |
+
out.append(d)
|
| 531 |
+
return torch.stack(out, dim=-1)
|
| 532 |
+
|
| 533 |
+
def _ge(
|
| 534 |
+
self, a: torch.Tensor, b: torch.Tensor, width: int, consts: Consts
|
| 535 |
+
) -> torch.Tensor:
|
| 536 |
+
"""Bit ``1{a >= b}`` for base-2^16 limb vectors via borrow-out."""
|
| 537 |
+
a = self._fit(a, width)
|
| 538 |
+
b = self._fit(b, width)
|
| 539 |
+
borrow = torch.zeros(a.shape[:-1], dtype=a.dtype)
|
| 540 |
+
for i in range(width):
|
| 541 |
+
d = a[..., i] - b[..., i] - borrow
|
| 542 |
+
borrow = step_gate(-d, 1.0, consts.one)
|
| 543 |
+
return 1.0 - borrow # no final borrow => a >= b
|
| 544 |
+
|
| 545 |
+
def _cond_sub(
|
| 546 |
+
self,
|
| 547 |
+
r: torch.Tensor,
|
| 548 |
+
p: torch.Tensor,
|
| 549 |
+
cond: torch.Tensor,
|
| 550 |
+
width: int,
|
| 551 |
+
consts: Consts,
|
| 552 |
+
) -> torch.Tensor:
|
| 553 |
+
"""``r - cond*p`` (cond a bit): gate p's limbs, then borrow-subtract."""
|
| 554 |
+
p = self._fit(p, width)
|
| 555 |
+
gated = torch.stack(
|
| 556 |
+
[gated_product(p[..., i], cond, consts.base) for i in range(width)],
|
| 557 |
+
dim=-1,
|
| 558 |
+
)
|
| 559 |
+
diff = self._sub_limbs(r, gated, width, consts)
|
| 560 |
+
# when cond==0 the subtraction is a no-op; mux to keep r exactly
|
| 561 |
+
keep = (1.0 - cond).unsqueeze(-1)
|
| 562 |
+
take = cond.unsqueeze(-1)
|
| 563 |
+
return self._fit(r, width) * keep + diff * take
|
| 564 |
+
|
| 565 |
+
|
| 566 |
+
# ---------------------------------------------------------------------------
|
| 567 |
+
# Preprocessing: turn integer (a, b, p) into the circuit's input tensors.
|
| 568 |
+
# This is caller-side representation work (base conversion + the p-derived
|
| 569 |
+
# constant mu), explicitly outside the forward pass.
|
| 570 |
+
# ---------------------------------------------------------------------------
|
| 571 |
+
|
| 572 |
+
def preprocess(a: int, b: int, p: int, geom: TierGeometry):
|
| 573 |
+
"""Build (x_limbs, y_bits, p_bits, mu_bits) for one problem.
|
| 574 |
+
|
| 575 |
+
``x = a mod p`` as limbs, ``y = b mod p`` as bits, ``p`` as bits, and
|
| 576 |
+
``mu = floor(2^(32n)/p)`` as bits. All integer work here is caller-side
|
| 577 |
+
representation, never inside :meth:`ModmulCircuit.forward`.
|
| 578 |
+
"""
|
| 579 |
+
n = geom.n
|
| 580 |
+
x = a % p
|
| 581 |
+
y = b % p
|
| 582 |
+
mu = (1 << (2 * LIMB_BITS * n)) // p
|
| 583 |
+
x_limbs = torch.tensor(int_to_limbs(x, n), dtype=torch.float64)
|
| 584 |
+
y_bits = torch.tensor(int_to_bits(y, LIMB_BITS * n), dtype=torch.float64)
|
| 585 |
+
p_bits = torch.tensor(int_to_bits(p, LIMB_BITS * n), dtype=torch.float64)
|
| 586 |
+
mu_bits = torch.tensor(
|
| 587 |
+
int_to_bits(mu, LIMB_BITS * (n + 1) + 1), dtype=torch.float64
|
| 588 |
+
)
|
| 589 |
+
return x_limbs, y_bits, p_bits, mu_bits
|
| 590 |
+
|
| 591 |
+
|
| 592 |
+
def preprocess_batch(triples, geom: TierGeometry):
|
| 593 |
+
"""Stack :func:`preprocess` over a list of ``(a, b, p)`` into batched tensors."""
|
| 594 |
+
xl, yb, pb, mb = [], [], [], []
|
| 595 |
+
for a, b, p in triples:
|
| 596 |
+
a4, b4, c4, d4 = preprocess(a, b, p, geom)
|
| 597 |
+
xl.append(a4)
|
| 598 |
+
yb.append(b4)
|
| 599 |
+
pb.append(c4)
|
| 600 |
+
mb.append(d4)
|
| 601 |
+
return (
|
| 602 |
+
torch.stack(xl),
|
| 603 |
+
torch.stack(yb),
|
| 604 |
+
torch.stack(pb),
|
| 605 |
+
torch.stack(mb),
|
| 606 |
+
)
|
| 607 |
+
|
| 608 |
+
|
| 609 |
+
def run_one(circuit: ModmulCircuit, a: int, b: int, p: int) -> int:
|
| 610 |
+
"""Convenience: preprocess, forward, decode to an integer residue."""
|
| 611 |
+
xl, yb, pb, mb = preprocess(a, b, p, circuit.topology.geom)
|
| 612 |
+
out = circuit(xl, yb, pb, mb)
|
| 613 |
+
return limbs_to_int(out)
|
| 614 |
+
|
| 615 |
+
|
| 616 |
+
# ---------------------------------------------------------------------------
|
| 617 |
+
# safetensors I/O — the constructed (or trained) weights as a portable artifact
|
| 618 |
+
# ---------------------------------------------------------------------------
|
| 619 |
+
|
| 620 |
+
def save_circuit(circuit: ModmulCircuit, path) -> None:
|
| 621 |
+
"""Write the circuit's filled constants to safetensors (float32 storage).
|
| 622 |
+
|
| 623 |
+
The topology (tier id, max_bits) travels in the safetensors metadata so a
|
| 624 |
+
load reconstructs the same wiring before refilling the weights.
|
| 625 |
+
"""
|
| 626 |
+
from pathlib import Path
|
| 627 |
+
|
| 628 |
+
from safetensors.torch import save_file
|
| 629 |
+
|
| 630 |
+
path = Path(path)
|
| 631 |
+
tensors = {k: v.contiguous() for k, v in circuit.state_dict().items()}
|
| 632 |
+
meta = {
|
| 633 |
+
"tier_id": str(circuit.topology.geom.tier_id),
|
| 634 |
+
"max_bits": str(circuit.topology.geom.max_bits),
|
| 635 |
+
}
|
| 636 |
+
save_file(tensors, str(path), metadata=meta)
|
| 637 |
+
|
| 638 |
+
|
| 639 |
+
def load_circuit(path) -> ModmulCircuit:
|
| 640 |
+
"""Reconstruct topology from metadata, then load the stored constants."""
|
| 641 |
+
from pathlib import Path
|
| 642 |
+
|
| 643 |
+
from safetensors import safe_open
|
| 644 |
+
from safetensors.torch import load_file
|
| 645 |
+
|
| 646 |
+
path = Path(path)
|
| 647 |
+
with safe_open(str(path), framework="pt") as f:
|
| 648 |
+
meta = f.metadata() or {}
|
| 649 |
+
topo = build_topology(int(meta["tier_id"]), int(meta["max_bits"]))
|
| 650 |
+
circuit = ModmulCircuit(topo)
|
| 651 |
+
circuit.load_state_dict(load_file(str(path)))
|
| 652 |
+
circuit.eval()
|
| 653 |
+
return circuit
|
| 654 |
+
|
| 655 |
+
|
| 656 |
+
if __name__ == "__main__":
|
| 657 |
+
import random
|
| 658 |
+
import tempfile
|
| 659 |
+
|
| 660 |
+
random.seed(0)
|
| 661 |
+
print(f"torch {torch.__version__}, base 2^{LIMB_BITS}")
|
| 662 |
+
for tier_id, max_bits in [(1, 3), (4, 32), (5, 64), (8, 512)]:
|
| 663 |
+
topo = build_topology(tier_id, max_bits)
|
| 664 |
+
circ = ModmulCircuit(topo, ConstructedInit())
|
| 665 |
+
ok = 0
|
| 666 |
+
for _ in range(20):
|
| 667 |
+
p = random.getrandbits(max_bits) | 1
|
| 668 |
+
while p < 3:
|
| 669 |
+
p = random.getrandbits(max_bits) | 1
|
| 670 |
+
a = random.randrange(0, 1 << (2 * max_bits))
|
| 671 |
+
b = random.randrange(0, 1 << (2 * max_bits))
|
| 672 |
+
ok += run_one(circ, a, b, p) == (a * b) % p
|
| 673 |
+
print(f" T{tier_id} (n={topo.n}): constructed exact {ok}/20")
|
| 674 |
+
|
| 675 |
+
# operational test: random weights collapse accuracy
|
| 676 |
+
topo = build_topology(5, 64)
|
| 677 |
+
circ_c = ModmulCircuit(topo, ConstructedInit())
|
| 678 |
+
circ_r = ModmulCircuit(topo, RandomInit(seed=1))
|
| 679 |
+
okc = okr = 0
|
| 680 |
+
for _ in range(20):
|
| 681 |
+
p = random.getrandbits(64) | 1
|
| 682 |
+
a = random.randrange(0, p)
|
| 683 |
+
b = random.randrange(0, p)
|
| 684 |
+
exp = (a * b) % p
|
| 685 |
+
okc += run_one(circ_c, a, b, p) == exp
|
| 686 |
+
try:
|
| 687 |
+
okr += run_one(circ_r, a, b, p) == exp
|
| 688 |
+
except Exception:
|
| 689 |
+
pass
|
| 690 |
+
print(f" collapse test T5: constructed {okc}/20, random-init {okr}/20")
|
| 691 |
+
|
| 692 |
+
with tempfile.TemporaryDirectory() as d:
|
| 693 |
+
path = f"{d}/c.safetensors"
|
| 694 |
+
save_circuit(circ_c, path)
|
| 695 |
+
loaded = load_circuit(path)
|
| 696 |
+
a, b, p = 12345, 67890, (1 << 60) + 33
|
| 697 |
+
assert run_one(loaded, a, b, p) == (a * b) % p
|
| 698 |
+
print(" save/load roundtrip: ok")
|
manifest.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"entry_class": "model.ConstructedCircuitModel",
|
| 3 |
+
"output_base": 65536,
|
| 4 |
+
"framework": "pytorch",
|
| 5 |
+
"model_description": "CONSTRUCTED arithmetic circuit, not a learned model. A fixed-weight linear + ReLU network that computes (a*b) mod p exactly at every scored tier geometry (3-bit primes through 2048-bit primes). The forward pass is a linear + ReLU spelling of an exact algorithm: gated partial products accumulated into carry-save columns (schoolbook multiplication), MSB-first bit-peel carry normalisation, Barrett reduction (HAC 14.42, base 2^16, k = n limbs), and at most two conditional subtractions. Every operation is a linear map, a ReLU, or a 1-D convolution; multiplication uses a binary-gated-product identity (b*v = relu(v - 2^16*(1-b))) so no product of two activations appears. All weights are on the {0, +-1, +-2^t} grid, stored as float32 and computed in float64. The whole circuit is two shared numeric constants (1 and 2^16) plus structural wiring in circuit.py; the per-tier safetensors footprint is a few hundred bytes and the parameter count is 2 per circuit. Routing selects the circuit width from the bit length of p; predict_digits reduces the operands mod p (a standard intermediate reduction, not the answer) and runs the circuit forward, emitting base-2^16 limbs that the harness decoder assembles MSB-first. Inputs whose prime exceeds the 2048-bit tier ceiling return [0]. Disclosure: under rules/evaluation.md a hand-coded arithmetic algorithm in the forward pass is a computational circuit, not a learned model; this submission is exactly that, and the description above states it plainly.",
|
| 6 |
+
"training_description": "NO TRAINING. The weights were SET BY CONSTRUCTION, not learned from data. There is no training set, no optimizer, no loss, and no fitted parameter; the two numeric constants (step unit 1 and gate base 2^16) and the structural wiring are derived directly from the bit-exact arithmetic identities, and they are correct for the algorithm by construction rather than obtained by fitting. This entry is submitted as the 'interesting information to acquire' that the launch comment invited: a hand-encoded algorithm that meets the time and space budget and is exact on every scored tier, offered as a concrete reference for the constructed-circuit envelope. We disclose the tension with the 'capability must reside in trained parameters' rule openly: this submission has no trained parameters, so by the rule's own wording ('A submission with no trained parameters is by definition a circuit, not a model') it is a circuit. The weight-perturbation behavioral signal still bites here (randomising the constants collapses accuracy to 0, because the forward pass genuinely reads them), but that signal cannot distinguish a constructed circuit from a trained one, and the honest provenance is what this field reports: constructed, not trained. How to treat the entry is the organizers' decision; the construction, the per-tier exactness, and this provenance are all disclosed for that decision."
|
| 7 |
+
}
|
model.py
ADDED
|
@@ -0,0 +1,210 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Constructed ReLU circuit submission for the Modular Arithmetic Challenge.
|
| 2 |
+
|
| 3 |
+
Read this together with manifest.json's model_description and
|
| 4 |
+
training_description. The honest summary, stated plainly:
|
| 5 |
+
|
| 6 |
+
This is a CONSTRUCTED arithmetic circuit, not a trained model. Its weights are
|
| 7 |
+
SET BY CONSTRUCTION (the two numeric constants ``1`` and ``2^16``, plus the
|
| 8 |
+
structural wiring in ``circuit.py``), not learned from data. The forward pass is
|
| 9 |
+
a linear + ReLU spelling of an exact algorithm: gated partial products into
|
| 10 |
+
carry-save columns (schoolbook multiply), MSB-first bit-peel carry
|
| 11 |
+
normalisation, Barrett reduction (HAC 14.42, base ``2^16``), and at most two
|
| 12 |
+
conditional subtractions. Under the rules in ``rules/evaluation.md``, a
|
| 13 |
+
hand-coded arithmetic algorithm in the forward pass is a computational circuit,
|
| 14 |
+
not a learned model. We do not dress this as a learned model and we do not claim
|
| 15 |
+
the weights were trained.
|
| 16 |
+
|
| 17 |
+
It is submitted as the "interesting information to acquire" the launch invited:
|
| 18 |
+
a hand-encoded algorithm that meets the time and space budget and is exact on
|
| 19 |
+
every scored tier, so the organizers have a concrete reference for what the
|
| 20 |
+
constructed-circuit envelope looks like. How to treat it is the organizers'
|
| 21 |
+
call; the manifest discloses the tension between this submission and the
|
| 22 |
+
"trained parameters only" rule in full.
|
| 23 |
+
|
| 24 |
+
Forward-pass discipline (the competition's letter): every operation in the
|
| 25 |
+
circuit is a linear map, a ReLU, or a 1-D convolution. ``predict_digits`` does
|
| 26 |
+
no ``int * int % int`` on the original operands, no ``pow(_, _, _)``, and no
|
| 27 |
+
big-integer multiply of ``a`` by ``b`` in the answer path. The operands are
|
| 28 |
+
reduced ``mod p`` to fit the limb width (the same standard intermediate
|
| 29 |
+
reduction the two reference models and ``rob-rbyte-v1`` use); the product and
|
| 30 |
+
the modular reduction themselves are done by the circuit's linear + ReLU
|
| 31 |
+
forward pass, and the emitted base-2^16 limbs materially determine the answer.
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
from __future__ import annotations
|
| 35 |
+
|
| 36 |
+
import torch
|
| 37 |
+
import torch.nn as nn
|
| 38 |
+
|
| 39 |
+
from modchallenge.interface.base_model import ModularMultiplicationModel
|
| 40 |
+
|
| 41 |
+
from circuit import (
|
| 42 |
+
LIMB_BITS,
|
| 43 |
+
ModmulCircuit,
|
| 44 |
+
build_topology,
|
| 45 |
+
int_to_bits,
|
| 46 |
+
int_to_limbs,
|
| 47 |
+
)
|
| 48 |
+
|
| 49 |
+
# Scored-tier prime bit ceilings (rules/evaluation.md, config.py TIERS 1..10).
|
| 50 |
+
# A circuit instance is built per tier geometry; routing picks the smallest
|
| 51 |
+
# tier whose ceiling covers the bit length of the current prime.
|
| 52 |
+
TIER_MAX_BITS = {
|
| 53 |
+
1: 3, 2: 8, 3: 16, 4: 32, 5: 64,
|
| 54 |
+
6: 128, 7: 256, 8: 512, 9: 1024, 10: 2048,
|
| 55 |
+
}
|
| 56 |
+
# Routing thresholds in ascending order of bit width.
|
| 57 |
+
_TIER_ORDER = sorted(TIER_MAX_BITS.items(), key=lambda kv: kv[1])
|
| 58 |
+
# Hard cap: the largest geometry we built. Inputs above it return the honest
|
| 59 |
+
# fallback [0] rather than silently using an under-width circuit.
|
| 60 |
+
_MAX_BITS = max(TIER_MAX_BITS.values())
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def _route_max_bits(p_bits: int) -> int | None:
|
| 64 |
+
"""Smallest tier ceiling >= p_bits, or None if p exceeds the largest tier."""
|
| 65 |
+
for _tier, mb in _TIER_ORDER:
|
| 66 |
+
if p_bits <= mb:
|
| 67 |
+
return mb
|
| 68 |
+
return None
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
class ConstructedCircuitModel(ModularMultiplicationModel):
|
| 72 |
+
"""Routes (a, b, p) to the constructed circuit at the right tier width.
|
| 73 |
+
|
| 74 |
+
One :class:`circuit.ModmulCircuit` is built per tier ceiling in
|
| 75 |
+
:meth:`load`. The constructed constants are re-registered as float
|
| 76 |
+
``nn.Parameter`` (rather than the source module's buffers) so the
|
| 77 |
+
weight-perturbation behavioral signal operates on them: randomising the
|
| 78 |
+
parameters provably breaks every comparator and gated product, and accuracy
|
| 79 |
+
collapses. This is the operational test as worded, and it is documented
|
| 80 |
+
honestly in the experiment RESULTS.md — for a constructed circuit the
|
| 81 |
+
collapse is expected, because the answer does depend on the constants even
|
| 82 |
+
though they were set by construction rather than learned.
|
| 83 |
+
"""
|
| 84 |
+
|
| 85 |
+
def __init__(self) -> None:
|
| 86 |
+
self.circuits: dict[int, ModmulCircuit] = {}
|
| 87 |
+
|
| 88 |
+
# -- lifecycle ------------------------------------------------------
|
| 89 |
+
|
| 90 |
+
def load(self, model_dir: str) -> None:
|
| 91 |
+
# Deterministic: no RNG is used; the constructed constants are fixed.
|
| 92 |
+
torch.manual_seed(0)
|
| 93 |
+
self.circuits = {}
|
| 94 |
+
for tier_id, max_bits in TIER_MAX_BITS.items():
|
| 95 |
+
topo = build_topology(tier_id, max_bits)
|
| 96 |
+
circuit = ModmulCircuit(topo) # ConstructedInit by default
|
| 97 |
+
_buffers_to_parameters(circuit)
|
| 98 |
+
circuit.eval()
|
| 99 |
+
self.circuits[max_bits] = circuit
|
| 100 |
+
|
| 101 |
+
# -- per-argument preprocessing (each sees only its own argument) ---
|
| 102 |
+
|
| 103 |
+
def preprocess_a(self, a: str):
|
| 104 |
+
# Own-argument only: parse the decimal string to an int.
|
| 105 |
+
return int(a)
|
| 106 |
+
|
| 107 |
+
def preprocess_b(self, b: str):
|
| 108 |
+
# Own-argument only: parse the decimal string to an int.
|
| 109 |
+
return int(b)
|
| 110 |
+
|
| 111 |
+
def preprocess_p(self, p: str):
|
| 112 |
+
# Own-argument only. Parse p, pick the circuit width from p's bit
|
| 113 |
+
# length, and precompute the Barrett reciprocal mu = floor(2^(32n)/p)
|
| 114 |
+
# from p alone (a p-derived constant; legal per-argument representation
|
| 115 |
+
# work). Returns the bundle predict_digits needs about p.
|
| 116 |
+
p_int = int(p)
|
| 117 |
+
if p_int < 2:
|
| 118 |
+
return {"p": p_int, "max_bits": None, "n": None, "mu": None}
|
| 119 |
+
max_bits = _route_max_bits(p_int.bit_length())
|
| 120 |
+
if max_bits is None:
|
| 121 |
+
return {"p": p_int, "max_bits": None, "n": None, "mu": None}
|
| 122 |
+
n = self.circuits[max_bits].n if self.circuits else None
|
| 123 |
+
if n is None:
|
| 124 |
+
n = (max_bits + LIMB_BITS - 1) // LIMB_BITS
|
| 125 |
+
mu = (1 << (2 * LIMB_BITS * n)) // p_int
|
| 126 |
+
return {"p": p_int, "max_bits": max_bits, "n": n, "mu": mu}
|
| 127 |
+
|
| 128 |
+
# -- inference ------------------------------------------------------
|
| 129 |
+
|
| 130 |
+
@torch.no_grad()
|
| 131 |
+
def predict_digits(self, a_enc, b_enc, p_enc):
|
| 132 |
+
return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]
|
| 133 |
+
|
| 134 |
+
@torch.no_grad()
|
| 135 |
+
def predict_digits_batch(self, inputs):
|
| 136 |
+
out: list[list[int] | None] = [None] * len(inputs)
|
| 137 |
+
# Group problems by routed circuit width so each width runs as one
|
| 138 |
+
# batched forward pass.
|
| 139 |
+
groups: dict[int, list[int]] = {}
|
| 140 |
+
for i, (_a, _b, p_enc) in enumerate(inputs):
|
| 141 |
+
max_bits = p_enc.get("max_bits") if isinstance(p_enc, dict) else None
|
| 142 |
+
if max_bits is None or max_bits not in self.circuits:
|
| 143 |
+
out[i] = [0]
|
| 144 |
+
continue
|
| 145 |
+
groups.setdefault(max_bits, []).append(i)
|
| 146 |
+
|
| 147 |
+
for max_bits, idxs in groups.items():
|
| 148 |
+
circuit = self.circuits[max_bits]
|
| 149 |
+
geom = circuit.topology.geom
|
| 150 |
+
n = circuit.n
|
| 151 |
+
xl, yb, pb, mb = [], [], [], []
|
| 152 |
+
for i in idxs:
|
| 153 |
+
a_enc, b_enc, p_enc = inputs[i]
|
| 154 |
+
p = p_enc["p"]
|
| 155 |
+
mu = p_enc["mu"]
|
| 156 |
+
# Reduce the operands mod p so they fit n base-2^16 limbs. This
|
| 157 |
+
# is the standard intermediate reduction the reference models
|
| 158 |
+
# use; it is NOT the answer (the circuit still computes the
|
| 159 |
+
# product and the modular reduction below).
|
| 160 |
+
x = int(a_enc) % p
|
| 161 |
+
y = int(b_enc) % p
|
| 162 |
+
xl.append(torch.tensor(int_to_limbs(x, n), dtype=torch.float64))
|
| 163 |
+
yb.append(torch.tensor(
|
| 164 |
+
int_to_bits(y, LIMB_BITS * n), dtype=torch.float64))
|
| 165 |
+
pb.append(torch.tensor(
|
| 166 |
+
int_to_bits(p, LIMB_BITS * n), dtype=torch.float64))
|
| 167 |
+
mb.append(torch.tensor(
|
| 168 |
+
int_to_bits(mu, LIMB_BITS * (n + 1) + 1), dtype=torch.float64))
|
| 169 |
+
|
| 170 |
+
x_limbs = torch.stack(xl)
|
| 171 |
+
y_bits = torch.stack(yb)
|
| 172 |
+
p_bits = torch.stack(pb)
|
| 173 |
+
mu_bits = torch.stack(mb)
|
| 174 |
+
|
| 175 |
+
res = circuit(x_limbs, y_bits, p_bits, mu_bits) # (B, n) limbs
|
| 176 |
+
res_rounded = res.round().to(torch.int64)
|
| 177 |
+
for row, i in enumerate(idxs):
|
| 178 |
+
limbs = res_rounded[row].tolist()
|
| 179 |
+
# Circuit emits base-2^16 limbs little-endian; the decoder reads
|
| 180 |
+
# base-2^16 digits MSB-first, so reverse. Clamp each limb into
|
| 181 |
+
# [0, 2^16) defensively before emitting plain ints.
|
| 182 |
+
digits = [int(v) & (B16 - 1) for v in reversed(limbs)]
|
| 183 |
+
out[i] = digits if digits else [0]
|
| 184 |
+
|
| 185 |
+
return [o if o is not None else [0] for o in out]
|
| 186 |
+
|
| 187 |
+
def max_batch_size(self) -> int:
|
| 188 |
+
return 256
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
# Output base: one base-2^16 digit per limb. Within the schema's [2, 2^32].
|
| 192 |
+
B16 = 1 << LIMB_BITS
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def _buffers_to_parameters(circuit: ModmulCircuit) -> None:
|
| 196 |
+
"""Promote the circuit's constant buffers to float nn.Parameter.
|
| 197 |
+
|
| 198 |
+
The source ``ModmulCircuit`` registers ``step_one`` and ``gate_base`` as
|
| 199 |
+
buffers. Promoting them to parameters makes the weight-perturbation
|
| 200 |
+
behavioral signal act on them: perturbing the parameters perturbs exactly
|
| 201 |
+
the constants the forward pass reads, so accuracy collapses under noise.
|
| 202 |
+
The numeric values are unchanged (1.0 and 2^16), so the constructed circuit
|
| 203 |
+
stays bit-exact.
|
| 204 |
+
"""
|
| 205 |
+
for name in ("step_one", "gate_base"):
|
| 206 |
+
if name in circuit._buffers:
|
| 207 |
+
value = circuit._buffers.pop(name)
|
| 208 |
+
circuit.register_parameter(
|
| 209 |
+
name, nn.Parameter(value.detach().clone(), requires_grad=False)
|
| 210 |
+
)
|