tirex-mirror 2025.10.8__py3-none-any.whl → 2025.10.17__py3-none-any.whl

tirex/models/slstm/cell.py

@@ -100,7 +100,7 @@ class sLSTMCell(nn.Module):
 
     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)}."
+            f"Input size mismatch: Expected input size {self.config.embedding_dim * self.config.num_gates}, but got {x.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)
 
@@ -128,7 +128,7 @@ class sLSTMCellTorch:
         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]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         num_gates = 4
         num_heads = R.shape[0]
         S, B, _ = x.shape
@@ -167,7 +167,7 @@ class sLSTMCellTorch:
         iraw, fraw, zraw, oraw = torch.unbind(raw.view(raw.shape[0], 4, -1), dim=1)
 
         # Equations reference the xlstm paper on page 4: https://arxiv.org/pdf/2405.04517
-        logfplusm = m + F.logsigmoid(fraw)  # eq 15
+        logfplusm = m + F.logsigmoid(torch.clamp(fraw, max=15))  # eq 15 # Clamp to avoid subnomals
         mnew = torch.where(torch.all(n == 0.0), iraw, torch.max(iraw, logfplusm))  # eq 15
         ogate = torch.sigmoid(oraw)  # eq 14
         igate = torch.minimum(torch.exp(iraw - mnew), torch.ones_like(iraw))  # eq 16

tirex/models/slstm/layer.py

@@ -20,7 +20,7 @@ class sLSTMLayer(nn.Module):
         self.ogate = LinearHeadwiseExpand(in_features, num_heads)
 
         self.slstm_cell = sLSTMCell(self.config, backend)
-        self.group_norm = MultiHeadLayerNorm(ndim=in_features)
+        self.group_norm = MultiHeadLayerNorm(ndim=in_features, num_heads=num_heads)
 
     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)
@@ -50,18 +50,20 @@ class LinearHeadwiseExpand(nn.Module):
 
 
 class MultiHeadLayerNorm(nn.Module):
-    def __init__(self, ndim: int):
+    def __init__(self, ndim: int, num_heads: int):
         super().__init__()
         self.weight = nn.Parameter(torch.zeros(ndim))
+        self.num_heads = num_heads
 
     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
 
+        assert NH == self.num_heads
         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)
+        out = F.group_norm(gn_in_2, num_groups=self.num_heads, 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

@@ -79,12 +79,18 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
         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]]
+            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)
 
         # median as mean
-        mean = predictions[:, :, training_quantile_levels.index(0.5)]
+        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)
         return quantiles, mean
 
     @torch.inference_mode()
@@ -105,24 +111,8 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
 
         context = context.to(dtype=torch.float32)
         while remaining > 0:
-            if context.shape[-1] > max_context:
-                context = context[..., -max_context:]
-            if context.shape[-1] < min_context:
-                pad = torch.full(
-                    (context.shape[0], min_context - context.shape[-1]),
-                    fill_value=torch.nan,
-                    device=context.device,
-                    dtype=context.dtype,
-                )
-                context = torch.concat((pad, context), dim=1)
-            tokenized_tensor, tokenizer_state = self.tokenizer.context_input_transform(context)
             fut_rollouts = min(remaining, max_accelerated_rollout_steps)
-            with torch.no_grad():
-                prediction, _ = self._forward_model_tokenized(input_token=tokenized_tensor, rollouts=fut_rollouts)
-                prediction = prediction[:, :, -fut_rollouts:, :].to(tokenized_tensor)  # predicted token
-                # [bs, num_quantiles, num_predicted_token, output_patch_size]
-            prediction = self.tokenizer.output_transform(prediction, tokenizer_state)
-            prediction = prediction.flatten(start_dim=2)
+            prediction, fut_rollouts = self._forecast_single_step(context, max_context, min_context, fut_rollouts)
 
             predictions.append(prediction)
             remaining -= fut_rollouts
@@ -134,6 +124,33 @@ 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]:
+        if context.shape[-1] > max_context:
+            context = context[..., -max_context:]
+        if context.shape[-1] < min_context:
+            pad = torch.full(
+                (context.shape[0], min_context - context.shape[-1]),
+                fill_value=torch.nan,
+                device=context.device,
+                dtype=context.dtype,
+            )
+            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
+        # Shape: [bs, num_quantiles, num_predicted_token, output_patch_size]
+        prediction = self.tokenizer.output_transform(prediction, tokenizer_state)
+        prediction = prediction.flatten(start_dim=2)
+
+        return prediction, new_patch_count
+
     def _forward_model_tokenized(
         self,
         input_token: torch.Tensor,
@@ -165,21 +182,7 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
 
         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]
-
-        quantile_preds = self._forward_model(torch.cat((input_token, input_mask), dim=2))
+        quantile_preds, hidden_states = 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)
@@ -196,7 +199,7 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
 
         hidden_states = self.out_norm(hidden_states)
 
-        return self.output_patch_embedding(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

{tirex_mirror-2025.10.8.dist-info → tirex_mirror-2025.10.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.8
+Version: 2025.10.17
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -63,39 +63,37 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: LICENSE_MIRROR.txt
 License-File: NOTICE.txt
-Requires-Dist: torch==2.8.0
-Requires-Dist: einops==0.8.1
-Requires-Dist: huggingface-hub==0.34.4
-Requires-Dist: numpy==2.1.2
+Requires-Dist: torch
+Requires-Dist: einops
+Requires-Dist: huggingface-hub
+Requires-Dist: numpy
 Provides-Extra: cuda
-Requires-Dist: xlstm==2.0.4; extra == "cuda"
-Requires-Dist: ninja==1.13.0; extra == "cuda"
+Requires-Dist: xlstm; extra == "cuda"
+Requires-Dist: ninja; extra == "cuda"
 Provides-Extra: notebooks
-Requires-Dist: ipykernel==6.30.1; extra == "notebooks"
-Requires-Dist: matplotlib==3.10.6; extra == "notebooks"
-Requires-Dist: pandas==2.3.1; extra == "notebooks"
-Requires-Dist: python-dotenv==1.1.1; extra == "notebooks"
+Requires-Dist: ipykernel; extra == "notebooks"
+Requires-Dist: matplotlib; extra == "notebooks"
+Requires-Dist: pandas; extra == "notebooks"
+Requires-Dist: python-dotenv; extra == "notebooks"
 Provides-Extra: gluonts
-Requires-Dist: gluonts==0.16.2; extra == "gluonts"
-Requires-Dist: pandas==2.3.1; extra == "gluonts"
+Requires-Dist: gluonts; extra == "gluonts"
+Requires-Dist: pandas; extra == "gluonts"
 Provides-Extra: hfdataset
-Requires-Dist: datasets==3.6.0; extra == "hfdataset"
+Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
 Requires-Dist: fev>=0.6.0; extra == "test"
-Requires-Dist: pytest==8.4.2; extra == "test"
-Requires-Dist: requests==2.32.3; extra == "test"
+Requires-Dist: pytest; extra == "test"
 Provides-Extra: all
-Requires-Dist: xlstm==2.0.4; extra == "all"
-Requires-Dist: ninja==1.13.0; extra == "all"
-Requires-Dist: ipykernel==6.30.1; extra == "all"
-Requires-Dist: matplotlib==3.10.6; extra == "all"
-Requires-Dist: pandas==2.3.1; extra == "all"
-Requires-Dist: python-dotenv==1.1.1; extra == "all"
-Requires-Dist: gluonts==0.16.2; extra == "all"
-Requires-Dist: datasets==3.6.0; extra == "all"
-Requires-Dist: pytest==8.4.2; extra == "all"
-Requires-Dist: fev==0.6.0; extra == "all"
-Requires-Dist: requests==2.32.3; extra == "all"
+Requires-Dist: xlstm; extra == "all"
+Requires-Dist: ninja; extra == "all"
+Requires-Dist: ipykernel; extra == "all"
+Requires-Dist: matplotlib; extra == "all"
+Requires-Dist: pandas; extra == "all"
+Requires-Dist: python-dotenv; extra == "all"
+Requires-Dist: gluonts; extra == "all"
+Requires-Dist: datasets; extra == "all"
+Requires-Dist: pytest; extra == "all"
+Requires-Dist: fev>=0.6.0; extra == "all"
 Dynamic: license-file
 
 # tirex-mirror

tirex_mirror-2025.10.17.dist-info/RECORD

@@ -0,0 +1,21 @@
+tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
+tirex/base.py,sha256=JZuWzezuDenrlDekMaaPLBWzB9mw6vJ_qF2GWjou8ws,4229
+tirex/util.py,sha256=AQ8eM-7IpGeAqegVqflrjPQITzIEKltRcJOBbJmikcA,22758
+tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
+tirex/api_adapter/forecast.py,sha256=CpX9YfiQ1X6nvzODklCPzuKPlYidJDha8FnEU-6zr1Q,13919
+tirex/api_adapter/gluon.py,sha256=CbqL8ZgSRvxWAHK6TWqKdxUrMVWmGskRWOHNF84Lh1U,1819
+tirex/api_adapter/hf_data.py,sha256=TRyys2xKIGZS0Yhq2Eb61lWCMg5CWWn1yRlLIN1mU7o,1369
+tirex/api_adapter/standard_adapter.py,sha256=vdlxNs8mTUtPgK_5WMqYqNdMj8W44igqWsAgtggt_xk,2809
+tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
+tirex/models/patcher.py,sha256=EOXFkHsPkq0nuxRNLAbnrgJtcYq0IMC3YIg_16WArg4,3213
+tirex/models/tirex.py,sha256=R21JEa6dJDZ8epGgOJcNIqDoCt12uLlg31JnBPed0Do,9815
+tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
+tirex/models/slstm/cell.py,sha256=mZwR9qHMgRwsy4vLxAwd0SOEpq0Wo9MorPKBjHD9O9U,7259
+tirex/models/slstm/layer.py,sha256=hrDydQJIAHf5W0A0Rt0hXG4yKXrOSY-HPL0UbigR6Q8,2867
+tirex_mirror-2025.10.17.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
+tirex_mirror-2025.10.17.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
+tirex_mirror-2025.10.17.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
+tirex_mirror-2025.10.17.dist-info/METADATA,sha256=ULMjrGf-LNOyvpm5LsGauifIR-vrrVpj1pX8YolXeBg,11444
+tirex_mirror-2025.10.17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tirex_mirror-2025.10.17.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
+tirex_mirror-2025.10.17.dist-info/RECORD,,

tirex_mirror-2025.10.8.dist-info/RECORD

@@ -1,21 +0,0 @@
-tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
-tirex/base.py,sha256=JZuWzezuDenrlDekMaaPLBWzB9mw6vJ_qF2GWjou8ws,4229
-tirex/util.py,sha256=AQ8eM-7IpGeAqegVqflrjPQITzIEKltRcJOBbJmikcA,22758
-tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
-tirex/api_adapter/forecast.py,sha256=CpX9YfiQ1X6nvzODklCPzuKPlYidJDha8FnEU-6zr1Q,13919
-tirex/api_adapter/gluon.py,sha256=CbqL8ZgSRvxWAHK6TWqKdxUrMVWmGskRWOHNF84Lh1U,1819
-tirex/api_adapter/hf_data.py,sha256=TRyys2xKIGZS0Yhq2Eb61lWCMg5CWWn1yRlLIN1mU7o,1369
-tirex/api_adapter/standard_adapter.py,sha256=vdlxNs8mTUtPgK_5WMqYqNdMj8W44igqWsAgtggt_xk,2809
-tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
-tirex/models/patcher.py,sha256=EOXFkHsPkq0nuxRNLAbnrgJtcYq0IMC3YIg_16WArg4,3213
-tirex/models/tirex.py,sha256=JKNuCzTI6B9_yCbcmTf2UFjAQXulLNEmloqtAhKJKjQ,9830
-tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
-tirex/models/slstm/cell.py,sha256=JfCs1aUy9IHuz9RwExhUwiUtbg8WmbEg4upcO7hA5Rg,7229
-tirex/models/slstm/layer.py,sha256=93CAYuG-HmUpF7mBAQ-z1S1u2__W10EW5jPToR57qqM,2747
-tirex_mirror-2025.10.8.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
-tirex_mirror-2025.10.8.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
-tirex_mirror-2025.10.8.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
-tirex_mirror-2025.10.8.dist-info/METADATA,sha256=rhFTxXEJQofcI2KRgTyxq6YuGF4TkjwQ7HWo-XsDFUQ,11710
-tirex_mirror-2025.10.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tirex_mirror-2025.10.8.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
-tirex_mirror-2025.10.8.dist-info/RECORD,,