Skip to content

API reference

This reference is generated from the docstrings in src/tirex2/. It covers the public API surface exported from the top-level tirex2 package.

  • This overview — load_model and TimeseriesType.
  • Forecasting — ForecastModel and its forecasting methods.
  • Demo — Demo and plot_demo_forecast, used in the Quickstart.
  • Plotting — plot_multivariate, plot_forecast, plot_covariate.

tirex2.load_model 

load_model(
    ckpt_path: str | Path = "NX-AI/TiRex-2",
    device: str = "cuda",
    *,
    hf_kwargs: dict[str, Any] | None = None,
    use_flex_attention: bool | None = None,
    compile: bool = False,
    cuda_graphs: bool | None = None,
) -> ForecastModel

Load an inference-ready :class:TiRex2 from a checkpoint directory or HF repo.

Parameters:

Name Type Description Default
ckpt_path str or Path

Hugging Face repo id (org/repo or hf://org/repo), downloaded with huggingface_hub, or a local directory containing config.json and model.safetensors. Legacy model-config.yaml / model.ckpt checkpoints need tirex-2<=0.4.0.

 'NX-AI/TiRex-2'
device (cpu, cuda, mps)

Runtime device and recurrent-kernel family to use. This overrides any device/backend stored in the checkpoint config. "mps" runs on Apple Metal using the same pure-PyTorch (native) kernels as "cpu".

 "cpu"
hf_kwargs dict

Extra keyword arguments forwarded to snapshot_download for Hugging Face paths, e.g. {"revision": "main", "local_files_only": True}.

 None
use_flex_attention bool

If True, enable block-sparse FlexAttention in variate mixers; otherwise, use the default dense attention. FlexAttention may speed up large grouped batches on CUDA but adds compilation overhead.

 None
compile bool

If True, torch.compile the recurrent layers: mLSTM always, and sLSTM only on CPU/MPS.

 False
cuda_graphs bool

Replay the forward pass from a CUDA graph per input shape, which removes the kernel launch overhead. Defaults to True on CUDA.

 None

Returns:

Type Description
ForecastModel

The instantiated backbone (with the checkpoint weights loaded, set to evaluation mode) wrapped in a :class:ForecastModel that exposes the high-level forecast / forecast_gluon API.

Examples:

>>> import torch
>>> from tirex2 import TimeseriesType, load_model
>>> model = load_model("NX-AI/TiRex-2", device="cpu")
>>> ts = TimeseriesType(target=torch.randn(1, 128), past_covariates=None, future_covariates=None)
>>> forecast = model.forecast([ts], prediction_length=32, output_type="numpy")[0]
>>> forecast.shape
(1, 9, 32)
Source code in src/tirex2/base.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def load_model(
    ckpt_path: str | Path = "NX-AI/TiRex-2",
    device: str = "cuda",
    *,
    hf_kwargs: dict[str, Any] | None = None,
    use_flex_attention: bool | None = None,
    compile: bool = False,
    cuda_graphs: bool | None = None,
) -> ForecastModel:
    """Load an inference-ready :class:`TiRex2` from a checkpoint directory or HF repo.

    Parameters
    ----------
    ckpt_path : str or pathlib.Path
        Hugging Face repo id (``org/repo`` or ``hf://org/repo``), downloaded with
        ``huggingface_hub``, or a local directory containing ``config.json`` and
        ``model.safetensors``. Legacy ``model-config.yaml`` / ``model.ckpt``
        checkpoints need ``tirex-2<=0.4.0``.
    device : {"cpu", "cuda", "mps"}
        Runtime device and recurrent-kernel family to use. This overrides any
        device/backend stored in the checkpoint config. ``"mps"`` runs on Apple
        Metal using the same pure-PyTorch (native) kernels as ``"cpu"``.
    hf_kwargs : dict, optional
        Extra keyword arguments forwarded to ``snapshot_download`` for Hugging
        Face paths, e.g. ``{"revision": "main", "local_files_only": True}``.
    use_flex_attention : bool, optional
        If ``True``, enable block-sparse FlexAttention in variate mixers;
        otherwise, use the default dense attention. FlexAttention may speed up
        large grouped batches on CUDA but adds compilation overhead.
    compile : bool
        If ``True``, ``torch.compile`` the recurrent layers: mLSTM always, and sLSTM
        only on CPU/MPS.
    cuda_graphs : bool, optional
        Replay the forward pass from a CUDA graph per input shape, which removes
        the kernel launch overhead. Defaults to ``True`` on CUDA.

    Returns
    -------
    ForecastModel
        The instantiated backbone (with the checkpoint weights loaded, set to
        evaluation mode) wrapped in a :class:`ForecastModel` that exposes the
        high-level ``forecast`` / ``forecast_gluon`` API.

    Examples
    --------
    >>> import torch
    >>> from tirex2 import TimeseriesType, load_model
    >>> model = load_model("NX-AI/TiRex-2", device="cpu")
    >>> ts = TimeseriesType(target=torch.randn(1, 128), past_covariates=None, future_covariates=None)
    >>> forecast = model.forecast([ts], prediction_length=32, output_type="numpy")[0]
    >>> forecast.shape
    (1, 9, 32)
    """
    if device.startswith("cuda") and not torch.cuda.is_available():
        raise RuntimeError("Execution on CUDA was requested but is not available.")
    if device == "mps" and not torch.backends.mps.is_available():
        raise RuntimeError("Execution on MPS was requested but is not available.")

    ckpt_dir = _resolve_ckpt_dir(ckpt_path, hf_kwargs=hf_kwargs)
    config_file = _resolve_file(ckpt_dir, CONFIG_FILENAME, "config")
    weights_file = _resolve_file(ckpt_dir, WEIGHTS_FILENAME, "weights")

    config = json.loads(config_file.read_text())
    config["device"] = device
    if use_flex_attention is not None:
        for template in config["stack_config"]["templates"].values():
            template["variate_mixer"]["use_flex_attention"] = use_flex_attention

    model = TiRex2(**config)

    state_dict = load_safetensors(weights_file, device="cpu")
    model.load_state_dict(state_dict, strict=True)
    model.use_cuda_graphs = device.startswith("cuda") if cuda_graphs is None else cuda_graphs
    model.eval()
    if compile:
        # sLSTM compiled on cpu/mps only.
        targets = (mLSTMLayer,) if device.startswith("cuda") else (mLSTMLayer, sLSTMLayer)
        for module in model.modules():
            if isinstance(module, targets):
                module.compile()

    return ForecastModel(model)

tirex2.TimeseriesType dataclass

A single (possibly multivariate) series with optional covariates, as passed to :meth:~tirex2.api_adapter.forecast.ForecastModel.forecast.

Parameters:

Name Type Description Default
target Tensor

Target history, shape [V_t, T] (V_t target variates, context length T). A univariate series is still 2D, with V_t == 1.

 required
past_covariates Tensor or None

Covariates known only up to the current time, shape [V_p, T], matching the target's context length. None if there are no past covariates.

 required
future_covariates Tensor or None

Covariates known ahead of time for the whole forecast horizon, shape [V_f, >=T+H] (H is the requested prediction_length); extra trailing steps beyond T+H are ignored. None if there are no future covariates.

 required

Examples:

>>> import torch
>>> from tirex2 import TimeseriesType
>>> ts = TimeseriesType(
...     target=torch.randn(1, 128),
...     past_covariates=None,
...     future_covariates=None,
... )
>>> ts.past_length
128
Source code in src/tirex2/model/types.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@dataclass
class TimeseriesType:
    """A single (possibly multivariate) series with optional covariates, as passed to
    :meth:`~tirex2.api_adapter.forecast.ForecastModel.forecast`.

    Parameters
    ----------
    target : torch.Tensor
        Target history, shape ``[V_t, T]`` (``V_t`` target variates, context length ``T``).
        A univariate series is still 2D, with ``V_t == 1``.
    past_covariates : torch.Tensor or None
        Covariates known only up to the current time, shape ``[V_p, T]``, matching the
        target's context length. ``None`` if there are no past covariates.
    future_covariates : torch.Tensor or None
        Covariates known ahead of time for the whole forecast horizon, shape
        ``[V_f, >=T+H]`` (``H`` is the requested ``prediction_length``); extra trailing
        steps beyond ``T+H`` are ignored. ``None`` if there are no future covariates.

    Examples
    --------
    >>> import torch
    >>> from tirex2 import TimeseriesType
    >>> ts = TimeseriesType(
    ...     target=torch.randn(1, 128),
    ...     past_covariates=None,
    ...     future_covariates=None,
    ... )
    >>> ts.past_length
    128
    """

    target: torch.Tensor  # [V_t, T]
    past_covariates: torch.Tensor | None  # [V_p, T]
    future_covariates: torch.Tensor | None  # [V_f, >=T+H]; extra future steps are ignored

    @property
    def n_past_covariates(self) -> int:
        return 0 if self.past_covariates is None else len(self.past_covariates)

    @property
    def n_future_covariates(self) -> int:
        return 0 if self.future_covariates is None else len(self.future_covariates)

    @property
    def past_length(self) -> int:
        return self.target.shape[-1]

    @property
    def future_length(self) -> int:
        return 0 if self.future_covariates is None else self.future_covariates.shape[-1] - self.past_length