tirex-mirror 2025.10.1__py3-none-any.whl → 2025.10.2__py3-none-any.whl

tirex/base.py

@@ -84,6 +84,7 @@ def load_model(
     Args:
         path (str): Hugging Face path to the model (e.g. NX-AI/TiRex)
         device (str, optional): The device on which to load the model (e.g., "cuda:0", "cpu").
+        backend (torch | cuda): What backend to use, torch or the custom CUDA kernels. Defaults to cuda when xlstm is installed, else torch.
         hf_kwargs (dict, optional): Keyword arguments to pass to the Hugging Face Hub download method.
         ckp_kwargs (dict, optional): Keyword arguments to pass when loading the checkpoint.
 

tirex/models/slstm/block.py

@@ -21,14 +21,8 @@ class sLSTMBlock(nn.Module):
         self.ffn = FeedForward(config.embedding_dim, up_proj_dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x_slstm = self.norm_slstm(x)
-
-        x_slstm = self.slstm_layer(x_slstm, slstm_state=None)
-        x = x + x_slstm
-
-        x_ffn = self.norm_ffn(x)
-        x_ffn = self.ffn(x_ffn)
-        x = x + x_ffn
+        x = x + self.slstm_layer(self.norm_slstm(x), slstm_state=None)
+        x = x + self.ffn(self.norm_ffn(x))
         return x
 
 
@@ -41,8 +35,8 @@ class FeedForward(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = F.silu(self.proj_up_gate(x)) * self.proj_up(x)
-        y = self.proj_down(x)
-        return y
+        x = self.proj_down(x)
+        return x
 
 
 class RMSNorm(nn.Module):

tirex/models/slstm/cell.py

@@ -168,21 +168,24 @@ def slstm_forward_pointwise(
     states: torch.Tensor,  # dim [4, B, H]
 ) -> tuple[torch.Tensor, torch.Tensor]:
     raw = Wx + Ry + b
-    y, c, n, m = torch.unbind(states.view(4, states.shape[1], -1), dim=0)
 
     iraw, fraw, zraw, oraw = torch.unbind(raw.view(raw.shape[0], 4, -1), dim=1)
+    y, c, n, m = torch.unbind(states.view(4, states.shape[1], -1), dim=0)
+
     # with torch.no_grad():  # THE difference to maxg aka max_gradient (here max / max_static)
-    logfplusm = m + F.logsigmoid(fraw)
+    # Equations reference the xlstm paper on page 4: https://arxiv.org/pdf/2405.04517
+    logfplusm = m + F.logsigmoid(fraw)  # eq 15
     if torch.all(n == 0.0):
         mnew = iraw
     else:
-        mnew = torch.max(iraw, logfplusm)
-    ogate = torch.sigmoid(oraw)
-    igate = torch.minimum(torch.exp(iraw - mnew), torch.ones_like(iraw))
-    fgate = torch.minimum(torch.exp(logfplusm - mnew), torch.ones_like(iraw))
-    cnew = fgate * c + igate * torch.tanh(zraw)
-    nnew = fgate * n + igate
-    ynew = ogate * cnew / nnew
-
-    # shapes ([B,H], [B,H], [B,H]), ([B,H],[B,H],[B,H],[B,H])
-    return torch.stack((ynew, cnew, nnew, mnew), dim=0), torch.stack((igate, fgate, zraw, ogate), dim=0)
+        mnew = torch.max(iraw, logfplusm)  # eq 15
+    ogate = torch.sigmoid(oraw)  # eq 14
+    igate = torch.minimum(torch.exp(iraw - mnew), torch.ones_like(iraw))  # eq 16
+    fgate = torch.minimum(torch.exp(logfplusm - mnew), torch.ones_like(iraw))  # eq 17
+    zgate = torch.tanh(zraw)  # eq 11
+    cnew = fgate * c + igate * zgate  # eq 8
+    nnew = fgate * n + igate  # eq 9
+    hnew = ogate * cnew / nnew  # eq 10
+
+    # y (4, B, H), state (4, B, H)
+    return torch.stack((hnew, cnew, nnew, mnew), dim=0), torch.stack((igate, fgate, zraw, ogate), dim=0)

{tirex_mirror-2025.10.1.dist-info → tirex_mirror-2025.10.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.1
+Version: 2025.10.2
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -67,16 +67,17 @@ Requires-Dist: torch
 Requires-Dist: einops
 Requires-Dist: huggingface-hub
 Requires-Dist: numpy
-Requires-Dist: pandas
-Requires-Dist: tqdm
 Provides-Extra: cuda
 Requires-Dist: xlstm; extra == "cuda"
 Requires-Dist: ninja; extra == "cuda"
 Provides-Extra: notebooks
 Requires-Dist: ipykernel; extra == "notebooks"
 Requires-Dist: matplotlib; extra == "notebooks"
+Requires-Dist: pandas; extra == "notebooks"
+Requires-Dist: python-dotenv; extra == "notebooks"
 Provides-Extra: gluonts
 Requires-Dist: gluonts; extra == "gluonts"
+Requires-Dist: pandas; extra == "gluonts"
 Provides-Extra: hfdataset
 Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
@@ -87,6 +88,8 @@ Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"
 Requires-Dist: ipykernel; extra == "all"
 Requires-Dist: matplotlib; extra == "all"
+Requires-Dist: pandas; extra == "all"
+Requires-Dist: python-dotenv; extra == "all"
 Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
 Requires-Dist: pytest; extra == "all"

{tirex_mirror-2025.10.1.dist-info → tirex_mirror-2025.10.2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
-tirex/base.py,sha256=u_fcwaIKEzq9aAt3UWqH8QvaqXG7qEykLNaP_opY26M,3366
+tirex/base.py,sha256=fwyUTGL103kK5jgK5MoLSIHQcZb4lrox_D9fNbY1W1k,3507
 tirex/util.py,sha256=7DFVBXwGQA4niT9VhYbt8iKMBINJVW4LfwwpggFS0Us,469
 tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
 tirex/api_adapter/forecast.py,sha256=snv0sT1_1WzjkhP1YV-I7CMQmSChl93qFc3b6fwUAS0,8502
@@ -9,13 +9,13 @@ tirex/api_adapter/standard_adapter.py,sha256=bI3XGYlWQu5EDyhDZyYqOJMbwi5h1aovPQv
 tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
 tirex/models/patcher.py,sha256=EOXFkHsPkq0nuxRNLAbnrgJtcYq0IMC3YIg_16WArg4,3213
 tirex/models/tirex.py,sha256=Kglea86t_f3nXXHSjFgssxxrd1Qbwfr1eB_5gKfWYxM,9098
-tirex/models/slstm/block.py,sha256=DCOxmLQUb7HRO6wXTZMK4ICUI5LFpo7NC5a28oM-Vsc,2104
-tirex/models/slstm/cell.py,sha256=XWsn8I7HrUoMrUrfRCpl6Q88xbBz67bKEkdZ8gXE3hY,7444
+tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
+tirex/models/slstm/cell.py,sha256=ippaAPKI83j3_1l3pu9ks-iBGO641Elm1W4HsHgVu-c,7601
 tirex/models/slstm/layer.py,sha256=93CAYuG-HmUpF7mBAQ-z1S1u2__W10EW5jPToR57qqM,2747
-tirex_mirror-2025.10.1.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
-tirex_mirror-2025.10.1.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
-tirex_mirror-2025.10.1.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
-tirex_mirror-2025.10.1.dist-info/METADATA,sha256=HHaDkSlIQIHE4aIhG1EG3S9vK_tfuWT_HuLxy-6st6s,11265
-tirex_mirror-2025.10.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tirex_mirror-2025.10.1.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
-tirex_mirror-2025.10.1.dist-info/RECORD,,
+tirex_mirror-2025.10.2.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
+tirex_mirror-2025.10.2.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
+tirex_mirror-2025.10.2.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
+tirex_mirror-2025.10.2.dist-info/METADATA,sha256=Aq9VAU0pojVsrxwbfFLCmwmk1Gfl_Z_49G4yB-Z9eLY,11443
+tirex_mirror-2025.10.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tirex_mirror-2025.10.2.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
+tirex_mirror-2025.10.2.dist-info/RECORD,,