tests: un-xfail torch↔JAX goldens under the embed-internal token contract

[ocg-goodfire]

Summary

Fast-follow to the residual-start removal (3883362) and the #923 terminology sweep: the frozen torch↔JAX equivalence goldens (param_decomp/tests/equivalence/) and stacked-parity fixtures were generated under the old residual-fed contract and have been xfailed since the model went embed-internal (token ids in, embed[tokens] inside). This regenerates the fixtures — not the goldens — under the token contract and un-xfails the 10 pending tests.

Key move: the goldens are preserved bit-for-bit, no torch-oracle re-run needed. The loss terms are a pure function of the residual entering the first decomposed layer. The fixtures now carry tokens + embed where embed[tokens] reproduces the committed residual bit-exactly (one distinct token per position, its embed row set to that position's residual — a pure integer gather, no arithmetic). The JAX side now exercises the real token→embed forward, and its outputs still pin against the frozen torch-oracle values — which is exactly the parity these tests assert. Every pre-existing array in fixtures.npz is byte-identical (verified programmatically); only tokens + embed were added.

equivalence/ (6 tests un-xfailed)

• test_jax_matches_torch_reference[faith/imp/stoch/ppgd]: all four terms match the frozen torch_reference.json at fp32 tolerance (rel err ≤ 9e-7). SPEC S1/S3.
• test_sc_source_broadcasts_over_batch_in_masked_forward: token-fed, passes (S1/S16).
• test_chunk_plan_static_live_gate deleted, not un-xfailed: it pinned SPEC S2 under a layer-splitting chunk plan — a capability deliberately dropped by the static segmented masked forward (2f9af23: "Drops partial-layer (per_site / layerwise) recon support"). The residual-fed xfail masked that real blocker (the test now trips the layer-aligned assertion, not a numeric mismatch). The layer-aligned S2 realization remains pinned by the whole-layer stoch term and stacked test_chunk_plan_static_live_set_matches. Dead helpers (_forward_with_split_mlp, chunk_plan_static_gate_kl) removed with it.

stacked_parity/ (4 tests un-xfailed)

The fixtures pinned a residual-fed suffix (absolute layers 3..7 of an 8-layer tiny cfg, residual injected at layer 3), and site names carry absolute indices. _load now rebuilds the target as the full 8-layer embed-internal model with identity prefix layers (zeroed attn + zeroed MLP ⇒ block(x) == x): embed[tokens] (== the fixtures' residual) passes through the prefix unchanged, so layer 3 sees the identical residual and all CI-independent goldens hold. Calls migrated to the current API (prepare_compute_weights before masked_output).

• un-xfailed & passing: test_clean_output_matches, test_site_inputs_and_weight_deltas_match, test_masked_output_match, test_chunk_plan_static_live_set_matches (SPEC S2/S4/N2).
• untouched: test_train_trajectory_matches stays xfailed (strict) — its staleness is CI-numerics (old per-site-concat CIArch → chunkwise transformer), genuinely torch-oracle-dependent and out of scope here.

Prose sweep

gen_fixtures.py / jax_equivalence.py / test_equivalence.py / gen_stacked_fixtures.py no longer speak in "suffix weights" / "suffix forward" terms — the last live remnant of the residual-fed design's terminology in the repo. (References to the torch oracle's residual-fed design remain where they describe that historical contract accurately.)

Validation

• pytest param_decomp/tests/equivalence/ param_decomp/tests/stacked_parity/: 15 passed, 1 xfailed at default device count AND XLA_FLAGS=--xla_force_host_platform_device_count=4.
• python -m param_decomp.tests.equivalence.jax_equivalence: all 4 terms PASS vs the frozen golden.
• gen_fixtures.py re-run: byte-stable (all arrays identical across runs); pre-existing arrays byte-identical to the previously committed npz.
• ruff + basedpyright clean (pre-commit passed).

🤖 Generated with Claude Code