tirex-mirror 2025.9.9__tar.gz → 2025.9.10__tar.gz

{tirex_mirror-2025.9.9/src/tirex_mirror.egg-info → tirex_mirror-2025.9.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.9.9
+Version: 2025.9.10
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -64,7 +64,6 @@ License-File: LICENSE
 License-File: LICENSE_MIRROR.txt
 License-File: NOTICE.txt
 Requires-Dist: torch
-Requires-Dist: torchvision
 Requires-Dist: einops
 Requires-Dist: huggingface-hub
 Requires-Dist: numpy
@@ -82,6 +81,7 @@ Provides-Extra: hfdataset
 Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
 Requires-Dist: fev; extra == "test"
+Requires-Dist: pytest; extra == "test"
 Provides-Extra: all
 Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"
@@ -89,6 +89,7 @@ Requires-Dist: ipykernel; extra == "all"
 Requires-Dist: matplotlib; extra == "all"
 Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
+Requires-Dist: pytest; extra == "all"
 Requires-Dist: fev; extra == "all"
 Dynamic: license-file
 

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.10}/pyproject.toml

@@ -1,18 +1,18 @@
 [project]
 name = "tirex-mirror"
-version = "2025.09.09"
+version = "2025.09.10"
 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", "torchvision", "einops", "huggingface-hub", "numpy", "pandas", "tqdm",]
+dependencies = [ "torch", "einops", "huggingface-hub", "numpy", "pandas", "tqdm",]
 [[project.authors]]
 name = "Arpad Rozsas"
 email = "rozsasarpi@gmail.com"
 
 [build-system]
-requires = [ "setuptools>=42", "wheel",]
+requires = [ "setuptools>=77.0.3", "wheel",]
 build-backend = "setuptools.build_meta"
 
 [project.license]
@@ -27,8 +27,8 @@ cuda = [ "xlstm", "ninja",]
 notebooks = [ "ipykernel", "matplotlib",]
 gluonts = [ "gluonts",]
 hfdataset = [ "datasets",]
-test = [ "fev",]
-all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "gluonts", "datasets", "fev",]
+test = [ "fev", "pytest",]
+all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "gluonts", "datasets", "pytest", "fev",]
 
 [tool.docformatter]
 diff = false

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.10}/src/tirex/base.py

@@ -15,6 +15,15 @@ def skip_cuda():
     return os.getenv("TIREX_NO_CUDA", "False").lower() in ("true", "1", "t")
 
 
+def xlstm_available():
+    try:
+        from xlstm.blocks.slstm.cell import sLSTMCellConfig, sLSTMCellFuncGenerator
+
+        return True
+    except ModuleNotFoundError:
+        return False
+
+
 def parse_hf_repo_id(path):
     parts = path.split("/")
     return "/".join(parts[0:2])
@@ -86,8 +95,7 @@ def load_model(
     """
 
     if backend is None:
-        backend = "torch" if skip_cuda() else "cuda"
-    assert backend in ["torch", "cuda"], f"Backend can either be torch or cuda, not {backend}!"
+        backend = "torch" if skip_cuda() or not xlstm_available() else "cuda"
 
     try:
         _, model_id = parse_hf_repo_id(path).split("/")

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.10}/src/tirex/models/slstm/cell.py

@@ -3,6 +3,7 @@
 
 import warnings
 from dataclasses import asdict, dataclass
+from typing import Literal
 
 import torch
 import torch.nn as nn
@@ -15,10 +16,8 @@ from tirex.util import dataclass_from_dict
 class sLSTMBlockConfig:
     embedding_dim: int
     num_heads: int
-    num_blocks: int
     ffn_proj_factor: float = 2.6667
-
-    num_states: int = 4  # this is for the sLSTM, a standard LSTM  has 2
+    num_states: int = 4
     num_gates: int = 4
 
     @property
@@ -27,8 +26,9 @@ class sLSTMBlockConfig:
 
 
 class sLSTMCell(nn.Module):
-    def __init__(self, config: sLSTMBlockConfig, backend: str):
+    def __init__(self, config: sLSTMBlockConfig, backend: Literal["torch", "cuda"]):
         super().__init__()
+        assert backend in ["torch", "cuda"], f"Backend can either be torch or cuda, not {backend}!"
         self.config = config
         self.backend = backend
 

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.10}/src/tirex/models/tirex.py

@@ -33,10 +33,10 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
         super().__init__()
         self.config = TiRexZeroConfig(**model_config, train_ctx_len=train_ctx_len, nan_mask_value=0)
         assert self.config.input_patch_size == self.config.output_patch_size
-        self.backend = backend
 
         self.tokenizer = PatchedUniTokenizer(patch_size=self.config.input_patch_size)
 
+        num_blocks = self.config.block_kwargs["num_blocks"]
         block_config = dataclass_from_dict(sLSTMBlockConfig, self.config.block_kwargs)
         self.input_patch_embedding = ResidualBlock(
             in_dim=self.config.input_patch_size * 2,
@@ -44,9 +44,7 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
             out_dim=block_config.embedding_dim,
         )
 
-        self.blocks = nn.ModuleList(
-            [sLSTMBlock(block_config, backend=self.backend) for i in range(block_config.num_blocks)]
-        )
+        self.blocks = nn.ModuleList([sLSTMBlock(block_config, backend) for i in range(num_blocks)])
 
         self.out_norm = RMSNorm(block_config.embedding_dim)
 

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.10/src/tirex_mirror.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.9.9
+Version: 2025.9.10
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -64,7 +64,6 @@ License-File: LICENSE
 License-File: LICENSE_MIRROR.txt
 License-File: NOTICE.txt
 Requires-Dist: torch
-Requires-Dist: torchvision
 Requires-Dist: einops
 Requires-Dist: huggingface-hub
 Requires-Dist: numpy
@@ -82,6 +81,7 @@ Provides-Extra: hfdataset
 Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
 Requires-Dist: fev; extra == "test"
+Requires-Dist: pytest; extra == "test"
 Provides-Extra: all
 Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"
@@ -89,6 +89,7 @@ Requires-Dist: ipykernel; extra == "all"
 Requires-Dist: matplotlib; extra == "all"
 Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
+Requires-Dist: pytest; extra == "all"
 Requires-Dist: fev; extra == "all"
 Dynamic: license-file
 

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.10}/src/tirex_mirror.egg-info/SOURCES.txt

@@ -26,4 +26,5 @@ src/tirex_mirror.egg-info/top_level.txt
 tests/test_chronos_zs.py
 tests/test_forecast.py
 tests/test_forecast_adapter.py
+tests/test_slstm_torch_vs_cuda.py
 tests/test_standard_adapter.py

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.10}/src/tirex_mirror.egg-info/requires.txt

@@ -1,5 +1,4 @@
 torch
-torchvision
 einops
 huggingface-hub
 numpy
@@ -13,6 +12,7 @@ ipykernel
 matplotlib
 gluonts
 datasets
+pytest
 fev
 
 [cuda]
@@ -31,3 +31,4 @@ matplotlib
 
 [test]
 fev
+pytest

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.10}/tests/test_chronos_zs.py

@@ -57,7 +57,10 @@ def test_chronos_single(tirex_model, benchmark):
     assert evaluation_summary["MASE"] < 0.99, "MASE on the electricity task needs to be less than 0.99"
 
 
-@pytest.mark.skipif(os.getenv("CI"), reason="Skip full chromos benchmarking in the CI")
+@pytest.mark.skipif(
+    os.getenv("CI") is not None and os.getenv("CI_RUN_BENCHMARKS") is None,
+    reason="Skip Chronos benchmarks in CI",
+)
 def test_chronos_all(tirex_model, benchmark):
     tasks_wql = []
     tasks_mase = []

tirex_mirror-2025.9.10/tests/test_slstm_torch_vs_cuda.py

@@ -0,0 +1,82 @@
+import copy
+
+import pytest
+import torch
+
+from tirex.models.slstm.cell import sLSTMBlockConfig, sLSTMCell
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="This test needs CUDA.")
+
+
+@pytest.mark.parametrize("with_in_state", [True, False])
+def test_with_in_state(with_in_state):
+    run_slstm_torch_vs_cuda(with_in_state=with_in_state)
+
+
+@pytest.mark.parametrize("sequence_length", [1, 2, 4])
+def test_sequence_length(sequence_length):
+    run_slstm_torch_vs_cuda(sequence_length=sequence_length)
+
+
+@pytest.mark.parametrize("batch_size", [1, 2, 4])
+def test_batch_size(batch_size):
+    run_slstm_torch_vs_cuda(batch_size=batch_size)
+
+
+@pytest.mark.parametrize("num_heads", [4, 1])
+def test_num_heads(num_heads):
+    run_slstm_torch_vs_cuda(num_heads=num_heads, with_in_state=True, atol=5e-5)
+
+
+@pytest.mark.parametrize("hidden_size", [64, 8])
+def test_hidden_size(hidden_size):
+    run_slstm_torch_vs_cuda(hidden_size=hidden_size, with_in_state=True)
+
+
+def test_complex():
+    run_slstm_torch_vs_cuda(
+        hidden_size=128, batch_size=2, sequence_length=8, num_heads=4, with_in_state=True, atol=1e-5
+    )
+
+
+def test_long_sequence():
+    run_slstm_torch_vs_cuda(sequence_length=128, atol=1e-5)
+
+
+def set_seed(seed):
+    torch.use_deterministic_algorithms(True)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def run_slstm_torch_vs_cuda(
+    batch_size=1, sequence_length=1, with_in_state=False, num_heads=4, hidden_size=64, rtol=1.3e-6, atol=1e-6
+):
+    device_cuda = "cuda"
+    config = sLSTMBlockConfig(embedding_dim=hidden_size, num_heads=num_heads)
+
+    set_seed(42)
+    recurrent_kernel_weight = torch.randn(
+        (config.num_heads, config.head_dim, config.num_gates * config.head_dim), dtype=torch.bfloat16
+    )
+    bias_weight = torch.randn((config.num_heads * config.num_gates * config.head_dim), dtype=torch.bfloat16)
+
+    cell_torch = sLSTMCell(copy.deepcopy(config), backend="torch")
+    cell_torch._recurrent_kernel_.data = recurrent_kernel_weight
+    cell_torch._bias_.data = bias_weight
+
+    cell_cuda = sLSTMCell(copy.deepcopy(config), backend="cuda").to(device_cuda)
+    cell_cuda._recurrent_kernel_.data = recurrent_kernel_weight.to(device_cuda)
+    cell_cuda._bias_.data = bias_weight.to(device_cuda)
+
+    set_seed(42)
+    current_input = torch.randn((batch_size, sequence_length, 4 * config.embedding_dim))
+    state = torch.randn((4, batch_size, hidden_size)) if with_in_state else None
+
+    output_torch, state_torch = cell_torch.forward(current_input, state)
+    output_cuda, state_cuda = cell_cuda.forward(
+        current_input.to(device_cuda), state.to(device_cuda) if state is not None else state
+    )
+
+    torch.testing.assert_close(output_torch, output_cuda.cpu(), rtol=rtol, atol=atol)
+    torch.testing.assert_close(state_torch, state_cuda.cpu(), rtol=rtol, atol=atol)