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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.9.9
+Version: 2025.9.24
 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
@@ -81,7 +80,8 @@ Requires-Dist: gluonts; extra == "gluonts"
 Provides-Extra: hfdataset
 Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
-Requires-Dist: fev; extra == "test"
+Requires-Dist: fev>=0.6.0; extra == "test"
+Requires-Dist: pytest; extra == "test"
 Provides-Extra: all
 Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"
@@ -89,7 +89,8 @@ Requires-Dist: ipykernel; extra == "all"
 Requires-Dist: matplotlib; extra == "all"
 Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
-Requires-Dist: fev; extra == "all"
+Requires-Dist: pytest; extra == "all"
+Requires-Dist: fev>=0.6.0; extra == "all"
 Dynamic: license-file
 
 # tirex-mirror

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

@@ -1,18 +1,18 @@
 [project]
 name = "tirex-mirror"
-version = "2025.09.09"
+version = "2025.09.24"
 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>=0.6.0", "pytest",]
+all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "gluonts", "datasets", "pytest", "fev>=0.6.0",]
 
 [tool.docformatter]
 diff = false

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.24}/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.24}/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.24}/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.24/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.24
 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
@@ -81,7 +80,8 @@ Requires-Dist: gluonts; extra == "gluonts"
 Provides-Extra: hfdataset
 Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
-Requires-Dist: fev; extra == "test"
+Requires-Dist: fev>=0.6.0; extra == "test"
+Requires-Dist: pytest; extra == "test"
 Provides-Extra: all
 Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"
@@ -89,7 +89,8 @@ Requires-Dist: ipykernel; extra == "all"
 Requires-Dist: matplotlib; extra == "all"
 Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
-Requires-Dist: fev; extra == "all"
+Requires-Dist: pytest; extra == "all"
+Requires-Dist: fev>=0.6.0; extra == "all"
 Dynamic: license-file
 
 # tirex-mirror

{tirex_mirror-2025.9.9 → tirex_mirror-2025.9.24}/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.24}/src/tirex_mirror.egg-info/requires.txt

@@ -1,5 +1,4 @@
 torch
-torchvision
 einops
 huggingface-hub
 numpy
@@ -13,7 +12,8 @@ ipykernel
 matplotlib
 gluonts
 datasets
-fev
+pytest
+fev>=0.6.0
 
 [cuda]
 xlstm
@@ -30,4 +30,5 @@ ipykernel
 matplotlib
 
 [test]
-fev
+fev>=0.6.0
+pytest

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

@@ -14,22 +14,30 @@ def geometric_mean(s):
 
 
 def eval_task(model, task):
-    past_data, _ = task.get_input_data(trust_remote_code=True)
-    quantile_levels = task.quantile_levels
-    past_data = past_data.with_format("torch").cast_column(
-        task.target_column, datasets.Sequence(datasets.Value("float32"))
-    )[task.target_column]
-    loaded_data = [t for t in past_data]
+    inference_time = 0.0
+    predictions_per_window = []
+    for window in task.iter_windows(trust_remote_code=True):
+        past_data, _ = fev.convert_input_data(window, adapter="datasets", as_univariate=True)
+        past_data = past_data.with_format("torch").cast_column("target", datasets.Sequence(datasets.Value("float32")))
+        loaded_targets = [t for t in past_data["target"]]
+
+        start_time = time.monotonic()
+        quantiles, means = model.forecast(
+            loaded_targets, quantile_levels=task.quantile_levels, prediction_length=task.horizon
+        )
+        inference_time += time.monotonic() - start_time
+
+        predictions_dict = {"predictions": means}
+        for idx, level in enumerate(task.quantile_levels):
+            predictions_dict[str(level)] = quantiles[:, :, idx]
 
-    start_time = time.monotonic()
-    quantiles, means = model.forecast(loaded_data, quantile_levels=quantile_levels, prediction_length=task.horizon)
-    inference_time = time.monotonic() - start_time
-    predictions_dict = {"predictions": means}
-    for idx, level in enumerate(quantile_levels):
-        predictions_dict[str(level)] = quantiles[:, :, idx]  # [num_items, horizon]
+        predictions_per_window.append(
+            fev.combine_univariate_predictions_to_multivariate(
+                datasets.Dataset.from_dict(predictions_dict), target_columns=task.target_columns
+            )
+        )
 
-    predictions = datasets.Dataset.from_dict(predictions_dict)
-    return predictions, inference_time
+    return predictions_per_window, inference_time
 
 
 @pytest.fixture
@@ -57,7 +65,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.24/tests/test_forecast.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from tirex import ForecastModel, load_model
+
+
+def load_tensor_from_txt_file(path):
+    base_path = Path(__file__).parent.resolve() / "data"
+    return torch.from_numpy(np.genfromtxt(base_path / path, dtype=np.float32))
+
+
+def load_tensor_from_pt_file(path):
+    base_path = Path(__file__).parent.resolve() / "data"
+    return torch.load(base_path / path)
+
+
+@pytest.fixture
+def tirex_model() -> ForecastModel:
+    return load_model("NX-AI/TiRex")
+
+
+def test_forecast_air_traffic(tirex_model):
+    context = load_tensor_from_txt_file("air_passengers.csv")[:-12]
+
+    quantiles, mean = tirex_model.forecast(context, prediction_length=24)
+
+    ref_mean = load_tensor_from_txt_file("air_passengers_forecast_ref.csv").unsqueeze(0)
+    ref_quantiles = load_tensor_from_pt_file("air_passengers_quantiles_ref.pt")
+
+    # default rtol & atol for bfloat16
+    torch.testing.assert_close(mean, ref_mean, rtol=1.6e-2, atol=1e-5)
+    torch.testing.assert_close(quantiles, ref_quantiles, rtol=1.6e-2, atol=1e-5)
+
+
+def test_forecast_seattle_5T(tirex_model):
+    context = load_tensor_from_txt_file("loop_seattle_5T.csv")[:-512]
+
+    quantiles, mean = tirex_model.forecast(context, prediction_length=768)
+
+    ref_mean = load_tensor_from_txt_file("loop_seattle_5T_forecast_ref.csv").unsqueeze(0)
+    ref_quantiles = load_tensor_from_pt_file("loop_seattle_5T_quantiles_ref.pt")
+
+    # default rtol & atol for bfloat16
+    torch.testing.assert_close(mean, ref_mean, rtol=1.6e-2, atol=1e-5)
+    torch.testing.assert_close(quantiles, ref_quantiles, rtol=1.6e-2, atol=1e-5)

tirex_mirror-2025.9.24/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)

tirex_mirror-2025.9.9/tests/test_forecast.py

@@ -1,37 +0,0 @@
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-from tirex import ForecastModel, load_model
-
-
-def load_tensor_from_file(path):
-    base_path = Path(__file__).parent.resolve() / "data"
-    return torch.from_numpy(np.genfromtxt(base_path / path, dtype=np.float32))
-
-
-@pytest.fixture
-def tirex_model() -> ForecastModel:
-    return load_model("NX-AI/TiRex")
-
-
-def test_forecast_air_traffic(tirex_model):
-    context = load_tensor_from_file("air_passengers.csv")[:-12]
-
-    quantiles, mean = tirex_model.forecast(context, prediction_length=24)
-
-    ref_data = load_tensor_from_file("air_passengers_forecast_ref.csv").unsqueeze(0)
-
-    torch.testing.assert_close(mean, ref_data, rtol=1.6e-2, atol=1e-5)  # default rtol & atol for bfloat16
-
-
-def test_forecast_seattle_5T(tirex_model):
-    context = load_tensor_from_file("loop_seattle_5T.csv")[:-512]
-
-    quantiles, mean = tirex_model.forecast(context, prediction_length=768)
-
-    ref_data = load_tensor_from_file("loop_seattle_5T_forecast_ref.csv").unsqueeze(0)
-
-    torch.testing.assert_close(mean, ref_data, rtol=1.6e-2, atol=1e-5)  # default rtol & atol for bfloat16