initium-deeplearning 0.1.0__py3-none-any.whl

initium/LNN_controller/cfc_backbone_parallel.py

@@ -0,0 +1,265 @@
+"""
+file: LNN_controller/cfc_backbone_parallel.py -- v1
+
+CfC as a whole-sequence backbone for SplitGraphDNC, plus a factory that
+builds ANY parallel backbone (or a stack of them) from a variant string.
+
+CfCBackboneParallel: (B, L, in_dim) -> (B, L, d_model), same contract as
+mamba_backbone_parallel.MambaBackboneParallel, and like it has NO notion of
+DNC/Memory/read vectors (it only ever sees the raw input stream X).
+NOTE: ncps' CfC loops over time internally (no parallel scan) and its backbone
+MLP consumes the previous hidden state, so this is sequential inside. The
+split-graph benefit is structural (backbone can never see reads), not a
+Mamba-style scan speedup.
+Block = pre-norm residual (Add -> LN -> CfC). Runs in real fp32 by default
+(force_fp32=True), same policy as mamba_backbone_parallel._ParallelBlock.
+
+build_parallel_backbone(variant=...): variant is "mamba1" | "mamba2" |
+"mamba3" | "cfc", or a "+"-joined stack such as "mamba2+cfc" (stage 0 maps
+in_dim -> d_model, later stages d_model -> d_model). num_blocks applies PER
+STAGE. This is the single dispatch point for new parallel backbones.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from initium.LNN_controller.cfc_controller import CfC, _require_ncps
+from initium.mamba_controller.mamba_backbone_parallel import MambaBackboneParallel
+from initium.MoE.moe_layer import MoEBlock
+
+_BACKBONE_KINDS = ("mamba1", "mamba2", "mamba3", "cfc")
+
+
+class _CfCParallelBlock(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        units,
+        mode,
+        backbone_units,
+        backbone_layers,
+        backbone_dropout,
+        activation,
+        mixed_memory,
+        residual,
+        force_fp32,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        _require_ncps()
+        units = d_model if units is None else units
+        self.residual, self.force_fp32 = residual, force_fp32
+        self.norm = nn.LayerNorm(d_model, device=device, dtype=dtype)
+        self.cfc = CfC(
+            input_size=d_model,
+            units=units,
+            proj_size=None if units == d_model else d_model,
+            return_sequences=True,
+            batch_first=True,
+            mixed_memory=mixed_memory,
+            mode=mode,
+            activation=activation,
+            backbone_units=backbone_units,
+            backbone_layers=backbone_layers,
+            backbone_dropout=backbone_dropout,
+        )
+        if device is not None or dtype is not None:
+            self.cfc.to(device=device, dtype=dtype)
+
+    def forward(self, x):  # x: (B, L, D)
+        if self.force_fp32:
+            with torch.autocast(device_type=x.device.type, enabled=False):
+                out, _ = self.cfc(self.norm(x.float()))
+        else:
+            out, _ = self.cfc(self.norm(x))
+        out = out.to(x.dtype)
+        return x + out if self.residual else out
+
+
+class CfCBackboneParallel(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        d_model,
+        num_blocks=2,
+        units=None,
+        mode="default",
+        backbone_units=512,
+        backbone_layers=1,
+        backbone_dropout=0.0,
+        activation="lecun_tanh",
+        mixed_memory=False,
+        residual=True,
+        force_fp32=True,
+        moe_enabled: bool = False,
+        moe_num_experts: int = 8,
+        moe_expert_dim: int | None = None,
+        moe_top_k: int = 1,
+        moe_capacity_factor: float = 1.5,
+        moe_load_balance_alpha: float = 0.01,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.in_adapter: nn.Module = (
+            nn.Identity()
+            if in_dim == d_model
+            else nn.Linear(in_dim, d_model, device=device, dtype=dtype)
+        )
+        self.blocks: nn.ModuleList = nn.ModuleList(
+            [
+                _CfCParallelBlock(
+                    d_model,
+                    units,
+                    mode,
+                    backbone_units,
+                    backbone_layers,
+                    backbone_dropout,
+                    activation,
+                    mixed_memory,
+                    residual,
+                    force_fp32,
+                    device=device,
+                    dtype=dtype,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+        # External interleave, applied to the WHOLE (B, L, d_model) sequence
+        # after each block (SwitchMoE reshapes any leading dims to a flat
+        # token batch, so this costs nothing extra to support here).
+        self.moe_enabled = moe_enabled
+        self.moe_blocks: nn.ModuleList | None = None
+        if moe_enabled:
+            self.moe_blocks = nn.ModuleList(
+                [
+                    MoEBlock(
+                        d_model,
+                        num_experts=moe_num_experts,
+                        expert_dim=moe_expert_dim,
+                        top_k=moe_top_k,
+                        capacity_factor=moe_capacity_factor,
+                        load_balance_alpha=moe_load_balance_alpha,
+                        device=device,
+                        dtype=dtype,
+                    )
+                    for _ in range(num_blocks)
+                ]
+            )
+
+    def forward(self, x):  # (B, L, in_dim) -> (B, L, d_model), called ONCE per training step
+        h = self.in_adapter(x)
+        for i, block in enumerate(self.blocks):
+            h = block(h)
+            if self.moe_enabled:
+                assert self.moe_blocks is not None
+                h = self.moe_blocks[i](h)
+        return h
+
+
+def build_parallel_backbone(
+    in_dim,
+    d_model,
+    num_blocks=2,
+    variant="mamba1",
+    d_state=16,
+    d_conv=4,
+    expand=2,
+    headdim=64,
+    cfc_kwargs=None,
+    moe_enabled: bool = False,
+    moe_num_experts: int = 8,
+    moe_expert_dim: int | None = None,
+    moe_top_k: int = 1,
+    moe_capacity_factor: float = 1.5,
+    moe_load_balance_alpha: float = 0.01,
+    device=None,
+    dtype=None,
+):
+    kinds = [k.strip().lower() for k in variant.split("+")]
+    bad = [k for k in kinds if k not in _BACKBONE_KINDS]
+    if bad:
+        raise ValueError(
+            f"build_parallel_backbone: unknown variant part(s) {bad} in {variant!r}, "
+            f"expected '+'-joined parts of {_BACKBONE_KINDS}"
+        )
+    stages: list[nn.Module] = []
+    cur = in_dim
+    for kind in kinds:
+        if kind == "cfc":
+            stages.append(
+                CfCBackboneParallel(
+                    cur,
+                    d_model,
+                    num_blocks=num_blocks,
+                    device=device,
+                    dtype=dtype,
+                    moe_enabled=moe_enabled,
+                    moe_num_experts=moe_num_experts,
+                    moe_expert_dim=moe_expert_dim,
+                    moe_top_k=moe_top_k,
+                    moe_capacity_factor=moe_capacity_factor,
+                    moe_load_balance_alpha=moe_load_balance_alpha,
+                    **(cfc_kwargs or {}),
+                )
+            )
+        else:
+            stages.append(
+                MambaBackboneParallel(
+                    in_dim=cur,
+                    d_model=d_model,
+                    num_blocks=num_blocks,
+                    variant=kind,
+                    d_state=d_state,
+                    d_conv=d_conv,
+                    expand=expand,
+                    headdim=headdim,
+                    device=device,
+                    dtype=dtype,
+                    moe_enabled=moe_enabled,
+                    moe_num_experts=moe_num_experts,
+                    moe_expert_dim=moe_expert_dim,
+                    moe_top_k=moe_top_k,
+                    moe_capacity_factor=moe_capacity_factor,
+                    moe_load_balance_alpha=moe_load_balance_alpha,
+                )
+            )
+        cur = d_model
+    return stages[0] if len(stages) == 1 else nn.Sequential(*stages)
+
+
+def collect_backbone_moe_layers(backbone: nn.Module) -> list:
+    """Flat list of every installed MoE sublayer on a backbone built by
+    build_parallel_backbone -- a single stage or a '+'-stacked nn.Sequential
+    of stages, each of which self-reports via .moe_enabled/.moe_blocks (same
+    convention every controller wrapper in this project uses)."""
+    stages = list(backbone) if isinstance(backbone, nn.Sequential) else [backbone]
+    moe_layers: list[nn.Module] = []
+    for stage in stages:
+        if getattr(stage, "moe_enabled", False):
+            blocks = stage.moe_blocks
+            if not isinstance(blocks, nn.ModuleList):
+                raise TypeError(
+                    "MoE-enabled backbone stages must expose a ModuleList of moe_blocks"
+                )
+            moe_layers.extend(blocks)
+    return moe_layers
+
+
+if (
+    __name__ == "__main__"
+):  # smoke test: python -m LNN_controller.cfc_backbone_parallel (from project root)
+    B, L, in_dim, d = 4, 12, 40, 32
+    for variant in ("cfc", "cfc+cfc"):
+        m = build_parallel_backbone(
+            in_dim, d, num_blocks=1, variant=variant, cfc_kwargs=dict(backbone_units=64)
+        )
+        y = m(torch.randn(B, L, in_dim))
+        assert y.shape == (B, L, d), y.shape
+        y.pow(2).mean().backward()
+        print(f"{variant}: OK")
+    print("For mamba1/2/3+cfc stacks, run the pilot below on a GPU (Mamba kernels need CUDA).")

initium/LNN_controller/cfc_controller.py

@@ -0,0 +1,325 @@
+"""
+file: LNN_controller/cfc_controller.py -- v1
+
+Standalone Closed-form Continuous-time (CfC) DNC controller (Hasani et al.,
+"Closed-form Continuous-time Neural Networks", arXiv 2106.13898):
+MambaDNC's rnn_type='cfc'. Uses the `ncps` package's CfC module as-is
+(pip install ncps -> from ncps.torch import CfC), driven ONE DNC timestep at
+a time as a length-1 sequence, interleaved with Memory reads/writes exactly
+like mamba_controller.py's controllers. ncps' CfC forward is out-of-place
+(returns new h), so it is BPTT-safe with no custom step() needed. Its state
+is bounded (convex mix of tanh heads), so no fp32-forcing/clamp workarounds.
+timespans are not used (ts=1.0 per step): this task has no timestamps.
+
+CONTROLLER-WRAPPER PROTOCOL (what makes this composable -- see also
+chained_controller.py):
+    wrapper.init_state(batch_size, device=None, dtype=None) -> state
+    wrapper(x: (B,1,in_dim), state)                         -> (out: (B,1,d_model), new_state)
+    wrapper.d_model, wrapper.in_adapter, wrapper.moe_enabled, wrapper.moe_blocks
+Any object following this can be a DNC controller, a split-graph combiner,
+or a stage in a ChainedControllerWrapper.
+
+Block = pre-norm residual (Add -> LN -> CfC), same pattern as the Mamba blocks.
+State per block: h (B, units) fp32, or (h, c) if mixed_memory=True (CfC-mmRNN).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import cast
+
+import torch
+import torch.nn as nn
+
+from initium.MoE.moe_layer import MoEBlock, MultiSourceMoEBlock
+
+_NCPS_IMPORT_ERROR: str | None
+
+try:
+    from ncps.torch import CfC
+except ImportError as _e:  # pragma: no cover - environment-dependent
+    CfC = None
+    _NCPS_IMPORT_ERROR = (
+        "cfc_controller.py requires the `ncps` package (`pip install ncps`). "
+        f"Original import error: {_e}"
+    )
+else:
+    _NCPS_IMPORT_ERROR = None
+
+
+def _require_ncps() -> None:
+    if CfC is None:
+        raise ImportError(_NCPS_IMPORT_ERROR)
+
+
+def _state_to_fp32(state):
+    # Keep the recurrent state fp32 between steps (ncps' cell returns fp16 under autocast).
+    if isinstance(state, tuple):
+        return tuple(s.float() for s in state)
+    return state.float()
+
+
+class CfCControllerBlock(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        units=None,
+        mode="default",
+        backbone_units=512,
+        backbone_layers=1,
+        backbone_dropout=0.0,
+        activation="lecun_tanh",
+        mixed_memory=False,
+        residual=True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        _require_ncps()
+        units = d_model if units is None else units
+        self.d_model, self.units = d_model, units
+        self.mixed_memory, self.residual = mixed_memory, residual
+        self.norm = nn.LayerNorm(d_model, device=device, dtype=dtype)
+        self.cfc = CfC(
+            input_size=d_model,
+            units=units,
+            proj_size=None
+            if units == d_model
+            else d_model,  # project back to d_model if state is wider/narrower
+            return_sequences=True,
+            batch_first=True,
+            mixed_memory=mixed_memory,
+            mode=mode,
+            activation=activation,
+            backbone_units=backbone_units,
+            backbone_layers=backbone_layers,
+            backbone_dropout=backbone_dropout,
+        )
+        if device is not None or dtype is not None:
+            self.cfc.to(device=device, dtype=dtype)
+
+    def init_state(self, batch_size, device=None, dtype=None):  # dtype ignored: always fp32
+        device = device if device is not None else next(self.parameters()).device
+        h = torch.zeros(batch_size, self.units, device=device, dtype=torch.float32)
+        return (h, torch.zeros_like(h)) if self.mixed_memory else h
+
+    def step(self, x, state):
+        with torch.autocast(device_type=x.device.type, enabled=False):
+            out, new_state = self.cfc(
+                self.norm(x.float()).unsqueeze(1),
+                _state_to_fp32(state) if state is not None else None,
+            )
+        out = out.squeeze(1).clamp(min=-1e4, max=1e4)
+        y = x + out.to(x.dtype) if self.residual else out.to(x.dtype)
+        return y, _state_to_fp32(new_state)
+
+
+class CfCControllerWrapper(nn.Module):
+    """Stack of `num_blocks` CfCControllerBlocks with the nn.LSTM-compatible
+    call convention DNC._layer_forward expects (see module docstring)."""
+
+    def __init__(
+        self,
+        in_dim,
+        d_model,
+        num_blocks=2,
+        units=None,
+        mode="default",
+        backbone_units=512,
+        backbone_layers=1,
+        backbone_dropout=0.0,
+        activation="lecun_tanh",
+        mixed_memory=False,
+        residual=True,
+        moe_enabled: bool = False,
+        moe_num_experts: int = 8,
+        moe_expert_dim: int | None = None,
+        moe_top_k: int = 1,
+        moe_capacity_factor: float = 1.5,
+        moe_load_balance_alpha: float = 0.01,
+        moe_source_dims: list[int] | None = None,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.d_model, self.num_blocks = d_model, num_blocks
+
+        # moe_source_dims: when given (>=1 entries, each an input SOURCE's
+        # OWN raw width -- e.g. [hidden_size, read_vectors_size] for
+        # SplitGraphDNC's [backbone_output, prev_read_vector] combiner
+        # input), this wrapper is driven via forward_multi_source() instead
+        # of forward(); widths need not match each other, which is exactly
+        # the "variable-length inputs" property CfC must keep once MoE
+        # sits in front of it. sum(moe_source_dims) must still equal
+        # in_dim, since in_dim is what every other bookkeeping path in this
+        # project (checkpoints, dim-compat checks) reads as one number.
+        self.moe_source_dims = list(moe_source_dims) if moe_source_dims else None
+        if self.moe_source_dims is not None:
+            if sum(self.moe_source_dims) != in_dim:
+                raise ValueError(
+                    f"CfCControllerWrapper: sum(moe_source_dims)="
+                    f"{sum(self.moe_source_dims)} != in_dim={in_dim}"
+                )
+            if not moe_enabled:
+                raise ValueError("CfCControllerWrapper: moe_source_dims requires moe_enabled=True")
+
+        self.in_adapter: nn.Module = (
+            nn.Identity()
+            if in_dim == d_model
+            else nn.Linear(in_dim, d_model, device=device, dtype=dtype)
+        )
+        self.blocks: nn.ModuleList = nn.ModuleList(
+            [
+                CfCControllerBlock(
+                    d_model,
+                    units=units,
+                    mode=mode,
+                    backbone_units=backbone_units,
+                    backbone_layers=backbone_layers,
+                    backbone_dropout=backbone_dropout,
+                    activation=activation,
+                    mixed_memory=mixed_memory,
+                    residual=residual,
+                    device=device,
+                    dtype=dtype,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+        self.moe_enabled = moe_enabled
+        self.moe_blocks: nn.ModuleList | None = None
+        self.source_in_adapters: nn.ModuleList | None = None
+        if moe_enabled and self.moe_source_dims is not None:
+            # Multi-source front-end: each raw source is projected to
+            # d_model, routed+combined by ONE source-aware Top-K MoE bank
+            # BEFORE entering the recurrent CfC stack (see
+            # forward_multi_source). self.moe_blocks holds this single
+            # MultiSourceMoEBlock so every existing "extend moe_layers from
+            # layer_controller.moe_blocks" call site (MambaDNC,
+            # ChainedControllerWrapper) keeps working with zero changes.
+            self.source_in_adapters = nn.ModuleList(
+                [
+                    nn.Identity()
+                    if w == d_model
+                    else nn.Linear(w, d_model, device=device, dtype=dtype)
+                    for w in self.moe_source_dims
+                ]
+            )
+            self.moe_blocks = nn.ModuleList(
+                [
+                    MultiSourceMoEBlock(
+                        d_model,
+                        num_sources=len(self.moe_source_dims),
+                        num_experts=moe_num_experts,
+                        expert_dim=moe_expert_dim,
+                        top_k=moe_top_k,
+                        capacity_factor=moe_capacity_factor,
+                        load_balance_alpha=moe_load_balance_alpha,
+                        device=device,
+                        dtype=dtype,
+                    )
+                ]
+            )
+        elif moe_enabled:
+            # Plain per-block external interleave, same convention as
+            # MambaControllerWrapper / Mamba2ControllerWrapper.
+            self.moe_blocks = nn.ModuleList(
+                [
+                    MoEBlock(
+                        d_model,
+                        num_experts=moe_num_experts,
+                        expert_dim=moe_expert_dim,
+                        top_k=moe_top_k,
+                        capacity_factor=moe_capacity_factor,
+                        load_balance_alpha=moe_load_balance_alpha,
+                        device=device,
+                        dtype=dtype,
+                    )
+                    for _ in range(num_blocks)
+                ]
+            )
+
+    def init_state(self, batch_size, device=None, dtype=None):
+        if device is None:
+            device = next(self.parameters()).device
+        blocks = cast(Iterable[CfCControllerBlock], self.blocks)
+        return [blk.init_state(batch_size, device=device) for blk in blocks]
+
+    def forward(self, input, hx):
+        assert input.dim() == 3 and input.size(1) == 1, (
+            "CfCControllerWrapper only supports single-timestep calls "
+            f"(got shape {tuple(input.shape)})."
+        )
+        if self.moe_source_dims is not None:
+            raise RuntimeError(
+                "CfCControllerWrapper was configured with moe_source_dims -- "
+                "call forward_multi_source(sources, hx) instead of forward()."
+            )
+        x = self.in_adapter(input.squeeze(1))
+        if hx is None:
+            hx = self.init_state(x.size(0), device=x.device)
+        new_hx = []
+        blocks = cast(Iterable[CfCControllerBlock], self.blocks)
+        for i, (block, state) in enumerate(zip(blocks, hx)):
+            x, new_state = block.step(x, state)
+            if self.moe_enabled:
+                assert self.moe_blocks is not None
+                x = self.moe_blocks[i](x)
+            new_hx.append(new_state)
+        return x.unsqueeze(1), new_hx
+
+    def forward_multi_source(self, sources: list[torch.Tensor], hx):
+        """Multi-source entry point (requires moe_source_dims). `sources[i]`
+        is a (B, moe_source_dims[i]) tensor -- e.g. SplitGraphDNC's
+        controller combiner passes [h_t, read_vec] directly instead of
+        pre-concatenating them. Each source is projected to d_model by its
+        own adapter, then routed+combined by the shared, source-aware
+        Top-K MoE bank BEFORE entering the recurrent CfC stack -- this is
+        what preserves CfC's "many inputs, specialized handling per input"
+        property once MoE sits in front of it, while the recurrent stack
+        itself still only ever sees one fused (B, d_model) vector per step
+        (CfC block internals are completely unchanged)."""
+        if self.moe_source_dims is None:
+            raise RuntimeError("forward_multi_source requires moe_source_dims to have been set")
+        assert self.source_in_adapters is not None and self.moe_blocks is not None
+        projected = [adapter(s) for adapter, s in zip(self.source_in_adapters, sources)]
+        fused_per_source = self.moe_blocks[0](projected)  # list[Tensor], one per source
+        x = torch.stack(fused_per_source, dim=0).sum(dim=0)  # combine into one fused input
+        if hx is None:
+            hx = self.init_state(x.size(0), device=x.device)
+        new_hx = []
+        blocks = cast(Iterable[CfCControllerBlock], self.blocks)
+        for block, state in zip(blocks, hx):
+            x, new_state = block.step(x, state)
+            new_hx.append(new_state)
+        return x.unsqueeze(1), new_hx
+
+
+if (
+    __name__ == "__main__"
+):  # smoke test: python -m LNN_controller.cfc_controller (from project root)
+    from initium.LNN_controller.chained_controller import ChainedControllerWrapper
+
+    B, T, in_dim, d = 4, 6, 40, 32
+    for mm in (False, True):
+        w = CfCControllerWrapper(in_dim, d, num_blocks=2, backbone_units=64, mixed_memory=mm)
+        hx, loss = w.init_state(B), torch.zeros(())
+        for _ in range(T):
+            out, hx = w(torch.randn(B, 1, in_dim), hx)
+            assert out.shape == (B, 1, d), out.shape
+            loss = loss + out.pow(2).mean()
+        loss.backward()  # BPTT across chained steps must not raise
+        print(f"standalone mixed_memory={mm}: OK (loss {float(loss):.4f})")
+    ch = ChainedControllerWrapper(
+        [
+            CfCControllerWrapper(in_dim, d, 1, backbone_units=64),
+            CfCControllerWrapper(d, d, 1, backbone_units=64),
+        ]
+    )
+    hx, loss = ch.init_state(B), torch.zeros(())
+    for _ in range(T):
+        out, hx = ch(torch.randn(B, 1, in_dim), hx)
+        loss = loss + out.pow(2).mean()
+    loss.backward()
+    print("chained: OK")

initium/LNN_controller/chained_controller.py

@@ -0,0 +1,83 @@
+"""
+file: LNN_controller/chained_controller.py -- v1
+
+Generic composer: runs several controller wrappers back-to-back per DNC
+timestep, each feeding the next (stage i's d_model must equal stage i+1's
+input width). Works with ANY wrapper following the controller-wrapper
+protocol in cfc_controller.py's docstring: MambaControllerWrapper,
+Mamba2ControllerWrapper, Mamba3ControllerWrapper, CfCControllerWrapper, or
+future ones. Example (Mamba backbone -> CfC head):
+
+    ChainedControllerWrapper([
+        MambaControllerWrapper(in_dim=in_dim, d_model=d, num_blocks=2),
+        CfCControllerWrapper(in_dim=d, d_model=d, num_blocks=1),
+    ])
+
+State is a list with one entry per stage. Also usable as a SplitGraphDNC
+combiner_wrapper (same protocol).
+"""
+
+from __future__ import annotations
+
+import torch.nn as nn
+
+
+def _in_features(stage) -> int:
+    a = stage.in_adapter
+    return a.in_features if isinstance(a, nn.Linear) else stage.d_model
+
+
+class ChainedControllerWrapper(nn.Module):
+    def __init__(self, stages, stage_kinds=None):
+        super().__init__()
+        stages = list(stages)
+        assert len(stages) >= 1, "ChainedControllerWrapper needs at least one stage"
+        for prev, nxt in zip(stages[:-1], stages[1:]):
+            assert _in_features(nxt) == prev.d_model, (
+                f"stage width mismatch: previous d_model={prev.d_model}, next stage input={_in_features(nxt)}"
+            )
+        self.stages = nn.ModuleList(stages)
+        self.d_model = stages[-1].d_model
+        self.in_adapter = stages[
+            0
+        ].in_adapter  # same object (not re-registered under a new path issue: shared reference)
+        # plain python list (NOT ModuleList): the MoE blocks are already registered inside their stages
+        self.moe_blocks = [
+            b for s in stages if getattr(s, "moe_enabled", False) for b in s.moe_blocks
+        ]
+        self.moe_enabled = len(self.moe_blocks) > 0
+        # Names each stage (e.g. ["mamba", "cfc"] for "mamba+cfc") so the
+        # training loop can log which stage a given ablation index refers
+        # to. None only for the [f"stage{i}"...] fallback below when a
+        # caller doesn't pass kinds.
+        self.stage_kinds = (
+            list(stage_kinds)
+            if stage_kinds is not None
+            else [f"stage{i}" for i in range(len(self.stages))]
+        )
+
+    def init_state(self, batch_size, device=None, dtype=None):
+        return [s.init_state(batch_size, device=device, dtype=dtype) for s in self.stages]
+
+    def forward(self, input, hx, skip_stages=None):
+        """skip_stages: optional set/list of stage indices to bypass, for
+        measuring each stage's contribution the same way pass_through_memory
+        measures Memory's contribution (see evaluate_traversal's
+        ablate_memory / SplitGraphDNC's combiner_skip_stages). Stage 0 still
+        gets its in_adapter applied (so downstream stages see the right
+        width) but skips its actual recurrent computation; every later index
+        is skipped entirely (its d_model-in == d_model-out, so identity is
+        exact)."""
+        if hx is None:
+            hx = [None] * len(self.stages)
+        skip_stages = set(skip_stages) if skip_stages else set()
+        x, new_hx = input, []
+        for i, (stage, state) in enumerate(zip(self.stages, hx)):
+            if i in skip_stages:
+                if i == 0:
+                    x = stage.in_adapter(x.squeeze(1)).unsqueeze(1)
+                new_hx.append(state)
+                continue
+            x, new_state = stage(x, state)
+            new_hx.append(new_state)
+        return x, new_hx