# Copyright the ShellElastic authors.
"""Reference PyTorch implementation of the MatFormer and Fully elastic supernets.

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

One set of weights holds five nested sub-networks (tiers Low, Medium, High, Extra, Max). Every tier
reads the leading block of each shared matrix, and its weights follow the μ-Elastic rules of the
paper (Tables 5 and 6). Every hidden matrix is stored as the two tensors of Eq. 23: the frozen
initialization ``init`` (W^init) and the trained ``delta``. Tier s uses

    W_s = r^(-1/2) * init[:rows, :cols] + r^(-1) * delta[:rows, :cols],

where r is the fan-in of the tier divided by the base fan-in (r = 1 for a matrix whose fan-in is
the same in every tier). This file combines these tensors when the model runs, in the dtype of the
parameters.

* Fully elastic (``variant = "fully_elastic"``). With base width g and supported set S
  (``widths``), tier s has model width s*g, s*g/d_head attention heads and FFN width s*h, where
  h = e*g (``mlp_base_width``) for the FFN expansion e. Every hidden matrix takes its top-left block
  with r = s; the token embedding and every norm gain are read up to the first s*g features, and the
  RMSNorms take their statistics over the active width. The shared readout is split in the same way
  and scaled by the learned per-tier scalar alpha_s (``lm_head.logit_scale``):
      U_s = alpha_s * (s^(-1/2) * U^init[:, :s*g] + s^(-1) * delta_U[:, :s*g]).
* MatFormer (``variant = "matformer"``). The model width D (``hidden_size``) is fixed and tier s
  keeps the first s*h0 hidden units of every FFN (base FFN width h0 = ``mlp_base_width``). Only the
  FFN down projection has a fan-in that changes with the tier (r = s); every other matrix has the
  fixed fan-in D and the multiplier 1 on both parts (Table 6), so W = W^init + delta in every tier.
  The released weights keep both parts of every split matrix (``split_matrices`` lists the matrices
  stored split; any other matrix is a single ``weight``). One readout ``lm_head.weight`` is shared by
  all tiers and scaled by one learned scalar alpha.
* In both, the output of the FFN down projection is divided by e (``mlp_output_divisor``; for
  MatFormer e = h0 / D takes the place of the FFN expansion).

Common architecture: 32 pre-norm layers with residual multiplier 1; RMSNorm with eps 1e-5 (this file
computes the statistics in float32 and casts the result back to the activation dtype before the gain
is applied); causal self-attention with a per-head RMSNorm of queries and keys (one gain per feature),
rotary position embeddings in the rotate-half form (base 10,000), softmax scale 1/sqrt(d_head) and an
elementwise sigmoid output gate, sigmoid(x W_G^T); SwiGLU FFN; no biases; untied embedding and readout
over the 50,432-token GPT-NeoX vocabulary (Table 8 of the paper; token ids of the GPT-NeoX-20B
tokenizer, Section 4.1); context 2,048 tokens.

Numerics. ``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)``. The learned scalar alpha and the divisor e are
applied when the model runs; ``extract`` folds them into the dense weights.

Example::

    model = ElasticForCausalLM.from_pretrained("path/to/checkpoint", device="cuda")
    logits = model(input_ids, tier="medium")              # [batch, tokens, vocab], float32
    logits, loss = model(input_ids[:, :-1], tier="low", labels=input_ids[:, 1:])
    weights = model.extract("low")                        # standalone dense Low model (state dict)
    model.save_subnetwork("low", "path/to/low")           # config.json + model.safetensors
"""

from __future__ import annotations

import json
import os
from collections import OrderedDict
from dataclasses import asdict, dataclass, field
from typing import List, Optional, Tuple, Union

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


DEFAULT_TIERS = ["low", "medium", "high", "extra", "max"]
VARIANTS = ("fully_elastic", "matformer")
# Every shared matrix of the model, by its name inside a layer (plus the readout).
MATRICES = ["attn.q_proj", "attn.k_proj", "attn.v_proj", "attn.gate_proj", "attn.o_proj",
            "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj", "lm_head"]


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


@dataclass
class ElasticConfig:
    """Hyperparameters of a MatFormer or Fully elastic supernet.

    Attributes:
        variant: ``"fully_elastic"`` (model width, heads and FFN width nest) or ``"matformer"``
            (only the FFN hidden width nests).
        vocab_size: Size of the vocabulary.
        num_hidden_layers: Number of transformer layers.
        hidden_size: Full model width (the width of the Max tier; the width of every MatFormer tier).
        head_dim: Attention head dimension (the same in every tier).
        base_width: Fully elastic only: model width per unit of s (the base width g).
        mlp_base_width: FFN hidden width per unit of s (e*g for Fully elastic, h0 for MatFormer).
        widths: Supported set S, strictly increasing, one entry per tier.
        tiers: Tier names, one per entry of ``widths``.
        mlp_output_divisor: e; the output of every FFN down projection is divided by it.
        split_matrices: Matrices stored as ``init`` + ``delta`` (Eq. 23); every other matrix is stored
            as one ``weight``. Names as in ``MATRICES``.
        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 readout are separate).
        tokenizer: Tokenizer the token ids refer to.
    """

    variant: str = "fully_elastic"
    vocab_size: int = 50432
    num_hidden_layers: int = 32
    hidden_size: int = 640
    head_dim: int = 64
    base_width: Optional[int] = 64
    mlp_base_width: int = 256
    widths: List[int] = field(default_factory=lambda: [6, 7, 8, 9, 10])
    tiers: List[str] = field(default_factory=lambda: list(DEFAULT_TIERS))
    mlp_output_divisor: float = 4.0
    split_matrices: List[str] = field(default_factory=lambda: list(MATRICES))
    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: ["ElasticForCausalLM"])
    model_type: str = "elastic"

    def __post_init__(self) -> None:
        self.widths = [int(w) for w in self.widths]
        self.tiers = [str(t) for t in self.tiers]
        self.split_matrices = [str(m) for m in self.split_matrices]
        if self.variant not in VARIANTS:
            raise ValueError(f"variant must be one of {VARIANTS}, got {self.variant!r}")
        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")
        unknown = sorted(set(self.split_matrices) - set(MATRICES))
        if unknown:
            raise ValueError(f"unknown matrices in split_matrices: {unknown}")
        if self.variant == "fully_elastic":
            # every matrix of a Fully elastic model changes with the tier, so it must be stored split
            if sorted(self.split_matrices) != sorted(MATRICES):
                raise ValueError("a Fully elastic model stores every matrix as init + delta")
            if self.base_width is None or self.base_width * self.widths[-1] != self.hidden_size:
                raise ValueError("hidden_size must equal base_width * widths[-1]")
            if self.base_width % self.head_dim != 0:
                raise ValueError("head_dim must divide base_width so that every tier has whole heads")
        else:
            if "mlp.down_proj" not in self.split_matrices or "lm_head" in self.split_matrices:
                raise ValueError("a MatFormer model stores the FFN down projection split and the readout whole")
            if self.hidden_size % self.head_dim != 0:
                raise ValueError("head_dim must divide hidden_size")
        if self.rope_precision not in ("bfloat16", "float32"):
            raise ValueError("rope_precision must be 'bfloat16' or 'float32'")
        if self.tie_word_embeddings:
            raise ValueError("these models do not tie the embedding and the readout")

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

    def ratio(self, tier: str) -> float:
        """Width ratio s of ``tier`` (its entry of the supported set)."""
        return float(self.widths[self.tiers.index(tier)])

    def width(self, tier: str) -> int:
        """Model width of ``tier``."""
        if self.variant == "matformer":
            return self.hidden_size
        return self.widths[self.tiers.index(tier)] * self.base_width

    def num_heads(self, tier: str) -> int:
        return self.width(tier) // self.head_dim

    def intermediate_size(self, tier: str) -> int:
        """FFN hidden width of ``tier``."""
        return self.widths[self.tiers.index(tier)] * self.mlp_base_width

    @property
    def max_intermediate_size(self) -> int:
        return self.widths[-1] * self.mlp_base_width

    def fan_in_ratio(self, matrix: str, tier: str) -> float:
        """Ratio r of Eq. 23 for ``matrix`` in ``tier``: the tier's fan-in over the base fan-in."""
        if self.variant == "fully_elastic" or matrix == "mlp.down_proj":
            return self.ratio(tier)
        return 1.0

    def num_logit_scales(self) -> int:
        """One learned readout scalar per tier (Fully elastic), or one shared scalar (MatFormer)."""
        return len(self.tiers) if self.variant == "fully_elastic" else 1

    def logit_scale_index(self, tier: str) -> int:
        return self.tiers.index(tier) if self.variant == "fully_elastic" else 0

    def subnetwork_config(self, tier: str) -> dict:
        """Configuration of the standalone dense model that ``ElasticForCausalLM.extract`` returns."""
        return OrderedDict(
            architectures=["DenseDecoder"], vocab_size=self.vocab_size, hidden_size=self.width(tier),
            intermediate_size=self.intermediate_size(tier), num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_heads(tier), head_dim=self.head_dim,
            max_position_embeddings=self.max_position_embeddings, rope_theta=self.rope_theta,
            rope_precision=self.rope_precision, rms_norm_eps=self.rms_norm_eps, hidden_act="silu",
            attention_output_gate="sigmoid", qk_norm="rmsnorm_per_head", tie_word_embeddings=False,
            tokenizer=self.tokenizer,
        )

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

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

    @classmethod
    def from_dict(cls, values: dict) -> "ElasticConfig":
        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) -> "ElasticConfig":
        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: Optional[int] = None) -> torch.Tensor:
    """RMSNorm over the last dimension (or within consecutive groups of ``group`` features).

    The statistics are computed in float32 and the result is cast back to the input dtype before the
    gain is applied.
    """
    h = x.float()
    if group is not None:
        lead = h.shape[:-1]
        h = h.reshape(*lead, -1, group)
        h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + eps)
        h = h.reshape(*lead, -1)
    else:
        h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + eps)
    return h.to(x.dtype) * weight


class RMSNorm(nn.Module):
    """Holds the full-width gain; a tier of width n uses its first n entries."""

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


class SharedMatrix(nn.Module):
    """A full-width ``[out_features, in_features]`` matrix whose tiers use its leading block.

    Stored either as one tensor ``weight``, or split (Eq. 23) as the frozen initialization ``init``
    and the trained ``delta``, which ``block`` combines with the multipliers of the tier.
    """

    def __init__(self, out_features: int, in_features: int, split: bool) -> None:
        super().__init__()
        self.split = split
        if split:
            self.init = nn.Parameter(torch.zeros(out_features, in_features))
            self.delta = nn.Parameter(torch.zeros(out_features, in_features))
        else:
            self.weight = nn.Parameter(torch.zeros(out_features, in_features))

    def block(self, rows: int, cols: int, ratio: float) -> torch.Tensor:
        """Weight of a tier: the leading ``[rows, cols]`` block, at fan-in ratio ``ratio``."""
        if self.split:
            return ratio ** -0.5 * self.init[:rows, :cols] + ratio ** -1.0 * self.delta[:rows, :cols]
        return self.weight[:rows, :cols]


class Readout(SharedMatrix):
    """Output projection with the learned readout scalars alpha (``logit_scale``)."""

    def __init__(self, vocab_size: int, width: int, split: bool, num_scales: int) -> None:
        super().__init__(vocab_size, width, split)
        self.logit_scale = nn.Parameter(torch.ones(num_scales))


def rotary_cos_sin(length: int, head_dim: int, theta: float, precision: str,
                   device: torch.device, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
    """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))  # 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, tokens, heads, head_dim]``: the two halves (a, b) of
    the last dimension are rotated to (-b, a)."""
    half = x.shape[-1] // 2
    rotated = torch.cat((-x[..., half:], x[..., :half]), dim=-1)
    return x * cos[None, :, None, :] + rotated * sin[None, :, None, :]


class Attention(nn.Module):
    """Causal self-attention with nested heads, per-head QK RMSNorm, RoPE and a sigmoid output gate."""

    def __init__(self, config: ElasticConfig) -> None:
        super().__init__()
        d, split = config.hidden_size, config.split_matrices
        self.config = config
        self.q_proj = SharedMatrix(d, d, "attn.q_proj" in split)
        self.k_proj = SharedMatrix(d, d, "attn.k_proj" in split)
        self.v_proj = SharedMatrix(d, d, "attn.v_proj" in split)
        self.gate_proj = SharedMatrix(d, d, "attn.gate_proj" in split)
        self.o_proj = SharedMatrix(d, d, "attn.o_proj" in split)
        self.q_norm = RMSNorm(d)
        self.k_norm = RMSNorm(d)

    def forward(self, x: torch.Tensor, tier: str, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        cfg = self.config
        batch, length, width = x.shape
        heads, head_dim = width // cfg.head_dim, cfg.head_dim
        shape = (batch, length, heads, head_dim)

        def weight(name: str) -> torch.Tensor:
            return getattr(self, name).block(width, width, cfg.fan_in_ratio(f"attn.{name}", tier))

        gate = F.linear(x, weight("gate_proj")).view(shape)
        # Queries, keys and values come from one matrix product with the three blocks stacked.
        qkv = torch.cat([weight("q_proj"), weight("k_proj"), weight("v_proj")], dim=0)
        q, k, v = F.linear(x, qkv).chunk(3, dim=-1)
        q = rms_norm(q, self.q_norm.weight[:width], cfg.rms_norm_eps, group=head_dim).view(shape)
        k = rms_norm(k, self.k_norm.weight[:width], cfg.rms_norm_eps, group=head_dim).view(shape)
        v = v.view(shape)
        q, k = apply_rotary(q, cos, sin), apply_rotary(k, cos, sin)
        out = F.scaled_dot_product_attention(
            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
            is_causal=length > 1, scale=head_dim ** -0.5,
        ).transpose(1, 2).contiguous()
        out = out * torch.sigmoid(gate)
        return F.linear(out.view(batch, length, width), weight("o_proj"))


class MLP(nn.Module):
    """SwiGLU FFN; tier s keeps the first s * mlp_base_width hidden units."""

    def __init__(self, config: ElasticConfig) -> None:
        super().__init__()
        d, h, split = config.hidden_size, config.max_intermediate_size, config.split_matrices
        self.config = config
        self.gate_proj = SharedMatrix(h, d, "mlp.gate_proj" in split)
        self.up_proj = SharedMatrix(h, d, "mlp.up_proj" in split)
        self.down_proj = SharedMatrix(d, h, "mlp.down_proj" in split)

    def forward(self, x: torch.Tensor, tier: str) -> torch.Tensor:
        cfg = self.config
        width, hidden = x.shape[-1], cfg.intermediate_size(tier)
        gate = F.linear(x, self.gate_proj.block(hidden, width, cfg.fan_in_ratio("mlp.gate_proj", tier)))
        up = F.linear(x, self.up_proj.block(hidden, width, cfg.fan_in_ratio("mlp.up_proj", tier)))
        down = self.down_proj.block(width, hidden, cfg.fan_in_ratio("mlp.down_proj", tier))
        return F.linear(F.silu(gate) * up, down) / cfg.mlp_output_divisor


class DecoderLayer(nn.Module):
    """Pre-norm transformer layer."""

    def __init__(self, config: ElasticConfig) -> None:
        super().__init__()
        self.eps = config.rms_norm_eps
        self.attn_norm = RMSNorm(config.hidden_size)
        self.attn = Attention(config)
        self.mlp_norm = RMSNorm(config.hidden_size)
        self.mlp = MLP(config)

    def forward(self, x: torch.Tensor, tier: str, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        width = x.shape[-1]
        h = x + self.attn(rms_norm(x, self.attn_norm.weight[:width], self.eps), tier, cos, sin)
        return h + self.mlp(rms_norm(h, self.mlp_norm.weight[:width], self.eps), tier)


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


class ElasticForCausalLM(nn.Module):
    """MatFormer or Fully elastic supernet: all tiers share one set of full-width weights."""

    def __init__(self, config: ElasticConfig) -> None:
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size)
        self.lm_head = Readout(config.vocab_size, config.hidden_size, "lm_head" in config.split_matrices,
                               config.num_logit_scales())

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

    @classmethod
    def from_pretrained(cls, path: str, device: Union[str, torch.device] = "cpu",
                        dtype: Optional[torch.dtype] = None) -> "ElasticForCausalLM":
        """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 = ElasticConfig.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``."""
        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 _tier(self, tier: Optional[str]) -> str:
        tier = tier or self.config.tiers[-1]
        if tier not in self.config.tiers:
            raise ValueError(f"unknown tier {tier!r}; this model has {self.config.tiers}")
        return tier

    def hidden_states(self, input_ids: torch.Tensor, tier: Optional[str] = None) -> torch.Tensor:
        """Output of the terminal norm of ``tier``: ``[batch, tokens, width of the tier]``."""
        cfg = self.config
        tier = self._tier(tier)
        length = input_ids.shape[1]
        if length > cfg.max_position_embeddings:
            raise ValueError(f"{length} tokens exceed the context of {cfg.max_position_embeddings}")
        width = cfg.width(tier)
        x = self.embed_tokens(input_ids)[..., :width]
        cos, sin = rotary_cos_sin(length, cfg.head_dim, cfg.rope_theta, cfg.rope_precision, x.device, x.dtype)
        for layer in self.layers:
            x = layer(x, tier, cos, sin)
        return rms_norm(x, self.norm.weight[:width], cfg.rms_norm_eps)

    def readout_weight(self, tier: Optional[str] = None) -> torch.Tensor:
        """Readout of ``tier``, ``[vocab_size, width]``, with its learned scalar alpha applied."""
        cfg = self.config
        tier = self._tier(tier)
        block = self.lm_head.block(cfg.vocab_size, cfg.width(tier), cfg.fan_in_ratio("lm_head", tier))
        return block * self.lm_head.logit_scale[cfg.logit_scale_index(tier)]

    def forward(self, input_ids: torch.Tensor, tier: Optional[str] = None,
                labels: Optional[torch.Tensor] = None):
        """Next-token logits of ``tier`` (default: the widest), as float32.

        Args:
            input_ids: ``[batch, tokens]`` token ids (no padding; causal attention over the sequence).
            tier: One of ``config.tiers``.
            labels: Optional ``[batch, tokens]`` targets; the mean cross-entropy is returned as well
                (label -100 is ignored).
        """
        logits = F.linear(self.hidden_states(input_ids, tier), self.readout_weight(tier)).float()
        if labels is None:
            return logits
        loss = F.cross_entropy(logits.flatten(0, 1), labels.flatten(), ignore_index=-100)
        return logits, loss

    @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]
            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 --------------------------------------------------------------------

    @torch.no_grad()
    def extract(self, tier: str) -> "OrderedDict[str, torch.Tensor]":
        """Weights of ``tier`` as a standalone dense model, in the format of the released dense models
        (Standalone / Sampled standalone): every split matrix combined with the multipliers of the
        tier, the FFN divisor e folded into ``down_proj``, alpha folded into ``lm_head``, and the
        embedding and norm gains cut to the width of the tier. The weights are computed in the dtype of
        this model. ``config.subnetwork_config(tier)`` describes the result."""
        cfg = self.config
        tier = self._tier(tier)
        width, hidden = cfg.width(tier), cfg.intermediate_size(tier)
        out: "OrderedDict[str, torch.Tensor]" = OrderedDict()
        out["embed_tokens.weight"] = self.embed_tokens.weight[:, :width]
        for i, layer in enumerate(self.layers):
            p = f"layers.{i}."
            attn, mlp = layer.attn, layer.mlp
            out[p + "attn_norm.weight"] = layer.attn_norm.weight[:width]
            for name in ("q_proj", "k_proj", "v_proj", "gate_proj", "o_proj"):
                out[p + f"attn.{name}.weight"] = getattr(attn, name).block(
                    width, width, cfg.fan_in_ratio(f"attn.{name}", tier))
            out[p + "attn.q_norm.weight"] = attn.q_norm.weight[:width]
            out[p + "attn.k_norm.weight"] = attn.k_norm.weight[:width]
            out[p + "mlp_norm.weight"] = layer.mlp_norm.weight[:width]
            out[p + "mlp.gate_proj.weight"] = mlp.gate_proj.block(hidden, width, cfg.fan_in_ratio("mlp.gate_proj", tier))
            out[p + "mlp.up_proj.weight"] = mlp.up_proj.block(hidden, width, cfg.fan_in_ratio("mlp.up_proj", tier))
            out[p + "mlp.down_proj.weight"] = mlp.down_proj.block(
                width, hidden, cfg.fan_in_ratio("mlp.down_proj", tier)) / cfg.mlp_output_divisor
        out["norm.weight"] = self.norm.weight[:width]
        out["lm_head.weight"] = self.readout_weight(tier)
        return OrderedDict((k, v.detach().clone().contiguous()) for k, v in out.items())

    def save_subnetwork(self, tier: str, path: str) -> None:
        """Write ``extract(tier)`` as ``model.safetensors`` with its ``config.json`` to ``path``."""
        from safetensors.torch import save_file

        os.makedirs(path, exist_ok=True)
        state = self.extract(tier)
        values = self.config.subnetwork_config(tier)
        values["torch_dtype"] = str(state["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")
        save_file(dict(state), os.path.join(path, "model.safetensors"), metadata={"format": "pt"})

    def num_parameters(self, tier: Optional[str] = None) -> int:
        """Parameters stored for the whole supernet, or the parameters of the dense model of ``tier``."""
        if tier is None:
            return sum(t.numel() for t in self.state_dict().values())
        cfg = self.config
        tier = self._tier(tier)
        n, h, v = cfg.width(tier), cfg.intermediate_size(tier), cfg.vocab_size
        per_layer = 5 * n * n + 3 * n * h + 4 * n
        return 2 * v * n + cfg.num_hidden_layers * per_layer + n


__all__ = [
    "ElasticConfig",
    "ElasticForCausalLM",
    "SharedMatrix",
    "Readout",
    "rms_norm",
    "rotary_cos_sin",
    "DEFAULT_TIERS",
    "MATRICES",
]