From f2315c78cd51fc33757cd365fd396c6a5c5b5e2e Mon Sep 17 00:00:00 2001
From: Antovigo <antvig@protonmail.com>
Date: Fri, 24 Jul 2026 19:24:52 +0000
Subject: [PATCH 1/5] feat(init): coupled component initialization
 (weight_init: kaiming | coupled)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

JAX port of the torch-lineage coupled init (experiment/8B_targeted
e5cdcbc80): unit-norm seed on the narrow side, wide side its raw W-image
(d_in <= d_out: v_c ~ unit norm, U_c <- (W v_c)^T; else u_c ~ unit norm,
V_c <- W^T u_c). No C-dependent rescale — components sit at W's natural
scale and the component sum equals W restricted to the seed span, the
delta carrying the complement.

pd.weight_init selects the init inside init_decomposition (match
dispatch, kaiming default keeps the eval_shape restore path and stored
configs untouched). W is recovered protocol-only via weight_deltas on
zero V/U — no new DecomposedModel method, no per-target changes. The
draw is vmapped per shape group like init_stack_arrays so the compiled
init doesn't scale with site count, and the placed init grew an optional
target_weights arg (riding as jit arguments, not closure constants)
instead of a second placement helper.

Verified: unit tests (seed determinism, W-image/span properties, stacked
same-shape sites), make check, and TMS 5-2 engine smokes for both modes.
The full core suite segfaults at test_llama8b under this box's memory
pressure with or without this change; the test passes in isolation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UVg2f69eq8mhDrpZZ3ZEmM
---
 param_decomp/components.py                    | 35 +++++++++
 param_decomp/configs.py                       |  8 ++
 param_decomp/run_state.py                     | 20 ++++-
 .../targets/glu_transformer_sharding.py       | 20 +++--
 param_decomp/tests/test_component_init.py     | 78 +++++++++++++++++++
 5 files changed, 153 insertions(+), 8 deletions(-)
 create mode 100644 param_decomp/tests/test_component_init.py

diff --git a/param_decomp/components.py b/param_decomp/components.py
index 2fac68c75..bcc67cc7e 100644
--- a/param_decomp/components.py
+++ b/param_decomp/components.py
@@ -169,6 +169,41 @@ def init_stack_arrays(
     return stacked
 
 
+def init_coupled_component_stacks(
+    sites: tuple[SiteSpec, ...], target_weights: dict[str, Array], key: Array
+) -> ComponentStacks:
+    """Coupled init: unit-norm seed on the narrow side, wide side its raw W-image.
+
+    `d_in <= d_out`: `v_c ~ unit norm`, `U_c <- (W v_c)^T`; else `u_c ~ unit norm`,
+    `V_c <- W^T u_c`. No C-dependent rescale — components sit at W's natural scale and
+    the component sum equals W restricted to the seed span (`W V V^T` when V is seeded),
+    the delta carrying the complement. One key per site, split in site order; vmapped
+    per shape group like `init_stack_arrays` so the compiled init doesn't scale with
+    site count."""
+    keys = jax.random.split(key, len(sites))
+    site_index = {spec.name: idx for idx, spec in enumerate(sites)}
+    stacked: dict[VUShape, tuple[Array, Array]] = {}
+    for (d_in, d_out, c), specs in vu_shape_groups(sites).items():
+        ws = jnp.stack([target_weights[spec.name] for spec in specs]).astype(jnp.float32)
+        assert ws.shape == (len(specs), d_out, d_in), (ws.shape, (d_in, d_out, c))
+        ks = keys[jnp.array([site_index[spec.name] for spec in specs])]
+        if d_in <= d_out:
+
+            def seed_v(k: Array, w: Array, s: tuple[int, int] = (d_in, c)) -> tuple[Array, Array]:
+                v = jax.random.normal(k, s)
+                v = v / jnp.linalg.norm(v, axis=0, keepdims=True)
+                return v, (w @ v).T
+        else:
+
+            def seed_v(k: Array, w: Array, s: tuple[int, int] = (c, d_out)) -> tuple[Array, Array]:
+                u = jax.random.normal(k, s)
+                u = u / jnp.linalg.norm(u, axis=1, keepdims=True)
+                return w.T @ u.T, u
+
+        stacked[(d_in, d_out, c)] = jax.vmap(seed_v)(ks, ws)
+    return ComponentStacks(stacks=stacked, site_slots=site_slots_for(sites))
+
+
 def component_stacks_from_sites(vu: dict[str, tuple[Array, Array]]) -> ComponentStacks:
     """Build the stacked `ComponentStacks` from a per-site `{name: (V, U)}` dict (site order =
     dict order). The explicit-arrays constructor for toys and tests; the trainer inits
diff --git a/param_decomp/configs.py b/param_decomp/configs.py
index eea12ab23..223f2b3d8 100644
--- a/param_decomp/configs.py
+++ b/param_decomp/configs.py
@@ -1041,6 +1041,14 @@ def _strip_removed_jax_unsupported_fields(cls, data: object) -> object:
         default=0,
         description="Random seed for reproducibility, including LM dataset shuffling.",
     )
+    weight_init: Literal["kaiming", "coupled"] = Field(
+        default="kaiming",
+        description="How component V/U are initialized. 'kaiming': iid normal "
+        "(V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5)), ignores W. 'coupled': unit-norm seed on "
+        "the narrow side, wide side its raw W-image (U_c from (W v_c)^T when d_in <= d_out, "
+        "else V_c from W^T u_c) — W-natural scale, component sum ~ W on a rank-C random "
+        "subspace (the delta carries the complement).",
+    )
     loss_metrics: list[AnyLossMetricConfig] = Field(
         default_factory=list,
         description=(
diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py
index a0c898743..238a62a8f 100644
--- a/param_decomp/run_state.py
+++ b/param_decomp/run_state.py
@@ -8,6 +8,7 @@
 """
 
 from collections.abc import Callable
+from typing import Literal
 
 import equinox as eqx
 import jax
@@ -20,6 +21,7 @@
 
 from param_decomp.adversary import PersistentAdversary, init_sources_adam_state
 from param_decomp.ci_fn import ChunkwiseTransformerCIArch, CIFnArch
+from param_decomp.components import component_stacks_from_sites
 from param_decomp.configs import (
     AdamPGDConfig,
     AdamWOptimizerConfig,
@@ -163,14 +165,26 @@ def init_decomposition(
     init_key: PRNGKeyArray,
     mesh: Mesh,
     rules: PlacementRules,
+    weight_init: Literal["kaiming", "coupled"] = "kaiming",
 ) -> Decomposition:
     """The trained-product half of `init_train_state`, factored out so a consumer can
     `jax.eval_shape` it to recover the saved `decomposition` item's tree structure
-    without building (or knowing about) the optimizers/adversaries."""
+    without building (or knowing about) the optimizers/adversaries (the `weight_init`
+    default is safe there — both inits produce the same tree)."""
     ci_key = random.fold_in(init_key, 1)
     # V/U placement derives from the rules table; the CI fn still declares its own
     # per-leaf shardings (the staged placement migration hasn't reached it).
-    components = init_component_stacks_placed(model.sites, init_key, rules)
+    match weight_init:
+        case "kaiming":
+            components = init_component_stacks_placed(model.sites, init_key, rules)
+        case "coupled":
+            # The protocol exposes W only through `weight_deltas`; on zero V/U the delta IS W.
+            zero_vu = component_stacks_from_sites(
+                {s.name: (jnp.zeros((s.d_in, s.C)), jnp.zeros((s.C, s.d_out))) for s in model.sites}
+            )
+            components = init_component_stacks_placed(
+                model.sites, init_key, rules, target_weights=model.weight_deltas(zero_vu)
+            )
     ci_fn = init_ci_fn_placed(ci_fn_arch, model.sites, ci_key, mesh)
     assert ci_fn.has_position_axis == model.has_position_axis, (
         f"CI fn has_position_axis={ci_fn.has_position_axis} but model declares "
@@ -195,7 +209,7 @@ def init_train_state(
     assert isinstance(positions, Positioned) == model.has_position_axis, (
         f"{positions} does not match the model's has_position_axis={model.has_position_axis}"
     )
-    decomposition = init_decomposition(model, ci_fn_arch, init_key, mesh, rules)
+    decomposition = init_decomposition(model, ci_fn_arch, init_key, mesh, rules, pd.weight_init)
     components, ci_fn = decomposition.components, decomposition.ci_fn
     losses = build_loss_terms(pd.loss_metrics, model.site_names)
     persistent = persistent_configs(losses.recon)
diff --git a/param_decomp/targets/glu_transformer_sharding.py b/param_decomp/targets/glu_transformer_sharding.py
index 08b334f5b..24da04337 100644
--- a/param_decomp/targets/glu_transformer_sharding.py
+++ b/param_decomp/targets/glu_transformer_sharding.py
@@ -64,6 +64,7 @@
     ComponentStacks,
     SiteSpec,
     init_component_stacks,
+    init_coupled_component_stacks,
 )
 from param_decomp.configs import SourceShape
 from param_decomp.model import PositionAxis, Positioned, Positionless
@@ -89,15 +90,24 @@ def place_target(tgt: GLUDecomposedModel, mesh: Mesh) -> GLUDecomposedModel:
 
 
 def init_component_stacks_placed(
-    sites: tuple[SiteSpec, ...], key: PRNGKeyArray, rules: PlacementRules
+    sites: tuple[SiteSpec, ...],
+    key: PRNGKeyArray,
+    rules: PlacementRules,
+    target_weights: dict[str, Array] | None = None,
 ) -> ComponentStacks:
     """Seeded V/U init placed by `component_stacks_shardings(_, rules)` (the run's placement
     policy), values bit-identical to the retired per-site init (pinned by `test_sharding`).
     One jit, 2×n_shapes sharded outputs — the persistence layout IS the stacked layout, so
-    the old two-stage stack-then-unstack fan-out (and its transient extra copy) is gone."""
-    abstract = eqx.filter_eval_shape(partial(init_component_stacks, sites), key)
-    placement = component_stacks_shardings(abstract, rules)
-    return jax.jit(partial(init_component_stacks, sites), out_shardings=placement)(key)
+    the old two-stage stack-then-unstack fan-out (and its transient extra copy) is gone.
+    `target_weights` selects the coupled init (`init_coupled_component_stacks`), riding as
+    jit ARGUMENTS (not closure constants) so the frozen W arrays are not baked into the
+    compiled init; None is the kaiming default."""
+    if target_weights is None:
+        init, args = partial(init_component_stacks, sites), (key,)
+    else:
+        init, args = partial(init_coupled_component_stacks, sites), (target_weights, key)
+    placement = component_stacks_shardings(eqx.filter_eval_shape(init, *args), rules)
+    return jax.jit(init, out_shardings=placement)(*args)
 
 
 def init_ci_fn_placed(
diff --git a/param_decomp/tests/test_component_init.py b/param_decomp/tests/test_component_init.py
new file mode 100644
index 000000000..000f58008
--- /dev/null
+++ b/param_decomp/tests/test_component_init.py
@@ -0,0 +1,78 @@
+"""Tests for the `coupled` component initialization (`init_coupled_component_stacks`)."""
+
+import jax
+import jax.numpy as jnp
+
+from param_decomp.components import SiteSpec, init_coupled_component_stacks
+
+
+def _random_weight(d_out: int, d_in: int) -> jax.Array:
+    return jax.random.normal(jax.random.PRNGKey(0), (d_out, d_in))
+
+
+def _component_sum(V: jax.Array, U: jax.Array) -> jax.Array:
+    return (V @ U).T
+
+
+def _col_space_residual(u: jax.Array, w: jax.Array) -> jax.Array:
+    q_out, _, _ = jnp.linalg.svd(w, full_matrices=False)
+    return u - (u @ q_out) @ q_out.T
+
+
+def _row_space_residual(v: jax.Array, w: jax.Array) -> jax.Array:
+    _, _, vh = jnp.linalg.svd(w, full_matrices=False)
+    q_in = vh.T
+    return v - q_in @ (q_in.T @ v)
+
+
+def _init_single(w: jax.Array, c: int, seed: int) -> tuple[jax.Array, jax.Array]:
+    d_out, d_in = w.shape
+    sites = (SiteSpec(name="m", d_in=d_in, d_out=d_out, C=c),)
+    vu = init_coupled_component_stacks(sites, {"m": w}, jax.random.PRNGKey(seed))
+    return vu.site("m")
+
+
+def test_coupled_seeds_are_unit_norm_and_derived_side_is_raw_w_image():
+    w = _random_weight(10, 6)  # d_in < d_out: V seeded, U derived
+    V, U = _init_single(w, c=8, seed=0)
+    assert jnp.allclose(jnp.linalg.norm(V, axis=0), jnp.ones(8), atol=1e-5)
+    assert jnp.allclose(U, (w @ V).T, atol=1e-6)
+    assert jnp.abs(_col_space_residual(U, w)).max() < 1e-5
+
+    w_t = w.T  # d_in > d_out: U seeded, V derived
+    V_t, U_t = _init_single(w_t, c=8, seed=0)
+    assert jnp.allclose(jnp.linalg.norm(U_t, axis=1), jnp.ones(8), atol=1e-5)
+    assert jnp.allclose(V_t, w_t.T @ U_t.T, atol=1e-6)
+    assert jnp.abs(_row_space_residual(V_t, w_t)).max() < 1e-5
+
+
+def test_coupled_sum_is_w_restricted_to_seed_span():
+    w = _random_weight(10, 6)
+    V, U = _init_single(w, c=4, seed=0)
+    # sum_c u_c v_c^T = W V V^T exactly.
+    assert jnp.allclose(_component_sum(V, U), w @ V @ V.T, atol=1e-5)
+
+
+def test_coupled_is_seed_deterministic():
+    w = _random_weight(10, 6)
+    V_a, U_a = _init_single(w, c=8, seed=7)
+    V_b, U_b = _init_single(w, c=8, seed=7)
+    assert jnp.array_equal(V_a, V_b) and jnp.array_equal(U_a, U_b)
+
+
+def test_coupled_stacked_sites_get_independent_draws_coupled_to_their_own_w():
+    # Two sites sharing one (d_in, d_out, C) shape group stack on one axis; each slice
+    # must be coupled to ITS OWN W with an independent seed draw.
+    key = jax.random.PRNGKey(3)
+    w_a = jax.random.normal(jax.random.fold_in(key, 0), (10, 6))
+    w_b = jax.random.normal(jax.random.fold_in(key, 1), (10, 6))
+    sites = (
+        SiteSpec(name="a", d_in=6, d_out=10, C=8),
+        SiteSpec(name="b", d_in=6, d_out=10, C=8),
+    )
+    vu = init_coupled_component_stacks(sites, {"a": w_a, "b": w_b}, jax.random.PRNGKey(0))
+    V_a, U_a = vu.site("a")
+    V_b, U_b = vu.site("b")
+    assert not jnp.allclose(V_a, V_b)
+    assert jnp.allclose(U_a, (w_a @ V_a).T, atol=1e-6)
+    assert jnp.allclose(U_b, (w_b @ V_b).T, atol=1e-6)

From 13587a67d5e5106ba3262fa73cca9a8348a64177 Mon Sep 17 00:00:00 2001
From: Antovigo <antvig@protonmail.com>
Date: Fri, 24 Jul 2026 20:31:27 +0000
Subject: [PATCH 2/5] refactor(init): explicit per-mode placed inits over a
 shared placement helper
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The optional-target_weights overload made init_component_stacks_placed
select the init implicitly (None => kaiming) — too implicit, and a dead
end for future weight_init options. Back to named entry points
(init_component_stacks_placed / init_coupled_component_stacks_placed),
each a thin wrapper over _init_component_stacks_via, which owns the one
eval_shape -> component_stacks_shardings -> jit(out_shardings) dance.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UVg2f69eq8mhDrpZZ3ZEmM
---
 param_decomp/run_state.py                     |  5 ++-
 .../targets/glu_transformer_sharding.py       | 42 ++++++++++++-------
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py
index 238a62a8f..0561d4f5c 100644
--- a/param_decomp/run_state.py
+++ b/param_decomp/run_state.py
@@ -42,6 +42,7 @@
 from param_decomp.targets.glu_transformer_sharding import (
     init_ci_fn_placed,
     init_component_stacks_placed,
+    init_coupled_component_stacks_placed,
     init_sources_sharded,
 )
 from param_decomp.train import Decomposition, TrainingItem, TrainState
@@ -182,8 +183,8 @@ def init_decomposition(
             zero_vu = component_stacks_from_sites(
                 {s.name: (jnp.zeros((s.d_in, s.C)), jnp.zeros((s.C, s.d_out))) for s in model.sites}
             )
-            components = init_component_stacks_placed(
-                model.sites, init_key, rules, target_weights=model.weight_deltas(zero_vu)
+            components = init_coupled_component_stacks_placed(
+                model.sites, model.weight_deltas(zero_vu), init_key, rules
             )
     ci_fn = init_ci_fn_placed(ci_fn_arch, model.sites, ci_key, mesh)
     assert ci_fn.has_position_axis == model.has_position_axis, (
diff --git a/param_decomp/targets/glu_transformer_sharding.py b/param_decomp/targets/glu_transformer_sharding.py
index 24da04337..1a8cb97d2 100644
--- a/param_decomp/targets/glu_transformer_sharding.py
+++ b/param_decomp/targets/glu_transformer_sharding.py
@@ -45,6 +45,7 @@
 placement construction / inside `.shardings` (fail-fast), never a silent replicate.
 """
 
+from collections.abc import Callable
 from functools import partial
 
 import equinox as eqx
@@ -77,6 +78,7 @@
     "hsdp_mesh",
     "place_target",
     "init_component_stacks_placed",
+    "init_coupled_component_stacks_placed",
     "init_ci_fn_placed",
     "init_sources_sharded",
     "shard_batch",
@@ -89,25 +91,37 @@ def place_target(tgt: GLUDecomposedModel, mesh: Mesh) -> GLUDecomposedModel:
     return place_via_shardings(tgt, tgt.shardings(mesh))
 
 
+def _init_component_stacks_via(
+    init: "Callable[..., ComponentStacks]",
+    args: tuple[PRNGKeyArray | dict[str, Array], ...],
+    rules: PlacementRules,
+) -> ComponentStacks:
+    """Run a `ComponentStacks` init placed by `component_stacks_shardings(_, rules)` (the
+    run's placement policy). One jit, 2×n_shapes sharded outputs — the persistence layout
+    IS the stacked layout, so no host-side full tree ever exists; `args` ride as jit
+    ARGUMENTS (not closure constants) so nothing large is baked into the compiled init."""
+    placement = component_stacks_shardings(eqx.filter_eval_shape(init, *args), rules)
+    return jax.jit(init, out_shardings=placement)(*args)
+
+
 def init_component_stacks_placed(
+    sites: tuple[SiteSpec, ...], key: PRNGKeyArray, rules: PlacementRules
+) -> ComponentStacks:
+    """Placed kaiming V/U init, values bit-identical to the retired per-site init (pinned
+    by `test_sharding`)."""
+    return _init_component_stacks_via(partial(init_component_stacks, sites), (key,), rules)
+
+
+def init_coupled_component_stacks_placed(
     sites: tuple[SiteSpec, ...],
+    target_weights: dict[str, Array],
     key: PRNGKeyArray,
     rules: PlacementRules,
-    target_weights: dict[str, Array] | None = None,
 ) -> ComponentStacks:
-    """Seeded V/U init placed by `component_stacks_shardings(_, rules)` (the run's placement
-    policy), values bit-identical to the retired per-site init (pinned by `test_sharding`).
-    One jit, 2×n_shapes sharded outputs — the persistence layout IS the stacked layout, so
-    the old two-stage stack-then-unstack fan-out (and its transient extra copy) is gone.
-    `target_weights` selects the coupled init (`init_coupled_component_stacks`), riding as
-    jit ARGUMENTS (not closure constants) so the frozen W arrays are not baked into the
-    compiled init; None is the kaiming default."""
-    if target_weights is None:
-        init, args = partial(init_component_stacks, sites), (key,)
-    else:
-        init, args = partial(init_coupled_component_stacks, sites), (target_weights, key)
-    placement = component_stacks_shardings(eqx.filter_eval_shape(init, *args), rules)
-    return jax.jit(init, out_shardings=placement)(*args)
+    """Placed coupled V/U init (`init_coupled_component_stacks`)."""
+    return _init_component_stacks_via(
+        partial(init_coupled_component_stacks, sites), (target_weights, key), rules
+    )
 
 
 def init_ci_fn_placed(

From 602ece037f6211573613e344929fdda84763bee6 Mon Sep 17 00:00:00 2001
From: Antovigo <antvig@protonmail.com>
Date: Fri, 24 Jul 2026 23:25:53 +0000
Subject: [PATCH 3/5] refactor(init): one universal init_component_stacks with
 an explicit weight_init selector
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Per review: per-mode entry points (init_kaiming_* / init_coupled_*) made
every call site import one name per option and left the bare
init_component_stacks name falsely suggesting universality. Now there is
ONE public init_component_stacks(sites, key, weight_init, target_weights)
— mode is an explicit argument (match-dispatched; target_weights required
iff coupled, asserted), the per-mode draws are private _kaiming/_coupled
stack-array helpers, and the placed variant carries the same signature.
The WeightInit literal lives in components.py and PDConfig reuses it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UVg2f69eq8mhDrpZZ3ZEmM
---
 param_decomp/components.py                    | 41 +++++++++++++------
 param_decomp/configs.py                       |  3 +-
 param_decomp/run.py                           |  3 +-
 param_decomp/run_state.py                     | 19 ++++-----
 .../targets/glu_transformer_sharding.py       | 41 ++++++-------------
 param_decomp/tests/test_component_init.py     |  8 ++--
 6 files changed, 59 insertions(+), 56 deletions(-)

diff --git a/param_decomp/components.py b/param_decomp/components.py
index bcc67cc7e..36397b2df 100644
--- a/param_decomp/components.py
+++ b/param_decomp/components.py
@@ -12,7 +12,7 @@
 from collections.abc import Iterator
 from dataclasses import dataclass
 from functools import cache
-from typing import ClassVar, Generic
+from typing import ClassVar, Generic, Literal
 
 import equinox as eqx
 import jax
@@ -72,6 +72,9 @@ def dequantize_fp8(q: Array, scale: Array) -> Array:
 
 VUShape = tuple[int, int, int]  # (d_in, d_out, C)
 
+# How `init_component_stacks` seeds V/U; the config schema (`PDConfig.weight_init`) reuses it.
+WeightInit = Literal["kaiming", "coupled"]
+
 # site name -> (shape group, slot on the group's stack axis); static, canonical site order
 SiteSlots = tuple[tuple[str, VUShape, int], ...]
 
@@ -150,7 +153,7 @@ def group_lengths(self) -> dict[VUShape, int]:
         return lengths
 
 
-def init_stack_arrays(
+def _kaiming_stack_arrays(
     sites: tuple[SiteSpec, ...], key: Array
 ) -> dict[VUShape, tuple[Array, Array]]:
     """Seeded stack init: `{(d_in, d_out, C): (V [g, d_in, C], U [g, C, d_out])}`, vmapped
@@ -169,16 +172,16 @@ def init_stack_arrays(
     return stacked
 
 
-def init_coupled_component_stacks(
+def _coupled_stack_arrays(
     sites: tuple[SiteSpec, ...], target_weights: dict[str, Array], key: Array
-) -> ComponentStacks:
-    """Coupled init: unit-norm seed on the narrow side, wide side its raw W-image.
+) -> dict[VUShape, tuple[Array, Array]]:
+    """Coupled stack init: unit-norm seed on the narrow side, wide side its raw W-image.
 
     `d_in <= d_out`: `v_c ~ unit norm`, `U_c <- (W v_c)^T`; else `u_c ~ unit norm`,
     `V_c <- W^T u_c`. No C-dependent rescale — components sit at W's natural scale and
     the component sum equals W restricted to the seed span (`W V V^T` when V is seeded),
     the delta carrying the complement. One key per site, split in site order; vmapped
-    per shape group like `init_stack_arrays` so the compiled init doesn't scale with
+    per shape group like `_kaiming_stack_arrays` so the compiled init doesn't scale with
     site count."""
     keys = jax.random.split(key, len(sites))
     site_index = {spec.name: idx for idx, spec in enumerate(sites)}
@@ -201,7 +204,7 @@ def seed_v(k: Array, w: Array, s: tuple[int, int] = (c, d_out)) -> tuple[Array,
                 return w.T @ u.T, u
 
         stacked[(d_in, d_out, c)] = jax.vmap(seed_v)(ks, ws)
-    return ComponentStacks(stacks=stacked, site_slots=site_slots_for(sites))
+    return stacked
 
 
 def component_stacks_from_sites(vu: dict[str, tuple[Array, Array]]) -> ComponentStacks:
@@ -222,11 +225,25 @@ def component_stacks_from_sites(vu: dict[str, tuple[Array, Array]]) -> Component
     return ComponentStacks(stacks=stacks, site_slots=site_slots_for(sites))
 
 
-def init_component_stacks(sites: tuple[SiteSpec, ...], key: Array) -> ComponentStacks:
-    """Small random fp32 V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5) per site, built directly in
-    the stacked persistence layout; the weight-delta channel carries the faithfulness
-    residual at init (before faithfulness warmup)."""
-    return ComponentStacks(stacks=init_stack_arrays(sites, key), site_slots=site_slots_for(sites))
+def init_component_stacks(
+    sites: tuple[SiteSpec, ...],
+    key: Array,
+    weight_init: WeightInit = "kaiming",
+    target_weights: dict[str, Array] | None = None,
+) -> ComponentStacks:
+    """Seeded V/U masters in the stacked persistence layout, per `weight_init`:
+    'kaiming' — small random fp32 V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5), ignores W
+    (the weight-delta channel carries the faithfulness residual at init); 'coupled' —
+    `_coupled_stack_arrays`, seeded from `target_weights` (required there, refused
+    otherwise)."""
+    match weight_init:
+        case "kaiming":
+            assert target_weights is None, "kaiming init draws blind — no target_weights"
+            stacks = _kaiming_stack_arrays(sites, key)
+        case "coupled":
+            assert target_weights is not None, "coupled init seeds from W — pass target_weights"
+            stacks = _coupled_stack_arrays(sites, target_weights, key)
+    return ComponentStacks(stacks=stacks, site_slots=site_slots_for(sites))
 
 
 def site_out(
diff --git a/param_decomp/configs.py b/param_decomp/configs.py
index 223f2b3d8..b4c363511 100644
--- a/param_decomp/configs.py
+++ b/param_decomp/configs.py
@@ -31,6 +31,7 @@
 )
 
 from param_decomp.base_config import BaseConfig, Probability
+from param_decomp.components import WeightInit
 from param_decomp.schedule import ScheduleConfig
 
 # ---------------------------------------------------------------------------
@@ -1041,7 +1042,7 @@ def _strip_removed_jax_unsupported_fields(cls, data: object) -> object:
         default=0,
         description="Random seed for reproducibility, including LM dataset shuffling.",
     )
-    weight_init: Literal["kaiming", "coupled"] = Field(
+    weight_init: WeightInit = Field(
         default="kaiming",
         description="How component V/U are initialized. 'kaiming': iid normal "
         "(V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5)), ignores W. 'coupled': unit-norm seed on "
diff --git a/param_decomp/run.py b/param_decomp/run.py
index 6116dd524..d2681cf76 100644
--- a/param_decomp/run.py
+++ b/param_decomp/run.py
@@ -516,7 +516,8 @@ def run_decomposition_training(
     rules = placement_rules
     if is_main:
         audit = component_stacks_audit(
-            eqx.filter_eval_shape(_partial(init_component_stacks, model.sites), init_key), rules
+            eqx.filter_eval_shape(_partial(init_component_stacks, model.sites), init_key),
+            rules,
         )
         print(
             rules.describe(
diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py
index 0561d4f5c..ddd04d061 100644
--- a/param_decomp/run_state.py
+++ b/param_decomp/run_state.py
@@ -8,7 +8,6 @@
 """
 
 from collections.abc import Callable
-from typing import Literal
 
 import equinox as eqx
 import jax
@@ -21,7 +20,7 @@
 
 from param_decomp.adversary import PersistentAdversary, init_sources_adam_state
 from param_decomp.ci_fn import ChunkwiseTransformerCIArch, CIFnArch
-from param_decomp.components import component_stacks_from_sites
+from param_decomp.components import WeightInit, component_stacks_from_sites
 from param_decomp.configs import (
     AdamPGDConfig,
     AdamWOptimizerConfig,
@@ -42,7 +41,6 @@
 from param_decomp.targets.glu_transformer_sharding import (
     init_ci_fn_placed,
     init_component_stacks_placed,
-    init_coupled_component_stacks_placed,
     init_sources_sharded,
 )
 from param_decomp.train import Decomposition, TrainingItem, TrainState
@@ -166,26 +164,27 @@ def init_decomposition(
     init_key: PRNGKeyArray,
     mesh: Mesh,
     rules: PlacementRules,
-    weight_init: Literal["kaiming", "coupled"] = "kaiming",
+    weight_init: WeightInit = "kaiming",
 ) -> Decomposition:
     """The trained-product half of `init_train_state`, factored out so a consumer can
     `jax.eval_shape` it to recover the saved `decomposition` item's tree structure
     without building (or knowing about) the optimizers/adversaries (the `weight_init`
     default is safe there — both inits produce the same tree)."""
     ci_key = random.fold_in(init_key, 1)
-    # V/U placement derives from the rules table; the CI fn still declares its own
-    # per-leaf shardings (the staged placement migration hasn't reached it).
     match weight_init:
         case "kaiming":
-            components = init_component_stacks_placed(model.sites, init_key, rules)
+            target_weights = None
         case "coupled":
             # The protocol exposes W only through `weight_deltas`; on zero V/U the delta IS W.
             zero_vu = component_stacks_from_sites(
                 {s.name: (jnp.zeros((s.d_in, s.C)), jnp.zeros((s.C, s.d_out))) for s in model.sites}
             )
-            components = init_coupled_component_stacks_placed(
-                model.sites, model.weight_deltas(zero_vu), init_key, rules
-            )
+            target_weights = model.weight_deltas(zero_vu)
+    # V/U placement derives from the rules table; the CI fn still declares its own
+    # per-leaf shardings (the staged placement migration hasn't reached it).
+    components = init_component_stacks_placed(
+        model.sites, init_key, rules, weight_init, target_weights
+    )
     ci_fn = init_ci_fn_placed(ci_fn_arch, model.sites, ci_key, mesh)
     assert ci_fn.has_position_axis == model.has_position_axis, (
         f"CI fn has_position_axis={ci_fn.has_position_axis} but model declares "
diff --git a/param_decomp/targets/glu_transformer_sharding.py b/param_decomp/targets/glu_transformer_sharding.py
index 1a8cb97d2..ec5e47d87 100644
--- a/param_decomp/targets/glu_transformer_sharding.py
+++ b/param_decomp/targets/glu_transformer_sharding.py
@@ -45,7 +45,6 @@
 placement construction / inside `.shardings` (fail-fast), never a silent replicate.
 """
 
-from collections.abc import Callable
 from functools import partial
 
 import equinox as eqx
@@ -64,8 +63,8 @@
 from param_decomp.components import (
     ComponentStacks,
     SiteSpec,
+    WeightInit,
     init_component_stacks,
-    init_coupled_component_stacks,
 )
 from param_decomp.configs import SourceShape
 from param_decomp.model import PositionAxis, Positioned, Positionless
@@ -78,7 +77,6 @@
     "hsdp_mesh",
     "place_target",
     "init_component_stacks_placed",
-    "init_coupled_component_stacks_placed",
     "init_ci_fn_placed",
     "init_sources_sharded",
     "shard_batch",
@@ -91,37 +89,24 @@ def place_target(tgt: GLUDecomposedModel, mesh: Mesh) -> GLUDecomposedModel:
     return place_via_shardings(tgt, tgt.shardings(mesh))
 
 
-def _init_component_stacks_via(
-    init: "Callable[..., ComponentStacks]",
-    args: tuple[PRNGKeyArray | dict[str, Array], ...],
-    rules: PlacementRules,
-) -> ComponentStacks:
-    """Run a `ComponentStacks` init placed by `component_stacks_shardings(_, rules)` (the
-    run's placement policy). One jit, 2×n_shapes sharded outputs — the persistence layout
-    IS the stacked layout, so no host-side full tree ever exists; `args` ride as jit
-    ARGUMENTS (not closure constants) so nothing large is baked into the compiled init."""
-    placement = component_stacks_shardings(eqx.filter_eval_shape(init, *args), rules)
-    return jax.jit(init, out_shardings=placement)(*args)
-
-
 def init_component_stacks_placed(
-    sites: tuple[SiteSpec, ...], key: PRNGKeyArray, rules: PlacementRules
-) -> ComponentStacks:
-    """Placed kaiming V/U init, values bit-identical to the retired per-site init (pinned
-    by `test_sharding`)."""
-    return _init_component_stacks_via(partial(init_component_stacks, sites), (key,), rules)
-
-
-def init_coupled_component_stacks_placed(
     sites: tuple[SiteSpec, ...],
-    target_weights: dict[str, Array],
     key: PRNGKeyArray,
     rules: PlacementRules,
+    weight_init: WeightInit = "kaiming",
+    target_weights: dict[str, Array] | None = None,
 ) -> ComponentStacks:
-    """Placed coupled V/U init (`init_coupled_component_stacks`)."""
-    return _init_component_stacks_via(
-        partial(init_coupled_component_stacks, sites), (target_weights, key), rules
+    """`init_component_stacks` placed by `component_stacks_shardings(_, rules)` (the run's
+    placement policy); kaiming values bit-identical to the retired per-site init (pinned
+    by `test_sharding`). One jit, 2×n_shapes sharded outputs — the persistence layout IS
+    the stacked layout, so no host-side full tree ever exists; `key`/`target_weights` ride
+    as jit ARGUMENTS (not closure constants) so the frozen W arrays are not baked into the
+    compiled init."""
+    init = partial(init_component_stacks, sites, weight_init=weight_init)
+    placement = component_stacks_shardings(
+        eqx.filter_eval_shape(init, key, target_weights=target_weights), rules
     )
+    return jax.jit(init, out_shardings=placement)(key, target_weights=target_weights)
 
 
 def init_ci_fn_placed(
diff --git a/param_decomp/tests/test_component_init.py b/param_decomp/tests/test_component_init.py
index 000f58008..c8cca260d 100644
--- a/param_decomp/tests/test_component_init.py
+++ b/param_decomp/tests/test_component_init.py
@@ -1,9 +1,9 @@
-"""Tests for the `coupled` component initialization (`init_coupled_component_stacks`)."""
+"""Tests for the `coupled` mode of `init_component_stacks`."""
 
 import jax
 import jax.numpy as jnp
 
-from param_decomp.components import SiteSpec, init_coupled_component_stacks
+from param_decomp.components import SiteSpec, init_component_stacks
 
 
 def _random_weight(d_out: int, d_in: int) -> jax.Array:
@@ -28,7 +28,7 @@ def _row_space_residual(v: jax.Array, w: jax.Array) -> jax.Array:
 def _init_single(w: jax.Array, c: int, seed: int) -> tuple[jax.Array, jax.Array]:
     d_out, d_in = w.shape
     sites = (SiteSpec(name="m", d_in=d_in, d_out=d_out, C=c),)
-    vu = init_coupled_component_stacks(sites, {"m": w}, jax.random.PRNGKey(seed))
+    vu = init_component_stacks(sites, jax.random.PRNGKey(seed), "coupled", {"m": w})
     return vu.site("m")
 
 
@@ -70,7 +70,7 @@ def test_coupled_stacked_sites_get_independent_draws_coupled_to_their_own_w():
         SiteSpec(name="a", d_in=6, d_out=10, C=8),
         SiteSpec(name="b", d_in=6, d_out=10, C=8),
     )
-    vu = init_coupled_component_stacks(sites, {"a": w_a, "b": w_b}, jax.random.PRNGKey(0))
+    vu = init_component_stacks(sites, jax.random.PRNGKey(0), "coupled", {"a": w_a, "b": w_b})
     V_a, U_a = vu.site("a")
     V_b, U_b = vu.site("b")
     assert not jnp.allclose(V_a, V_b)

From ca8a86c579ee5b631e3a15a18a6f4b9025d81646 Mon Sep 17 00:00:00 2001
From: Antovigo <antvig@protonmail.com>
Date: Fri, 24 Jul 2026 23:45:24 +0000
Subject: [PATCH 4/5] test(init): keep only the essential coupled-init tests

Drop the component-sum test (an algebraic corollary of the exact
U = (W V)^T assertion) and the seed-determinism test (a property of
JAX's keyed PRNG, not of this code). What remains: the defining W-image
property on both orientation branches, and the multi-site shape-group
stacking wiring.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UVg2f69eq8mhDrpZZ3ZEmM
---
 param_decomp/tests/test_component_init.py | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/param_decomp/tests/test_component_init.py b/param_decomp/tests/test_component_init.py
index c8cca260d..6f1b52cc2 100644
--- a/param_decomp/tests/test_component_init.py
+++ b/param_decomp/tests/test_component_init.py
@@ -10,10 +10,6 @@ def _random_weight(d_out: int, d_in: int) -> jax.Array:
     return jax.random.normal(jax.random.PRNGKey(0), (d_out, d_in))
 
 
-def _component_sum(V: jax.Array, U: jax.Array) -> jax.Array:
-    return (V @ U).T
-
-
 def _col_space_residual(u: jax.Array, w: jax.Array) -> jax.Array:
     q_out, _, _ = jnp.linalg.svd(w, full_matrices=False)
     return u - (u @ q_out) @ q_out.T
@@ -46,20 +42,6 @@ def test_coupled_seeds_are_unit_norm_and_derived_side_is_raw_w_image():
     assert jnp.abs(_row_space_residual(V_t, w_t)).max() < 1e-5
 
 
-def test_coupled_sum_is_w_restricted_to_seed_span():
-    w = _random_weight(10, 6)
-    V, U = _init_single(w, c=4, seed=0)
-    # sum_c u_c v_c^T = W V V^T exactly.
-    assert jnp.allclose(_component_sum(V, U), w @ V @ V.T, atol=1e-5)
-
-
-def test_coupled_is_seed_deterministic():
-    w = _random_weight(10, 6)
-    V_a, U_a = _init_single(w, c=8, seed=7)
-    V_b, U_b = _init_single(w, c=8, seed=7)
-    assert jnp.array_equal(V_a, V_b) and jnp.array_equal(U_a, U_b)
-
-
 def test_coupled_stacked_sites_get_independent_draws_coupled_to_their_own_w():
     # Two sites sharing one (d_in, d_out, C) shape group stack on one axis; each slice
     # must be coupled to ITS OWN W with an independent seed draw.

From eb3b5bcf2c9b819aeb73a33338898584f170e8c5 Mon Sep 17 00:00:00 2001
From: Antovigo <antvig@protonmail.com>
Date: Sat, 25 Jul 2026 00:21:41 +0000
Subject: [PATCH 5/5] feat(model): target_weight protocol accessor; drop the
 zero-V/U weight recovery
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

DecomposedModel grows target_weight(name) -> frozen W in stored dtype —
the torch lineage's ComponentModel.target_weight, restored. weight_deltas
in every target is rewritten on top of it, and the coupled init reads
W directly instead of the weight_deltas(zeros) trick, which cost a full
zero V/U tree, an eager fp32 copy of every site's W, and all-site zero
matmuls at startup. The fp32 upcast now happens inside the jitted init,
per shape group.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UVg2f69eq8mhDrpZZ3ZEmM
---
 param_decomp/model.py                           | 4 ++++
 param_decomp/run_state.py                       | 8 ++------
 param_decomp/targets/glu_transformer.py         | 9 +++++++--
 param_decomp/targets/llama_simple_mlp.py        | 8 ++++++--
 param_decomp/tests/test_attn_patterns_eval.py   | 4 ++++
 param_decomp/tests/test_eval.py                 | 4 ++++
 param_decomp/tests/test_generic_model_io.py     | 4 ++++
 param_decomp_lab/experiments/resid_mlp/model.py | 3 +++
 param_decomp_lab/experiments/tms/model.py       | 3 +++
 9 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/param_decomp/model.py b/param_decomp/model.py
index 8cd95d759..8f4b7dfdf 100644
--- a/param_decomp/model.py
+++ b/param_decomp/model.py
@@ -217,6 +217,10 @@ def masked_site_outputs(
         the recon grid, which stays KL-on-final-logits."""
         ...
 
+    def target_weight(self, name: str) -> Float[Array, "d_out d_in"]:
+        """The frozen target weight W for site `name`, in its stored dtype."""
+        ...
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Float[Array, "d_out d_in"]]:
         """fp32 `W − V@U` per site from the fp32 master `vu`."""
         ...
diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py
index ddd04d061..2ec8615f0 100644
--- a/param_decomp/run_state.py
+++ b/param_decomp/run_state.py
@@ -20,7 +20,7 @@
 
 from param_decomp.adversary import PersistentAdversary, init_sources_adam_state
 from param_decomp.ci_fn import ChunkwiseTransformerCIArch, CIFnArch
-from param_decomp.components import WeightInit, component_stacks_from_sites
+from param_decomp.components import WeightInit
 from param_decomp.configs import (
     AdamPGDConfig,
     AdamWOptimizerConfig,
@@ -175,11 +175,7 @@ def init_decomposition(
         case "kaiming":
             target_weights = None
         case "coupled":
-            # The protocol exposes W only through `weight_deltas`; on zero V/U the delta IS W.
-            zero_vu = component_stacks_from_sites(
-                {s.name: (jnp.zeros((s.d_in, s.C)), jnp.zeros((s.C, s.d_out))) for s in model.sites}
-            )
-            target_weights = model.weight_deltas(zero_vu)
+            target_weights = {name: model.target_weight(name) for name in model.site_names}
     # V/U placement derives from the rules table; the CI fn still declares its own
     # per-leaf shardings (the staged placement migration hasn't reached it).
     components = init_component_stacks_placed(
diff --git a/param_decomp/targets/glu_transformer.py b/param_decomp/targets/glu_transformer.py
index d285406a7..23a3c7594 100644
--- a/param_decomp/targets/glu_transformer.py
+++ b/param_decomp/targets/glu_transformer.py
@@ -1136,12 +1136,17 @@ def masked_component_activations(
         assert set(collect_activations) == set(live), (sorted(collect_activations), sorted(live))
         return collect_activations
 
+    def target_weight(self, name: str) -> Array:
+        """The frozen W for site `name` (stored dtype): the site's slice of the stacked
+        per-layer weights."""
+        layer, kind = parse_site_name(name)
+        return _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], self.stacked), kind)
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         """fp32 `W − V@U` per site from fp32 masters (faithfulness input)."""
         out: dict[str, Array] = {}
         for spec in self.sites:
-            layer, kind = parse_site_name(spec.name)
-            W = _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], self.stacked), kind)
+            W = self.target_weight(spec.name)
             V, U = vu.site(spec.name)
             out[spec.name] = (
                 W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T
diff --git a/param_decomp/targets/llama_simple_mlp.py b/param_decomp/targets/llama_simple_mlp.py
index 408237bdc..3b663a653 100644
--- a/param_decomp/targets/llama_simple_mlp.py
+++ b/param_decomp/targets/llama_simple_mlp.py
@@ -488,12 +488,16 @@ def masked_site_outputs(
         assert set(collect) == set(live), (sorted(collect), sorted(live))
         return collect
 
+    def target_weight(self, name: str) -> Array:
+        """The frozen W for site `name`, in its stored dtype."""
+        layer_idx, kind = parse_site_name(name)
+        return _frozen_site_weight(self.layers[layer_idx], kind)
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         """fp32 `W − V@U` per site from fp32 masters (faithfulness input)."""
         out: dict[str, Array] = {}
         for spec in self.sites:
-            layer_idx, kind = parse_site_name(spec.name)
-            W = _frozen_site_weight(self.layers[layer_idx], kind)
+            W = self.target_weight(spec.name)
             V, U = vu.site(spec.name)
             out[spec.name] = (
                 W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T
diff --git a/param_decomp/tests/test_attn_patterns_eval.py b/param_decomp/tests/test_attn_patterns_eval.py
index 460adc2cf..976b4a7a2 100644
--- a/param_decomp/tests/test_attn_patterns_eval.py
+++ b/param_decomp/tests/test_attn_patterns_eval.py
@@ -275,6 +275,10 @@ def masked_site_outputs(
         del vu, resid, masks, delta_masks, routes, live, has_delta
         raise AssertionError("positionless stub fn must not be called")
 
+    def target_weight(self, name: str) -> jax.Array:
+        del name
+        raise AssertionError("positionless stub fn must not be called")
+
     def weight_deltas(self, vu: Any) -> dict[str, jax.Array]:
         del vu
         raise AssertionError("positionless stub fn must not be called")
diff --git a/param_decomp/tests/test_eval.py b/param_decomp/tests/test_eval.py
index 9fd5e790a..9327e18b5 100644
--- a/param_decomp/tests/test_eval.py
+++ b/param_decomp/tests/test_eval.py
@@ -131,6 +131,10 @@ def masked_site_outputs(
         del vu, resid, masks, delta_masks, routes, live, has_delta
         raise AssertionError("positionless stub fn must not be called")
 
+    def target_weight(self, name: str) -> jax.Array:
+        del name
+        raise AssertionError("positionless stub fn must not be called")
+
     def weight_deltas(self, vu: Any) -> dict[str, jax.Array]:
         del vu
         raise AssertionError("positionless stub fn must not be called")
diff --git a/param_decomp/tests/test_generic_model_io.py b/param_decomp/tests/test_generic_model_io.py
index 281f25d46..f87fca0e7 100644
--- a/param_decomp/tests/test_generic_model_io.py
+++ b/param_decomp/tests/test_generic_model_io.py
@@ -166,6 +166,10 @@ def masked_site_outputs(
             hidden = hidden + delta_masks[SITE][..., None] * (resid @ delta.T)
         return {SITE: hidden}
 
+    def target_weight(self, name: str) -> Array:
+        assert name == SITE, name
+        return self.W
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         V, U = vu.site(SITE)
         return {
diff --git a/param_decomp_lab/experiments/resid_mlp/model.py b/param_decomp_lab/experiments/resid_mlp/model.py
index baffcb854..2854918d1 100644
--- a/param_decomp_lab/experiments/resid_mlp/model.py
+++ b/param_decomp_lab/experiments/resid_mlp/model.py
@@ -457,6 +457,9 @@ def masked_site_outputs(
             self.target, prepared, resid, masks, delta_masks, routes, live, has_delta
         )
 
+    def target_weight(self, name: str) -> Array:
+        return _frozen_site_weight(self.target, name)
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         return weight_deltas_fp32(self.target, vu, self.sites)
 
diff --git a/param_decomp_lab/experiments/tms/model.py b/param_decomp_lab/experiments/tms/model.py
index 615b76822..fdbdbffad 100644
--- a/param_decomp_lab/experiments/tms/model.py
+++ b/param_decomp_lab/experiments/tms/model.py
@@ -389,6 +389,9 @@ def masked_site_outputs(
             self.target, prepared, resid, masks, delta_masks, routes, live, has_delta
         )
 
+    def target_weight(self, name: str) -> Array:
+        return _frozen_site_weight(self.target, name)
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         return weight_deltas_fp32(self.target, vu, self.sites)