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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.9.24
+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.9.24 → tirex_mirror-2025.10.2}/pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "tirex-mirror"
-version = "2025.09.24"
+version = "2025.10.02"
 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", "einops", "huggingface-hub", "numpy", "pandas", "tqdm",]
+dependencies = [ "torch", "einops", "huggingface-hub", "numpy",]
 [[project.authors]]
 name = "Arpad Rozsas"
 email = "rozsasarpi@gmail.com"
@@ -24,11 +24,11 @@ Issues = "https://github.com/rozsasarpi/tirex-mirror/issues"
 
 [project.optional-dependencies]
 cuda = [ "xlstm", "ninja",]
-notebooks = [ "ipykernel", "matplotlib",]
-gluonts = [ "gluonts",]
+notebooks = [ "ipykernel", "matplotlib", "pandas", "python-dotenv",]
+gluonts = [ "gluonts", "pandas",]
 hfdataset = [ "datasets",]
 test = [ "fev>=0.6.0", "pytest",]
-all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "gluonts", "datasets", "pytest", "fev>=0.6.0",]
+all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "pandas", "python-dotenv", "gluonts", "datasets", "pytest", "fev>=0.6.0",]
 
 [tool.docformatter]
 diff = false

{tirex_mirror-2025.9.24 → tirex_mirror-2025.10.2}/src/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_mirror-2025.9.24 → tirex_mirror-2025.10.2}/src/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_mirror-2025.9.24 → tirex_mirror-2025.10.2}/src/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.9.24 → tirex_mirror-2025.10.2/src/tirex_mirror.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.9.24
+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.9.24 → tirex_mirror-2025.10.2}/src/tirex_mirror.egg-info/requires.txt

@@ -2,14 +2,14 @@ torch
 einops
 huggingface-hub
 numpy
-pandas
-tqdm
 
 [all]
 xlstm
 ninja
 ipykernel
 matplotlib
+pandas
+python-dotenv
 gluonts
 datasets
 pytest
@@ -21,6 +21,7 @@ ninja
 
 [gluonts]
 gluonts
+pandas
 
 [hfdataset]
 datasets
@@ -28,6 +29,8 @@ datasets
 [notebooks]
 ipykernel
 matplotlib
+pandas
+python-dotenv
 
 [test]
 fev>=0.6.0

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

@@ -1,3 +1,6 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
 import math
 import os
 import time

{tirex_mirror-2025.9.24 → tirex_mirror-2025.10.2}/tests/test_forecast.py

@@ -1,3 +1,6 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
 from pathlib import Path
 
 import numpy as np

{tirex_mirror-2025.9.24 → tirex_mirror-2025.10.2}/tests/test_slstm_torch_vs_cuda.py

@@ -1,3 +1,6 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
 import copy
 
 import pytest