Docs/mkdocs site

.github/workflows/docs.yml

@@ -0,0 +1,38 @@
+name: Deploy Documentation
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - "docs/**"
+      - "src/tirex2/**"
+      - "mkdocs.yml"
+      - ".github/workflows/docs.yml"
+  workflow_dispatch:
+
+permissions:
+  contents: write
+
+jobs:
+  build-and-deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install documentation dependencies
+        run: pip install -r docs/requirements.txt
+
+      - name: Build documentation
+        run: mkdocs build --strict
+
+      - name: Deploy to GitHub Pages
+        uses: peaceiris/actions-gh-pages@v4
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          publish_dir: ./site
+          publish_branch: gh-pages

.github/workflows/test-inference.yaml

@@ -9,6 +9,21 @@ on:
 jobs:
   inference-tests:
     runs-on: ubuntu-latest
+    env:
+      # Local Broker
+      TEST_MQTT_BROKER_HOST: localhost
+      TEST_MQTT_BROKER_PORT: "1883"
+    services:
+      mqtt:
+        image: emqx/emqx:5.10.4
+        ports:
+          - 1883:1883
+        options: >-
+          --health-cmd "emqx ctl status"
+          --health-interval 5s
+          --health-timeout 5s
+          --health-retries 20
+          --health-start-period 10s
     steps:
       - uses: actions/checkout@v5
       - name: Setup Python
@@ -28,11 +43,4 @@ jobs:
 
       - name: Run inference tests
         working-directory: ./inference
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
-        run: |
-          if [ -z "$HF_TOKEN" ]; then
-            echo "HF_TOKEN secret is not configured; skipping gated inference tests."
-            exit 0
-          fi
-          pytest tests -s
+        run: pytest tests -s

.gitignore

@@ -11,6 +11,7 @@ output/
 *.csv
 build/
 dist/
+/site/
 
 .ipynb_checkpoints/
 .idea/

CONTRIBUTING.md

@@ -0,0 +1,58 @@
+# Contributing
+
+Development and documentation contribution guidelines.
+
+## Environment Setup
+
+- **Pixi** (recommended): install [Pixi](https://pixi.prefix.dev/latest/), then run `pixi install`.
+  Environments are defined in [`pyproject.toml`](pyproject.toml) under `tool.pixi.environments`
+  (e.g. `cuda128`, `cuda126`, `test-cu128`, `test-cu126`, `example`, `example-cu128`, `example-cu126`).
+- **pip**: create a virtual environment and install the package in editable mode:
+  `python -m venv .venv && source .venv/bin/activate && pip install -e ".[examples,fev,gluonts]"`
+- **Tooling**: run `pre-commit install` once, then `pre-commit run --all-files` and `pixi run test`
+  (or `pytest test/` in a pip environment) before opening a PR.
+
+## Workflow Overview
+
+1. Branch from `main` and keep changes focused (docs versus code versus tooling).
+2. Run pre-commit and tests locally before pushing.
+3. Open a PR with a clear summary, test notes, and any follow-up TODOs.
+4. Address CI feedback — red checks block review.
+
+Commit messages are linted by [`conventional-pre-commit`](.pre-commit-config.yaml) and must
+follow `type(scope): summary` with one of `chore`, `ci`, `docs`, `feat`, `fix`, `test`
+(a scope is required).
+
+## Documentation Specifics
+
+The documentation site lives under [`docs/`](docs/) and is built with
+[MkDocs](https://www.mkdocs.org/) + [Material for MkDocs](https://squidfunk.github.io/mkdocs-material/)
++ [mkdocstrings](https://mkdocstrings.github.io/).
+
+- Install docs dependencies: `pip install -r docs/requirements.txt`
+- Preview locally: `mkdocs serve`
+- Build (as CI does): `mkdocs build --strict`
+- The [API reference](docs/api/) is generated automatically from docstrings in
+  `src/tirex2/` via mkdocstrings — update the docstring, not the generated page, and add a
+  runnable usage example to any public function or class that doesn't already have one.
+- Add new guides under `docs/` and register them in the `nav` section of [`mkdocs.yml`](mkdocs.yml).
+
+## Commit & Review Etiquette
+
+- Avoid committing generated artifacts (e.g. `.pixi/`, `output/`, `model`, `*.csv`,
+  `*.egg-info`, `site/` — see [`.gitignore`](.gitignore)) unless they are intended changes.
+- Rebase (don't merge) when syncing from `main`.
+- Respond to every review comment; clarify disagreements rather than ignoring them.
+
+## Getting Help
+
+- Open a draft PR early for directional feedback.
+- Use GitHub Issues/Discussions for larger proposals.
+
+## NXAI Contributor License Agreement
+
+Read the full CLA for Individual Contributors here: [CLA](https://github.com/NX-AI/CLA/blob/main/CLA.md)
+
+### Contact
+
+If you have any question about the CLA, feel free to reach out to [contact@nx-ai.com](mailto:contact@nx-ai.com)

README.md

@@ -8,10 +8,13 @@
 [![Hugging Face](https://img.shields.io/badge/HuggingFace-TiRex--2-yellow?logo=huggingface)](https://huggingface.co/NX-AI/TiRex-2)
 [![PyPI](https://img.shields.io/pypi/v/tirex-2?color=blue)](https://pypi.org/project/tirex-2/)
 [![PyPI Downloads](https://static.pepy.tech/personalized-badge/tirex-2?period=total&units=INTERNATIONAL_SYSTEM&left_color=GREY&right_color=BLUE&left_text=downloads)](https://pepy.tech/projects/tirex-2)
+[![Docker](https://img.shields.io/badge/GHCR-tirex2--cpu%20%2F%20tirex2--gpu-2496ED?logo=docker&logoColor=white)](https://github.com/NX-AI/tirex-2/pkgs/container/tirex2-cpu)
 [![License](https://img.shields.io/badge/license-Apache--2.0-green)](LICENSE)
+[![docs](https://img.shields.io/website?url=https%3A%2F%2Fnx-ai.github.io%2Ftirex-2%2F&label=docs&up_message=online&up_color=green&down_message=offline&down_color=red)](https://nx-ai.github.io/tirex-2/)
 [![Tests](https://github.com/NX-AI/tirex-2/actions/workflows/test.yaml/badge.svg)](https://github.com/NX-AI/tirex-2/actions/workflows/test.yaml)
 [![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit)](https://github.com/pre-commit/pre-commit)
 [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NX-AI/tirex-2/blob/main/examples/getting_started.ipynb)
+[![TiRex-2 Demo](https://img.shields.io/badge/HuggingFace-TiRex--2%20Demo-yellow?logo=huggingface)](https://huggingface.co/spaces/NX-AI/TiRex-2-demo)
 
 </div>
 
@@ -120,22 +123,6 @@ fig.show()
 ```
 ![output of plot_multivariate function visualizing context and forecast of multivariate input](/resources/multivariate-prediction.png)
 
-### Optional FlexAttention for large multivariate batches
-
-The default dense attention backend avoids compilation overhead and is suitable for small batches. For large
-CUDA batches containing many independent multivariate series, opt into the block-sparse FlexAttention backend
-when loading the checkpoint:
-
-```python
-model = load_model("NX-AI/TiRex-2", device="cuda", use_flex_attention=True)
-forecasts = model.forecast(timeseries, prediction_length=64, batch_size=64)
-```
-
-FlexAttention compiles its CUDA kernel on first use and can be slower for small batches, so benchmark both the
-backend and `batch_size` on the target GPU. Leave `use_flex_attention` unset to preserve the checkpoint setting
-and package default, or pass `False` to force dense attention. The two CUDA kernels are numerically close but
-not bit-identical; re-evaluate forecast metrics when changing backends in reproducible benchmarks.
-
 
 
 ### Benchmarking

docs/about.md

@@ -0,0 +1,21 @@
+# About NXAI
+
+![NXAI](images/nxai-logo-light.svg#only-light){ width="160" }
+![NXAI](images/nxai-logo-dark.svg#only-dark){ width="160" }
+
+TiRex-2 is developed by [NXAI](https://nx-ai.com) — *"Minds for Machines."* NXAI builds AI
+that masters any cognitive task, developing proprietary xLSTM memory architectures designed
+to outperform Transformer architectures. TiRex-2 is part of their time-series product line,
+focused on deploying efficient, edge-optimized models for real-world industrial
+applications.
+
+- Website: [nx-ai.com](https://nx-ai.com)
+- GitHub: [github.com/NX-AI](https://github.com/NX-AI)
+- Hugging Face: [huggingface.co/NX-AI](https://huggingface.co/NX-AI)
+- LinkedIn: [linkedin.com/company/nxai](https://www.linkedin.com/company/nxai/)
+- Contact: [contact@nx-ai.com](mailto:contact@nx-ai.com)
+
+TiRex-2 builds on [xLSTM](https://arxiv.org/abs/2405.04517) and generalizes NXAI's original
+univariate forecasting model, [TiRex](https://github.com/NX-AI/tirex) (see the
+[Introduction](introduction.md)). For additional capabilities beyond this open-source
+release, see [TiRex-2 Pro](pro.md).

docs/api/demo.md

@@ -0,0 +1,8 @@
+# Demo utilities
+
+Used to build the runnable examples in the [Quickstart](../getting-started/quickstart.md)
+and [Covariates](../how-to/covariates.md) pages.
+
+::: tirex2.demo.Demo
+
+::: tirex2.demo.plot_demo_forecast

docs/api/forecast.md

@@ -0,0 +1,3 @@
+# Forecasting API
+
+::: tirex2.ForecastModel

docs/api/index.md

@@ -0,0 +1,14 @@
+# API reference
+
+This reference is generated from the docstrings in
+[`src/tirex2/`](https://github.com/NX-AI/tirex-2/tree/main/src/tirex2). It covers the public
+API surface exported from the top-level `tirex2` package.
+
+- [Forecasting](forecast.md) — `load_model`, `TimeseriesType`, `ForecastModel`.
+- [Demo](demo.md) — `Demo` and `plot_demo_forecast`, used in the
+  [Quickstart](../getting-started/quickstart.md).
+- [Plotting](plotting.md) — `plot_multivariate`, `plot_forecast`, `plot_covariate`.
+
+::: tirex2.load_model
+
+::: tirex2.TimeseriesType

docs/api/plotting.md

@@ -0,0 +1,10 @@
+# Plotting
+
+Requires either `matplotlib` or `plotly` to be installed, e.g. via
+`pip install "tirex-2[examples]"`.
+
+::: tirex2.plotting.plot_multivariate
+
+::: tirex2.plotting.plot_forecast
+
+::: tirex2.plotting.plot_covariate

docs/benchmarks.md

@@ -0,0 +1,59 @@
+# Benchmarks
+
+TiRex-2's reported results are reproducible against two public leaderboards. Full
+instructions live in the repository next to the benchmark code:
+
+- [GIFT-Eval](https://github.com/NX-AI/tirex-2/blob/main/examples/gifteval/README.md) —
+  reproduce results on the
+  [GIFT-Eval](https://huggingface.co/spaces/Salesforce/GIFT-Eval) leaderboard.
+- [fev-bench](https://github.com/NX-AI/tirex-2/blob/main/examples/fevbench/README.md) —
+  reproduce results on the [fev-bench](https://huggingface.co/spaces/autogluon/fev-bench)
+  leaderboard.
+
+## GIFT-Eval
+
+Download the dataset once:
+
+```bash
+pixi run -e example-cu128 huggingface-cli download Salesforce/GiftEval --repo-type=dataset --local-dir PATH_TO_SAVE
+```
+
+Run the benchmark, choosing a model type:
+
+```bash
+pixi run -e example-cu128 python examples/gifteval/run_gifteval.py </path/to/gifteval_storage> pretrained
+```
+
+- `pretrained` loads `NX-AI/TiRex-2-gifteval-pretrain`.
+- `zero-shot` loads `NX-AI/TiRex-2-gifteval-zs`.
+
+By default (`--eval-mode multivariate`) the native multivariate target is kept intact and
+scored jointly, which exercises TiRex-2's cross-variate path but is **not** directly
+comparable to the public GIFT-Eval leaderboard. Pass `--eval-mode univariate` to split every
+multivariate dataset into independent univariate channels, matching the official leaderboard
+protocol:
+
+```bash
+pixi run -e example-cu126 python examples/gifteval/run_gifteval.py \
+    </path/to/gifteval_storage> <ckpt_dir> --eval-mode univariate
+```
+
+An interactive notebook is also available: start `pixi run notebook` and open
+`examples/gifteval/gifteval.ipynb`.
+
+## fev-bench
+
+Optionally pre-download the data:
+
+```bash
+pixi run -e example-cu128 huggingface-cli download autogluon/fev_datasets --repo-type=dataset --local-dir </path/to/fevbench/store>
+```
+
+Run the benchmark — this always loads `NX-AI/TiRex-2-fevbench` from Hugging Face:
+
+```bash
+pixi run fevbench [/path/to/fevbench_storage] [--tasks examples/fevbench/tasks.yaml]
+```
+
+If the storage path is omitted, the dataset is downloaded at runtime and cached under
+`$HOME/.cache`.

docs/contributing.md

@@ -0,0 +1,7 @@
+# Contributing
+
+Development and documentation contribution guidelines live in
+[`CONTRIBUTING.md`](https://github.com/NX-AI/tirex-2/blob/main/CONTRIBUTING.md) at the
+repository root — covering environment setup (Pixi/pip), pre-commit and testing, how to
+preview and build this documentation site, commit/review etiquette, and the NXAI Contributor
+License Agreement.