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.

• Forecasting — load_model, TimeseriesType, ForecastModel.
• 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,
) -> ForecastModel

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

Parameters:

Name	Type	Description	Default
ckpt_path	str or Path	Local directory holding model-config.yaml and model.ckpt. Values of the form hf://org/repo or org/repo are treated as Hugging Face model repo ids and downloaded with :func:huggingface_hub.snapshot_download.	'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	Override every variate mixer's checkpoint setting. True enables block-sparse FlexAttention, which can reduce the cost of large grouped multivariate batches on CUDA but adds first-call compilation overhead. False forces dense attention. Leave as None to preserve the checkpoint configuration and package defaults.	None
compile	bool	If True, torch.compile the recurrent layers: mLSTM always, and sLSTM only on CPU/MPS.	False

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

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,
) -> ForecastModel:
    """Load an inference-ready :class:`TiRex2` from a checkpoint directory or HF repo.

    Parameters
    ----------
    ckpt_path : str or pathlib.Path
        Local directory holding ``model-config.yaml`` and ``model.ckpt``. Values
        of the form ``hf://org/repo`` or ``org/repo`` are treated as Hugging Face
        model repo ids and downloaded with :func:`huggingface_hub.snapshot_download`.
    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
        Override every variate mixer's checkpoint setting. ``True`` enables
        block-sparse FlexAttention, which can reduce the cost of large grouped
        multivariate batches on CUDA but adds first-call compilation overhead.
        ``False`` forces dense attention. Leave as ``None`` to preserve the
        checkpoint configuration and package defaults.
    compile : bool
        If True, ``torch.compile`` the recurrent layers: mLSTM always, and sLSTM
        only on CPU/MPS.

    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 = ckpt_dir / CONFIG_FILENAME
    weights_file = ckpt_dir / CKPT_FILENAME
    if not config_file.is_file():
        raise FileNotFoundError(f"Expected model config at {config_file}")
    if not weights_file.is_file():
        raise FileNotFoundError(f"Expected model checkpoint at {weights_file}")

    with config_file.open() as f:
        config: dict[str, Any] = yaml.safe_load(f)

    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)

    checkpoint = torch.load(weights_file, map_location="cpu", weights_only=True)
    state_dict = checkpoint.get("state_dict", checkpoint) if isinstance(checkpoint, dict) else checkpoint
    model.load_state_dict(state_dict, strict=True)
    model.eval()
    if compile:
        # slSTM compiled on cpu/mps only.
        targets = (mLSTMLayer,) if device == "cuda" else (mLSTMLayer, _FlashRNNLayer)
        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

@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