tirex-mirror 2025.10.25__tar.gz → 2025.10.29__tar.gz

{tirex_mirror-2025.10.25/src/tirex_mirror.egg-info → tirex_mirror-2025.10.29}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.25
+Version: 2025.10.29
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT

{tirex_mirror-2025.10.25 → tirex_mirror-2025.10.29}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tirex-mirror"
-version = "2025.10.25"
+version = "2025.10.29"
 description = "Unofficial mirror of NX-AI/tirex for packaging"
 readme = "README.md"
 requires-python = ">=3.11"

tirex_mirror-2025.10.29/src/tirex/models/patcher.py

@@ -0,0 +1,47 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from typing import NamedTuple
+
+import torch
+
+
+class StandardScalerState(NamedTuple):
+    loc: torch.Tensor
+    scale: torch.Tensor
+
+
+class StandardScaler:
+    def scale(self, x: torch.Tensor) -> tuple[torch.Tensor, StandardScalerState]:
+        state = self.get_loc_scale(x)
+        return ((x - state.loc) / state.scale), state
+
+    def re_scale(self, x: torch.Tensor, state: StandardScalerState) -> torch.Tensor:
+        return x * state.scale + state.loc
+
+    def get_loc_scale(self, x: torch.Tensor, eps=1e-5):
+        loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=0.0)
+        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) + eps, scale)
+        return StandardScalerState(loc=loc, scale=scale)
+
+
+class PatchedTokenizer:
+    def __init__(self, patch_size: int):
+        self.patch_size = patch_size
+        self.scaler = StandardScaler()
+
+    def input_transform(self, data: torch.Tensor) -> tuple[torch.Tensor, StandardScalerState]:
+        assert data.ndim == 2
+        assert data.shape[1] % self.patch_size == 0, "Length of data has to be a multiple of patch_size!"
+
+        scaled_data, scale_state = self.scaler.scale(data)
+        patched_data = scaled_data.unfold(dimension=-1, size=self.patch_size, step=self.patch_size)
+        return patched_data, scale_state
+
+    def output_transform(self, data: torch.Tensor, scaler_state: StandardScalerState) -> torch.Tensor:
+        assert data.shape[-1] == self.patch_size
+
+        rescaled_data = self.scaler.re_scale(data.reshape(data.shape[0], -1), scaler_state)
+        unpatched_data = rescaled_data.view(*data.shape[:-2], data.shape[-2] * self.patch_size)
+        return unpatched_data

{tirex_mirror-2025.10.25 → tirex_mirror-2025.10.29}/src/tirex/models/tirex.py

@@ -11,7 +11,7 @@ import torch.nn.functional as F
 from ..api_adapter.forecast import ForecastModel
 from ..base import PretrainedModel
 from ..util import dataclass_from_dict
-from .patcher import PatchedUniTokenizer
+from .patcher import PatchedTokenizer
 from .slstm.block import RMSNorm, sLSTMBlock, sLSTMBlockConfig
 
 LOGGER = logging.getLogger()
@@ -34,7 +34,7 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
         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.tokenizer = PatchedUniTokenizer(patch_size=self.config.input_patch_size)
+        self.tokenizer = PatchedTokenizer(patch_size=self.config.input_patch_size)
 
         num_blocks = self.config.block_kwargs["num_blocks"]
         block_config = dataclass_from_dict(sLSTMBlockConfig, self.config.block_kwargs)
@@ -58,64 +58,41 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
     def register_name(cls):
         return "TiRex"
 
+    @torch.inference_mode()
     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,
+        max_accelerated_rollout_steps: int = 1,
     ) -> 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)):
-            quantile_indices = torch.tensor(
-                [training_quantile_levels.index(q) for q in quantile_levels],
-                dtype=torch.long,
-                device=predictions.device,
-            )
-            quantiles = torch.index_select(predictions, dim=-1, index=quantile_indices)
-        else:
-            quantiles = self._interpolate_quantiles(predictions, quantile_levels)
+        quantiles = self._forecast_tensor(context, prediction_length, new_patch_count=max_accelerated_rollout_steps)
+        quantiles = quantiles.to(torch.device(output_device)).swapaxes(1, 2)
 
-        # median as mean
-        median_idx = torch.tensor([training_quantile_levels.index(0.5)], dtype=torch.long, device=predictions.device)
-        mean = torch.index_select(predictions, dim=-1, index=median_idx).squeeze(-1)
+        mean = quantiles[:, :, self.config.quantiles.index(0.5)].squeeze(-1)  # median as mean
         return quantiles, mean
 
-    @torch.inference_mode()
     def _forecast_tensor(
         self,
         context: torch.Tensor,
         prediction_length: int | None = None,
-        max_context: int | None = None,
-        max_accelerated_rollout_steps: int = 1,
+        new_patch_count: int = 1,
     ) -> torch.Tensor:
         predictions = []
         if prediction_length is None:
             prediction_length = self.tokenizer.patch_size
         remaining = -(prediction_length // -self.tokenizer.patch_size)
-        if max_context is None:
-            max_context = self.config.train_ctx_len
-        min_context = max(self.config.train_ctx_len, max_context)
 
         context = context.to(dtype=torch.float32)
         while remaining > 0:
-            fut_rollouts = min(remaining, max_accelerated_rollout_steps)
-            prediction, fut_rollouts = self._forecast_single_step(context, max_context, min_context, fut_rollouts)
+            new_patch_count = min(remaining, new_patch_count)
+            prediction = self._forecast_single_step(context, new_patch_count)
 
             predictions.append(prediction)
-            remaining -= fut_rollouts
+            remaining -= new_patch_count
 
             if remaining <= 0:
                 break
@@ -124,13 +101,9 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
 
         return torch.cat(predictions, dim=-1)[..., :prediction_length].to(dtype=torch.float32)
 
-    def _forecast_single_step(
-        self,
-        context: torch.Tensor,
-        max_context: int,
-        min_context: int,
-        new_patch_count: int = 1,
-    ) -> tuple[torch.Tensor, int]:
+    def _forecast_single_step(self, context: torch.Tensor, new_patch_count: int = 1) -> torch.Tensor:
+        max_context, min_context = self.config.train_ctx_len, self.config.train_ctx_len
+
         if context.shape[-1] > max_context:
             context = context[..., -max_context:]
         if context.shape[-1] < min_context:
@@ -142,38 +115,32 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
             )
             context = torch.concat((pad, context), dim=1)
 
-        tokenized_tensor, tokenizer_state = self.tokenizer.context_input_transform(context)
-        prediction, _ = self._forward_model_tokenized(input_token=tokenized_tensor, rollouts=new_patch_count)
-        prediction = prediction[:, :, -new_patch_count:, :].to(tokenized_tensor)  # predicted token
+        input_token, tokenizer_state = self.tokenizer.input_transform(context)
+        prediction = self._forward_model_tokenized(input_token=input_token, new_patch_count=new_patch_count)
+        predicted_token = prediction[:, :, -new_patch_count:, :].to(input_token)  # predicted token
         # Shape: [bs, num_quantiles, num_predicted_token, output_patch_size]
-        prediction = self.tokenizer.output_transform(prediction, tokenizer_state)
-        prediction = prediction.flatten(start_dim=2)
+        predicted_token = self.tokenizer.output_transform(predicted_token, tokenizer_state)
 
-        return prediction, new_patch_count
+        return predicted_token
 
-    def _forward_model_tokenized(
-        self,
-        input_token: torch.Tensor,
-        input_mask=None,
-        rollouts=1,
-    ):
+    def _forward_model_tokenized(self, input_token: torch.Tensor, input_mask=None, new_patch_count=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
+        assert new_patch_count >= 1
         bs, numb_ctx_token, token_dim = input_token.shape
-        if rollouts > 1:
+        if new_patch_count > 1:
             input_token_rollout_pad = torch.full(
-                (bs, rollouts - 1, token_dim),
+                (bs, new_patch_count - 1, token_dim),
                 fill_value=torch.nan,
                 device=input_token.device,
                 dtype=input_token.dtype,
             )
             input_token = torch.cat((input_token, input_token_rollout_pad), dim=1)
             input_mask_rollout_pad = torch.full(
-                (bs, rollouts - 1, token_dim),
+                (bs, new_patch_count - 1, token_dim),
                 fill_value=False,
                 device=input_mask.device,
                 dtype=input_mask.dtype,
@@ -182,16 +149,16 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
 
         input_token = torch.nan_to_num(input_token, nan=self.config.nan_mask_value)
 
-        quantile_preds, hidden_states = self._forward_model(torch.cat((input_token, input_mask), dim=2))
+        quantile_preds = self._forward_model(torch.cat((input_token, input_mask), dim=2))
 
         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
+        return quantile_preds
 
-    def _forward_model(self, input: torch.Tensor):
+    def _forward_model(self, input: torch.Tensor) -> torch.Tensor:
         hidden_states = self.input_patch_embedding(input)
 
         for block in self.blocks:
@@ -199,29 +166,7 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
 
         hidden_states = self.out_norm(hidden_states)
 
-        return self.output_patch_embedding(hidden_states), 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
+        return self.output_patch_embedding(hidden_states)
 
     def on_load_checkpoint(self, checkpoint: dict) -> None:
         # rename keys of state_dict, because the block_stack was moved directly into the tirex model

{tirex_mirror-2025.10.25 → tirex_mirror-2025.10.29/src/tirex_mirror.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.25
+Version: 2025.10.29
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT

{tirex_mirror-2025.10.25 → tirex_mirror-2025.10.29}/src/tirex_mirror.egg-info/SOURCES.txt

@@ -27,6 +27,7 @@ tests/test_chronos_zs.py
 tests/test_compile.py
 tests/test_forecast.py
 tests/test_forecast_adapter.py
+tests/test_patcher.py
 tests/test_slstm_torch_vs_cuda.py
 tests/test_standard_adapter.py
 tests/test_util_freq.py

tirex_mirror-2025.10.29/tests/test_patcher.py

@@ -0,0 +1,40 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+
+from tirex.models.patcher import PatchedTokenizer
+
+
+def rms(x: torch.Tensor):
+    return torch.nanmean(x.square(), dim=-1, keepdim=True).sqrt()
+
+
+def test_patcher_decode_encode():
+    patcher = PatchedTokenizer(patch_size=32)
+
+    input = torch.randn((2, 256))
+
+    patched_context, state = patcher.input_transform(input)
+    output = patcher.output_transform(patched_context, state)
+
+    assert patched_context.shape == (2, 8, 32)
+    assert input.shape == output.shape
+    torch.testing.assert_close(input, output)
+
+    context_rms = rms(patched_context.view(2, -1) - state.loc)
+    context_mean = torch.nanmean(patched_context.view(2, -1), dim=-1, keepdim=True)
+    torch.testing.assert_close(context_rms, torch.ones((2, 1)), rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(context_mean, torch.zeros((2, 1)))
+
+
+def test_patcher_nan():
+    patcher = PatchedTokenizer(patch_size=32)
+
+    input = torch.randn((2, 256))
+    input[0, 0:64] = torch.nan
+
+    patched_context, state = patcher.input_transform(input)
+    output = patcher.output_transform(patched_context, state)
+
+    torch.testing.assert_close(input, output, equal_nan=True)

tirex_mirror-2025.10.25/src/tirex/models/patcher.py

@@ -1,84 +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
-
-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