fix(ops): cancellation-free table-lookup builders (min-clamp + saturating steps)

The map_select table-lookup builders _table_lookup_row_vector /
_table_lookup_column_mask / _table_lookup_index_staircase built a single
piecewise_linear over 2*(n-1)+1 breakpoints with input_scale=sharpness -- the
same single-projection sum delta*relu(s*x - b) floor_int had. The partial sums
grow to ~s*x*n and overflow fp32's 2^24 exact-integer limit, so the
alternating-sign terms cancel into garbage. Two triggers: (a) a scaled index far
outside [0,n-1] pushes sharpness*index past 2^24; (b) a tall table's top
breakpoint position passes 2^24. _table_lookup_index_staircase(n=256, s=100) at
x=100000 returned 0.0 (should clamp to 255 -- the index was never clamped before
the staircase); n=20000, s=1000 collapsed an in-range x=19999 to 0.0. Dormant at
current call sites (small tables, in-range pre-rounded indices) but a landmine
for the DOOM texture/palette lookup track's larger tables.

Rebuild all three on a shared _saturating_step_select helper, the table analogue
of floor_int's saturating-step staircase:
  out = top - sum_k relu(1 - step_k) * deltas[k-1],
  step_k = relu(t_k) - relu(t_k - W),  t_k = s*(x - (k-0.5)) + 0.5,
rows reconstructed from consecutive differences (deltas[k-1] = value[k] -
value[k-1], value[n-1] = top). Each step is a 2-term per-output difference, so
the accumulation's partial sums are bounded by the table's total variation, not
by s*x -- removing the cancellation regardless of n (trigger b). A leading
min(index, n-1) clamp (one sublayer, (n-1) - relu((n-1) - x), a single ReLU of a
magnitude, no cancellation) bounds the out-of-range case so W -- sized for the
in-range span -- always resolves the saturating difference (trigger a); the
lower edge needs no clamp. The +0.5 centers the ramp on the boundary,
reproducing the old two-row linear blend in the transition band, so off-grid
values are unchanged.

Cost: one min sublayer + the two chained ReLU sublayers of the staircase, vs the
old single sublayer. Reproducer tests at the smallest layer
(test_table_lookup_2d_out_of_range_index_clamps_to_edge, _tall_table_stays_exact,
_index_staircase_clamps_and_rounds_at_any_magnitude, _3d_large_flattened_rows_exact);
existing table_lookup tests (incl transition band + compile-fidelity) stay green.
table_lookup_2d/_3d are not measured ops so op_noise_data.json is unchanged;
docs/numerical_noise_findings.md "Known gap" -> "Resolved". Drops now-dead
_lookup_breakpoints / _lookup_integer_index_at and the unused piecewise_linear /
Assert module imports.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: physicsrob

docs/numerical_noise_findings.md

@@ -171,26 +171,56 @@ Exact at any magnitude/sharpness AND compiled-precise; pinned by
 slope-change-sum structure but is correct for its documented **integer-only**
 contract (integer inputs land in flat zones); it is not reworked.
 
-### Known gap: the `map_select` table-lookup staircases have the same 2^24 cancellation
+### Resolved: the `map_select` table-lookup staircases had the same 2^24 cancellation
 
 An audit for the `floor_int` cancellation class found one untriggered sibling.
 `_table_lookup_index_staircase` / `_table_lookup_row_vector` /
-`_table_lookup_column_mask` (`torchwright/ops/map_select.py`) build a
+`_table_lookup_column_mask` (`torchwright/ops/map_select.py`) built a
 `piecewise_linear` with `2·(n−1)+1` breakpoints and `input_scale=sharpness`
 (default 100), feeding `table_lookup_2d`/`table_lookup_3d`. Same single
 alternating-sign ReLU sum, so the same overflow:
-`_table_lookup_index_staircase(n=256, sharpness=100)` at `x=100000` returns
-`0.0` (should clamp to 255 — the index is never clamped before the staircase),
-and `n=20000, sharpness=1000` collapses an in-range `x=19999` to `0.0`
-(`s·max_breakpoint ≈ 2e7 > 2^24`). **Untriggered today**: current table
-consumers use small tables with in-range, pre-rounded integer indices, and the
-DOOM texture/palette lookup track is still a Phase-F stub. Fix before that
-track lands large tables — clamp the index into `[0, n−1]` *before* the
-staircase, and/or apply the saturating-step reformulation. The audit also
-flagged `compare` as a milder 2-term variant that fails silently for `|x| ≳
-2^24/sharpness` (default-sharpness boundary ~1.68e6; live callers stay far
-under it), and confirmed `square`/`exp`/`multiply_2d`/`low_rank_2d` safe
-(monotone slopes or `input_scale=1`).
+`_table_lookup_index_staircase(n=256, sharpness=100)` at `x=100000` returned
+`0.0` (should clamp to 255 — the index was never clamped before the staircase),
+and `n=20000, sharpness=1000` collapsed an in-range `x=19999` to `0.0`
+(`s·max_breakpoint ≈ 2e7 > 2^24`). Two triggers: **(a)** a scaled index far
+outside `[0, n−1]` pushes `sharpness·x` past `2^24`; **(b)** a tall table's top
+breakpoint position passes `2^24`.
+
+**Fixed** by rebuilding all three on the shared `_saturating_step_select`
+helper, the table analogue of `floor_int`'s saturating-step staircase. Two
+parts, one per trigger:
+
+- **(b) The staircase.** Each half-integer boundary `k − 0.5` contributes a
+  *bounded* step `step_k = relu(t_k) − relu(t_k − W)`
+  (`t_k = s·(x − (k−0.5)) + 0.5`, `W = max(2, 8·ulp(s·(n−1)+1))`), and the
+  output accumulates `out = top − Σ_k relu(1 − step_k)·deltas[k−1]` (rows
+  reconstructed from consecutive differences: `deltas[k−1] = value[k] −
+  value[k−1]`, `value[n−1] = top`). Every partial sum is bounded by the table's
+  total variation `Σ_k |deltas[k−1]|`, not by `s·x`, so it never overflows
+  regardless of `n`. The `+ 0.5` centers the ramp on the boundary
+  (indicator `= 0.5` there), reproducing the old two-row linear blend in the
+  transition band — off-grid values are unchanged.
+- **(a) The clamp.** A leading `min(index, n−1)` (one sublayer,
+  `(n−1) − relu((n−1) − x)` — a single ReLU of a magnitude, no cancellation)
+  bounds the out-of-range case so `W`, sized for the in-range span, always
+  resolves the saturating difference. The lower edge needs no clamp: a hugely
+  negative index leaves every step off (ReLU of a negative is exactly 0) and
+  selects row 0.
+
+Cost: one `min` sublayer plus the two chained ReLU sublayers of the staircase,
+versus the old single sublayer. Exact at any magnitude/sharpness and
+compiled-precise; pinned by `test_table_lookup_2d_out_of_range_index_clamps_to_edge`,
+`test_table_lookup_2d_tall_table_stays_exact`,
+`test_table_lookup_index_staircase_clamps_and_rounds_at_any_magnitude`, and
+`test_table_lookup_3d_large_flattened_rows_exact` in
+`tests/ops/test_resampling_primitives.py`. (`table_lookup_2d`/`_3d` are not in
+the measured-op set in `scripts/measure_op_noise.py`; adding them there for a
+tracked noise number is a reasonable follow-on.)
+
+The same audit flagged `compare` as a milder 2-term variant that fails silently
+for `|x| ≳ 2^24/sharpness` (default-sharpness boundary ~1.68e6; live callers
+stay far under it) and confirmed `square`/`exp`/`multiply_2d`/`low_rank_2d` safe
+(monotone slopes or `input_scale=1`); those are unchanged.
 
 ### Rel-error reporting: ramp-zone artefacts
 

tests/ops/test_resampling_primitives.py

@@ -1130,3 +1130,120 @@ def test_linear_bin_index_inv_range_probe():
     assert (
         report.first_divergent is None
     ), f"probe divergence on hoisted inv_range:\n{report.format_short()}"
+
+
+# ---------------------------------------------------------------------------
+# Cancellation-free index staircase (regression)
+#
+# The old single-projection ``piecewise_linear`` builders summed ~n ReLU terms
+# of magnitude ~sharpness*index into one accumulator, overflowing fp32's 2^24
+# exact-integer limit so the alternating-sign terms cancelled into garbage.
+# Two triggers: (a) a scaled index far outside [0, n-1] (sharpness*index past
+# 2^24), and (b) a tall table whose top breakpoint position passes 2^24.  See
+# docs/numerical_noise_findings.md, "the map_select table-lookup staircases".
+# Mirrors test_floor_int_wide_range_large_magnitude; the cancellation lives in
+# the exact-math fp32 sum that node.compute / reference_eval reproduces.
+# ---------------------------------------------------------------------------
+
+
+def test_table_lookup_2d_out_of_range_index_clamps_to_edge():
+    """Trigger (a): a row index far outside [0, rows-1] clamps to the table
+    edge instead of collapsing (index 1e5 at sharpness 100 -> sharpness*index
+    1e7 > 2^24 used to return ~0)."""
+    rows = 256
+    table = torch.arange(rows * 2, dtype=torch.float32).reshape(rows, 2)
+    out_node = _build_table_lookup_2d_graph(table)  # default sharpness=100
+
+    i_val = torch.tensor([[1.0e5], [-50.0], [0.0], [128.0]], dtype=torch.float32)
+    j_val = torch.zeros_like(i_val)
+    inputs = {"i": i_val, "j": j_val}
+    expected = torch.tensor(
+        [
+            [float(table[255, 0])],
+            [float(table[0, 0])],
+            [float(table[0, 0])],
+            [float(table[128, 0])],
+        ]
+    )
+    cache = reference_eval(out_node, inputs, i_val.shape[0])
+    assert torch.allclose(cache[out_node], expected, atol=5e-3), cache[out_node]
+
+    # ...and the min-clamp + saturating-step chain compiles precisely.
+    report = probe_graph(
+        out_node,
+        pos_encoding=None,
+        input_values=inputs,
+        n_pos=i_val.shape[0],
+        d=1024,
+        d_head=16,
+        verbose=False,
+        atol=5e-3,
+    )
+    assert report.first_divergent is None, report.format_short()
+
+
+def test_table_lookup_2d_tall_table_stays_exact():
+    """Trigger (b): a 20000-row table at sharpness 100 and 1000 keeps in-range
+    indices exact (an in-range index used to collapse to 0).  Oracle-only:
+    compiling a 20000-row table is needlessly expensive and the cancellation is
+    in the exact-math fp32 sum, like the floor_int wide-range regression."""
+    rows = 20000
+    table = torch.arange(rows, dtype=torch.float32).reshape(rows, 1)
+    idxs = [0.0, 1.0, 9999.0, 12345.0, 19998.0, 19999.0]
+    for s in (100.0, 1000.0):
+        i = create_input("i", 1)
+        j = create_input("j", 1)
+        out_node = table_lookup_2d(i, j, table, sharpness=s)
+        inputs = {
+            "i": torch.tensor([[v] for v in idxs]),
+            "j": torch.zeros(len(idxs), 1),
+        }
+        got = reference_eval(out_node, inputs, len(idxs))[out_node].squeeze(-1)
+        expected = torch.tensor(idxs)
+        assert torch.allclose(got, expected, atol=5e-3), (s, got, expected)
+
+
+def test_table_lookup_index_staircase_clamps_and_rounds_at_any_magnitude():
+    """The scalar staircase computes clamp(round(x), 0, n-1) exactly at any
+    magnitude — the smallest layer that reproduces both triggers (out-of-range
+    x and large n)."""
+    from torchwright.ops.map_select import _table_lookup_index_staircase
+
+    for n, s in [(256, 100.0), (256, 1000.0), (20000, 1000.0)]:
+        x = create_input("x", 1)
+        node = _table_lookup_index_staircase(
+            x, n, sharpness=s, d_max=1024, name="st"
+        )
+        cases = [
+            (-1.0e6, 0.0),
+            (-0.4, 0.0),
+            (0.0, 0.0),
+            (0.6, 1.0),
+            (float(n // 3), float(n // 3)),
+            (float(n - 1), float(n - 1)),
+            (1.0e6, float(n - 1)),
+        ]
+        xs = torch.tensor([[c[0]] for c in cases])
+        out = reference_eval(node, {"x": xs}, len(cases))[node].squeeze(-1)
+        exp = torch.tensor([c[1] for c in cases])
+        assert torch.allclose(out, exp, atol=5e-3), (n, s, out, exp)
+
+
+def test_table_lookup_3d_large_flattened_rows_exact():
+    """The 3D lookup flattens (A, B) into A*B rows; the motivating 16x128x128
+    shape reshapes to 2048 rows — well into the tall-table regime — and stays
+    exact on the integer grid through the cancellation-free row builder."""
+    t3 = torch.arange(16 * 128 * 128, dtype=torch.float32).reshape(16, 128, 128)
+    i = create_input("i", 1)
+    j = create_input("j", 1)
+    k = create_input("k", 1)
+    node = table_lookup_3d(i, j, k, t3)
+    cells = [(0, 0, 0), (15, 127, 127), (7, 64, 33), (3, 5, 120)]
+    inputs = {
+        "i": torch.tensor([[float(a)] for a, _, _ in cells]),
+        "j": torch.tensor([[float(b)] for _, b, _ in cells]),
+        "k": torch.tensor([[float(c)] for _, _, c in cells]),
+    }
+    got = reference_eval(node, inputs, len(cells))[node].squeeze(-1)
+    expected = torch.tensor([float(t3[a, b, c]) for a, b, c in cells])
+    assert torch.allclose(got, expected, atol=5e-3), (got, expected)