Address review findings on the PosEncoding redesign

Follow-up to b192865, from an adversarial review of that commit.

- attend_most_recent_matching(exclude_self=True): drop the stale
  len(value)/len(key_vector) <= d_pos asserts. The pre-shift uses
  attend_to_offset, whose V/O auto-splits across heads, so the only real
  bound (d_head >= W+1 for the head's Q/K) is already enforced by the
  compiler. The old asserts rejected valid wide inputs.
- get_prev_value: rewrite the numerical-limit docstring. The max_pos
  bound is exact logit ordering — a recent false key outscores an early
  true once the runtime position exceeds ~2*max_pos ({0,1} cond) or
  ~4*max_pos ({-1,+1}) — NOT TF32 precision; the prior "needs fp32" note
  was wrong. Also document the cond conventions and the all-false window.
- attend_argmin/argmax/argmin_where: drop stale "len(value) <= d_pos"
  docstrings (wide V/O auto-splits).
- PosEncoding: reject d_pos < 3 (trig_width 0 divides by zero on use).

Tests:
- compiled round-trips for attend_argmin/argmax/_where (previously only
  oracle-tested), incl. a wide-value (40 > d_head) auto-split case.
- get_prev_value latch-within-budget and adjacent-true recency.
- _current_pos_attn_matrices zeroes the counter column across d_head.

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

Author: physicsrob

tests/compile/forward/test_forward_compile.py

@@ -27,6 +27,12 @@
 )
 from torchwright.ops.logic_ops import cond_gate, cond_add_vector
 from torchwright.ops.map_select import select, map_to_table
+from torchwright.ops.attention_ops import (
+    attend_argmin,
+    attend_argmax,
+    attend_argmin_where,
+    attend_argmax_where,
+)
 
 D = 256
 D_HEAD = 16
@@ -623,6 +629,50 @@ def test_compile_split_vo_large_ratio():
     ), f"Max diff: {(actual.cpu() - expected).abs().max().item():.6f}"
 
 
+def test_compile_selection_primitives():
+    """Compiled path for the simple selection primitives, whose d_qk shrank
+    to 1-2 and whose query is now a bare LiteralValue([1.0])."""
+    pos = create_pos_encoding()  # PosEncoding(17), trig_width 16 == D_HEAD
+    score = create_input("score", 1)
+    validity = create_input("valid", 1)
+    value = create_input("value", 3)
+
+    n_pos = 6
+    inputs = {
+        "score": torch.tensor([[3.0], [1.0], [4.0], [0.0], [2.0], [5.0]]),
+        # position 0 is always valid, so every causal window has a valid key
+        "valid": torch.tensor([[1.0], [-1.0], [1.0], [-1.0], [1.0], [-1.0]]),
+        "value": torch.arange(n_pos * 3, dtype=torch.float32).reshape(n_pos, 3),
+    }
+
+    for out in (
+        attend_argmin(pos, score, value),
+        attend_argmax(pos, score, value),
+        attend_argmin_where(pos, score, validity, value),
+        attend_argmax_where(pos, score, validity, value),
+    ):
+        _verify(out, n_pos=n_pos, input_values=inputs, pos_encoding=pos)
+
+
+def test_compile_selection_wide_value():
+    """Selection value wider than d_head must auto-split across heads (the
+    len(value) <= d_pos cap was removed)."""
+    pos = create_pos_encoding()
+    score = create_input("score", 1)
+    value = create_input("value", 40)  # > D_HEAD = 16
+    n_pos = 5
+    inputs = {
+        "score": torch.tensor([[3.0], [1.0], [2.0], [0.0], [4.0]]),
+        "value": torch.arange(n_pos * 40, dtype=torch.float32).reshape(n_pos, 40),
+    }
+    _verify(
+        attend_argmin(pos, score, value),
+        n_pos=n_pos,
+        input_values=inputs,
+        pos_encoding=pos,
+    )
+
+
 def test_compile_rejects_d_head_below_trig_width():
     """forward_compile rejects d_head < pos_encoding.trig_width early, before
     the copy heads would silently match on only part of the trig grid."""

tests/compile/forward/test_weight_writer.py

@@ -24,6 +24,7 @@
     MLPOp,
     write_attn_sublayer,
     write_mlp_sublayer,
+    _current_pos_attn_matrices,
 )
 from torchwright.compiler.groups.transformer_layer import TransformerLayer
 from torchwright.graph import Linear, ReLU, Attn, Add, Concatenate
@@ -191,6 +192,20 @@ def test_attn_compute():
     assert torch.allclose(result.cpu(), expected, atol=1e-4)
 
 
+def test_current_pos_matrices_exclude_counter():
+    """Same-position copy heads must match on the trig block only — the raw
+    counter column is zeroed regardless of how d_head compares to d_pos."""
+    pe = PosEncoding(17)  # counter_col == 16
+    counter_col = pe.counter_col
+    for d_head in (8, 16, 17, 20):
+        q, k = _current_pos_attn_matrices(pe, d_head)
+        assert torch.all(q[counter_col] == 0.0), f"q counter not zero at d_head={d_head}"
+        assert torch.all(k[counter_col] == 0.0), f"k counter not zero at d_head={d_head}"
+    # A trig row IS matched (non-degenerate diagonal).
+    q16, k16 = _current_pos_attn_matrices(pe, 16)
+    assert q16[0, 0] != 0.0 and k16[0, 0] != 0.0
+
+
 def test_attn_compute_small_d_head():
     """Attn node with d_head smaller than layer d_head — needs padding."""
     pos = _make_pos_encoding()

tests/graph/test_pos_encoding.py

@@ -44,6 +44,46 @@ def test_get_prev_value():
     assert torch.allclose(output, expected_prev_values)
 
 
+def test_get_prev_value_latch_within_budget():
+    """A value latched at an early trigger is held across a long-ish sequence
+    as long as the runtime length stays within the max_pos budget (< 2*max_pos
+    for a {0,1} cond)."""
+    n_pos = 400  # < 2 * default max_pos (512)
+    value_input = InputNode("value", 1, value_range=(-1000.0, 1000.0))
+    cond_input = InputNode("cond", 1, value_range=(-100.0, 100.0))
+    pos_encoding = PosEncoding(17)
+    out = pos_encoding.get_prev_value(value_input, cond_input)
+
+    cond = torch.zeros(n_pos, 1)
+    cond[3, 0] = 1.0  # single true at position 3
+    value = torch.arange(n_pos, dtype=torch.float32).reshape(n_pos, 1)
+    value[3, 0] = 42.0
+    res = out.compute(n_pos=n_pos, input_values={"value": value, "cond": cond})
+    # Every query at or after the trigger returns the latched value.
+    assert torch.allclose(res[3:], torch.full((n_pos - 3, 1), 42.0))
+
+
+def test_get_prev_value_adjacent_trues_pick_most_recent():
+    """Among adjacent true positions, the most recent (largest counter) wins."""
+    n_pos = 10
+    value_input = InputNode("value", 1, value_range=(-100.0, 100.0))
+    cond_input = InputNode("cond", 1, value_range=(-100.0, 100.0))
+    pos_encoding = PosEncoding(17)
+    out = pos_encoding.get_prev_value(value_input, cond_input)
+
+    cond = torch.zeros(n_pos, 1)
+    cond[5, 0] = 1.0
+    cond[6, 0] = 1.0  # adjacent trues at 5 and 6
+    value = torch.zeros(n_pos, 1)
+    value[5, 0] = 50.0
+    value[6, 0] = 60.0
+    res = out.compute(n_pos=n_pos, input_values={"value": value, "cond": cond})
+    # Most recent true (pos 6 -> 60) dominates; ~exp(-recency_gap) contamination
+    # from the adjacent true (50) keeps it just under 60.
+    assert torch.allclose(res[6], torch.tensor([60.0]), atol=0.05)
+    assert torch.allclose(res[9], torch.tensor([60.0]), atol=0.05)
+
+
 def test_attend_to_offset():
     input_values = torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0]])
     value_input = InputNode("value", 1, value_range=(-100.0, 100.0))

tests/ops/test_attention_ops.py

@@ -1250,31 +1250,24 @@ def test_most_recent_matching_exclude_self():
     assert abs(result_excl[3].item() - 30.0) < 1e-2, result_excl[3]
 
 
-def test_most_recent_matching_exclude_self_width_constraint():
-    """``exclude_self=True`` requires ``len(value) <= d_pos`` and
-    ``len(key_vector) <= d_pos`` because each pre-shift compiles to a
-    single attention head whose width is bounded by ``d_pos``.  The
-    default mode has no such limit.
+def test_most_recent_matching_exclude_self_accepts_wide_width():
+    """``exclude_self=True`` no longer caps key/value width: the pre-shift
+    uses ``attend_to_offset``, whose V/O auto-splits across heads.  Both modes
+    build for widths wider than the encoding (regression for the removed caps).
     """
-    import pytest
-
-    pe = _pe()  # d_pos = 16
+    pe = _pe()  # PosEncoding(17), trig_width 16
     query = InputNode("query", 4, value_range=(-1.0, 1.0))
     key = InputNode("key", 4, value_range=(-1.0, 1.0))
     wide_value = InputNode("value", 24, value_range=(-100.0, 100.0))
 
-    # Default mode accepts width 24 > d_pos = 16.
+    # Wide value (24 > trig_width 16) builds in both modes.
     attend_most_recent_matching(pe, query, key, wide_value)
+    attend_most_recent_matching(pe, query, key, wide_value, exclude_self=True)
 
-    # exclude_self mode rejects it.
-    with pytest.raises(AssertionError, match="value width"):
-        attend_most_recent_matching(pe, query, key, wide_value, exclude_self=True)
-
-    # Symmetric check on key width.
+    # Wide key_vector likewise.
     wide_query = InputNode("query", 24, value_range=(-1.0, 1.0))
     wide_key = InputNode("key", 24, value_range=(-1.0, 1.0))
     narrow_value = InputNode("value", 1, value_range=(-100.0, 100.0))
-    with pytest.raises(AssertionError, match="key_vector width"):
-        attend_most_recent_matching(
-            pe, wide_query, wide_key, narrow_value, exclude_self=True
-        )
+    attend_most_recent_matching(
+        pe, wide_query, wide_key, narrow_value, exclude_self=True
+    )

torchwright/graph/pos_encoding.py

@@ -31,6 +31,13 @@ def __init__(self, d_pos: int):
                 f"even d_pos would leave an incomplete sin/cos pair. Use "
                 f"{d_pos + 1} or {d_pos - 1}."
             )
+        if d_pos < 3:
+            raise ValueError(
+                f"PosEncoding requires d_pos >= 3 (got {d_pos}): at least one "
+                f"complete sin/cos pair (trig_width >= 2) plus the counter "
+                f"column. trig_width = 0 has no position discrimination and "
+                f"divides by zero when building the frequency grid."
+            )
         self.d_pos = d_pos
         super().__init__(d_pos, [])
 
@@ -130,18 +137,42 @@ def get_prev_value(
         """Most-recent previous ``value`` at a position where ``cond`` is true.
 
         For each query position, selects ``value`` at the most recent causal
-        position where ``cond == 1.0`` (false positions must be ``<= 0``).
-        Recency is ranked by the raw position counter: among true positions
-        the largest counter (most recent) wins by ``recency_logit_gap`` of
-        logit per position.  A condition gate that dominates the full recency
-        span makes any true key beat any false key.
-
-        Numerical limit.  The gate forces peak logits near
-        ``3 * recency_logit_gap * max_pos``.  On A100 TF32 (~1e-3 relative
-        precision) the per-position recency gap stays resolvable only while
-        ``max_pos`` is a few hundred; the default 256 is safe, ~512 is the
-        ceiling.  Larger reach needs an fp32 matmul for this head.  Runtime
-        positions must be ``<= max_pos``.
+        position where ``cond`` is true.  ``cond`` may be either torchwright
+        boolean convention: ``{0, 1}`` (false = 0) or ``{-1, +1}`` (false = -1);
+        true must be ``+1`` and false must be ``<= 0``.  Recency is ranked by
+        the raw position counter: among true positions the largest counter
+        (most recent) wins by ``recency_logit_gap`` of logit per position.  A
+        condition gate (``2 * recency_logit_gap * max_pos``) makes a true key
+        outscore a false key.
+
+        **Selection contract — this is exact arithmetic, not precision.**  The
+        gate is frozen at graph-build time from ``max_pos``, but the raw
+        counter grows with the *actual* runtime sequence length.  A true key at
+        counter ``t`` beats a false key at counter ``f`` only while
+        ``gate > recency_logit_gap * (f - t)``, i.e. while the runtime position
+        stays below ``2 * max_pos`` (false = 0) or ``4 * max_pos``
+        (false = -1).  **Past that bound a recent false key silently outscores
+        an early true** (e.g. a value latched at an early trigger and held
+        across a long rollout) — no error, no tie, just the wrong value.  So
+        ``max_pos`` must be set to at least the longest sequence this head will
+        ever see (the default 256 covers sequences up to ~512 / ~1024).  This
+        mirrors the caller-supplied ``match_gain`` invariant documented on
+        :func:`~torchwright.ops.attention_ops.attend_most_recent_matching`.
+        (fp32 does **not** move this bound — it is logit ordering, not mantissa
+        width.)
+
+        **Tiebreak precision — TF32, soft.**  Among *multiple* true positions,
+        resolving the most recent needs the ``recency_logit_gap`` (8) to exceed
+        the TF32 noise floor at the peak logit (~``3 * recency_logit_gap *
+        max_pos``); this holds for a few hundred positions and then degrades
+        gracefully to "approximately the most recent true" (a blend of the last
+        few trues), not a wrong false.  The common single-true latch never hits
+        this.
+
+        **All-false window.**  If no causal position has ``cond`` true, the
+        gate term is uniform and the head degrades to pure recency, returning
+        the most recent position's value.  Callers must ensure at least one
+        true key exists in every consumed window, or gate the result.
         """
         assert len(cond) == 1, "get_prev_value expects a 1-D boolean cond"
         d_v = len(value)

torchwright/ops/attention_ops.py

@@ -295,7 +295,8 @@ def attend_argmin(
         pos_encoding: The graph's positional encoding node.
         score: 1D scalar node (``len(score) == 1``).
         value: Node whose value to read at the winning position.
-            Must satisfy ``len(value) <= pos_encoding.d_pos``.
+            No width constraint — wide V/O is auto-split across physical
+            heads by the compiler.
 
     Returns:
         A new ``Attn`` node of width ``len(value)`` equal to ``value`` at
@@ -328,7 +329,8 @@ def attend_argmax(
     Args:
         pos_encoding: The graph's positional encoding node.
         score: 1D scalar node.
-        value: Node whose value to read. ``len(value) <= pos_encoding.d_pos``.
+        value: Node whose value to read. No width constraint (wide V/O
+            auto-splits across heads).
 
     Returns:
         Attn node of width ``len(value)`` equal to ``value`` at the
@@ -384,7 +386,8 @@ def attend_argmin_where(
         pos_encoding: The graph's positional encoding node.
         score: 1D scalar node.
         validity: 1D boolean node (+1 valid, −1 invalid).
-        value: Node to read. ``len(value) <= pos_encoding.d_pos``.
+        value: Node to read. No width constraint (wide V/O auto-splits
+            across heads).
 
     Returns:
         Attn node of width ``len(value)``.
@@ -1279,12 +1282,11 @@ def attend_most_recent_matching(
             Implemented by pre-shifting ``key_vector`` and ``value`` by
             one position via :meth:`PosEncoding.attend_to_offset`, so
             that at slot ``i`` the head sees the predecessor's key and
-            value.  Costs two extra attention heads (one for the key
-            shift, one for the value shift) and adds the constraint
-            ``len(key_vector) <= d_pos`` and ``len(value) <= d_pos``
-            (the shift's per-head width limit).  At query position 0
-            the result is degenerate (no prior position exists) — the
-            caller must not consume it there.
+            value.  Costs two extra attention heads for the key and
+            value shifts (wide key/value auto-split across more heads;
+            no width cap).  At query position 0 the result is degenerate
+            (no prior position exists) — the caller must not consume it
+            there.
         assert_hardness_gt: If set, wraps the output in a softmax
             hardness assertion verifying the max attention weight per
             query exceeds this threshold during ``debug=True`` forward
@@ -1307,16 +1309,10 @@ def attend_most_recent_matching(
     )
     d_pos = pos_encoding.d_pos
     if exclude_self:
-        assert len(key_vector) <= d_pos, (
-            f"attend_most_recent_matching(exclude_self=True): "
-            f"key_vector width ({len(key_vector)}) must be <= d_pos ({d_pos}); "
-            "the position-shift step uses one attention head per shift."
-        )
-        assert len(value) <= d_pos, (
-            f"attend_most_recent_matching(exclude_self=True): "
-            f"value width ({len(value)}) must be <= d_pos ({d_pos}); "
-            "the position-shift step uses one attention head per shift."
-        )
+        # The shift uses attend_to_offset, whose Q/K width is trig_width and
+        # whose V/O auto-splits across heads, so key_vector / value may be any
+        # width (the only binding constraint, d_head >= W + 1 for this head's
+        # Q/K, is enforced by the compiler).
         # Shift key and value back by one position so that at slot i the
         # head sees key_vector[i-1] / value[i-1].  The causal mask still
         # admits i in [0, j], but the rightmost slot i = j now carries