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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.17
+Version: 2025.10.21
 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.17 → tirex_mirror-2025.10.21}/pyproject.toml

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

{tirex_mirror-2025.10.17 → tirex_mirror-2025.10.21}/src/tirex/base.py

@@ -76,7 +76,9 @@ class PretrainedModel(ABC):
         model = model.to(device)
 
         if compile and backend == "torch":
-            sLSTMCellTorch.slstm_forward = torch.compile(sLSTMCellTorch.slstm_forward, mode="max-autotune")
+            compiled_slstm_forward = torch.compile(sLSTMCellTorch.slstm_forward)
+            for block in model.blocks:
+                block.slstm_layer.slstm_cell._impl_forward_torch = compiled_slstm_forward
         return model
 
     @classmethod

{tirex_mirror-2025.10.17 → tirex_mirror-2025.10.21}/src/tirex/models/slstm/cell.py

@@ -38,6 +38,8 @@ class sLSTMCell(nn.Module):
 
         self._bias_ = nn.Parameter(torch.empty((config.num_heads * config.num_gates * config.head_dim), dtype=None))
 
+        self._impl_forward_torch = sLSTMCellTorch.slstm_forward
+
     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)
@@ -62,7 +64,7 @@ class sLSTMCell(nn.Module):
             .reshape(-1)
         )
 
-        return sLSTMCellTorch.slstm_forward(input, state, recurrent_kernel, bias)
+        return self._impl_forward_torch(input, state, recurrent_kernel, bias)
 
     def _impl_cuda(self, input: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
         if input.device.type != "cuda":

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.17
+Version: 2025.10.21
 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.17 → tirex_mirror-2025.10.21}/src/tirex_mirror.egg-info/SOURCES.txt

@@ -24,6 +24,7 @@ src/tirex_mirror.egg-info/dependency_links.txt
 src/tirex_mirror.egg-info/requires.txt
 src/tirex_mirror.egg-info/top_level.txt
 tests/test_chronos_zs.py
+tests/test_compile.py
 tests/test_forecast.py
 tests/test_forecast_adapter.py
 tests/test_jupyterlab.py

tirex_mirror-2025.10.21/tests/test_compile.py

@@ -0,0 +1,39 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import time
+
+import torch
+
+from tirex import load_model
+
+
+def measure_model_execution_time(model):
+    context = torch.randn((1, 128))
+
+    _, __ = model.forecast(context, prediction_length=32)  # warmup
+
+    start = time.time()
+    _, __ = model.forecast(context, prediction_length=32)
+    end = time.time()
+
+    return end - start
+
+
+def test_compileable():
+    model = load_model("NX-AI/TiRex", backend="torch", compile=True)
+
+    context = torch.randn((1, 128))
+    _, mean = model.forecast(context, prediction_length=32)
+
+    assert mean.shape == (1, 32)
+
+
+def test_compiled_faster():
+    model_uncompiled = load_model("NX-AI/TiRex", backend="torch", compile=False)
+    model_compiled = load_model("NX-AI/TiRex", backend="torch", compile=True)
+
+    time_uncompiled = measure_model_execution_time(model_uncompiled)
+    time_compiled = measure_model_execution_time(model_compiled)
+
+    assert time_compiled < time_uncompiled, "Compiled model has to be faster than uncompiled one!"