tirex-mirror 2025.9.9__py3-none-any.whl → 2025.9.10__py3-none-any.whl

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/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/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.dist-info → tirex_mirror-2025.9.10.dist-info}/METADATA

@@ -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.dist-info → tirex_mirror-2025.9.10.dist-info}/RECORD

@@ -1,5 +1,5 @@
 tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
-tirex/base.py,sha256=ODUyhYFR33ZYffu7dxDwsb9m2IiZAnGHIXvA81crbjQ,3245
+tirex/base.py,sha256=u_fcwaIKEzq9aAt3UWqH8QvaqXG7qEykLNaP_opY26M,3366
 tirex/util.py,sha256=7DFVBXwGQA4niT9VhYbt8iKMBINJVW4LfwwpggFS0Us,469
 tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
 tirex/api_adapter/forecast.py,sha256=snv0sT1_1WzjkhP1YV-I7CMQmSChl93qFc3b6fwUAS0,8502
@@ -8,14 +8,14 @@ tirex/api_adapter/hf_data.py,sha256=T1eaxqC3OO9yOzIvw4sr55x6iA2AHKJTZd36rROM4fQ,
 tirex/api_adapter/standard_adapter.py,sha256=bI3XGYlWQu5EDyhDZyYqOJMbwi5h1aovPQvfHuWETJk,2618
 tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
 tirex/models/patcher.py,sha256=EOXFkHsPkq0nuxRNLAbnrgJtcYq0IMC3YIg_16WArg4,3213
-tirex/models/tirex.py,sha256=dclEckb6CmvESeX_LwT2kaCNTB7deTFovIOQUIFF5J8,9117
+tirex/models/tirex.py,sha256=Kglea86t_f3nXXHSjFgssxxrd1Qbwfr1eB_5gKfWYxM,9098
 tirex/models/slstm/block.py,sha256=DCOxmLQUb7HRO6wXTZMK4ICUI5LFpo7NC5a28oM-Vsc,2104
-tirex/models/slstm/cell.py,sha256=4_pQcXOOT16aEpKIi4A-yEnj4qKK6pFyFADD2nGPzGc,7366
+tirex/models/slstm/cell.py,sha256=XWsn8I7HrUoMrUrfRCpl6Q88xbBz67bKEkdZ8gXE3hY,7444
 tirex/models/slstm/layer.py,sha256=93CAYuG-HmUpF7mBAQ-z1S1u2__W10EW5jPToR57qqM,2747
-tirex_mirror-2025.9.9.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
-tirex_mirror-2025.9.9.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
-tirex_mirror-2025.9.9.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
-tirex_mirror-2025.9.9.dist-info/METADATA,sha256=u9C_cIb8FtaHUep1XrFTeI7UAsVRtNJt2VSQo7420Vo,11200
-tirex_mirror-2025.9.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tirex_mirror-2025.9.9.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
-tirex_mirror-2025.9.9.dist-info/RECORD,,
+tirex_mirror-2025.9.10.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
+tirex_mirror-2025.9.10.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
+tirex_mirror-2025.9.10.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
+tirex_mirror-2025.9.10.dist-info/METADATA,sha256=-hCSFjZcCbFNmc1rRJiVWXdKN7YSD6Uwi1QQkSmbOVc,11251
+tirex_mirror-2025.9.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tirex_mirror-2025.9.10.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
+tirex_mirror-2025.9.10.dist-info/RECORD,,