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

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

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

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

{tirex_mirror-2025.10.18 → tirex_mirror-2025.10.25}/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.18 → tirex_mirror-2025.10.25}/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.18 → tirex_mirror-2025.10.25/src/tirex_mirror.egg-info}/PKG-INFO

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

@@ -24,9 +24,9 @@ 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
 tests/test_slstm_torch_vs_cuda.py
 tests/test_standard_adapter.py
 tests/test_util_freq.py

tirex_mirror-2025.10.25/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!"

tirex_mirror-2025.10.18/tests/test_jupyterlab.py

@@ -1,52 +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
-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
-        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:
-        pytest.fail(f"❌ Could not connect to {cpu_url} (connection refused or server not running)")
-
-    except requests.exceptions.Timeout:
-        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}")