fix(toys): pd-tms/pd-resid-mlp eval crash — missing remat kwarg

[ocg-goodfire]

What

CIFn.__call__ grew a keyword-only remat: bool, but the toy composition roots' eval closures were never updated — both pd-tms and pd-resid-mlp crashed at the first eval step:

TypeError: LayerwiseMLPCIFn.__call__() missing 1 required keyword-only argument: 'remat'

Reproduced with stock tms_5-5_SMOKE.yaml on f423233fb. (Found by bridge crew-try-muon while using the TMS smoke as a muon testbed.)

Why the tests didn't catch it

Each run.py carried a near-duplicate of its domain's model.single_feature_ci helper. The tests exercise the helper (which passes remat=False correctly); the run.py duplicate rotted unseen.

Fix

Deduplicate rather than patch: model.single_feature_ci now returns the full CI bundle (run.py needs upper for the UV figure, tests take .lower), and each run.py jit-wraps the helper instead of re-implementing it. The remat=False decision now lives in exactly one place per domain, on the code path the tests cover. remat is a no-op for the MLP CI fns, so remat=False is semantics-neutral.

Verification

• Crash reproduced verbatim on unfixed tree, gone after fix
• pd-tms --config tms_5-5_SMOKE.yaml and pd-resid-mlp --config resid_mlp_1l_SMOKE.yaml run end-to-end on CPU; eval/identity_ci_error/* metrics identical between the minimal kwarg-only fix and the dedup version
• Toy suites: 42 passed + 2 slow passed (--runslow)
• basedpyright + ruff clean

Remaining gap (not addressed here)

The toy composition roots (run.py: YAML→BuiltRun→engine + eval_fn wiring) still have no automated coverage — the SMOKE configs are run manually. A ~10s --runslow test invoking the SMOKE config with a tmpdir PARAM_DECOMP_OUT_DIR would close this class of rot entirely; left out to keep this PR minimal.

🤖 Generated with Claude Code

[ocg-goodfire]

Base corrected main → feature/jax: the remat kwarg this fixes exists only on feature/jax (main's lab has no remat call sites), so the bug doesn't exist on main. Branch still carries foreign commit f423233 (checkpoint fix, delivered by #927); a rebase dropping it is staged at bridge/task-fix-bridge-pr-base-branches/clean-928 pending force-push approval.

[ocg-goodfire]

Branch rewritten (with lease) per the base-branch cleanup task: dropped foreign commit f423233 (carried by #927) and rebased the toys-remat fix onto current feature/jax (now 1b192fc, patch-identical to 5e120e1). Verified: tms + resid_mlp test suites pass on the new head (42 passed, 2 skipped).