tirex-mirror 2025.9.2__py3-none-any.whl → 2025.9.9__py3-none-any.whl

tirex/api_adapter/forecast.py

@@ -8,7 +8,6 @@ import torch
 
 from .standard_adapter import ContextType, get_batches
 
-
 DEF_TARGET_COLUMN = "target"
 DEF_META_COLUMNS = ("start", "item_id")
 

tirex/base.py

@@ -3,13 +3,18 @@
 
 import os
 from abc import ABC, abstractmethod
-from typing import TypeVar
+from typing import Literal, TypeVar
 
+import torch
 from huggingface_hub import hf_hub_download
 
 T = TypeVar("T", bound="PretrainedModel")
 
 
+def skip_cuda():
+    return os.getenv("TIREX_NO_CUDA", "False").lower() in ("true", "1", "t")
+
+
 def parse_hf_repo_id(path):
     parts = path.split("/")
     return "/".join(parts[0:2])
@@ -23,19 +28,30 @@ class PretrainedModel(ABC):
         cls.REGISTRY[cls.register_name()] = cls
 
     @classmethod
-    def from_pretrained(cls: type[T], path, device: str = "cuda:0", hf_kwargs=None, ckp_kwargs=None) -> T:
+    def from_pretrained(
+        cls: type[T], path: str, backend: str, device: str | None = None, hf_kwargs=None, ckp_kwargs=None
+    ) -> T:
         if hf_kwargs is None:
             hf_kwargs = {}
         if ckp_kwargs is None:
             ckp_kwargs = {}
+        if device is None:
+            device = "cuda:0" if backend == "cuda" else "cpu"
         if os.path.exists(path):
             print("Loading weights from local directory")
             checkpoint_path = path
         else:
             repo_id = parse_hf_repo_id(path)
             checkpoint_path = hf_hub_download(repo_id=repo_id, filename="model.ckpt", **hf_kwargs)
-        model = cls.load_from_checkpoint(checkpoint_path, map_location=device, **ckp_kwargs)
-        model.after_load_from_checkpoint()
+
+        # load lightning checkpoint
+        checkpoint = torch.load(checkpoint_path, map_location=device, **ckp_kwargs, weights_only=True)
+        model: T = cls(backend=backend, **checkpoint["hyper_parameters"])
+        model.on_load_checkpoint(checkpoint)
+        model.load_state_dict(checkpoint["state_dict"])
+
+        if backend == "cuda":
+            model = model.to(device)
         return model
 
     @classmethod
@@ -43,17 +59,22 @@ class PretrainedModel(ABC):
     def register_name(cls) -> str:
         pass
 
-    def after_load_from_checkpoint(self):
+    def on_load_checkpoint(self):
         pass
 
 
-def load_model(path: str, device: str = "cuda:0", hf_kwargs=None, ckp_kwargs=None) -> PretrainedModel:
+def load_model(
+    path: str,
+    device: str | None = None,
+    backend: Literal["torch", "cuda"] | None = None,
+    hf_kwargs=None,
+    ckp_kwargs=None,
+) -> PretrainedModel:
     """Loads a TiRex model. This function attempts to load the specified model.
 
     Args:
         path (str): Hugging Face path to the model (e.g. NX-AI/TiRex)
         device (str, optional): The device on which to load the model (e.g., "cuda:0", "cpu").
-                                If you want to use "cpu" you need to deactivate the sLSTM CUDA kernels (check repository FAQ!).
         hf_kwargs (dict, optional): Keyword arguments to pass to the Hugging Face Hub download method.
         ckp_kwargs (dict, optional): Keyword arguments to pass when loading the checkpoint.
 
@@ -63,6 +84,11 @@ def load_model(path: str, device: str = "cuda:0", hf_kwargs=None, ckp_kwargs=Non
     Examples:
         model: ForecastModel = load_model("NX-AI/TiRex")
     """
+
+    if backend is None:
+        backend = "torch" if skip_cuda() else "cuda"
+    assert backend in ["torch", "cuda"], f"Backend can either be torch or cuda, not {backend}!"
+
     try:
         _, model_id = parse_hf_repo_id(path).split("/")
     except:
@@ -70,4 +96,5 @@ def load_model(path: str, device: str = "cuda:0", hf_kwargs=None, ckp_kwargs=Non
     model_cls = PretrainedModel.REGISTRY.get(model_id, None)
     if model_cls is None:
         raise ValueError(f"Invalid model id {model_id}")
-    return model_cls.from_pretrained(path, device=device, hf_kwargs=hf_kwargs, ckp_kwargs=ckp_kwargs)
+
+    return model_cls.from_pretrained(path, device=device, backend=backend, hf_kwargs=hf_kwargs, ckp_kwargs=ckp_kwargs)

tirex/models/patcher.py

@@ -0,0 +1,84 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from dataclasses import dataclass
+
+import torch
+
+
+class StandardScaler:
+    def __init__(self, eps: float = 1e-5, nan_loc: float = 0.0):
+        self.eps = eps
+        self.nan_loc = nan_loc
+
+    def scale(
+        self,
+        x: torch.Tensor,
+        loc_scale: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        if loc_scale is None:
+            loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=self.nan_loc)
+            scale = torch.nan_to_num(torch.nanmean((x - loc).square(), dim=-1, keepdim=True).sqrt(), nan=1.0)
+            scale = torch.where(scale == 0, torch.abs(loc) + self.eps, scale)
+        else:
+            loc, scale = loc_scale
+
+        return ((x - loc) / scale), (loc, scale)
+
+    def re_scale(self, x: torch.Tensor, loc_scale: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+        loc, scale = loc_scale
+        return x * scale + loc
+
+
+class Patcher:
+    def __init__(self, patch_size: int, patch_stride: int, left_pad: bool):
+        self.patch_size = patch_size
+        self.patch_stride = patch_stride
+        self.left_pad = left_pad
+        assert self.patch_size % self.patch_stride == 0
+
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.ndim == 2
+        length = x.shape[-1]
+
+        if length < self.patch_size or (length % self.patch_stride != 0):
+            if length < self.patch_size:
+                padding_size = (
+                    *x.shape[:-1],
+                    self.patch_size - (length % self.patch_size),
+                )
+            else:
+                padding_size = (
+                    *x.shape[:-1],
+                    self.patch_stride - (length % self.patch_stride),
+                )
+            padding = torch.full(size=padding_size, fill_value=torch.nan, dtype=x.dtype, device=x.device)
+            if self.left_pad:
+                x = torch.concat((padding, x), dim=-1)
+            else:
+                x = torch.concat((x, padding), dim=-1)
+
+        return x.unfold(dimension=-1, size=self.patch_size, step=self.patch_stride)
+
+
+@dataclass
+class PatchedUniTokenizerState:
+    scale_state: float
+
+
+class PatchedUniTokenizer:
+    def __init__(self, patch_size: int, patch_stride: int | None = None, scaler: StandardScaler | None = None):
+        self.patch_size = patch_size
+        self.patch_stride = patch_size if patch_stride is None else patch_stride
+        self.scaler = StandardScaler() if scaler is None else scaler
+        self.patcher = Patcher(self.patch_size, self.patch_stride, left_pad=True)
+
+    def context_input_transform(self, data: torch.Tensor):
+        assert data.ndim == 2
+        data, scale_state = self.scaler.scale(data)
+        return self.patcher(data), PatchedUniTokenizerState(scale_state)
+
+    def output_transform(self, data: torch.Tensor, tokenizer_state: PatchedUniTokenizerState):
+        data_shape = data.shape
+        data = self.scaler.re_scale(data.reshape(data_shape[0], -1), tokenizer_state.scale_state).view(*data_shape)
+        return data

tirex/models/slstm/block.py

@@ -0,0 +1,60 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from tirex.models.slstm.layer import sLSTMBlockConfig, sLSTMLayer
+from tirex.util import round_up_to_next_multiple_of
+
+
+class sLSTMBlock(nn.Module):
+    def __init__(self, config: sLSTMBlockConfig, backend: str):
+        super().__init__()
+        self.config = config
+        self.norm_slstm = RMSNorm(config.embedding_dim)
+        self.slstm_layer = sLSTMLayer(config, backend)
+        self.norm_ffn = RMSNorm(config.embedding_dim)
+
+        up_proj_dim = round_up_to_next_multiple_of(config.embedding_dim * config.ffn_proj_factor, 64)
+        self.ffn = FeedForward(config.embedding_dim, up_proj_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_slstm = self.norm_slstm(x)
+
+        x_slstm = self.slstm_layer(x_slstm, slstm_state=None)
+        x = x + x_slstm
+
+        x_ffn = self.norm_ffn(x)
+        x_ffn = self.ffn(x_ffn)
+        x = x + x_ffn
+        return x
+
+
+class FeedForward(nn.Module):
+    def __init__(self, embedding_dim: int, up_proj_dim: int):
+        super().__init__()
+        self.proj_up_gate = nn.Linear(embedding_dim, up_proj_dim, bias=False)
+        self.proj_up = nn.Linear(embedding_dim, up_proj_dim, bias=False)
+        self.proj_down = nn.Linear(up_proj_dim, embedding_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.silu(self.proj_up_gate(x)) * self.proj_up(x)
+        y = self.proj_down(x)
+        return y
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, num_features: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(num_features))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self._rms_normalize(x.float()).to(x.dtype)
+        x = x * self.weight
+        return x
+
+    def _rms_normalize(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)

tirex/models/slstm/cell.py

@@ -0,0 +1,188 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import warnings
+from dataclasses import asdict, dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from tirex.util import dataclass_from_dict
+
+
+@dataclass
+class sLSTMBlockConfig:
+    embedding_dim: int
+    num_heads: int
+    num_blocks: int
+    ffn_proj_factor: float = 2.6667
+
+    num_states: int = 4  # this is for the sLSTM, a standard LSTM  has 2
+    num_gates: int = 4
+
+    @property
+    def head_dim(self):
+        return self.embedding_dim // self.num_heads
+
+
+class sLSTMCell(nn.Module):
+    def __init__(self, config: sLSTMBlockConfig, backend: str):
+        super().__init__()
+        self.config = config
+        self.backend = backend
+
+        self._recurrent_kernel_ = nn.Parameter(
+            torch.empty((config.num_heads, config.head_dim, config.num_gates * config.head_dim), dtype=None)
+        )
+
+        self._bias_ = nn.Parameter(torch.empty((config.num_heads * config.num_gates * config.head_dim), dtype=None))
+
+    def forward(self, input: torch.Tensor, state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        input = self._get_input(input)
+        state = self._get_state(input, state)
+
+        if self.backend == "torch":
+            all_states = self._impl_torch(input, state)
+        elif self.backend == "cuda":
+            all_states = self._impl_cuda(input, state)
+
+        state = all_states[:, -1]
+        output = self._permute_output(all_states[0][1:])
+        return output.to(input.dtype), state.to(input.dtype)
+
+    def _impl_torch(self, input: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
+        input = input.to(dtype=torch.bfloat16)
+        state = state.to(dtype=torch.bfloat16)
+        recurrent_kernel = self._recurrent_kernel_.to(dtype=torch.bfloat16)
+        bias = self._bias_.to(dtype=torch.float32)
+
+        input = input.view(input.shape[0], input.shape[1], -1)
+        bias = (
+            bias.reshape(self.config.num_heads, self.config.num_gates, self.config.head_dim)
+            .permute(1, 0, 2)
+            .reshape(-1)
+        )
+
+        return slstm_forward(input, state, recurrent_kernel, bias)[0]
+
+    def _impl_cuda(self, input: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
+        if input.device.type != "cuda":
+            warnings.warn(
+                f"You use TiRex with sLSTM CUDA kernels BUT DO NOT LOAD THE DEVICE ON A CUDA DEVICE (device type is {input.device.type})!"
+                "This is not supported and calls to the model will likely lead to an error if you dont move your model to a CUDA device!"
+                "If you want to run TiRex on CPU you need to disable sLSTM CUDA kernels but be aware of the downsides (see FAQ)"
+            )
+
+        if not hasattr(self, "func"):
+            try:
+                from xlstm.blocks.slstm.cell import sLSTMCellConfig as sLSTMCellConfigCuda, sLSTMCellFuncGenerator
+            except ModuleNotFoundError:
+                raise ValueError(
+                    'xlstm package not found! To use the custom cuda backend, install the additional dependencies with: pip install -e ".[cuda]"'
+                )
+            cuda_config = dataclass_from_dict(
+                sLSTMCellConfigCuda, {**asdict(self.config), "hidden_size": self.config.embedding_dim}
+            )
+            self.func = sLSTMCellFuncGenerator(False, cuda_config)
+
+        input = input.permute(0, 1, 3, 2, 4).reshape(input.shape[0], input.shape[1], -1)
+
+        return self.func.apply(
+            False,
+            input.contiguous(),
+            state.contiguous(),
+            self._recurrent_kernel_.contiguous(),
+            self._bias_.contiguous(),
+        )
+
+    def _get_input(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-1] == self.config.embedding_dim * self.config.num_gates, (
+            f"Input size mismatch: Expected input size {self.config.embedding_dim * self.config.num_gates}, but got {input.size(-1)}."
+        )
+        return x.view(x.shape[0], x.shape[1], self.config.num_gates, self.config.num_heads, -1).permute(1, 0, 2, 3, 4)
+
+    def _get_state(self, input: torch.Tensor, state: torch.Tensor | None) -> torch.Tensor:
+        B = input.shape[1]
+        if state is None:
+            state = torch.zeros(
+                (self.config.num_states, B, self.config.embedding_dim),
+                dtype=input.dtype,
+                device=input.device,
+            )
+
+        assert state.shape == (self.config.num_states, B, self.config.embedding_dim)
+        return state
+
+    def _permute_output(self, output: torch.Tensor) -> torch.Tensor:
+        output = output.view(output.shape[0], output.shape[1], self.config.num_heads, self.config.head_dim)
+        return output.permute(1, 2, 0, 3)
+
+
+def slstm_forward(
+    x: torch.Tensor,  # [S, B, G*I]
+    states: torch.Tensor,  # [4, B, H] only the first is used for recurrence!
+    R: torch.Tensor,  # [K, R*H, H] - K num_heads
+    b: torch.Tensor,  # [T*H]
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    num_states = states.shape[0]
+    sequence_dim = x.shape[0]
+    # this only works for a fully-connected RNN, for a hin change this
+    num_gates_r = R.shape[2] // R.shape[1]
+    hidden_dim = R.shape[1] * R.shape[0]
+    batch_dim = x.shape[1]
+    num_heads = R.shape[0]
+
+    assert batch_dim == states.shape[1]
+    assert hidden_dim == states.shape[2]
+
+    states_all = torch.zeros(
+        [num_states, sequence_dim + 1, batch_dim, hidden_dim],
+        device=x.device,
+        dtype=x.dtype,
+    )
+    states_all[:, 0] = states
+    for i, Wx_t in enumerate(x.unbind(dim=0)):
+        Ry = (
+            states[0]
+            .reshape(batch_dim, num_heads, 1, -1)
+            .matmul(R.unsqueeze(0))
+            .reshape(batch_dim, num_heads, num_gates_r, -1)
+            .transpose(1, 2)
+            .reshape(batch_dim, -1)
+        )
+        sdtype = states.dtype
+        Wx_t, Ry, b, states = Wx_t.float(), Ry.float(), b.float(), states.float()
+        states, gates = slstm_forward_pointwise(Wx_t, Ry, b, states)
+        states = states.to(dtype=sdtype)
+        states_all[:, i + 1] = states
+
+    # shapes ([S, B, H], ([B,H], [B,H], [B,H])
+    return states_all, states
+
+
+def slstm_forward_pointwise(
+    Wx: torch.Tensor,  # dim [B, 4*H]
+    Ry: torch.Tensor,  # dim [B, 4*H]
+    b: torch.Tensor,  # dim [1, 4*H]
+    states: torch.Tensor,  # dim [4, B, H]
+) -> tuple[torch.Tensor, torch.Tensor]:
+    raw = Wx + Ry + b
+    y, c, n, m = torch.unbind(states.view(4, states.shape[1], -1), dim=0)
+
+    iraw, fraw, zraw, oraw = torch.unbind(raw.view(raw.shape[0], 4, -1), dim=1)
+    # with torch.no_grad():  # THE difference to maxg aka max_gradient (here max / max_static)
+    logfplusm = m + F.logsigmoid(fraw)
+    if torch.all(n == 0.0):
+        mnew = iraw
+    else:
+        mnew = torch.max(iraw, logfplusm)
+    ogate = torch.sigmoid(oraw)
+    igate = torch.minimum(torch.exp(iraw - mnew), torch.ones_like(iraw))
+    fgate = torch.minimum(torch.exp(logfplusm - mnew), torch.ones_like(iraw))
+    cnew = fgate * c + igate * torch.tanh(zraw)
+    nnew = fgate * n + igate
+    ynew = ogate * cnew / nnew
+
+    # shapes ([B,H], [B,H], [B,H]), ([B,H],[B,H],[B,H],[B,H])
+    return torch.stack((ynew, cnew, nnew, mnew), dim=0), torch.stack((igate, fgate, zraw, ogate), dim=0)

tirex/models/slstm/layer.py

@@ -0,0 +1,67 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .cell import sLSTMBlockConfig, sLSTMCell
+
+
+class sLSTMLayer(nn.Module):
+    def __init__(self, config: sLSTMBlockConfig, backend: str):
+        super().__init__()
+        self.config = config
+
+        in_features, num_heads = self.config.embedding_dim, self.config.num_heads
+        self.fgate = LinearHeadwiseExpand(in_features, num_heads)
+        self.igate = LinearHeadwiseExpand(in_features, num_heads)
+        self.zgate = LinearHeadwiseExpand(in_features, num_heads)
+        self.ogate = LinearHeadwiseExpand(in_features, num_heads)
+
+        self.slstm_cell = sLSTMCell(self.config, backend)
+        self.group_norm = MultiHeadLayerNorm(ndim=in_features)
+
+    def forward(self, x: torch.Tensor, slstm_state: torch.Tensor | None = None) -> torch.Tensor:
+        x_g = torch.cat((self.fgate(x), self.igate(x), self.zgate(x), self.ogate(x)), dim=-1)
+
+        y, slstm_state = self.slstm_cell(x_g, state=slstm_state)
+
+        return self.group_norm(y).transpose(1, 2).view(x.shape[0], x.shape[1], -1)
+
+
+class LinearHeadwiseExpand(nn.Module):
+    def __init__(self, in_features, num_heads, expand_factor_up: float = 1):
+        super().__init__()
+        assert num_heads <= in_features, "num_heads must be <= in_features"
+        assert in_features % num_heads == 0, "in_features must be a multiple of num_heads"
+        self.num_heads = num_heads
+
+        out_features = round(expand_factor_up * in_features)
+        out_features_per_head = out_features // num_heads
+        self.weight = nn.Parameter(torch.empty(num_heads, out_features_per_head, in_features // num_heads))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        x = x.view(*shape[:-1], self.num_heads, -1)
+        x = torch.einsum("...hd,hod->...ho", x, self.weight)
+        x = x.reshape(*shape[:-1], -1)
+        return x
+
+
+class MultiHeadLayerNorm(nn.Module):
+    def __init__(self, ndim: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(ndim))
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        assert input.dim() == 4, "Input must be 4D tensor (B, NH, S, DH)"
+        B, NH, S, DH = input.shape
+
+        gn_in_1 = input.transpose(1, 2)  # (B, S, NH, DH)
+        gn_in_2 = gn_in_1.reshape(B * S, NH * DH)  # (B * S, NH * DH)
+        residual_weight = 1.0 + self.weight
+        out = F.group_norm(gn_in_2, num_groups=NH, weight=residual_weight)
+        # (B * S), (NH * DH) -> (B, S, NH, DH) -> (B, NH, S, DH)
+        out = out.view(B, S, NH, DH).transpose(1, 2)
+        return out

tirex/models/tirex.py

@@ -2,18 +2,17 @@
 # This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
 import logging
-import warnings
-from contextlib import redirect_stdout
 from dataclasses import dataclass
 
-import lightning as L
 import torch
-from dacite import Config, from_dict
+import torch.nn as nn
+import torch.nn.functional as F
 
+from ..api_adapter.forecast import ForecastModel
 from ..base import PretrainedModel
-from .components import PatchedUniTokenizer, ResidualBlock, StreamToLogger
-from .mixed_stack import skip_cuda, xLSTMMixedLargeBlockStack, xLSTMMixedLargeConfig
-from .predict_utils import TensorQuantileUniPredictMixin
+from ..util import dataclass_from_dict
+from .patcher import PatchedUniTokenizer
+from .slstm.block import RMSNorm, sLSTMBlock, sLSTMBlockConfig
 
 LOGGER = logging.getLogger()
 
@@ -25,113 +24,70 @@ class TiRexZeroConfig:
     quantiles: list[float]
     block_kwargs: dict
     input_ff_dim: int
+    train_ctx_len: int
+    nan_mask_value: int = 0
 
 
-class TiRexZero(L.LightningModule, PretrainedModel, TensorQuantileUniPredictMixin):
-    def __init__(self, model_config: dict, train_ctx_len=None):
+class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
+    def __init__(self, backend, model_config: TiRexZeroConfig, train_ctx_len=None):
         super().__init__()
-        self.model_config: TiRexZeroConfig = from_dict(TiRexZeroConfig, model_config, config=Config(strict=True))
-        assert self.model_config.input_patch_size == self.model_config.output_patch_size
-        self.train_ctx_len = train_ctx_len
+        self.config = TiRexZeroConfig(**model_config, train_ctx_len=train_ctx_len, nan_mask_value=0)
+        assert self.config.input_patch_size == self.config.output_patch_size
+        self.backend = backend
 
-        # Block Stack
-        self.nan_mask_value = 0
-        self.block_stack, resolved_config = self.init_block(self.model_config.block_kwargs)
-        self.model_config.block_kwargs = resolved_config
+        self.tokenizer = PatchedUniTokenizer(patch_size=self.config.input_patch_size)
 
-        # Input Layer
+        block_config = dataclass_from_dict(sLSTMBlockConfig, self.config.block_kwargs)
         self.input_patch_embedding = ResidualBlock(
-            in_dim=self.model_config.input_patch_size * 2,
-            h_dim=self.model_config.input_ff_dim,
-            out_dim=self.model_config.block_kwargs.embedding_dim,
+            in_dim=self.config.input_patch_size * 2,
+            h_dim=self.config.input_ff_dim,
+            out_dim=block_config.embedding_dim,
         )
-        self.tokenizer = PatchedUniTokenizer(
-            patch_size=self.model_config.input_patch_size,
+
+        self.blocks = nn.ModuleList(
+            [sLSTMBlock(block_config, backend=self.backend) for i in range(block_config.num_blocks)]
         )
 
-        # Output Layer
-        self.num_quantiles = len(self.model_config.quantiles)
-        quantiles = torch.tensor(self.model_config.quantiles)
-        self.register_buffer("quantiles", quantiles, persistent=False)
+        self.out_norm = RMSNorm(block_config.embedding_dim)
 
         self.output_patch_embedding = ResidualBlock(
-            in_dim=self.model_config.block_kwargs.embedding_dim,
-            h_dim=self.model_config.input_ff_dim,
-            out_dim=self.num_quantiles * self.model_config.output_patch_size,
+            in_dim=block_config.embedding_dim,
+            h_dim=self.config.input_ff_dim,
+            out_dim=len(self.config.quantiles) * self.config.output_patch_size,
         )
 
-        self.save_hyperparameters()
-
     @classmethod
     def register_name(cls):
         return "TiRex"
 
-    def init_block(self, block_kwargs):
-        config = from_dict(xLSTMMixedLargeConfig, block_kwargs)
-        log_redirect = StreamToLogger(LOGGER, logging.INFO)
-        with redirect_stdout(log_redirect):  # avoid excessive print statements of sLSTM compile
-            model = xLSTMMixedLargeBlockStack(config)
-        return model, config
-
-    @property
-    def quantiles(self):
-        return self.model.quantiles
-
-    def _forward_model_tokenized(
+    def _forecast_quantiles(
         self,
-        input_token,
-        input_mask=None,
-        rollouts=1,
-    ):
-        input_mask = (
-            input_mask.to(input_token.dtype)
-            if input_mask is not None
-            else torch.isnan(input_token).logical_not().to(input_token.dtype)
-        )
-        assert rollouts >= 1
-        bs, numb_ctx_token, token_dim = input_token.shape
-        if rollouts > 1:
-            input_token = torch.cat(
-                (
-                    input_token,
-                    torch.full(
-                        (bs, rollouts - 1, token_dim),
-                        fill_value=torch.nan,
-                        device=input_token.device,
-                        dtype=input_token.dtype,
-                    ),
-                ),
-                dim=1,
-            )
-            input_mask = torch.cat(
-                (
-                    input_mask,
-                    torch.full(
-                        (bs, rollouts - 1, token_dim),
-                        fill_value=False,
-                        device=input_mask.device,
-                        dtype=input_mask.dtype,
-                    ),
-                ),
-                dim=1,
-            )
-        input_token = torch.nan_to_num(input_token, nan=self.nan_mask_value)
-        input_embeds = self.input_patch_embedding(torch.cat((input_token, input_mask), dim=2))
-
-        # hidden_states = []
-        # for rollout in range(rollout):
-        x = self.block_stack(input_embeds)
-        if isinstance(x, tuple):
-            hidden_states = x[0]
+        context: torch.Tensor,
+        prediction_length: int | None = None,
+        quantile_levels: list[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
+        output_device: str = "cpu",
+        auto_cast: bool = False,
+        **predict_kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = self.input_patch_embedding.hidden_layer.weight.device
+        context = context.to(device)
+
+        with torch.autocast(device_type=device.type, enabled=auto_cast):
+            predictions = self._forecast_tensor(
+                context=context, prediction_length=prediction_length, **predict_kwargs
+            ).detach()
+        predictions = predictions.to(torch.device(output_device)).swapaxes(1, 2)
+
+        training_quantile_levels = self.config.quantiles
+
+        if set(quantile_levels).issubset(set(training_quantile_levels)):
+            quantiles = predictions[..., [training_quantile_levels.index(q) for q in quantile_levels]]
         else:
-            hidden_states = x
+            quantiles = self._interpolate_quantiles(predictions, quantile_levels)
 
-        quantile_preds = self.output_patch_embedding(hidden_states)
-        quantile_preds = torch.unflatten(quantile_preds, -1, (self.num_quantiles, self.model_config.output_patch_size))
-        quantile_preds = torch.transpose(quantile_preds, 1, 2)  # switch quantile and num_token_dimension
-        # quantile_preds: [batch_size, num_quantiles, num_token, output_patch_size]
-
-        return quantile_preds, hidden_states
+        # median as mean
+        mean = predictions[:, :, training_quantile_levels.index(0.5)]
+        return quantiles, mean
 
     @torch.inference_mode()
     def _forecast_tensor(
@@ -146,13 +102,10 @@ class TiRexZero(L.LightningModule, PretrainedModel, TensorQuantileUniPredictMixi
             prediction_length = self.tokenizer.patch_size
         remaining = -(prediction_length // -self.tokenizer.patch_size)
         if max_context is None:
-            max_context = self.train_ctx_len
-        min_context = max(self.train_ctx_len, max_context)
+            max_context = self.config.train_ctx_len
+        min_context = max(self.config.train_ctx_len, max_context)
 
-        context = context.to(
-            device=self.device,
-            dtype=torch.float32,
-        )
+        context = context.to(dtype=torch.float32)
         while remaining > 0:
             if context.shape[-1] > max_context:
                 context = context[..., -max_context:]
@@ -181,51 +134,92 @@ class TiRexZero(L.LightningModule, PretrainedModel, TensorQuantileUniPredictMixi
 
             context = torch.cat([context, torch.full_like(prediction[:, 0, :], fill_value=torch.nan)], dim=-1)
 
-        return torch.cat(predictions, dim=-1)[..., :prediction_length].to(
-            dtype=torch.float32,
-        )
+        return torch.cat(predictions, dim=-1)[..., :prediction_length].to(dtype=torch.float32)
 
-    def on_load_checkpoint(self, checkpoint: dict) -> None:
-        state_dict = checkpoint["state_dict"]
-        load_vanilla_kernel = skip_cuda()
-        if load_vanilla_kernel:
-            warnings.warn(
-                "You use TiRex without sLSTM CUDA kernels! This might slow down the model considerably and might degrade forecasting results!"
-                "Set the environment variable TIREX_NO_CUDA to 0 to avoid this!"
+    def _forward_model_tokenized(
+        self,
+        input_token: torch.Tensor,
+        input_mask=None,
+        rollouts=1,
+    ):
+        input_mask = (
+            input_mask.to(input_token.dtype)
+            if input_mask is not None
+            else torch.isnan(input_token).logical_not().to(input_token.dtype)
+        )
+        assert rollouts >= 1
+        bs, numb_ctx_token, token_dim = input_token.shape
+        if rollouts > 1:
+            input_token_rollout_pad = torch.full(
+                (bs, rollouts - 1, token_dim),
+                fill_value=torch.nan,
+                device=input_token.device,
+                dtype=input_token.dtype,
             )
-            block_kwargs = self.model_config.block_kwargs
-            head_dim = block_kwargs.embedding_dim // block_kwargs.num_heads
-            num_gates = 4
-            new_state_dict = {}
-            for k, v in state_dict.items():
-                if "slstm_layer.slstm_cell._recurrent_kernel_" in k:
-                    new_state_dict[k] = (
-                        v.reshape(
-                            block_kwargs.num_heads,
-                            head_dim,
-                            num_gates,
-                            head_dim,
-                        )
-                        .permute(0, 2, 3, 1)
-                        .reshape(
-                            block_kwargs.num_heads,
-                            num_gates * head_dim,
-                            head_dim,
-                        )
-                    )
-                    # new_state_dict[k] = v.permute(0, 2, 1)
-                elif "slstm_layer.slstm_cell._bias_" in k:
-                    new_state_dict[k] = (
-                        v.reshape(block_kwargs.num_heads, num_gates, head_dim).permute(1, 0, 2).reshape(-1)
-                    )
-                else:
-                    new_state_dict[k] = v
-            checkpoint["state_dict"] = new_state_dict
-
-    def after_load_from_checkpoint(self):
-        if not skip_cuda() and self.device.type != "cuda":
-            warnings.warn(
-                f"You use TiRex with sLSTM CUDA kernels BUT DO NOT LOAD THE DEVICE ON A CUDA DEVICE (device type is {self.device.type})!"
-                "This is not supported and calls to the model will likely lead to an error if you dont move your model to a CUDA device!"
-                "If you want to run TiRex on CPU you need to disable sLSTM CUDA kernels but be aware of the downsides (see FAQ)"
+            input_token = torch.cat((input_token, input_token_rollout_pad), dim=1)
+            input_mask_rollout_pad = torch.full(
+                (bs, rollouts - 1, token_dim),
+                fill_value=False,
+                device=input_mask.device,
+                dtype=input_mask.dtype,
+            )
+            input_mask = torch.cat((input_mask, input_mask_rollout_pad), dim=1)
+
+        input_token = torch.nan_to_num(input_token, nan=self.config.nan_mask_value)
+
+        hidden_states = self.input_patch_embedding(torch.cat((input_token, input_mask), dim=2))
+
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+
+        hidden_states = self.out_norm(hidden_states)
+
+        quantile_preds = self.output_patch_embedding(hidden_states)
+        quantile_preds = torch.unflatten(
+            quantile_preds, -1, (len(self.config.quantiles), self.config.output_patch_size)
+        )
+        quantile_preds = torch.transpose(quantile_preds, 1, 2)  # switch quantile and num_token_dimension
+        # quantile_preds: [batch_size, num_quantiles, num_token, output_patch_size]
+
+        return quantile_preds, hidden_states
+
+    def _interpolate_quantiles(self, predictions: torch.Tensor, quantile_levels: list[float]):
+        training_quantile_levels = self.config.quantiles
+        if min(quantile_levels) < min(training_quantile_levels) or max(quantile_levels) > max(training_quantile_levels):
+            logging.warning(
+                f"Requested quantile levels ({quantile_levels}) fall outside the range of "
+                f"quantiles the model was trained on ({training_quantile_levels}). "
+                "Predictions for out-of-range quantiles will be clamped to the nearest "
+                "boundary of the trained quantiles (i.e., minimum or maximum trained level). "
+                "This can significantly impact prediction accuracy, especially for extreme quantiles. "
             )
+
+        augmented_predictions = torch.cat(
+            [predictions[..., [0]], predictions, predictions[..., [-1]]],
+            dim=-1,
+        )
+        quantiles = torch.quantile(
+            augmented_predictions,
+            q=torch.tensor(quantile_levels, dtype=augmented_predictions.dtype),
+            dim=-1,
+        ).permute(1, 2, 0)
+        return quantiles
+
+    def on_load_checkpoint(self, checkpoint: dict) -> None:
+        # rename keys of state_dict, because the block_stack was moved directly into the tirex model
+        checkpoint["state_dict"] = {k.replace("block_stack.", ""): v for k, v in checkpoint["state_dict"].items()}
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_dim: int, h_dim: int, out_dim: int) -> 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)
+
+    def forward(self, x: torch.Tensor):
+        hid = F.relu(self.hidden_layer(x))
+        out = self.output_layer(hid)
+        res = self.residual_layer(x)
+        out = out + res
+        return out

tirex/util.py

@@ -0,0 +1,13 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from dataclasses import fields
+
+
+def round_up_to_next_multiple_of(x: int, multiple_of: int) -> int:
+    return int(((x + multiple_of - 1) // multiple_of) * multiple_of)
+
+
+def dataclass_from_dict(cls, dict: dict):
+    class_fields = {f.name for f in fields(cls)}
+    return cls(**{k: v for k, v in dict.items() if k in class_fields})

{tirex_mirror-2025.9.2.dist-info → tirex_mirror-2025.9.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.9.2
+Version: 2025.9.9
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -65,17 +65,17 @@ License-File: LICENSE_MIRROR.txt
 License-File: NOTICE.txt
 Requires-Dist: torch
 Requires-Dist: torchvision
-Requires-Dist: xlstm
 Requires-Dist: einops
-Requires-Dist: ninja
 Requires-Dist: huggingface-hub
-Requires-Dist: lightning
 Requires-Dist: numpy
 Requires-Dist: pandas
-Requires-Dist: dacite
 Requires-Dist: tqdm
+Provides-Extra: cuda
+Requires-Dist: xlstm; extra == "cuda"
+Requires-Dist: ninja; extra == "cuda"
 Provides-Extra: notebooks
 Requires-Dist: ipykernel; extra == "notebooks"
+Requires-Dist: matplotlib; extra == "notebooks"
 Provides-Extra: gluonts
 Requires-Dist: gluonts; extra == "gluonts"
 Provides-Extra: hfdataset
@@ -83,7 +83,10 @@ Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
 Requires-Dist: fev; extra == "test"
 Provides-Extra: all
+Requires-Dist: xlstm; extra == "all"
+Requires-Dist: ninja; extra == "all"
 Requires-Dist: ipykernel; extra == "all"
+Requires-Dist: matplotlib; extra == "all"
 Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
 Requires-Dist: fev; extra == "all"

tirex_mirror-2025.9.9.dist-info/RECORD

@@ -0,0 +1,21 @@
+tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
+tirex/base.py,sha256=ODUyhYFR33ZYffu7dxDwsb9m2IiZAnGHIXvA81crbjQ,3245
+tirex/util.py,sha256=7DFVBXwGQA4niT9VhYbt8iKMBINJVW4LfwwpggFS0Us,469
+tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
+tirex/api_adapter/forecast.py,sha256=snv0sT1_1WzjkhP1YV-I7CMQmSChl93qFc3b6fwUAS0,8502
+tirex/api_adapter/gluon.py,sha256=faiYyn0kBBVQKbpWqrVoyylxZUrmr-qce66twpguVds,1827
+tirex/api_adapter/hf_data.py,sha256=T1eaxqC3OO9yOzIvw4sr55x6iA2AHKJTZd36rROM4fQ,1377
+tirex/api_adapter/standard_adapter.py,sha256=bI3XGYlWQu5EDyhDZyYqOJMbwi5h1aovPQvfHuWETJk,2618
+tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
+tirex/models/patcher.py,sha256=EOXFkHsPkq0nuxRNLAbnrgJtcYq0IMC3YIg_16WArg4,3213
+tirex/models/tirex.py,sha256=dclEckb6CmvESeX_LwT2kaCNTB7deTFovIOQUIFF5J8,9117
+tirex/models/slstm/block.py,sha256=DCOxmLQUb7HRO6wXTZMK4ICUI5LFpo7NC5a28oM-Vsc,2104
+tirex/models/slstm/cell.py,sha256=4_pQcXOOT16aEpKIi4A-yEnj4qKK6pFyFADD2nGPzGc,7366
+tirex/models/slstm/layer.py,sha256=93CAYuG-HmUpF7mBAQ-z1S1u2__W10EW5jPToR57qqM,2747
+tirex_mirror-2025.9.9.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
+tirex_mirror-2025.9.9.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
+tirex_mirror-2025.9.9.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
+tirex_mirror-2025.9.9.dist-info/METADATA,sha256=u9C_cIb8FtaHUep1XrFTeI7UAsVRtNJt2VSQo7420Vo,11200
+tirex_mirror-2025.9.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tirex_mirror-2025.9.9.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
+tirex_mirror-2025.9.9.dist-info/RECORD,,

tirex/models/components.py

@@ -1,147 +0,0 @@
-# Copyright (c) NXAI GmbH.
-# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
-
-
-from dataclasses import dataclass, field
-from typing import Any
-
-import torch
-
-SCALER_STATE = "scaler_state"
-
-
-class ResidualBlock(torch.nn.Module):
-    def __init__(
-        self,
-        in_dim: int,
-        h_dim: int,
-        out_dim: int,
-        dropout: float = 0,
-    ) -> None:
-        super().__init__()
-        self.dropout = torch.nn.Dropout(dropout)
-        self.hidden_layer = torch.nn.Linear(in_dim, h_dim)
-        self.output_layer = torch.nn.Linear(h_dim, out_dim)
-        self.residual_layer = torch.nn.Linear(in_dim, out_dim)
-        self.act = torch.nn.ReLU()
-
-    def forward(self, x: torch.Tensor):
-        hid = self.act(self.hidden_layer(x))
-        out = self.output_layer(hid)
-        res = self.residual_layer(x)
-        out = out + res
-        return out
-
-
-@dataclass
-class StandardScaler:
-    eps: float = 1e-5
-    nan_loc: float = 0.0
-
-    def scale(
-        self,
-        x: torch.Tensor,
-        loc_scale: tuple[torch.Tensor, torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
-        if loc_scale is None:
-            loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=self.nan_loc)
-            scale = torch.nan_to_num(torch.nanmean((x - loc).square(), dim=-1, keepdim=True).sqrt(), nan=1.0)
-            scale = torch.where(scale == 0, torch.abs(loc) + self.eps, scale)
-        else:
-            loc, scale = loc_scale
-
-        return ((x - loc) / scale), (loc, scale)
-
-    def re_scale(self, x: torch.Tensor, loc_scale: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
-        loc, scale = loc_scale
-        return x * scale + loc
-
-
-@dataclass
-class _Patcher:
-    patch_size: int
-    patch_stride: int
-    left_pad: bool
-
-    def __post_init__(self):
-        assert self.patch_size % self.patch_stride == 0
-
-    def __call__(self, x: torch.Tensor) -> torch.Tensor:
-        assert x.ndim == 2
-        length = x.shape[-1]
-
-        if length < self.patch_size or (length % self.patch_stride != 0):
-            if length < self.patch_size:
-                padding_size = (
-                    *x.shape[:-1],
-                    self.patch_size - (length % self.patch_size),
-                )
-            else:
-                padding_size = (
-                    *x.shape[:-1],
-                    self.patch_stride - (length % self.patch_stride),
-                )
-            padding = torch.full(size=padding_size, fill_value=torch.nan, dtype=x.dtype, device=x.device)
-            if self.left_pad:
-                x = torch.concat((padding, x), dim=-1)
-            else:
-                x = torch.concat((x, padding), dim=-1)
-
-        x = x.unfold(dimension=-1, size=self.patch_size, step=self.patch_stride)
-        return x
-
-
-@dataclass
-class PatchedUniTokenizer:
-    patch_size: int
-    scaler: Any = field(default_factory=StandardScaler)
-    patch_stride: int | None = None
-
-    def __post_init__(self):
-        if self.patch_stride is None:
-            self.patch_stride = self.patch_size
-        self.patcher = _Patcher(self.patch_size, self.patch_stride, left_pad=True)
-
-    def context_input_transform(self, data: torch.Tensor):
-        assert data.ndim == 2
-        data, scale_state = self.scaler.scale(data)
-        return self.patcher(data), {SCALER_STATE: scale_state}
-
-    def output_transform(self, data: torch.Tensor, tokenizer_state: dict):
-        data_shape = data.shape
-        data = self.scaler.re_scale(data.reshape(data_shape[0], -1), tokenizer_state[SCALER_STATE]).view(*data_shape)
-        return data
-
-
-class StreamToLogger:
-    """Fake file-like stream object that redirects writes to a logger
-    instance."""
-
-    def __init__(self, logger, log_level):
-        self.logger = logger
-        self.log_level = log_level
-        self.linebuf = ""  # Buffer for partial lines
-
-    def write(self, message):
-        # Filter out empty messages (often from just a newline)
-        if message.strip():
-            self.linebuf += message
-            # If the message contains a newline, process the full line
-            if "\n" in self.linebuf:
-                lines = self.linebuf.splitlines(keepends=True)
-                for line in lines:
-                    if line.endswith("\n"):
-                        # Log full lines without the trailing newline (logger adds its own)
-                        self.logger.log(self.log_level, line.rstrip("\n"))
-                    else:
-                        # Keep partial lines in buffer
-                        self.linebuf = line
-                        return
-                self.linebuf = ""  # All lines processed
-            # If no newline, keep buffering
-
-    def flush(self):
-        # Log any remaining buffered content when flush is called
-        if self.linebuf.strip():
-            self.logger.log(self.log_level, self.linebuf.rstrip("\n"))
-            self.linebuf = ""

tirex/models/mixed_stack.py

@@ -1,143 +0,0 @@
-# Copyright (c) NXAI GmbH.
-# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
-
-
-import os
-from dataclasses import dataclass, field
-
-import torch
-from torch import nn
-from xlstm.blocks.slstm.layer import sLSTMLayer, sLSTMLayerConfig
-from xlstm.xlstm_large import xLSTMLargeConfig
-from xlstm.xlstm_large.components import RMSNorm
-from xlstm.xlstm_large.model import FeedForward, mLSTMBlock, mLSTMStateType
-
-
-def skip_cuda():
-    return os.getenv("TIREX_NO_CUDA", "False").lower() in ("true", "1", "t")
-
-
-def init_cell(config: xLSTMLargeConfig, block_idx, num_blocks):
-    return sLSTMLayer(
-        sLSTMLayerConfig(
-            embedding_dim=config.embedding_dim,
-            num_heads=config.num_heads,
-            conv1d_kernel_size=0,  # 0 means no convolution included
-            group_norm_weight=True,
-            dropout=0,
-            # CellConfig
-            backend="vanilla" if skip_cuda() else "cuda",
-            bias_init="powerlaw_blockdependent",
-            recurrent_weight_init="zeros",
-            num_gates=4,
-            gradient_recurrent_cut=False,
-            gradient_recurrent_clipval=None,
-            forward_clipval=None,
-            batch_size=8,  # needed?
-            _block_idx=block_idx,
-            _num_blocks=num_blocks,
-        )
-    )
-
-
-sLSTMLayerStateType = tuple[torch.Tensor, torch.Tensor]
-sLSTMStateType = dict[int, sLSTMLayerStateType]
-
-
-class sLSTMBlock(nn.Module):
-    def __init__(self, config: xLSTMLargeConfig, block_idx: int, num_blocks: int):
-        super().__init__()
-        self.config = config
-        self.norm_slstm = RMSNorm(
-            num_features=config.embedding_dim,
-            eps=config.norm_eps,
-            use_weight=True,
-            use_bias=config.use_bias,
-            force_float32_reductions=config.norm_reduction_force_float32,
-        )
-        self.slstm_layer = init_cell(config, block_idx, num_blocks)
-
-        self.norm_ffn = RMSNorm(
-            num_features=config.embedding_dim,
-            eps=config.norm_eps,
-            use_weight=True,
-            use_bias=config.use_bias,
-            force_float32_reductions=config.norm_reduction_force_float32,
-        )
-        self.ffn = FeedForward(config)
-
-    def forward(
-        self, x: torch.Tensor, state: sLSTMLayerStateType | None = None
-    ) -> tuple[torch.Tensor, sLSTMLayerStateType]:
-        x_slstm = self.norm_slstm(x)
-        if state is None:
-            conv_state, slstm_state = None, None
-        else:
-            conv_state, slstm_state = state
-        x_slstm, state = self.slstm_layer(x_slstm, conv_state, slstm_state, return_last_state=True)
-        x = x + x_slstm
-
-        x_ffn = self.norm_ffn(x)
-        x_ffn = self.ffn(x_ffn)
-        x = x + x_ffn
-
-        return x, (state["conv_state"], state["slstm_state"])
-
-
-@dataclass
-class xLSTMMixedLargeConfig(xLSTMLargeConfig):
-    slstm_at: list[int] = field(default_factory=list)
-    all_slstm: bool = True
-
-    @property
-    def block_types(self):
-        return ["s" if i in self.slstm_at or self.all_slstm else "m" for i in range(self.num_blocks)]
-
-
-class xLSTMMixedLargeBlockStack(nn.Module):
-    config_class = xLSTMMixedLargeConfig
-
-    def __init__(self, config: xLSTMMixedLargeConfig):
-        super().__init__()
-        self.config = config
-
-        self.blocks = nn.ModuleList(
-            [
-                sLSTMBlock(config, block_idx=i, num_blocks=config.num_blocks) if t == "s" else mLSTMBlock(config)
-                for i, t in enumerate(config.block_types)
-            ]
-        )
-
-        if self.config.add_out_norm:
-            self.out_norm = RMSNorm(
-                num_features=config.embedding_dim,
-                eps=config.norm_eps,
-                use_weight=True,
-                use_bias=config.use_bias,
-                force_float32_reductions=config.norm_reduction_force_float32,
-            )
-        else:
-            self.out_norm = nn.Identity()
-
-    def forward(
-        self, x: torch.Tensor, state: mLSTMStateType | sLSTMStateType | None = None
-    ) -> tuple[torch.Tensor, mLSTMStateType]:
-        if state is None:
-            state = {i: None for i in range(len(self.blocks))}
-
-        for i, block in enumerate(self.blocks):
-            block_state = state[i]
-            x, block_state_new = block(x, block_state)
-
-            if block_state is None:
-                state[i] = block_state_new
-            else:
-                pass
-                ## layer state is a tuple of three tensors: c, n, m
-                ## we update the state in place in order to avoid creating new tensors
-                # for state_idx in range(len(block_state)):
-                #    state[i][state_idx].copy_(block_state_new[state_idx])
-
-        x = self.out_norm(x)
-
-        return x, state

tirex/models/predict_utils.py

@@ -1,72 +0,0 @@
-# Copyright (c) NXAI GmbH.
-# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
-
-
-import logging
-from abc import abstractmethod
-
-import torch
-
-from ..api_adapter.forecast import ForecastModel
-
-LOGGER = logging.getLogger()
-
-
-class TensorQuantileUniPredictMixin(ForecastModel):
-    @abstractmethod
-    def _forecast_tensor(
-        self,
-        context: torch.Tensor,
-        prediction_length: int | None = None,
-        **predict_kwargs,
-    ) -> torch.Tensor:
-        pass
-
-    @property
-    @abstractmethod
-    def quantiles(self):
-        pass
-
-    def _forecast_quantiles(
-        self,
-        context: torch.Tensor,
-        prediction_length: int | None = None,
-        quantile_levels: list[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
-        output_device: str = "cpu",
-        auto_cast: bool = False,
-        **predict_kwargs,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        with torch.autocast(device_type=self.device.type, enabled=auto_cast):
-            predictions = self._forecast_tensor(
-                context=context, prediction_length=prediction_length, **predict_kwargs
-            ).detach()
-        predictions = predictions.to(torch.device(output_device)).swapaxes(1, 2)
-
-        training_quantile_levels = list(self.quantiles)
-
-        if set(quantile_levels).issubset(set(training_quantile_levels)):
-            quantiles = predictions[..., [training_quantile_levels.index(q) for q in quantile_levels]]
-        else:
-            if min(quantile_levels) < min(training_quantile_levels) or max(quantile_levels) > max(
-                training_quantile_levels
-            ):
-                logging.warning(
-                    f"Requested quantile levels ({quantile_levels}) fall outside the range of "
-                    f"quantiles the model was trained on ({training_quantile_levels}). "
-                    "Predictions for out-of-range quantiles will be clamped to the nearest "
-                    "boundary of the trained quantiles (i.e., minimum or maximum trained level). "
-                    "This can significantly impact prediction accuracy, especially for extreme quantiles. "
-                )
-            # Interpolate quantiles
-            augmented_predictions = torch.cat(
-                [predictions[..., [0]], predictions, predictions[..., [-1]]],
-                dim=-1,
-            )
-            quantiles = torch.quantile(
-                augmented_predictions,
-                q=torch.tensor(quantile_levels, dtype=augmented_predictions.dtype),
-                dim=-1,
-            ).permute(1, 2, 0)
-        # median as mean
-        mean = predictions[:, :, training_quantile_levels.index(0.5)]
-        return quantiles, mean

tirex_mirror-2025.9.2.dist-info/RECORD

@@ -1,19 +0,0 @@
-tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
-tirex/base.py,sha256=F18v9tTbLH0-nX-PC6kBAkYQHkS1T_7OQD6_aN6EjMw,2623
-tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
-tirex/api_adapter/forecast.py,sha256=iOVP_L7fYlp1ZjyrQe2b8fwuEcxTYOszfZ5f9VDqKHU,8503
-tirex/api_adapter/gluon.py,sha256=faiYyn0kBBVQKbpWqrVoyylxZUrmr-qce66twpguVds,1827
-tirex/api_adapter/hf_data.py,sha256=T1eaxqC3OO9yOzIvw4sr55x6iA2AHKJTZd36rROM4fQ,1377
-tirex/api_adapter/standard_adapter.py,sha256=bI3XGYlWQu5EDyhDZyYqOJMbwi5h1aovPQvfHuWETJk,2618
-tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
-tirex/models/components.py,sha256=sluhMbV6KL3W1ESoC5Nyoxdge9WSNx98alc8NG85dv0,4991
-tirex/models/mixed_stack.py,sha256=ffpdhwCrPAbpp4_s1q8Z0Ei7iZ2TsqzVzOPe3BQPW9w,4790
-tirex/models/predict_utils.py,sha256=QUMZZ4_Sxa09UaHs1DG-MbfP8j_XwYt0x1zemdSEcFI,2749
-tirex/models/tirex.py,sha256=bFxtcpQB9-Hnayy_4bqif-o75DwO3-W0wJxelS8F_6c,9243
-tirex_mirror-2025.9.2.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
-tirex_mirror-2025.9.2.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
-tirex_mirror-2025.9.2.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
-tirex_mirror-2025.9.2.dist-info/METADATA,sha256=Ekx7wxImQuw0we5lCuz1fK7rULcY3q4K1IqKuJjZm_M,11028
-tirex_mirror-2025.9.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tirex_mirror-2025.9.2.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
-tirex_mirror-2025.9.2.dist-info/RECORD,,