refactor(ci): chunkwise CI fn — protocol, read_activations seam, scale-driven sharding

[ocg-goodfire]

Description

Replaces the single global-concat transformer CI fn with a CIFn protocol (dict[InputTap, Array] -> CI) and a chunkwise-transformer implementation, and untangles the concerns that were welded into the old arch.

• Protocol + bundle. CIFn: input keyspace (opaque tap keys) independent of output keyspace (decomposition sites); output sites must partition the model's sites. CI = {logits, lower, upper}, with the two squashings (S5/S6) centralized in CI.from_logits (was copy-pasted per impl).
• Model seam. DecomposedModel.site_inputs → read_activations(frozen, resid, wanted) — a general clean-path accessor keyed by opaque taps. The target is the sole key interpreter: the LM serves resid.{i} residual taps (chunkwise input) and per-site matrix inputs (harvest); toys serve per-site inputs. Deletes three copies of the per-site frozen-forward math.
• Chunkwise transformer. Sites partition into chunks; each chunk reads one or more taps (RMS-normed per tap → concatenated → input_dim), runs an independent pre-norm bidirectional-RoPE transformer, emits per-site CI. Per-chunk transformers are stacked along an n_chunks axis and run under one eqx.filter_vmap. input_dim is a generic linear width computed lab-side — no residual-dim/transformer concept leaks into core. layerwise/global are degenerate chunkings.
• One module, peers. All CI-fn arches (chunkwise transformer + the positionless MLPs, formerly ci_fn_mlp.py) live in ci_fn.py; CIFnArch + build_ci_fn unify construction.
• Sharding decoupled from arch. Scale-driven init_decomp_vu_placed / init_ci_fn_placed (shardable = mesh>1 and all C%n==0); the match ci_fn_arch in run_state is gone. A non-tileable CI matrix replicating is logged, not silent.
• Flat config. Single-type CI-config union (drops the nested mode × fn_type unions + the legacy null-drop shim).
• Lab chunk generator. blocks_per_chunk consecutive blocks → one residual tap in, block sites out.
• SPEC. §4.6 / S4 / S5 / S27 updated to the chunkwise realization (same read_activations → CI contract).

Related Issue

N/A

Motivation and Context

The old CI fn was a single transformer over every site's input concatenated — expensive, with sharding welded to arch identity (MLP→replicated / transformer→sharded, an incidental correlation) and a nested mode × fn_type config. The chunkwise design is the "no-regrets" generalization: it interpolates between layerwise (one chunk per site) and global (one chunk), reads the residual stream entering each chunk rather than every per-site activation, and makes the input/output keyspaces independent (enabling future read-from-elsewhere variants). Along the way the genuinely-orthogonal axes are separated: architecture (in ci_fn.py), construction (build_ci_fn), and placement (scale-driven), so adding a CI fn no longer touches a dispatch match.

How Has This Been Tested?

• make type: 0 errors (after make install-dev syncs the worktree venv).
• Core tests: 117 passed / 7 skipped / 1 xfailed at 1 device; sharding-sensitive subset 14 passed at 4 sim-devices (exercises the new C-sharded path).
• Equivalence suite: 12 passed, unchanged — confirms the CI-fn change is correctly isolated from the recon / faithfulness / source-term goldens.
• Toys (TMS / ResidMLP): 45 passed, end-to-end through read_activations + the training step.
• A real core bug surfaced only by running (a lazy-split StopIteration in per-chunk init, invisible to the type checker) and is fixed; per-site read_activations outputs are proven bit-identical to the deleted re-threads via the committed out::site_input stacked-parity goldens.

Does this PR introduce a breaking change?

Yes.

• CI-fn config schema changed to the flat type form (type: chunkwise_transformer / layerwise_mlp / global_mlp); old mode/fn_type YAMLs must be migrated (in-repo configs are migrated here).
• CI-fn numerics changed by design — the CI fn is now JAX-native and no longer bit-faithful to the torch oracle's global-concat CI fn. The rest of the spec stays oracle-grounded.
• Checkpoints with the old CI fn are not loadable; the stacked-parity CI-trajectory golden is xfail pending a torch-oracle regen.

🤖 Generated with Claude Code