[kotselidis]

Description

This PR adds TiRex-2 (NXAI) to AINode's built-in forecasting models, next to timer_xl, sundial, chronos2, moirai2, toto and moment. TiRex-2 is an xLSTM time series foundation model with about 83M parameters. It forecasts nine quantiles (0.1 to 0.9), up to 320 steps ahead, for univariate and multivariate series, and it can use past-only and future-known covariates. The model is registered as tirex2. AINode downloads its weights from the Hugging Face repository NX-AI/TiRex-2 at startup, as it does for the other built-in models.

Usage

SELECT * FROM FORECAST(
  model_id => 'tirex2',
  targets => (SELECT time, s FROM db.t WHERE device = 'd1' ORDER BY time DESC LIMIT 1024) ORDER BY time,
  output_length => 24)

Once the weights are downloaded, SHOW MODELS lists tirex2 with model type tirex2, category builtin and state active.

Example: the input is 1024 points of 50 + 10·sin(2πt/24) + 0.01·t. The request is a TForecastReq with outputLength = 24, sent to AINode's inference manager on CPU with the published checkpoint. The response is 24 values (the median forecast). The first six:

step	forecast	true value
1	51.636	51.580
2	50.628	50.591
3	50.294	50.260
4	50.641	50.611
5	51.603	51.620
6	53.277	53.219

The mean absolute error is 0.112 over 24 steps and 0.126 over 96 steps. If output_length is above the model's horizon, the request fails with a clear message: 1599: TiRex-2 forecasts at most 320 steps, but output_length is 400.

Model code

iotdb/ainode/core/model/tirex2/ contains the inference model from NX-AI/tirex-2 0.5.0 (Apache License 2.0), in the same way as chronos2/ contains the code from chronos-forecasting. The model code is unchanged except for these points:

• Imports use the iotdb.ainode.core.model.tirex2 package, and files carry the ASF header followed by NXAI's copyright notice.
• The recurrent kernels follow the device of their inputs. Upstream picks the kernels when the model is built. AINode builds the model on CPU and then moves it to the device it is loaded on, so the choice now happens at run time. CUDA inputs use the Triton sLSTM and mLSTM kernels when Triton is available; other inputs use the PyTorch kernels.
• Triton counts as available only if its kernel modules import. Triton compiles a kernel from its Python source. In a packaged AINode binary, Triton may be bundled without that source; the model then falls back to the PyTorch kernels instead of failing.
• The activation named in config.json must be a torch.nn module, such as torch.nn.SiLU. Upstream imports any dotted path from the config.
• Code that AINode does not need is left out: the high-level ForecastModel API, dataframe/GluonTS/fev adapters, plotting and the demo. The narwhals dependency goes with them.

TiRex-2's own loader replays the forward pass from CUDA graphs. This PR leaves CUDA graphs off (use_cuda_graphs = False), so every forecast runs eagerly.

As upstream does, TiRex-2 wraps its residual blocks, and FlexAttention where it is used, in torch.compile. This PR keeps that. PyTorch compiles these blocks on the first forecast, so the first request after loading is slower. On CPU, compilation needs a C++ compiler on the AINode host. This was tested with PyTorch 2.14.1 on CPU, but not in the packaged AINode binary.

AINode integration

• modeling_tirex2.TiRex2ForPrediction is the model, loaded through PyTorchModelHubMixin (hub_mixin_cls in ModelInfo). AINode already supports this loading path. The transformers PreTrainedModel path is not used, because the published model.safetensors stores one weight per gate. TiRex-2's load_state_dict hooks fuse those weights into the layout of its recurrent layers, and transformers' loader bypasses the hooks. The weights load with strict=True, so a checkpoint with a missing or extra weight is rejected.
• pipeline_tirex2.TiRex2Pipeline turns each input into one TiRex-2 series. The targets of an input, shape (target_count, input_length), are forecast jointly. A covariate that appears only in past_covariates becomes a past covariate. A covariate that also appears in future_covariates becomes a known covariate, built from its past and future values. All inputs of a request run in one batch. The pipeline returns the 0.5 quantile.
• output_length must be between 1 and 320. TiRex-2 has no autoregressive mode, so the pipeline rejects longer horizons instead of truncating them.

Dependencies

There are no new dependencies. The model needs PyTorch, NumPy, safetensors and huggingface_hub, which AINode already pins. On Linux, the CUDA wheels of PyTorch include Triton. TiRex-2 0.5.0 replaced its earlier flashrnn and mlstm_kernels dependencies, and their CUDA toolkit requirement, with these Triton kernels. pyproject.toml, ainode.spec and ainode.xml are unchanged. ainode.spec already collects every module under iotdb.ainode.core.

Licences

TiRex-2's code and its published weights are both under the Apache License 2.0 (code licence, weights licence). The weights are downloaded at run time and are not part of the distribution. This PR adds TiRex-2 to LICENSE, NOTICE and AINode's NOTICE-binary, with the copyright line from TiRex-2's NOTICE file.

NXAI's intended-use notice says the release is not intended for high-risk uses under the EU AI Act. The notice states that it does not change the licence.

Tests

• tests/test_tirex2.py (unittest, CPU only, no downloads) uses a two-block TiRex-2 with random weights. The tests:
  • save it in the published checkpoint's per-gate layout and check that it reloads with identical weights;
  • check that a checkpoint with a missing weight is rejected;
  • load it through AINode's built-in loader and the TiRex2Pipeline, and check the output shapes for 1-D and 2-D targets and that the result is the model's median;
  • check how covariates are assembled, that output_length above the horizon is rejected, and that the activation lookup is restricted.
• The tirex2 entry is added to AINodeTestUtils.BUILTIN_LTSM_MAP, so AINodeSharedClusterIT lists, loads and forecasts with it on CPU and GPU, like the other built-in foundation models.
• With the published checkpoint on CPU, the forecasts are bit-identical to the unmodified TiRex-2 0.5.0 code for a univariate series, a multivariate series with past and known covariates, and a batch of series of different lengths.

---

This PR has:

• been self-reviewed.
• added or updated version, license, or notice information
• added comments explaining the "why" and the intent of the code wherever would not be obvious
  for an unfamiliar reader.
• added unit tests or modified existing tests to cover new code paths, ensuring the threshold
  for code coverage.
• added integration tests.

---

Key changed/added classes (or packages if there are too many classes) in this PR

• iotdb/ainode/core/model/tirex2/ (new): the TiRex-2 model, TiRex2ForPrediction, TiRex2Pipeline
• iotdb/ainode/core/model/model_info.py: tirex2 entry in BUILTIN_HF_TRANSFORMERS_MODEL_MAP
• iotdb-core/ainode/tests/test_tirex2.py (new)
• integration-test/.../ainode/utils/AINodeTestUtils.java: tirex2 in BUILTIN_LTSM_MAP
• LICENSE, NOTICE, iotdb-core/ainode/src/assembly/resources/NOTICE-binary