# Copyright the ShellElastic authors.
"""Reference PyTorch implementation of the ShellElastic supernet.

This file implements the model of "ShellElastic: Scalable Elastic Pretraining with Simultaneous Maximal Updates and
Hyperparameter Transfer" (Jin et al.). It depends only on ``torch``; ``safetensors`` is needed to load released
weights.

Notation follows Section 3.1 of the paper.

* Width. With base width ``g`` and supported set ``S = {s_1 < ... < s_K = M}`` (stored in ``config.widths``, in
  units of ``g``), the residual stream has ``M * g`` features split into contiguous groups
  ``G_k = [s_{k-1} g, s_k g)``. ``G_1`` is the dense core, ``G_2 .. G_K`` are the shells. The sub-network of tier
  ``k`` (Low, Medium, High, Extra, Max) is the core plus the first ``k - 1`` shells, i.e. the first ``s_k g``
  features.
* Structured projections (Eq. 2). Every hidden matrix is block lower triangular: the block row ``W_k`` that
  produces group ``G_k`` reads only the first ``s_k g`` input features and has shape ``[|G_k|, s_k g]``. Each
  matrix is stored by block row: ``<proj>.core.weight`` is ``W_1`` (the rows of the core ``G_1``), and
  ``<proj>.shell.<j>.weight`` is ``W_{j+1}`` (the rows of the shell ``G_{j+1}``, ``j = 1 .. K-1``), of shape
  ``[|G_{j+1}|, s_{j+1} g]``. The zero blocks above the diagonal are never stored.
* Group-local normalization (Eq. 3). The two pre-norms of every layer are RMSNorms whose statistics are computed
  separately over the core and over each shell (one gain per feature).
* Attention (Eq. 4). Tier ``k`` uses the first ``s_k g / d_head`` heads; heads never straddle a group. Queries and
  keys get a per-head RMSNorm and rotary position embeddings, and the attention output is multiplied elementwise
  by ``sigmoid(x W_G^T)`` before the output projection.
* FFN (Eq. 5). SwiGLU with hidden width ``e`` times the model width; ``gate_proj`` / ``up_proj`` have block
  boundaries ``s_k g`` on the input axis and ``e s_k g`` on the output axis, ``down_proj`` the reverse.
* Readout (Eq. 6). Every tier has its own terminal RMSNorm over all of its ``s_k g`` features (``norm.<tier>``) and
  its own readout ``U_s`` (``lm_head.<tier>``). In the released weights the learned logit scale ``alpha_s`` is
  already multiplied into ``lm_head.<tier>.weight`` and the FFN multiplier ``1/e`` into ``down_proj``, so inference
  needs no extra scalars.

Because no operation lets the features of a tier depend on the groups outside it, one forward pass at width
``M g`` contains the exact hidden states of every smaller tier in its leading features (Proposition 1).
``forward(..., tiers=[...])`` uses this to return several tiers' logits from one trunk pass, and ``extract(tier)``
returns a standalone smaller model with identical outputs.

Numerics. This file computes the RMSNorm statistics in float32 and casts the result back to the activation dtype
before the gain is applied. ``config.rope_precision`` selects the rotary tables: ``"bfloat16"`` (the default)
rounds the positions, the frequencies, their product and the cosine and sine to bfloat16; ``"float32"`` computes
them in float32. The released weights are float32; run the model in float32, or cast it to bfloat16 with
``from_pretrained(..., dtype=torch.bfloat16)``.

Example::

    model = ShellElasticForCausalLM.from_pretrained("path/to/checkpoint", device="cuda")
    logits = model(input_ids, tier="medium")                 # [batch, seq, vocab]
    all_logits = model(input_ids, tiers=model.config.tiers)  # dict: tier -> logits
    logits, loss = model(input_ids[:, :-1], tier="low", labels=input_ids[:, 1:])
    low = model.extract("low")                               # standalone Low sub-network
"""

from __future__ import annotations

import json
import os
from dataclasses import asdict, dataclass, field
from typing import Dict, List, Optional, Sequence, Union

import torch
import torch.nn as nn
import torch.nn.functional as F


DEFAULT_TIERS = ["low", "medium", "high", "extra", "max"]


# --------------------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------------------


@dataclass
class ShellElasticConfig:
    """Hyperparameters of a ShellElastic supernet (or of a sub-network extracted from it).

    Attributes:
        vocab_size: Size of the vocabulary.
        num_hidden_layers: Number of transformer layers ``L``.
        base_width: Base width ``g``; every group is a multiple of it.
        widths: Supported set ``S`` in units of ``g``, strictly increasing. The last entry is ``M`` (the full width
            is ``M * g``). Group ``G_1`` (the core) has ``widths[0] * g`` features, every later group
            ``(widths[k] - widths[k-1]) * g``.
        tiers: One name per supported width (Low, Medium, High, Extra, Max).
        readout_tiers: Tiers that have their own terminal norm and readout in this checkpoint (default: all tiers;
            a sub-network returned by ``extract`` keeps only its own).
        head_dim: Attention head dimension ``d_head`` (must divide ``base_width``).
        mlp_base_width: FFN hidden features per base width, i.e. ``e * g`` with FFN expansion ``e``.
        rope_theta: Base of the rotary position embedding.
        rope_precision: Rotary tables, ``"bfloat16"`` (default: position, frequency and their product rounded to
            bfloat16, and the cosine and sine rounded to bfloat16, in any compute precision) or ``"float32"``
            (computed in float32).
        rms_norm_eps: Epsilon of every RMSNorm.
        max_position_embeddings: Context length (the sequence length of 2,048 tokens).
        tie_word_embeddings: Always false (the embedding and the readouts are separate).
        tokenizer: Tokenizer the token ids refer to.
    """

    vocab_size: int = 50432
    num_hidden_layers: int = 32
    base_width: int = 128
    widths: List[int] = field(default_factory=lambda: [4, 5, 6, 7, 8])
    tiers: List[str] = field(default_factory=lambda: list(DEFAULT_TIERS))
    readout_tiers: Optional[List[str]] = None
    head_dim: int = 128
    mlp_base_width: int = 128
    rope_theta: float = 10000.0
    rope_precision: str = "bfloat16"
    rms_norm_eps: float = 1e-5
    max_position_embeddings: int = 2048
    tie_word_embeddings: bool = False
    tokenizer: str = "EleutherAI/gpt-neox-20b"
    architectures: List[str] = field(default_factory=lambda: ["ShellElasticForCausalLM"])
    model_type: str = "shellelastic"

    def __post_init__(self) -> None:
        self.widths = [int(w) for w in self.widths]
        self.tiers = [str(t) for t in self.tiers]
        if self.readout_tiers is None:
            self.readout_tiers = list(self.tiers)
        if any(b <= a for a, b in zip([0] + self.widths, self.widths)):
            raise ValueError(f"widths must be positive and strictly increasing, got {self.widths}")
        if len(self.tiers) != len(self.widths):
            raise ValueError("tiers and widths must have the same length")
        if not self.readout_tiers or any(t not in self.tiers for t in self.readout_tiers):
            raise ValueError(f"readout_tiers {self.readout_tiers} must be a non-empty subset of tiers {self.tiers}")
        self.readout_tiers = sorted(set(self.readout_tiers), key=self.tiers.index)
        if self.base_width % self.head_dim != 0:
            raise ValueError("head_dim must divide base_width so that heads never straddle a group")
        if self.rope_precision not in ("bfloat16", "float32"):
            raise ValueError("rope_precision must be 'bfloat16' or 'float32'")
        if self.tie_word_embeddings:
            raise ValueError("ShellElastic models do not tie the embedding and the readouts")

    # ---- derived quantities --------------------------------------------------------------

    @property
    def hidden_size(self) -> int:
        """Full model width ``M * g``."""
        return self.widths[-1] * self.base_width

    @property
    def boundaries(self) -> List[int]:
        """Right edges of the residual groups, ``[s_1 g, ..., s_K g]``."""
        return [s * self.base_width for s in self.widths]

    @property
    def mlp_boundaries(self) -> List[int]:
        """Right edges of the FFN hidden groups, ``[e s_1 g, ..., e s_K g]``."""
        return [s * self.mlp_base_width for s in self.widths]

    @property
    def group_sizes(self) -> List[int]:
        """Sizes of the core and of every shell."""
        b = self.boundaries
        return [b[0]] + [hi - lo for lo, hi in zip(b, b[1:])]

    def num_groups(self, tier: str) -> int:
        """Number of groups (core + shells) used by ``tier``."""
        return self.tiers.index(tier) + 1

    def width(self, tier: str) -> int:
        """Model width of ``tier``."""
        return self.boundaries[self.tiers.index(tier)]

    def for_tier(self, tier: str) -> "ShellElasticConfig":
        """Configuration of the standalone sub-network ``tier`` (see ``extract``)."""
        k = self.num_groups(tier)
        values = asdict(self)
        values.update(widths=self.widths[:k], tiers=self.tiers[:k], readout_tiers=[tier])
        return ShellElasticConfig(**values)

    # ---- (de)serialization ---------------------------------------------------------------

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, values: dict) -> "ShellElasticConfig":
        known = {f for f in cls.__dataclass_fields__}
        return cls(**{k: v for k, v in values.items() if k in known})

    @classmethod
    def from_json_file(cls, path: str) -> "ShellElasticConfig":
        with open(path, "r", encoding="utf-8") as f:
            return cls.from_dict(json.load(f))


# --------------------------------------------------------------------------------------
# Building blocks
# --------------------------------------------------------------------------------------


def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float,
             group_sizes: Optional[Sequence[int]] = None) -> torch.Tensor:
    """RMSNorm over the last dimension, with separate statistics per feature group (Eq. 3).

    The statistics are computed in float32 and the result is cast back to the input dtype before the gain is
    applied. ``group_sizes=None`` normalizes over the whole last dimension.
    """
    xf = x.float()
    chunks = xf.split(list(group_sizes), dim=-1) if group_sizes is not None else (xf,)
    normed = [c * torch.rsqrt(c.pow(2).mean(dim=-1, keepdim=True) + eps) for c in chunks]
    return torch.cat(normed, dim=-1).to(x.dtype) * weight


class RMSNorm(nn.Module):
    """RMSNorm with one gain per feature. ``group_sizes`` selects group-local statistics.

    The gain has the full width of the module; a narrower input uses its leading entries.
    """

    def __init__(self, width: int, eps: float) -> None:
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(width))

    def forward(self, x: torch.Tensor, group_sizes: Optional[Sequence[int]] = None) -> torch.Tensor:
        return rms_norm(x, self.weight[: x.shape[-1]], self.eps, group_sizes)


class ShellLinear(nn.Module):
    """Block-lower-triangular linear map, stored as one weight per output group (block row).

    ``core.weight`` (the block row of the core ``G_1``) has shape ``[out_bounds[0], in_bounds[0]]``;
    ``shell[str(j)].weight`` (the block row of the shell ``G_{j+1}``, ``j = 1 .. K-1``) has shape
    ``[out_bounds[j] - out_bounds[j-1], in_bounds[j]]``: its rows read only the first ``in_bounds[j]`` input
    features. There is no bias.
    """

    def __init__(self, in_bounds: Sequence[int], out_bounds: Sequence[int]) -> None:
        super().__init__()
        self.in_bounds = list(in_bounds)
        self.out_bounds = list(out_bounds)
        self.core = nn.Linear(self.in_bounds[0], self.out_bounds[0], bias=False)
        self.shell = nn.ModuleDict({
            str(k): nn.Linear(self.in_bounds[k], self.out_bounds[k] - self.out_bounds[k - 1], bias=False)
            for k in range(1, len(self.out_bounds))
        })

    def block_rows(self, num_groups: int) -> List[torch.Tensor]:
        """Weights of the first ``num_groups`` block rows (core first)."""
        return [self.core.weight] + [self.shell[str(k)].weight for k in range(1, num_groups)]

    def forward(self, x: torch.Tensor, num_groups: int) -> torch.Tensor:
        # y_{G_k} = x[:, :s_k g] W_k^T  (Eq. 2), concatenated over the active groups.
        return torch.cat([F.linear(x[..., : w.shape[1]], w) for w in self.block_rows(num_groups)], dim=-1)

    def dense_weight(self, num_groups: int) -> torch.Tensor:
        """The equivalent dense ``[out, in]`` matrix of the first ``num_groups`` groups, with the blocks above the
        diagonal filled with zeros."""
        n_in = self.in_bounds[num_groups - 1]
        return torch.cat([F.pad(w, (0, n_in - w.shape[1])) for w in self.block_rows(num_groups)], dim=0)


def rotary_cos_sin(length: int, head_dim: int, theta: float, precision: str,
                   device: torch.device, dtype: torch.dtype):
    """Cosine and sine tables of shape ``[length, head_dim]`` for rotate-half RoPE."""
    frequency = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))  # computed on the CPU
    position = torch.arange(length, dtype=torch.float32)
    frequency, position = frequency.to(device), position.to(device)
    if precision == "bfloat16":
        # Position, frequency, angle (their product, which is exact in float32 before rounding) and the cosine /
        # sine are each rounded to bfloat16.
        angles = torch.outer(position.bfloat16().float(), frequency.bfloat16().float()).bfloat16().float()
        angles = torch.cat([angles, angles], dim=-1)
        return angles.cos().bfloat16().to(dtype), angles.sin().bfloat16().to(dtype)
    angles = torch.outer(position, frequency)
    angles = torch.cat([angles, angles], dim=-1)
    return angles.cos().to(dtype), angles.sin().to(dtype)


def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Rotate-half RoPE for ``x`` of shape ``[batch, seq, heads, head_dim]``."""
    x1, x2 = x.chunk(2, dim=-1)
    rotated = torch.cat([-x2, x1], dim=-1)
    return x * cos[None, :, None, :] + rotated * sin[None, :, None, :]


class ShellAttention(nn.Module):
    """Causal self-attention with nested heads, QK-norm, RoPE and an elementwise sigmoid gate."""

    def __init__(self, config: ShellElasticConfig) -> None:
        super().__init__()
        b, width = config.boundaries, config.hidden_size
        self.head_dim = config.head_dim
        self.q_proj = ShellLinear(b, b)
        self.k_proj = ShellLinear(b, b)
        self.v_proj = ShellLinear(b, b)
        self.gate_proj = ShellLinear(b, b)
        self.o_proj = ShellLinear(b, b)
        # Per-head RMSNorm of queries and keys, with one gain per feature (head-major order).
        self.q_norm = RMSNorm(width, config.rms_norm_eps)
        self.k_norm = RMSNorm(width, config.rms_norm_eps)

    def _per_head_norm(self, norm: RMSNorm, x: torch.Tensor) -> torch.Tensor:
        heads = x.shape[-2]
        gain = norm.weight[: heads * self.head_dim].view(heads, self.head_dim)
        return rms_norm(x, gain, norm.eps)

    def forward(self, x: torch.Tensor, num_groups: int, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        batch, seq, width = x.shape
        heads = width // self.head_dim
        shape = (batch, seq, heads, self.head_dim)
        q = self.q_proj(x, num_groups).view(shape)
        k = self.k_proj(x, num_groups).view(shape)
        v = self.v_proj(x, num_groups).view(shape)
        q = apply_rotary(self._per_head_norm(self.q_norm, q), cos, sin)
        k = apply_rotary(self._per_head_norm(self.k_norm, k), cos, sin)
        out = F.scaled_dot_product_attention(
            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
            is_causal=True, scale=self.head_dim ** -0.5,
        ).transpose(1, 2)
        out = out * torch.sigmoid(self.gate_proj(x, num_groups).view(shape))
        return self.o_proj(out.reshape(batch, seq, width), num_groups)


class ShellMLP(nn.Module):
    """SwiGLU feed-forward network with block-lower-triangular projections."""

    def __init__(self, config: ShellElasticConfig) -> None:
        super().__init__()
        b, h = config.boundaries, config.mlp_boundaries
        self.gate_proj = ShellLinear(b, h)
        self.up_proj = ShellLinear(b, h)
        self.down_proj = ShellLinear(h, b)  # the released weights already include the 1/e multiplier

    def forward(self, x: torch.Tensor, num_groups: int) -> torch.Tensor:
        hidden = F.silu(self.gate_proj(x, num_groups)) * self.up_proj(x, num_groups)
        return self.down_proj(hidden, num_groups)


class ShellElasticLayer(nn.Module):
    """Pre-norm transformer layer with group-local branch normalization."""

    def __init__(self, config: ShellElasticConfig) -> None:
        super().__init__()
        self.group_sizes = config.group_sizes
        self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = ShellAttention(config)
        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = ShellMLP(config)

    def forward(self, x: torch.Tensor, num_groups: int, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        groups = self.group_sizes[:num_groups]
        x = x + self.attn(self.attn_norm(x, groups), num_groups, cos, sin)
        x = x + self.mlp(self.mlp_norm(x, groups), num_groups)
        return x


# --------------------------------------------------------------------------------------
# Model
# --------------------------------------------------------------------------------------


class ShellElasticForCausalLM(nn.Module):
    """ShellElastic supernet: a shared trunk plus one terminal norm and readout per tier."""

    def __init__(self, config: ShellElasticConfig) -> None:
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([ShellElasticLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = nn.ModuleDict({t: RMSNorm(config.width(t), config.rms_norm_eps) for t in config.readout_tiers})
        self.lm_head = nn.ModuleDict({
            t: nn.Linear(config.width(t), config.vocab_size, bias=False) for t in config.readout_tiers
        })

    # ---- loading -------------------------------------------------------------------------

    @classmethod
    def from_pretrained(cls, path: str, device: Union[str, torch.device] = "cpu",
                        dtype: Optional[torch.dtype] = None) -> "ShellElasticForCausalLM":
        """Load ``config.json`` and ``model.safetensors`` from the directory ``path``; ``dtype`` optionally casts
        the (float32) weights, e.g. to ``torch.bfloat16``. The weights are read and cast on the CPU and then moved
        to ``device``, so the device holds only one copy of the model, in ``dtype``."""
        from safetensors.torch import load_file

        config = ShellElasticConfig.from_json_file(os.path.join(path, "config.json"))
        with torch.device("meta"):
            model = cls(config)
        state = load_file(os.path.join(path, "model.safetensors"), device="cpu")
        if dtype is not None:
            for name in list(state):  # cast on the CPU, one tensor at a time
                if state[name].is_floating_point():
                    state[name] = state[name].to(dtype)
        model.load_state_dict(state, strict=True, assign=True)
        return model.to(device).eval()

    def save_pretrained(self, path: str) -> None:
        """Write ``config.json`` and ``model.safetensors`` to the directory ``path`` (for example, a sub-network
        returned by ``extract``); ``from_pretrained`` reads them back."""
        from safetensors.torch import save_file

        os.makedirs(path, exist_ok=True)
        values = self.config.to_dict()
        values["torch_dtype"] = str(self.embed_tokens.weight.dtype).replace("torch.", "")
        with open(os.path.join(path, "config.json"), "w", encoding="utf-8") as f:
            json.dump(values, f, indent=2)
            f.write("\n")
        state = {name: t.detach().contiguous() for name, t in self.state_dict().items()}
        save_file(state, os.path.join(path, "model.safetensors"), metadata={"format": "pt"})

    # ---- forward -------------------------------------------------------------------------

    def hidden_states(self, input_ids: torch.Tensor, tier: Optional[str] = None) -> torch.Tensor:
        """Trunk output (before the terminal norm) of width ``config.width(tier)``.

        Its first ``config.width(t)`` features are the hidden states of every smaller tier ``t``.
        """
        tier = tier or self.config.tiers[-1]
        num_groups = self.config.num_groups(tier)
        width = self.config.width(tier)
        seq = input_ids.shape[1]
        if seq > self.config.max_position_embeddings:
            raise ValueError(f"{seq} tokens exceed the context of {self.config.max_position_embeddings}")
        x = self.embed_tokens(input_ids)[..., :width]
        cos, sin = rotary_cos_sin(seq, self.config.head_dim, self.config.rope_theta,
                                  self.config.rope_precision, x.device, x.dtype)
        for layer in self.layers:
            x = layer(x, num_groups, cos, sin)
        return x

    def logits(self, hidden: torch.Tensor, tier: str) -> torch.Tensor:
        """Terminal norm and readout of ``tier`` (Eq. 6) applied to a trunk output of width >= its own."""
        if tier not in self.lm_head:
            raise KeyError(f"tier {tier!r} has no readout in this model (available: {self.config.readout_tiers})")
        h = hidden[..., : self.config.width(tier)]
        return self.lm_head[tier](self.norm[tier](h))

    def forward(self, input_ids: torch.Tensor, tier: Optional[str] = None, tiers: Optional[Sequence[str]] = None,
                labels: Optional[torch.Tensor] = None):
        """Next-token logits (in the model's dtype).

        Args:
            input_ids: ``[batch, seq]`` token ids (no padding; causal attention over the sequence).
            tier: Tier to evaluate (default: the widest tier with a readout).
            tiers: If given, evaluate all of these tiers with one trunk pass at the widest of them and return a
                ``{tier: logits}`` dict.
            labels: Optional ``[batch, seq]`` targets (-100 is ignored). The mean cross-entropy (computed in float32)
                is returned as well: ``(logits, loss)``, or ``({tier: logits}, {tier: loss})`` with ``tiers``.
        """
        if tiers is None:
            tier = tier or self.config.readout_tiers[-1]
            logits = self.logits(self.hidden_states(input_ids, tier), tier)
            return logits if labels is None else (logits, self.loss(logits, labels))
        widest = max(tiers, key=self.config.tiers.index)
        hidden = self.hidden_states(input_ids, widest)
        out = {t: self.logits(hidden, t) for t in tiers}
        return out if labels is None else (out, {t: self.loss(v, labels) for t, v in out.items()})

    @staticmethod
    def loss(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        """Mean next-token cross-entropy in float32 (label -100 is ignored)."""
        return F.cross_entropy(logits.float().flatten(0, 1), labels.flatten(), ignore_index=-100)

    @torch.no_grad()
    def generate(self, input_ids: torch.Tensor, max_new_tokens: int, tier: Optional[str] = None,
                 temperature: float = 0.0, top_k: Optional[int] = None) -> torch.Tensor:
        """Simple sampling loop without a key/value cache (greedy when ``temperature == 0``)."""
        for _ in range(max_new_tokens):
            context = input_ids[:, -self.config.max_position_embeddings:]
            logits = self(context, tier=tier)[:, -1].float()
            if temperature <= 0:
                next_ids = logits.argmax(dim=-1, keepdim=True)
            else:
                logits = logits / temperature
                if top_k is not None:
                    kth = torch.topk(logits, top_k, dim=-1).values[:, -1:]
                    logits = logits.masked_fill(logits < kth, float("-inf"))
                next_ids = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
            input_ids = torch.cat([input_ids, next_ids], dim=1)
        return input_ids

    # ---- sub-networks --------------------------------------------------------------------

    def subnetwork_state_dict(self, tier: str) -> Dict[str, torch.Tensor]:
        """Weights of the standalone sub-network ``tier``: the embedding columns, the core and the shells it
        contains, the leading entries of every norm gain, and its own terminal norm and readout. Loads into
        ``ShellElasticForCausalLM(config.for_tier(tier))``."""
        if tier not in self.lm_head:
            raise KeyError(f"tier {tier!r} has no readout in this model (available: {self.config.readout_tiers})")
        k, width = self.config.num_groups(tier), self.config.width(tier)
        out = {}
        for name, tensor in self.state_dict().items():
            if name.startswith(("norm.", "lm_head.")):
                if name.split(".")[1] == tier:
                    out[name] = tensor
            elif name == "embed_tokens.weight":
                out[name] = tensor[:, :width]
            elif ".shell." in name:
                if int(name.split(".shell.")[1].split(".")[0]) < k:
                    out[name] = tensor
            elif name.endswith("norm.weight"):  # branch norms and query/key norms
                out[name] = tensor[:width]
            else:  # core block rows
                out[name] = tensor
        return out

    def extract(self, tier: str) -> "ShellElasticForCausalLM":
        """A standalone copy of the sub-network ``tier``; its logits equal ``self(x, tier=tier)``."""
        with torch.device("meta"):
            model = ShellElasticForCausalLM(self.config.for_tier(tier))
        state = {n: t.clone() for n, t in self.subnetwork_state_dict(tier).items()}
        model.load_state_dict(state, strict=True, assign=True)
        return model.train(self.training)

    def dense_state_dict(self, tier: str) -> Dict[str, torch.Tensor]:
        """Returns the weights of the sub-network ``tier`` as a state dict (nothing is written to disk), with every
        structured projection as one dense matrix (``...<proj>.weight``, zeros above the block diagonal). The
        group-local norms are still required to run it; this format is provided for inspection and porting."""
        k = self.config.num_groups(tier)
        out = {}
        for name, tensor in self.subnetwork_state_dict(tier).items():
            if ".core." in name or ".shell." in name:
                continue
            out[name] = tensor
        for name, module in self.named_modules():
            if isinstance(module, ShellLinear):
                out[f"{name}.weight"] = module.dense_weight(k)
        return out

    def num_parameters(self, tier: Optional[str] = None) -> int:
        """Parameters stored for the whole model, or used by ``tier`` (non-zero weights it reads, its embedding
        columns, and its own terminal norm and readout)."""
        tensors = self.state_dict() if tier is None else self.subnetwork_state_dict(tier)
        return sum(t.numel() for t in tensors.values())


__all__ = [
    "ShellElasticConfig",
    "ShellElasticForCausalLM",
    "ShellLinear",
    "RMSNorm",
    "rms_norm",
    "rotary_cos_sin",
]