Spaces:

hugging-apps
/

ci-proto

Running on Zero

ci-proto / continuousintelligence /ci_proto /_model.py

Upload folder using huggingface_hub

ba6c090 verified 2 days ago

21.1 kB

"""Continuous Intelligence's patch model and multivariate attention.
Includes model components adapted by Continuous Intelligence from TiRex-2:
https://github.com/NX-AI/tirex-2/tree/bb5427f4a97379a17756e67bc846c4f4b569eb48/src/tirex2/model
The gated feed-forward and RMS normalisation components also derive from xLSTM:
https://github.com/NX-AI/xlstm/tree/ab22eadbd245f293dd8dec38ed29963d73758a12/xlstm/xlstm_large
Copyright (c) NXAI GmbH.
Copyright 2024 NXAI GmbH and its affiliates 2024
Upstream portions are licensed under Apache-2.0; see NOTICE.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from functools import lru_cache
from types import SimpleNamespace
from typing import Any

import torch
from torch import Tensor, nn
from torch.nn import functional as F

from ._mamba import Mamba3, RMSNorm


class ResidualBlock(nn.Module):
    """Patch embedding and quantile projection with a linear residual."""

    def __init__(self, in_dim: int, h_dim: int, out_dim: int, act: nn.Module) -> None:
        super().__init__()
        self.hidden_layer = nn.Linear(in_dim, h_dim)
        self.output_layer = nn.Linear(h_dim, out_dim)
        self.residual_layer = nn.Linear(in_dim, out_dim)
        self.act = act

    def forward(self, x: Tensor) -> Tensor:
        return self.output_layer(self.act(self.hidden_layer(x))) + self.residual_layer(x)


class RMSLayerNorm(nn.Module):
    """Final norm preserving float32 normalisation for reduced-precision activations."""

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

    def forward(self, x: Tensor) -> Tensor:
        variance = x.float().pow(2).mean(-1, keepdim=True)
        return self.weight * (x * torch.rsqrt(variance + self.eps))


class GatedFFN(nn.Module):
    """The time mixer's gated feed-forward projection."""

    def __init__(self, dim: int, constants: dict[str, Any]) -> None:
        super().__init__()
        multiple = constants.get("ffn_round_up_to_multiple_of", 64)
        factor = constants.get("ffn_proj_factor", 2.6667)
        if type(multiple) is not int or multiple < 1:
            raise ValueError("ffn_round_up_to_multiple_of must be a positive integer")
        if not math.isfinite(factor) or factor <= 0:
            raise ValueError("ffn_proj_factor must be positive and finite")
        width = math.ceil(dim * factor / multiple) * multiple
        bias = constants.get("use_bias", False)
        self.proj_up_gate = nn.Linear(dim, width, bias=bias)
        self.proj_up = nn.Linear(dim, width, bias=bias)
        self.proj_down = nn.Linear(width, dim, bias=bias)

    def forward(self, x: Tensor) -> Tensor:
        return self.proj_down(F.silu(self.proj_up_gate(x)) * self.proj_up(x))


class MambaCell(nn.Module):
    """Keep the checkpoint's cell prefix around the selected Mamba recurrence."""

    def __init__(self, dim: int, cfg: dict[str, Any]) -> None:
        super().__init__()
        expand, heads = cfg.get("mamba_expand", 2), cfg.get("num_heads", 4)
        if any(type(value) is not int or value < 1 for value in (expand, heads)):
            raise ValueError("mamba_expand and time num_heads must be positive integers")
        if expand * dim % heads:
            raise ValueError("expanded embedding dimension must be divisible by time heads")
        if cfg.get("model_type") != "bi-mamba3" or cfg.get("mamba_conv_width") is not None:
            raise ValueError("checkpoint must use the convolution-free Mamba-3 cell")
        rank = cfg.get("mamba_rank")
        if rank is not None and (type(rank) is not int or rank < 1):
            raise ValueError("mamba_rank must be a positive integer")
        self.mamba = Mamba3(
            d_model=dim,
            d_state=cfg.get("mamba_state_dim", 128),
            expand=expand,
            headdim=expand * dim // heads,
            is_mimo=rank is not None,
            mimo_rank=rank if rank is not None else 1,
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.mamba(x)


class TimeMixer(nn.Module):
    """Causal target mixing and tied bidirectional mixing for known covariates."""

    def __init__(self, dim: int, cfg: dict[str, Any], constants: dict[str, Any]) -> None:
        super().__init__()
        norm = dict(
            eps=constants.get("norm_eps", 1e-6),
            force_float32_reductions=constants.get("norm_reduction_force_float32", True),
            use_bias=constants.get("use_bias", False),
        )
        self.norm_lstm = RMSNorm(dim, **norm)
        self.fwd_cell = MambaCell(dim, cfg)
        self.rev_cell = self.fwd_cell
        self.combination = nn.Linear(2 * dim, dim, bias=False)
        self.norm_ffn = RMSNorm(dim, **norm)
        self.ffn = GatedFFN(dim, constants)

    def forward(
        self, x: Tensor, future_mask: Tensor, reverse_indices: Tensor | None = None
    ) -> Tensor:
        normalised = self.norm_lstm(x)
        indices = torch.where(future_mask)[0] if reverse_indices is None else reverse_indices
        if indices.numel() and not self.training:
            joined = torch.cat([normalised, normalised[indices].flip(1)])
            mixed = self.fwd_cell(joined)
            out, reverse = mixed[: len(x)], mixed[len(x) :].flip(1)
        else:
            out = self.fwd_cell(normalised)
            if indices.numel():
                reverse = self.rev_cell(normalised[indices].flip(1).contiguous()).flip(1)
        if indices.numel():
            mixed = self.combination(torch.cat([out[indices], reverse], dim=-1))
            out = out.clone()
            out[indices] = mixed.to(out.dtype)
        out = x + out
        return out + self.ffn(self.norm_ffn(out))


def _pack_groups(groups: Tensor, targets: Tensor) -> SimpleNamespace | None:
    groups = groups.squeeze(-1) if groups.ndim == 2 else groups
    if groups.shape != targets.shape:
        raise ValueError("group_vector and target_mask must have the same shape")
    _, inverse, counts = torch.unique(groups, return_inverse=True, return_counts=True)
    keep = counts[inverse] > 1
    if not bool(keep.any()):
        return None
    _, group_index, sizes = torch.unique(inverse[keep], return_inverse=True, return_counts=True)
    count, n_groups, width = group_index.numel(), sizes.numel(), int(sizes.max())
    order = torch.argsort(group_index, stable=True)
    first = torch.cumsum(sizes, dim=0) - sizes
    slot = torch.empty_like(group_index)
    slot[order] = torch.arange(count, device=groups.device) - first[group_index[order]]
    flat_idx = group_index * width + slot
    take = torch.full((n_groups * width,), count, device=groups.device)
    take[flat_idx] = torch.arange(count, device=groups.device)
    valid = torch.zeros(n_groups, width, dtype=torch.bool, device=groups.device)
    valid[group_index, slot] = True
    target = torch.zeros_like(valid)
    target[group_index, slot] = targets[keep]
    allowed = valid[:, :, None] & valid[:, None, :]
    allowed &= ~(~target[:, :, None] & target[:, None, :])
    allowed |= ~valid[:, :, None] & torch.eye(width, dtype=torch.bool, device=groups.device)
    return SimpleNamespace(
        regular=False,
        n_groups=n_groups,
        keep=keep,
        take=take,
        flat_idx=flat_idx,
        width=width,
        bias=torch.where(allowed, 0.0, -torch.inf)[:, None],
    )


@dataclass(frozen=True)
class ForecastLayout:
    """Cached row roles and attention geometry; contained tensors are read-only."""

    groups: Tensor
    targets: Tensor
    future: Tensor
    past: Tensor
    reverse_indices: Tensor
    packing: SimpleNamespace | None


@lru_cache(maxsize=32)
def forecast_layout(
    samples: int, channels: int, targets: int, known: int, device: torch.device
) -> ForecastLayout:
    """Reuse geometry for independent samples with equal target and covariate counts."""
    if samples < 1 or not 0 < targets <= channels or not 0 <= known <= channels - targets:
        raise ValueError("forecast layout needs samples and valid channel roles")
    channel = torch.arange(channels, device=device)
    target = channel < targets
    future = (channel >= targets) & (channel < targets + known)
    allowed = ~(~target[:, None] & target[None, :])
    packing = (
        None
        if channels == 1
        else SimpleNamespace(
            regular=True,
            n_groups=samples,
            width=channels,
            bias=torch.where(allowed, 0.0, -torch.inf)[None, None],
        )
    )
    reverse = (
        torch.arange(samples, device=device)[:, None] * channels
        + torch.arange(targets, targets + known, device=device)[None, :]
    ).flatten()
    return ForecastLayout(
        torch.arange(samples, device=device).repeat_interleave(channels),
        target.repeat(samples),
        future.repeat(samples),
        (~(target | future)).repeat(samples),
        reverse,
        packing,
    )


class GroupedAttention(nn.Module):
    """Attention within each multivariate sample, with asymmetric covariate masking."""

    def __init__(self, dim: int, cfg: dict[str, Any]) -> None:
        super().__init__()
        self.heads = cfg.get("num_heads", 4)
        if type(self.heads) is not int or self.heads < 1:
            raise ValueError("variate num_heads must be a positive integer")
        if dim % self.heads:
            raise ValueError("embedding dimension must be divisible by variate heads")
        self.head_dim = dim // self.heads
        self.use_qk_norm = cfg.get("use_qk_norm", True)
        self.WK = nn.Linear(dim, dim, bias=False)
        self.WQ = nn.Linear(dim, dim, bias=False)
        self.WV = nn.Linear(dim, dim, bias=False)
        self.WO = nn.Linear(dim, dim, bias=False)

    def forward(self, x: Tensor, packing: SimpleNamespace) -> Tensor:
        length, rows, _ = x.shape
        heads, hd, width = self.heads, self.head_dim, packing.width
        n_groups = packing.n_groups
        q = self.WQ(x).view(length, rows, heads, hd)
        k = self.WK(x).view(length, rows, heads, hd)
        v = self.WV(x).view(length, rows, heads, hd)
        if self.use_qk_norm:
            q = F.rms_norm(q, (hd,)).to(v.dtype)
            k = F.rms_norm(k, (hd,)).to(v.dtype)
        qkv = torch.cat((q, k, v), dim=2).flatten(2)
        if not packing.regular:
            qkv = F.pad(qkv, (0, 0, 0, 1)).index_select(1, packing.take)
        packed = qkv.view(length, n_groups, width, 3, heads, hd)
        q = packed[:, :, :, 0].permute(0, 1, 3, 2, 4)
        k = packed[:, :, :, 1].permute(0, 1, 3, 2, 4)
        v = packed[:, :, :, 2].permute(0, 1, 3, 2, 4)
        scores = q @ k.transpose(-1, -2) * (1 / math.sqrt(hd)) + packing.bias.to(q.dtype)
        out = (torch.softmax(scores, dim=-1) @ v).permute(0, 1, 3, 2, 4)
        out = out.reshape(length, n_groups * width, heads * hd)
        return self.WO(out if packing.regular else out.index_select(1, packing.flat_idx))


class MLP(nn.Module):
    """The variate mixer's feed-forward projection."""

    def __init__(self, dim: int, act: nn.Module) -> None:
        super().__init__()
        self.wi = nn.Linear(dim, 4 * dim, bias=False)
        self.wo = nn.Linear(4 * dim, dim, bias=False)
        self.act = act

    def forward(self, x: Tensor) -> Tensor:
        return self.wo(self.act(self.wi(x)))


class VariateMixer(nn.Module):
    """Mix related rows; singleton groups pass through unchanged."""

    def __init__(self, dim: int, cfg: dict[str, Any], act: nn.Module) -> None:
        super().__init__()
        self.norm_attn = nn.RMSNorm(dim, eps=1e-6)
        self.attn = GroupedAttention(dim, cfg)
        self.norm_ffn = nn.RMSNorm(dim, eps=1e-6)
        self.ffn = MLP(dim, act)

    def forward(self, x: Tensor, packing: SimpleNamespace | None) -> Tensor:
        if packing is None:
            return x
        kept = x if packing.regular else x[:, packing.keep]
        kept = kept + self.attn(self.norm_attn(kept), packing)
        kept = kept + self.ffn(self.norm_ffn(kept))
        if packing.regular:
            return kept
        out = x.clone()
        out[:, packing.keep] = kept
        return out


class Block(nn.Module):
    """The time mixer followed by its variate mixer."""

    def __init__(
        self,
        dim: int,
        cfg: dict[str, Any],
        constants: dict[str, Any],
        act: nn.Module,
    ) -> None:
        super().__init__()
        if cfg.get("mixing_order", "time_first") != "time_first":
            raise ValueError("checkpoint must use time_first mixing")
        self.time_mixer = TimeMixer(dim, cfg["time_mixer"], constants)
        self.variate_mixer = VariateMixer(dim, cfg["variate_mixer"], act)

    def forward(
        self,
        x: Tensor,
        packing: SimpleNamespace | None,
        future_mask: Tensor,
        reverse_indices: Tensor,
    ) -> Tensor:
        x = self.time_mixer(x, future_mask, reverse_indices)
        return self.variate_mixer(x.transpose(0, 1), packing).transpose(0, 1)


class Model(nn.Module):
    """Inference-only model; dimensions come from the checkpoint configuration.
    ``forward`` maps flattened ``[variables, time]`` rows to
    ``[variables, quantiles, padded_time]``. Each output patch predicts the next input patch.
    """

    def __init__(self, raw: dict[str, Any]) -> None:
        super().__init__()
        self.config = SimpleNamespace(**raw)
        cfg = self.config
        if raw.get("input_embedding") is not None or raw.get("quantile_mlp", False):
            raise ValueError("checkpoint must use the patch embeddings")
        if any(raw["scaler_cfg"].get(key, False) for key in ("causal_patch", "known_whole_row")):
            raise ValueError("checkpoint must use whole-row observed-value scaling")
        for name in (
            "embedding_dim",
            "input_patch_size",
            "output_patch_size",
            "context_len",
            "future_len",
            "num_blocks",
            "input_ff_dim",
        ):
            if name == "context_len" and raw[name] is None:
                continue
            if type(raw[name]) is not int or raw[name] < 1:
                raise ValueError(f"{name} must be a positive integer")
        if not math.isfinite(cfg.h_expand) or cfg.h_expand <= 0:
            raise ValueError("h_expand must be positive and finite")
        if int(cfg.embedding_dim * cfg.h_expand) < 1:
            raise ValueError("h_expand must produce at least one hidden feature")
        if cfg.input_patch_size != cfg.output_patch_size:
            raise ValueError("input and output patch sizes must match")
        if (
            not isinstance(cfg.quantiles, (list, tuple))
            or not cfg.quantiles
            or any(not 0 < q < 1 for q in cfg.quantiles)
        ):
            raise ValueError("quantiles must be levels in (0, 1)")
        if any(b <= a for a, b in zip(cfg.quantiles, cfg.quantiles[1:], strict=False)):
            raise ValueError("quantiles must be strictly increasing")
        recipe = cfg.stack_config["recipe"]
        if len(recipe) != cfg.num_blocks:
            raise ValueError("recipe length must equal num_blocks")
        activation = raw.get("act_func", "SiLU").removeprefix("torch.nn.")
        if activation not in ("SiLU", "GELU", "ReLU"):
            raise ValueError(f"unsupported activation {activation!r}")
        act = getattr(nn, activation)()
        dim = cfg.embedding_dim
        self.input_patch_embedding = ResidualBlock(
            cfg.input_patch_size * 2, cfg.input_ff_dim, dim, act
        )
        self.stack = nn.ModuleList(
            Block(dim, cfg.stack_config["templates"][name], raw.get("xlstm", {}), act)
            for name in recipe
        )
        norm = raw.get("stack_out_norm_config")
        self.stack_out_norm = (
            nn.Identity()
            if norm is None
            else nn.RMSNorm(dim, eps=norm.get("eps", 1e-7))
            if norm.get("use_rmsnorm", False)
            else RMSLayerNorm(dim, norm.get("eps", 1e-7))
        )
        self.register_buffer("quantiles", torch.tensor(cfg.quantiles), persistent=False)
        self.output_patch_embedding = ResidualBlock(
            dim,
            int(dim * cfg.h_expand),
            len(cfg.quantiles) * cfg.output_patch_size,
            act,
        )

    def forward(
        self,
        batch: dict[str, Tensor],
        *,
        future_len: int | None = None,
        horizon: int | None = None,
        layout: ForecastLayout | None = None,
    ) -> Tensor:
        """Run inference, projecting only forecast tokens when ``horizon`` is supplied."""
        x = batch["x"]
        future_len = self.config.future_len if future_len is None else future_len
        groups, targets = batch.get("group_vector"), batch.get("target_mask")
        future, past = batch.get("future_covariate_mask"), batch.get("past_covariate_mask")
        if groups is None:
            if targets is not None:
                raise ValueError("target_mask requires group_vector")
            groups = torch.arange(x.shape[0], device=x.device)
        if targets is None:
            if future is not None or past is not None:
                raise ValueError("covariate masks require target_mask")
            targets = torch.ones(x.shape[0], dtype=torch.bool, device=x.device)
        declared = future is not None or past is not None
        future = torch.zeros_like(targets) if future is None else future
        past = ~(targets | future) if past is None else past
        if layout is not None:
            if (
                groups is not layout.groups
                or targets is not layout.targets
                or future is not layout.future
                or past is not layout.past
                or groups.shape != (x.shape[0],)
                or groups.device != x.device
            ):
                raise ValueError("forecast layout must match the batch masks and rows")
        else:
            if bool((future & targets).any()) or bool((past & (targets | future)).any()):
                raise ValueError("target and covariate masks must not overlap")
            if not bool((targets | future | past).all()):
                raise ValueError("target and covariate masks must cover every row")
        if horizon is not None and (type(horizon) is not int or not 0 < horizon <= future_len):
            raise ValueError("horizon must be a positive integer within the future window")
        if declared:
            if x.shape[-1] < future_len:
                raise ValueError("inputs with covariate masks must include the full future window")
            window = torch.zeros_like(x, dtype=torch.bool)
            window[:, -future_len:] = (targets | past)[:, None]
            x = torch.where(window, torch.nan, x)

        scaler = self.config.scaler_cfg
        wide = x.double()
        loc = torch.nan_to_num(torch.nanmean(wide, dim=-1, keepdim=True), nan=0.0)
        scale = torch.nan_to_num(
            torch.nanmean((wide - loc) ** 2, dim=-1, keepdim=True).sqrt(), nan=1.0
        )
        scale = scale.clamp_min(scaler.get("eps", 1e-8))
        binary = torch.zeros_like(loc, dtype=torch.bool)
        if scaler.get("binaryaware", False):
            binary = ((x == 0) | (x == 1) | torch.isnan(x)).all(dim=-1, keepdim=True)
            loc = torch.where(binary, 0.0, loc)
            scale = torch.where(binary, 1.0, scale)
        values = ((wide - loc) / scale).to(x.dtype)
        loc, scale = loc.to(x.dtype), scale.to(x.dtype)
        if scaler.get("use_arcsinh", False):
            values = torch.where(binary, values, values.asinh())
        patch = self.config.input_patch_size
        values = F.pad(values, ((-values.shape[-1]) % patch, 0), value=torch.nan)
        values = values.unfold(-1, patch, patch)
        present = (~torch.isnan(values)).to(values.dtype)
        x = self.input_patch_embedding(
            torch.cat((torch.nan_to_num(values, nan=0.0), present), dim=-1)
        )
        packing = _pack_groups(groups, targets) if layout is None else layout.packing
        reverse_indices = torch.where(future)[0] if layout is None else layout.reverse_indices
        for block in self.stack:
            x = block(x, packing, future, reverse_indices)
        if horizon is not None:
            start = x.shape[1] * patch - future_len - patch
            if start < 0:
                raise ValueError("forecast projection needs a context patch before the future")
            offset = start % patch
            x = x[:, start // patch : -(-(start + horizon) // patch)]
        x = self.output_patch_embedding(self.stack_out_norm(x))
        x = x.unflatten(-1, (len(self.config.quantiles), self.config.output_patch_size))
        x = x.transpose(1, 2).flatten(2)
        if horizon is not None:
            x = x[..., offset : offset + horizon]
        if scaler.get("use_arcsinh", False):
            x = torch.where(binary[:, None], x, torch.sinh(x.clamp(-20.0, 20.0)))
        return x * scale[:, None] + loc[:, None]

	"""Continuous Intelligence's patch model and multivariate attention.

	Includes model components adapted by Continuous Intelligence from TiRex-2:
	https://github.com/NX-AI/tirex-2/tree/bb5427f4a97379a17756e67bc846c4f4b569eb48/src/tirex2/model

	The gated feed-forward and RMS normalisation components also derive from xLSTM:
	https://github.com/NX-AI/xlstm/tree/ab22eadbd245f293dd8dec38ed29963d73758a12/xlstm/xlstm_large

	Copyright (c) NXAI GmbH.
	Copyright 2024 NXAI GmbH and its affiliates 2024
	Upstream portions are licensed under Apache-2.0; see NOTICE.
	"""

	from __future__ import annotations

	import math
	from dataclasses import dataclass
	from functools import lru_cache
	from types import SimpleNamespace
	from typing import Any

	import torch
	from torch import Tensor, nn
	from torch.nn import functional as F

	from ._mamba import Mamba3, RMSNorm


	class ResidualBlock(nn.Module):
	"""Patch embedding and quantile projection with a linear residual."""

	def __init__(self, in_dim: int, h_dim: int, out_dim: int, act: nn.Module) -> None:
	super().__init__()
	self.hidden_layer = nn.Linear(in_dim, h_dim)
	self.output_layer = nn.Linear(h_dim, out_dim)
	self.residual_layer = nn.Linear(in_dim, out_dim)
	self.act = act

	def forward(self, x: Tensor) -> Tensor:
	return self.output_layer(self.act(self.hidden_layer(x))) + self.residual_layer(x)


	class RMSLayerNorm(nn.Module):
	"""Final norm preserving float32 normalisation for reduced-precision activations."""

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

	def forward(self, x: Tensor) -> Tensor:
	variance = x.float().pow(2).mean(-1, keepdim=True)
	return self.weight * (x * torch.rsqrt(variance + self.eps))


	class GatedFFN(nn.Module):
	"""The time mixer's gated feed-forward projection."""

	def __init__(self, dim: int, constants: dict[str, Any]) -> None:
	super().__init__()
	multiple = constants.get("ffn_round_up_to_multiple_of", 64)
	factor = constants.get("ffn_proj_factor", 2.6667)
	if type(multiple) is not int or multiple < 1:
	raise ValueError("ffn_round_up_to_multiple_of must be a positive integer")
	if not math.isfinite(factor) or factor <= 0:
	raise ValueError("ffn_proj_factor must be positive and finite")
	width = math.ceil(dim * factor / multiple) * multiple
	bias = constants.get("use_bias", False)
	self.proj_up_gate = nn.Linear(dim, width, bias=bias)
	self.proj_up = nn.Linear(dim, width, bias=bias)
	self.proj_down = nn.Linear(width, dim, bias=bias)

	def forward(self, x: Tensor) -> Tensor:
	return self.proj_down(F.silu(self.proj_up_gate(x)) * self.proj_up(x))


	class MambaCell(nn.Module):
	"""Keep the checkpoint's cell prefix around the selected Mamba recurrence."""

	def __init__(self, dim: int, cfg: dict[str, Any]) -> None:
	super().__init__()
	expand, heads = cfg.get("mamba_expand", 2), cfg.get("num_heads", 4)
	if any(type(value) is not int or value < 1 for value in (expand, heads)):
	raise ValueError("mamba_expand and time num_heads must be positive integers")
	if expand * dim % heads:
	raise ValueError("expanded embedding dimension must be divisible by time heads")
	if cfg.get("model_type") != "bi-mamba3" or cfg.get("mamba_conv_width") is not None:
	raise ValueError("checkpoint must use the convolution-free Mamba-3 cell")
	rank = cfg.get("mamba_rank")
	if rank is not None and (type(rank) is not int or rank < 1):
	raise ValueError("mamba_rank must be a positive integer")
	self.mamba = Mamba3(
	d_model=dim,
	d_state=cfg.get("mamba_state_dim", 128),
	expand=expand,
	headdim=expand * dim // heads,
	is_mimo=rank is not None,
	mimo_rank=rank if rank is not None else 1,
	)

	def forward(self, x: Tensor) -> Tensor:
	return self.mamba(x)


	class TimeMixer(nn.Module):
	"""Causal target mixing and tied bidirectional mixing for known covariates."""

	def __init__(self, dim: int, cfg: dict[str, Any], constants: dict[str, Any]) -> None:
	super().__init__()
	norm = dict(
	eps=constants.get("norm_eps", 1e-6),
	force_float32_reductions=constants.get("norm_reduction_force_float32", True),
	use_bias=constants.get("use_bias", False),
	)
	self.norm_lstm = RMSNorm(dim, **norm)
	self.fwd_cell = MambaCell(dim, cfg)
	self.rev_cell = self.fwd_cell
	self.combination = nn.Linear(2 * dim, dim, bias=False)
	self.norm_ffn = RMSNorm(dim, **norm)
	self.ffn = GatedFFN(dim, constants)

	def forward(
	self, x: Tensor, future_mask: Tensor, reverse_indices: Tensor \| None = None
	) -> Tensor:
	normalised = self.norm_lstm(x)
	indices = torch.where(future_mask)[0] if reverse_indices is None else reverse_indices
	if indices.numel() and not self.training:
	joined = torch.cat([normalised, normalised[indices].flip(1)])
	mixed = self.fwd_cell(joined)
	out, reverse = mixed[: len(x)], mixed[len(x) :].flip(1)
	else:
	out = self.fwd_cell(normalised)
	if indices.numel():
	reverse = self.rev_cell(normalised[indices].flip(1).contiguous()).flip(1)
	if indices.numel():
	mixed = self.combination(torch.cat([out[indices], reverse], dim=-1))
	out = out.clone()
	out[indices] = mixed.to(out.dtype)
	out = x + out
	return out + self.ffn(self.norm_ffn(out))


	def _pack_groups(groups: Tensor, targets: Tensor) -> SimpleNamespace \| None:
	groups = groups.squeeze(-1) if groups.ndim == 2 else groups
	if groups.shape != targets.shape:
	raise ValueError("group_vector and target_mask must have the same shape")
	_, inverse, counts = torch.unique(groups, return_inverse=True, return_counts=True)
	keep = counts[inverse] > 1
	if not bool(keep.any()):
	return None
	_, group_index, sizes = torch.unique(inverse[keep], return_inverse=True, return_counts=True)
	count, n_groups, width = group_index.numel(), sizes.numel(), int(sizes.max())
	order = torch.argsort(group_index, stable=True)
	first = torch.cumsum(sizes, dim=0) - sizes
	slot = torch.empty_like(group_index)
	slot[order] = torch.arange(count, device=groups.device) - first[group_index[order]]
	flat_idx = group_index * width + slot
	take = torch.full((n_groups * width,), count, device=groups.device)
	take[flat_idx] = torch.arange(count, device=groups.device)
	valid = torch.zeros(n_groups, width, dtype=torch.bool, device=groups.device)
	valid[group_index, slot] = True
	target = torch.zeros_like(valid)
	target[group_index, slot] = targets[keep]
	allowed = valid[:, :, None] & valid[:, None, :]
	allowed &= ~(~target[:, :, None] & target[:, None, :])
	allowed \|= ~valid[:, :, None] & torch.eye(width, dtype=torch.bool, device=groups.device)
	return SimpleNamespace(
	regular=False,
	n_groups=n_groups,
	keep=keep,
	take=take,
	flat_idx=flat_idx,
	width=width,
	bias=torch.where(allowed, 0.0, -torch.inf)[:, None],
	)


	@dataclass(frozen=True)
	class ForecastLayout:
	"""Cached row roles and attention geometry; contained tensors are read-only."""

	groups: Tensor
	targets: Tensor
	future: Tensor
	past: Tensor
	reverse_indices: Tensor
	packing: SimpleNamespace \| None


	@lru_cache(maxsize=32)
	def forecast_layout(
	samples: int, channels: int, targets: int, known: int, device: torch.device
	) -> ForecastLayout:
	"""Reuse geometry for independent samples with equal target and covariate counts."""
	if samples < 1 or not 0 < targets <= channels or not 0 <= known <= channels - targets:
	raise ValueError("forecast layout needs samples and valid channel roles")
	channel = torch.arange(channels, device=device)
	target = channel < targets
	future = (channel >= targets) & (channel < targets + known)
	allowed = ~(~target[:, None] & target[None, :])
	packing = (
	None
	if channels == 1
	else SimpleNamespace(
	regular=True,
	n_groups=samples,
	width=channels,
	bias=torch.where(allowed, 0.0, -torch.inf)[None, None],
	)
	)
	reverse = (
	torch.arange(samples, device=device)[:, None] * channels
	+ torch.arange(targets, targets + known, device=device)[None, :]
	).flatten()
	return ForecastLayout(
	torch.arange(samples, device=device).repeat_interleave(channels),
	target.repeat(samples),
	future.repeat(samples),
	(~(target \| future)).repeat(samples),
	reverse,
	packing,
	)


	class GroupedAttention(nn.Module):
	"""Attention within each multivariate sample, with asymmetric covariate masking."""

	def __init__(self, dim: int, cfg: dict[str, Any]) -> None:
	super().__init__()
	self.heads = cfg.get("num_heads", 4)
	if type(self.heads) is not int or self.heads < 1:
	raise ValueError("variate num_heads must be a positive integer")
	if dim % self.heads:
	raise ValueError("embedding dimension must be divisible by variate heads")
	self.head_dim = dim // self.heads
	self.use_qk_norm = cfg.get("use_qk_norm", True)
	self.WK = nn.Linear(dim, dim, bias=False)
	self.WQ = nn.Linear(dim, dim, bias=False)
	self.WV = nn.Linear(dim, dim, bias=False)
	self.WO = nn.Linear(dim, dim, bias=False)

	def forward(self, x: Tensor, packing: SimpleNamespace) -> Tensor:
	length, rows, _ = x.shape
	heads, hd, width = self.heads, self.head_dim, packing.width
	n_groups = packing.n_groups
	q = self.WQ(x).view(length, rows, heads, hd)
	k = self.WK(x).view(length, rows, heads, hd)
	v = self.WV(x).view(length, rows, heads, hd)
	if self.use_qk_norm:
	q = F.rms_norm(q, (hd,)).to(v.dtype)
	k = F.rms_norm(k, (hd,)).to(v.dtype)
	qkv = torch.cat((q, k, v), dim=2).flatten(2)
	if not packing.regular:
	qkv = F.pad(qkv, (0, 0, 0, 1)).index_select(1, packing.take)
	packed = qkv.view(length, n_groups, width, 3, heads, hd)
	q = packed[:, :, :, 0].permute(0, 1, 3, 2, 4)
	k = packed[:, :, :, 1].permute(0, 1, 3, 2, 4)
	v = packed[:, :, :, 2].permute(0, 1, 3, 2, 4)
	scores = q @ k.transpose(-1, -2) * (1 / math.sqrt(hd)) + packing.bias.to(q.dtype)
	out = (torch.softmax(scores, dim=-1) @ v).permute(0, 1, 3, 2, 4)
	out = out.reshape(length, n_groups * width, heads * hd)
	return self.WO(out if packing.regular else out.index_select(1, packing.flat_idx))


	class MLP(nn.Module):
	"""The variate mixer's feed-forward projection."""

	def __init__(self, dim: int, act: nn.Module) -> None:
	super().__init__()
	self.wi = nn.Linear(dim, 4 * dim, bias=False)
	self.wo = nn.Linear(4 * dim, dim, bias=False)
	self.act = act

	def forward(self, x: Tensor) -> Tensor:
	return self.wo(self.act(self.wi(x)))


	class VariateMixer(nn.Module):
	"""Mix related rows; singleton groups pass through unchanged."""

	def __init__(self, dim: int, cfg: dict[str, Any], act: nn.Module) -> None:
	super().__init__()
	self.norm_attn = nn.RMSNorm(dim, eps=1e-6)
	self.attn = GroupedAttention(dim, cfg)
	self.norm_ffn = nn.RMSNorm(dim, eps=1e-6)
	self.ffn = MLP(dim, act)

	def forward(self, x: Tensor, packing: SimpleNamespace \| None) -> Tensor:
	if packing is None:
	return x
	kept = x if packing.regular else x[:, packing.keep]
	kept = kept + self.attn(self.norm_attn(kept), packing)
	kept = kept + self.ffn(self.norm_ffn(kept))
	if packing.regular:
	return kept
	out = x.clone()
	out[:, packing.keep] = kept
	return out


	class Block(nn.Module):
	"""The time mixer followed by its variate mixer."""

	def __init__(
	self,
	dim: int,
	cfg: dict[str, Any],
	constants: dict[str, Any],
	act: nn.Module,
	) -> None:
	super().__init__()
	if cfg.get("mixing_order", "time_first") != "time_first":
	raise ValueError("checkpoint must use time_first mixing")
	self.time_mixer = TimeMixer(dim, cfg["time_mixer"], constants)
	self.variate_mixer = VariateMixer(dim, cfg["variate_mixer"], act)

	def forward(
	self,
	x: Tensor,
	packing: SimpleNamespace \| None,
	future_mask: Tensor,
	reverse_indices: Tensor,
	) -> Tensor:
	x = self.time_mixer(x, future_mask, reverse_indices)
	return self.variate_mixer(x.transpose(0, 1), packing).transpose(0, 1)


	class Model(nn.Module):
	"""Inference-only model; dimensions come from the checkpoint configuration.

	``forward`` maps flattened ``[variables, time]`` rows to
	``[variables, quantiles, padded_time]``. Each output patch predicts the next input patch.
	"""

	def __init__(self, raw: dict[str, Any]) -> None:
	super().__init__()
	self.config = SimpleNamespace(**raw)
	cfg = self.config
	if raw.get("input_embedding") is not None or raw.get("quantile_mlp", False):
	raise ValueError("checkpoint must use the patch embeddings")
	if any(raw["scaler_cfg"].get(key, False) for key in ("causal_patch", "known_whole_row")):
	raise ValueError("checkpoint must use whole-row observed-value scaling")
	for name in (
	"embedding_dim",
	"input_patch_size",
	"output_patch_size",
	"context_len",
	"future_len",
	"num_blocks",
	"input_ff_dim",
	):
	if name == "context_len" and raw[name] is None:
	continue
	if type(raw[name]) is not int or raw[name] < 1:
	raise ValueError(f"{name} must be a positive integer")
	if not math.isfinite(cfg.h_expand) or cfg.h_expand <= 0:
	raise ValueError("h_expand must be positive and finite")
	if int(cfg.embedding_dim * cfg.h_expand) < 1:
	raise ValueError("h_expand must produce at least one hidden feature")
	if cfg.input_patch_size != cfg.output_patch_size:
	raise ValueError("input and output patch sizes must match")
	if (
	not isinstance(cfg.quantiles, (list, tuple))
	or not cfg.quantiles
	or any(not 0 < q < 1 for q in cfg.quantiles)
	):
	raise ValueError("quantiles must be levels in (0, 1)")
	if any(b <= a for a, b in zip(cfg.quantiles, cfg.quantiles[1:], strict=False)):
	raise ValueError("quantiles must be strictly increasing")
	recipe = cfg.stack_config["recipe"]
	if len(recipe) != cfg.num_blocks:
	raise ValueError("recipe length must equal num_blocks")
	activation = raw.get("act_func", "SiLU").removeprefix("torch.nn.")
	if activation not in ("SiLU", "GELU", "ReLU"):
	raise ValueError(f"unsupported activation {activation!r}")
	act = getattr(nn, activation)()
	dim = cfg.embedding_dim
	self.input_patch_embedding = ResidualBlock(
	cfg.input_patch_size * 2, cfg.input_ff_dim, dim, act
	)
	self.stack = nn.ModuleList(
	Block(dim, cfg.stack_config["templates"][name], raw.get("xlstm", {}), act)
	for name in recipe
	)
	norm = raw.get("stack_out_norm_config")
	self.stack_out_norm = (
	nn.Identity()
	if norm is None
	else nn.RMSNorm(dim, eps=norm.get("eps", 1e-7))
	if norm.get("use_rmsnorm", False)
	else RMSLayerNorm(dim, norm.get("eps", 1e-7))
	)
	self.register_buffer("quantiles", torch.tensor(cfg.quantiles), persistent=False)
	self.output_patch_embedding = ResidualBlock(
	dim,
	int(dim * cfg.h_expand),
	len(cfg.quantiles) * cfg.output_patch_size,
	act,
	)

	def forward(
	self,
	batch: dict[str, Tensor],
	*,
	future_len: int \| None = None,
	horizon: int \| None = None,
	layout: ForecastLayout \| None = None,
	) -> Tensor:
	"""Run inference, projecting only forecast tokens when ``horizon`` is supplied."""
	x = batch["x"]
	future_len = self.config.future_len if future_len is None else future_len
	groups, targets = batch.get("group_vector"), batch.get("target_mask")
	future, past = batch.get("future_covariate_mask"), batch.get("past_covariate_mask")
	if groups is None:
	if targets is not None:
	raise ValueError("target_mask requires group_vector")
	groups = torch.arange(x.shape[0], device=x.device)
	if targets is None:
	if future is not None or past is not None:
	raise ValueError("covariate masks require target_mask")
	targets = torch.ones(x.shape[0], dtype=torch.bool, device=x.device)
	declared = future is not None or past is not None
	future = torch.zeros_like(targets) if future is None else future
	past = ~(targets \| future) if past is None else past
	if layout is not None:
	if (
	groups is not layout.groups
	or targets is not layout.targets
	or future is not layout.future
	or past is not layout.past
	or groups.shape != (x.shape[0],)
	or groups.device != x.device
	):
	raise ValueError("forecast layout must match the batch masks and rows")
	else:
	if bool((future & targets).any()) or bool((past & (targets \| future)).any()):
	raise ValueError("target and covariate masks must not overlap")
	if not bool((targets \| future \| past).all()):
	raise ValueError("target and covariate masks must cover every row")
	if horizon is not None and (type(horizon) is not int or not 0 < horizon <= future_len):
	raise ValueError("horizon must be a positive integer within the future window")
	if declared:
	if x.shape[-1] < future_len:
	raise ValueError("inputs with covariate masks must include the full future window")
	window = torch.zeros_like(x, dtype=torch.bool)
	window[:, -future_len:] = (targets \| past)[:, None]
	x = torch.where(window, torch.nan, x)

	scaler = self.config.scaler_cfg
	wide = x.double()
	loc = torch.nan_to_num(torch.nanmean(wide, dim=-1, keepdim=True), nan=0.0)
	scale = torch.nan_to_num(
	torch.nanmean((wide - loc) ** 2, dim=-1, keepdim=True).sqrt(), nan=1.0
	)
	scale = scale.clamp_min(scaler.get("eps", 1e-8))
	binary = torch.zeros_like(loc, dtype=torch.bool)
	if scaler.get("binaryaware", False):
	binary = ((x == 0) \| (x == 1) \| torch.isnan(x)).all(dim=-1, keepdim=True)
	loc = torch.where(binary, 0.0, loc)
	scale = torch.where(binary, 1.0, scale)
	values = ((wide - loc) / scale).to(x.dtype)
	loc, scale = loc.to(x.dtype), scale.to(x.dtype)
	if scaler.get("use_arcsinh", False):
	values = torch.where(binary, values, values.asinh())
	patch = self.config.input_patch_size
	values = F.pad(values, ((-values.shape[-1]) % patch, 0), value=torch.nan)
	values = values.unfold(-1, patch, patch)
	present = (~torch.isnan(values)).to(values.dtype)
	x = self.input_patch_embedding(
	torch.cat((torch.nan_to_num(values, nan=0.0), present), dim=-1)
	)
	packing = _pack_groups(groups, targets) if layout is None else layout.packing
	reverse_indices = torch.where(future)[0] if layout is None else layout.reverse_indices
	for block in self.stack:
	x = block(x, packing, future, reverse_indices)
	if horizon is not None:
	start = x.shape[1] * patch - future_len - patch
	if start < 0:
	raise ValueError("forecast projection needs a context patch before the future")
	offset = start % patch
	x = x[:, start // patch : -(-(start + horizon) // patch)]
	x = self.output_patch_embedding(self.stack_out_norm(x))
	x = x.unflatten(-1, (len(self.config.quantiles), self.config.output_patch_size))
	x = x.transpose(1, 2).flatten(2)
	if horizon is not None:
	x = x[..., offset : offset + horizon]
	if scaler.get("use_arcsinh", False):
	x = torch.where(binary[:, None], x, torch.sinh(x.clamp(-20.0, 20.0)))
	return x * scale[:, None] + loc[:, None]