chaobrain/braintrace: braintrace 0.2.6

Upgrade if you use JAX 0.11. braintrace 0.2.5 does not work with jax/jaxlib 0.11 — every gradient path fails on the first transition evaluation. This release fixes that. Support for JAX 0.8 through 0.10 is unchanged, and there are no public API changes. A compatibility release. JAX 0.11 merged ClosedJaxpr into Jaxpr and, in doing so, silently changed what a compiler-built transition jaxpr reports about its own inputs. On 0.11.1 that broke 690 of the 2902 tests in the suite; all of them pass again. What broke A jaxpr now holds a single positional input list and derives the constvars / invars boundary from how many constant values are attached, instead of storing it alongside the symbols: constvars = all_invars[:len(consts)] # consts = attached VALUES invars = all_invars[len(consts):] ClosedJaxpr(jaxpr, consts) still round-trips, because attaching values restores the boundary. But the ETP compiler builds transition jaxprs — programs whose invars are the differentiated inputs (a hidden state at t-1, or an ETP primitive's output y) and whose constvars are surrounding intermediates bound from the forward pass at execution time. Those carry symbols but no attached values, so on 0.11 they began reporting constvars == [] and folding the constvars into invars. Every consumer that recovered the split by reading jaxpr.constvars back therefore passed too few values to jax.core.eval_jaxpr, which failed arity checking with ValueError: foreach() argument 2 is shorter than argument 1. Since that is the shared entry point for hidden→hidden and hidden→weight Jacobians, it took down essentially every gradient path — D-RTRL, ES-D-RTRL / pp_prop, EProp, OSTL, SnAp, UORO, DNI and all of the BPTT oracle cross-checks — across every ETP primitive family (dense, LoRA, sparse, convolutional, element-wise, embedding). The fix The split is now derived from the invar count, which every caller already knows since the invars are what it is about to feed in. Three helpers in braintrace._compatible_imports — jaxpr_all_invars, split_jaxpr_invars and jaxpr_constvars — do this with no version branching, so 0.8 → 0.11 all run the same code path. split_jaxpr_invars range-checks its argument, so a miscounted call raises at the compiler boundary rather than producing a misaligned argument list. Raising the jax floor to 0.11 was rejected: it would drop three supported versions for an internal representation change that is straightforward to abstract over. Also fixed: SnAp-n position analysis. build_snap_pattern let the adjacency walk seed from the transition jaxpr's own invars. Under the merged representation that set silently widened to include every constvar, so the derived neighbourhood was inflated — typically all the way to the conservative all-positions-couple fallback, which is correct but costs orders of magnitude in trace size. The group's hidden_invars are now passed explicitly. Verification Full suite: 2912 passed, 0 failed on jax/jaxlib 0.11.1 (was 690 failed / 2212 passed) mypy gate: clean across all 62 source files CI now pins 0.11.0 in the jax-version matrix alongside 0.8.0 / 0.9.0 / 0.10.0 and latest, so the merged-Jaxpr representation stays covered once latest moves past it New regression coverage at both levels: the helpers directly (round-trip, zero- and all-constvar edges, out-of-range rejection), and the compiler (a compiled group's recorded constvars match its jaxpr's leading inputs, and the transition still evaluates) The full analysis, including the alternatives considered, is in docs/specs/2026-08-18-jax-011-jaxpr-merge-compat.md. Full Changelog: https://github.com/chaobrain/braintrace/compare/v0.2.5...v0.2.6