P3: SnAp-n — generalise recurrence_scope to an n-step neighbourhood

[chaoming0625]

Summary

Adds recurrence_scope='sparse_n' and the braintrace.SnAp preset — SnAp-n
(Menick et al., 2021), the interpolation between the cheap and the exact end of
the recurrence-scope axis.

The point of doing this in a compiler rather than a library: the n-step
influence sparsity pattern is derived, not declared. position_graph.py
reads the hidden group's transition jaxpr, extracts the one-step position
adjacency from the recurrent mixing primitive, closes it n-1 times, and hands
back the retained neighbourhood. The trace's trailing num_state axis S
widens to a (neighbour, state) axis of width M = K·S, and not one of the
nine ETP primitive rules needed changing.

Both ends of the scale turned out to be in the codebase already: n=1
canonicalises to recurrence_scope='coupled' (OSTLRecurrent's rule), and any
n at or above a group's position-graph diameter saturates to full within-group
RTRL. diagonal sits below the scale rather than on it.

Measured

sparse_ring_rnn(n_rec=6) (diameter 5), chunk_size=3, T=8, against the BPTT
oracle on the ETP parameter:

n	K	trace scalars	rel. deviation from BPTT
1	1	12	2.32e-01
2	2	24	4.96e-03
3	3	36	5.93e-04
4	4	48	2.22e-04
5	5	60	8.03e-08
6	6	72	8.03e-08

Memory exactly linear in K and saturating with it; accuracy monotone, reaching
round-off at n = diameter.

Safety properties

• snap_anchor is declared per primitive and defaults to deny. A trace slot
  must have one well-defined hidden position its instantaneous term lands on.
  Nine primitives declare it; etp_emb and shared-axis etp_einsum are
  rejected with a message naming the legal scopes.
• Unreadable adjacency widens, and says so. An unregistered mixing
  primitive, more than one mixing equation, control flow, or a
  non-position-preserving tail all fall back to all-to-all with
  SNAP_PATTERN_CONSERVATIVE. K == 1 emits SNAP_PATTERN_DEGENERATE rather
  than quietly computing coupled under a sparse_n label.
• The reachability closure is boolean throughout. An integer closure counts
  paths and wraps, which would under-approximate the neighbourhood — the one
  error direction a conservative analysis may never have.
• The widened block Jacobian is budgeted at P·(K·S)² elements, with
  snap_max_jacobian_elements to raise it deliberately.

Tests

122 new tests; full suite 2388 passed, mypy braintrace clean.

Acceptance is a two-sided squeeze — n=1 equals OSTLRecurrent element-wise
and saturation equals BPTT, both through the finite-window oracle since the
full-window path is blind to every axis (F-23) — plus:

• an independent hand-written masked-RTRL reference for the interior orders
  (n = 2,3,4 match to <1e-6; a deliberately narrower reference misses by
  1.2e-1), because the endpoints are no evidence about the middle;
• all eleven anchored primitives executed at a saturating order on both the
  fast and legacy solve paths;
• a strictly directed fixture pinning that the neighbourhood runs downstream
  of its anchor — the orientation no size-based assertion can see.

Also

Documents F-31 in the known-limitations list: on a hidden group whose
y → hidden tail relabels positions, no within-group scope reaches BPTT,
saturated SnAp included, and the saturated rule is further from BPTT than the
cheaper ones (1.30 vs 1.03, against 1.3e-07 for the same model with the
relabelling removed). Pre-existing and structural — it hits diagonal and
coupled identically — detected and warned about rather than silent. The
conservative diagnostic previously read "the gradient stays correct", which is
true of the pattern but over-claims about the returned gradient; it now says the
approximation is not degraded and points at F-31.

Roadmap lessons 18–30 record what generalises, including the five findings an
adversarial review of the finished implementation turned up.

🤖 Generated with Claude Code

[sourcery-ai]

Sorry @chaoming0625, your pull request is larger than the review limit of 150000 diff characters

[codecov]

Codecov Report

❌ Patch coverage is 99.04306% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
braintrace/_op/sparse.py	84.61%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!