Skip to content

MLXPorts/tirex-2-mlx

 
 

Repository files navigation

TiRex emoji TiRex-2: Generalizing TiRex to Multivariate Data and Streaming

Paper Hugging Face PyPI PyPI Downloads License Tests pre-commit Open in Colab

This repository provides the pre-trained multivariate forecasting model TiRex-2 introduced in the paper TiRex-2: Generalizing TiRex to Multivariate Data and Streaming.

TiRex-2 Pro: This repository is our open-source release. Our pro version extends TiRex-2 with streaming, hardware-optimized inference (edge, embedded, and industrial PCs, among others), finetuning, and classification & regression support — see TiRex-2 Pro below or contact us at contact@nx-ai.com.

TiRex-2

TiRex-2 is a pretrained time series foundation model that forecasts one or many target variates directly from their history, optionally conditioned on past and future-known covariates. A single checkpoint serves both univariate and multivariate forecasting, built on a recurrent architecture designed for efficient streaming settings — all zero-shot, with no task-specific training or fine-tuning.

TiRex-2 generalizes our original univariate model, TiRex, to multivariate forecasting with past and future covariates.

Key facts

  • Zero-shot multivariate forecasting: TiRex-2 forecasts multiple target variates out of the box, without training or fine-tuning on your data.

  • Past and future-known covariates: TiRex-2 natively conditions on past covariates and future-known covariates, such as calendar features, holidays, promotions, or scheduled interventions.

  • Small active footprint: TiRex-2 activates 38.4M parameters in univariate mode and an additional 44.1M parameters for multivariate forecasting.

Installation

Via Pip

pip install tirex-2

Install with additional dependencies:

pip install "tirex-2[examples,fev,gluonts]"

The Python package installation is currently only tested on Linux and macOS. Docker usage is documented separately and includes Linux, macOS, and Windows Docker Desktop instructions.

Via Pixi

We use Pixi for our development and benchmarking environment to ensure that it is set up correctly. Run the following command to install it on your machine:

curl -fsSL https://pixi.sh/install.sh | sh

Access to Model Weights

TiRex-2's model weights are gated on Hugging Face. To be able to access them, either log in via the Hugging Face CLI, or generate yourself a Hugging Face access token (make sure to enable Read access to contents of all public gated repos you can access) and set it before loading the model:

CLI

# log in to Hugging Face and follow the prompts
huggingface-cli login

Access Token

# set access token as environment variable
import os
os.environ["HF_TOKEN"] = "<insert-hf-token>"

On Google Colab, you can instead store it as a Colab secret (via the key icon in the sidebar) named HF_TOKEN, which is picked up automatically without pasting the token into a cell.

Getting started

The most easy way for you to get started is by checking out our "Getting Started" notebook. Moreover, you can jump straight into testing out TiRex using Google Colab. If you have cloned this repository, you can also easily start the notebook via Pixi by running:

pixi run notebook

Note that for pixi, depending on your CUDA version and use-case, you may need to use another environment, e.g., -e example-cu128, that are defined in pyproject.toml under section tool.pixi.environments.

Minimal usage predicting a simple sine wave

import torch
from tirex2 import TimeseriesType, load_model
from tirex2.plotting import plot_multivariate  # requires matplotlib to be installed

# load model
model = load_model("NX-AI/TiRex-2", device="cpu")  # use `device="cuda"` if cuda is available

# generate data - target expects time series of shape (n_targets, context_length)
context = torch.sin(torch.arange(128).float() / 8)
ts = TimeseriesType(target=context.unsqueeze(0), past_covariates=None, future_covariates=None)

# perform forecast - each forecast is of shape (n_targets, 9 quantiles, prediction_length)
forecast = model.forecast([ts], prediction_length=32, output_type="numpy")[0]

# visualize result
fig = plot_multivariate(ts, forecast, engine="matplotlib")
fig.show()

output of plot_multivariate function visualizing context and forecast

Covariate example

This example originates from the "Getting Started" notebook, showing the value of additional covariates.

from tirex2 import load_model
from tirex2.demo import Demo, plot_demo_forecast

# load model
model = load_model("NX-AI/TiRex-2", device="cpu")  # use `device="cuda"` if cuda is available

# load data
demo = Demo.create_nonstationary_demo()
ts_univariate = demo.to_timeseries_type(include_covariates=False)
ts_multivariate = demo.to_timeseries_type(include_covariates=True)

# perform forecast - each forecast is of shape (n_targets, 9 quantiles, prediction_length)
forecasts = model.forecast(
    timeseries=[ts_univariate, ts_multivariate],
    prediction_length=demo.horizon,
    output_type="numpy",
)

# visualize result
fig = plot_demo_forecast(demo, *forecasts, engine="matplotlib")
fig.show()

output of plot_multivariate function visualizing context and forecast of multivariate input

Benchmarking

To reproduce our results for the GIFT-Eval and fev-bench leaderboards, follow the instructions in /examples/gifteval/ and /examples/fevbench/, respectively.

TiRex Docker image

For detailed instructions on building and running TiRex-2 in a Docker container, see the Docker README.

TiRex-2 Pro

TiRex-2 already provides state-of-the-art performance for zero-shot prediction, so you can use this open-source release without training on your own data.

Our pro version extends TiRex-2 with additional capabilities, including:

  • Streaming: incremental forecast updates as new observations arrive, without recomputing over the full history.
  • Speed: performance-optimized inference, including optimization for dedicated hardware such as edge, embedded, and industrial PC deployments.
  • Finetuning: models fine-tuned on your data or with different pretraining.
  • Classification & Regression: TiRex-2 adapted for classification and regression tasks.

If you are interested in any of these, please contact us at contact@nx-ai.com.

Cite

If you use TiRex in your research, please cite our work:

@misc{podest2026tirex2generalizingtirexmultivariate,
      title={TiRex-2: Generalizing TiRex to Multivariate Data and Streaming},
      author={Patrick Podest and Marco Pichler and Elias Bürger and Levente Zólyomi and Bernhard Voggenberger and Wilhelm Berghammer and Daniel Klotz and Sebastian Böck and Günter Klambauer and Sepp Hochreiter},
      year={2026},
      eprint={2607.01204},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2607.01204},
}

License

TiRex-2 is licensed under the Apache License 2.0.

About

TiRex-2 MLX: A Zero-Shot Timeseries Forecasting Model

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Packages

 
 
 

Contributors

Languages

  • Python 100.0%