tirex-mirror 2025.10.8__tar.gz → 2025.10.17__tar.gz

{tirex_mirror-2025.10.8/src/tirex_mirror.egg-info → tirex_mirror-2025.10.17}/PKG-INFO

@@ -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.8 → tirex_mirror-2025.10.17}/pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "tirex-mirror"
-version = "2025.10.08"
+version = "2025.10.17"
 description = "Unofficial mirror of NX-AI/tirex for packaging"
 readme = "README.md"
 requires-python = ">=3.11"
 classifiers = [ "Programming Language :: Python :: 3", "Operating System :: OS Independent",]
 keywords = [ "TiRex", "xLSTM", "Time Series", "Zero-shot", "Deep Learning",]
-dependencies = [ "torch==2.8.0", "einops==0.8.1", "huggingface-hub==0.34.4", "numpy==2.1.2",]
+dependencies = [ "torch", "einops", "huggingface-hub", "numpy",]
 [[project.authors]]
 name = "Arpad Rozsas"
 email = "rozsasarpi@gmail.com"
@@ -23,12 +23,12 @@ Repository = "https://github.com/rozsasarpi/tirex-mirror"
 Issues = "https://github.com/rozsasarpi/tirex-mirror/issues"
 
 [project.optional-dependencies]
-cuda = [ "xlstm==2.0.4", "ninja==1.13.0",]
-notebooks = [ "ipykernel==6.30.1", "matplotlib==3.10.6", "pandas==2.3.1", "python-dotenv==1.1.1",]
-gluonts = [ "gluonts==0.16.2", "pandas==2.3.1",]
-hfdataset = [ "datasets==3.6.0",]
-test = [ "fev>=0.6.0", "pytest==8.4.2", "requests==2.32.3",]
-all = [ "xlstm==2.0.4", "ninja==1.13.0", "ipykernel==6.30.1", "matplotlib==3.10.6", "pandas==2.3.1", "python-dotenv==1.1.1", "gluonts==0.16.2", "datasets==3.6.0", "pytest==8.4.2", "fev==0.6.0", "requests==2.32.3",]
+cuda = [ "xlstm", "ninja",]
+notebooks = [ "ipykernel", "matplotlib", "pandas", "python-dotenv",]
+gluonts = [ "gluonts", "pandas",]
+hfdataset = [ "datasets",]
+test = [ "fev>=0.6.0", "pytest",]
+all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "pandas", "python-dotenv", "gluonts", "datasets", "pytest", "fev>=0.6.0",]
 
 [tool.docformatter]
 diff = false

{tirex_mirror-2025.10.8 → tirex_mirror-2025.10.17}/src/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_mirror-2025.10.8 → tirex_mirror-2025.10.17}/src/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_mirror-2025.10.8 → tirex_mirror-2025.10.17}/src/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 → tirex_mirror-2025.10.17/src/tirex_mirror.egg-info}/PKG-INFO

@@ -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/src/tirex_mirror.egg-info/requires.txt

@@ -0,0 +1,37 @@
+torch
+einops
+huggingface-hub
+numpy
+
+[all]
+xlstm
+ninja
+ipykernel
+matplotlib
+pandas
+python-dotenv
+gluonts
+datasets
+pytest
+fev>=0.6.0
+
+[cuda]
+xlstm
+ninja
+
+[gluonts]
+gluonts
+pandas
+
+[hfdataset]
+datasets
+
+[notebooks]
+ipykernel
+matplotlib
+pandas
+python-dotenv
+
+[test]
+fev>=0.6.0
+pytest

{tirex_mirror-2025.10.8 → tirex_mirror-2025.10.17}/tests/test_jupyterlab.py

@@ -1,21 +1,45 @@
 # Copyright (c) NXAI GmbH.
 # This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
-import requests
+import logging
+import subprocess
+
 import pytest
+import requests
 
 cpu_url = "http://localhost:8889"
 # gpu_url = "http://localhost:8888" - will be added as soon as self-hosted gpu runner is available
 
+logger = logging.getLogger(__name__)
+
+
+def _docker_container_running() -> bool:
+    """Return True if any Docker container is currently running."""
+    try:
+        result = subprocess.run(
+            ["docker", "ps", "-q"],
+            check=True,
+            text=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+        )
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        return False
+
+    return bool(result.stdout.strip())
+
+
+pytestmark = pytest.mark.skipif(
+    not _docker_container_running(),
+    reason="requires Docker with a running container exposing JupyterLab",
+)
+
+
 def test_jupyterlab_running():
     """Check that the JupyterLab instance inside the container is reachable."""
     try:
         response = requests.get(cpu_url, timeout=5)  # timeout prevents hanging
-
-        print(f"✅ Connected to {cpu_url}")
-        print("Status Code:", response.status_code)
-
-        # Basic validation
+        logger.info("✅ Connected to %s, status code: %s", cpu_url, response.status_code)
         assert response.status_code in [200, 302], f"Unexpected status code: {response.status_code}"
 
     except requests.exceptions.ConnectionError:
@@ -25,4 +49,4 @@ def test_jupyterlab_running():
         pytest.fail(f"⏰ Connection to {cpu_url} timed out")
 
     except requests.exceptions.RequestException as e:
-        pytest.fail(f"⚠️ General error connecting to {cpu_url}: {e}")
+        pytest.fail(f"⚠️ General error connecting to {cpu_url}: {e}")

tirex_mirror-2025.10.8/src/tirex_mirror.egg-info/requires.txt

@@ -1,39 +0,0 @@
-torch==2.8.0
-einops==0.8.1
-huggingface-hub==0.34.4
-numpy==2.1.2
-
-[all]
-xlstm==2.0.4
-ninja==1.13.0
-ipykernel==6.30.1
-matplotlib==3.10.6
-pandas==2.3.1
-python-dotenv==1.1.1
-gluonts==0.16.2
-datasets==3.6.0
-pytest==8.4.2
-fev==0.6.0
-requests==2.32.3
-
-[cuda]
-xlstm==2.0.4
-ninja==1.13.0
-
-[gluonts]
-gluonts==0.16.2
-pandas==2.3.1
-
-[hfdataset]
-datasets==3.6.0
-
-[notebooks]
-ipykernel==6.30.1
-matplotlib==3.10.6
-pandas==2.3.1
-python-dotenv==1.1.1
-
-[test]
-fev>=0.6.0
-pytest==8.4.2
-requests==2.32.3