tirex-mirror 2025.11.13__tar.gz → 2025.11.15__tar.gz

{tirex_mirror-2025.11.13 → tirex_mirror-2025.11.15}/PKG-INFO

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

{tirex_mirror-2025.11.13 → tirex_mirror-2025.11.15}/pyproject.toml

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

{tirex_mirror-2025.11.13 → tirex_mirror-2025.11.15}/src/tirex/base.py

@@ -3,6 +3,7 @@
 
 import logging
 import os
+import warnings
 from abc import ABC, abstractmethod
 from typing import Literal, TypeVar
 
@@ -60,7 +61,7 @@ class PretrainedModel(ABC):
         if ckp_kwargs is None:
             ckp_kwargs = {}
         if device is None:
-            device = "cuda:0" if backend == "cuda" else "cpu"
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
         if os.path.exists(path):
             print("Loading weights from local directory")
             checkpoint_path = path
@@ -93,7 +94,7 @@ class PretrainedModel(ABC):
 def load_model(
     path: str,
     device: str | None = None,
-    backend: Literal["torch", "cuda"] | None = None,
+    backend: Literal["torch", "cuda"] = "torch",
     compile: bool = False,
     hf_kwargs=None,
     ckp_kwargs=None,
@@ -115,8 +116,22 @@ def load_model(
         model: ForecastModel = load_model("NX-AI/TiRex")
     """
 
-    if backend is None:
-        backend = "torch" if skip_cuda() or not xlstm_available() else "cuda"
+    if backend == "cuda" and not xlstm_available() and torch.cuda.is_available():
+        backend = "torch"
+        warnings.warn(
+            "Switching to 'torch' backend with device='gpu'. In oder to use a CUDA backend, please make sure that xlstm package is installed.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    if device is not None and device.startswith("cuda") and not torch.cuda.is_available():
+        raise ValueError(
+            "CUDA is not available! This could be because:\n"
+            "  - No GPU is present on this machine\n"
+            "  - GPU drivers are not installed or not functioning\n"
+            "  - PyTorch is installed without CUDA support (CPU-only version)\n"
+            "To resolve: use device='cpu' for CPU inference, or install CUDA-enabled PyTorch if a GPU is available."
+        )
 
     try:
         _, model_string = parse_hf_repo_id(path).split("/")

{tirex_mirror-2025.11.13 → tirex_mirror-2025.11.15}/src/tirex/models/slstm/cell.py

@@ -70,8 +70,8 @@ class sLSTMCell(nn.Module):
         if input.device.type != "cuda":
             warnings.warn(
                 f"You use TiRex with sLSTM CUDA kernels BUT DO NOT LOAD THE DEVICE ON A CUDA DEVICE (device type is {input.device.type})!"
-                "This is not supported and calls to the model will likely lead to an error if you dont move your model to a CUDA device!"
-                "If you want to run TiRex on CPU you need to disable sLSTM CUDA kernels but be aware of the downsides (see FAQ)"
+                "This is not supported and calls to the model will likely lead to an error if you don't move your model to a CUDA device!"
+                "If you want to run TiRex on CPU, please select backend='torch' and device='cpu'"
             )
 
         if not hasattr(self, "func"):

{tirex_mirror-2025.11.13 → tirex_mirror-2025.11.15}/src/tirex_mirror.egg-info/PKG-INFO

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

{tirex_mirror-2025.11.13 → tirex_mirror-2025.11.15}/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_load_model.py
 tests/test_patcher.py
 tests/test_slstm_torch_vs_cuda.py
 tests/test_standard_adapter.py

tirex_mirror-2025.11.15/tests/test_load_model.py

@@ -0,0 +1,67 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import warnings
+
+import pytest
+import torch
+
+from tirex import load_model
+
+
+def test_load_model_with_default_parameters():
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        model = load_model("NX-AI/TiRex")
+
+        assert model is not None
+        assert model.blocks[0].slstm_layer.slstm_cell.backend == "torch"
+
+        context = torch.randn(1, 64)
+        _, _ = model.forecast(context, prediction_length=32)
+        assert len(w) == 0  # no warnings check
+
+
+def test_load_model_with_cpu_device_no_warning():
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        model = load_model("NX-AI/TiRex", device="cpu")
+
+        assert model is not None
+        assert model.blocks[0].slstm_layer.slstm_cell.backend == "torch"
+        context = torch.randn(1, 64)
+        _, _ = model.forecast(context, prediction_length=32)
+        assert len(w) == 0  # no warnings check
+
+
+def test_load_model_cuda_device_without_cuda_raises_error():
+    if torch.cuda.is_available():
+        pytest.skip("CUDA is available, skipping no-CUDA test")
+
+    with pytest.raises(ValueError) as exc_info:
+        model = load_model("NX-AI/TiRex", device="cuda:0")
+
+    error_msg = str(exc_info.value)
+    assert "CUDA is not available" in error_msg
+    assert "No GPU is present" in error_msg
+    assert "PyTorch is installed without CUDA support (CPU-only version)" in error_msg
+    assert "To resolve: use device='cpu'" in error_msg
+
+
+def test_load_model_with_mps_device():
+    if not torch.backends.mps.is_available():
+        pytest.skip("MPS is not available, skipping MPS test")
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        model = load_model("NX-AI/TiRex", device="mps")
+
+        assert model is not None
+        assert model.blocks[0].slstm_layer.slstm_cell.backend == "torch"
+
+        device = next(model.parameters()).device
+        assert device.type == "mps"
+
+        context = torch.randn(1, 64).to("mps")
+        _, _ = model.forecast(context, prediction_length=32)
+        assert len(w) == 0  # no warnings check