feat: add tirex-2 model

.github/workflows/ci.yaml

@@ -161,4 +161,4 @@ jobs:
         uses: astral-sh/setup-uv@v6
 
       - name: Run pre-commit
-        run: uvx pre-commit run --hook-stage manual --all-files
+        run: uv run pre-commit run --all-files

docs/changelogs/index.md

@@ -2,6 +2,7 @@
 
 Welcome to the TimeCopilot Changelog. Here, you will find a comprehensive list of all the changes, updates, and improvements made to the TimeCopilot project. This section is designed to keep you informed about the latest features, bug fixes, and enhancements as we continue to develop and refine the TimeCopilot experience. Stay tuned for regular updates and feel free to explore the details of each release below.
 
+- [v0.0.30](v0.0.30.md)
 - [v0.0.29](v0.0.29.md)
 - [v0.0.28](v0.0.28.md)
 - [v0.0.27](v0.0.27.md)

docs/changelogs/v0.0.30.md

@@ -0,0 +1,30 @@
+### Features
+
+* **TiRex-2 foundation model**: Extended the [`TiRex`](https://github.com/NX-AI/tirex) integration to support [TiRex 2.0](https://github.com/NX-AI/tirex-2) alongside TiRex 1.0. Pass `repo_id="NX-AI/TiRex-2"` to use the new model through the same `unique_id` / `ds` / `y` API. TiRex 2.0 runs on CPU, CUDA, and Apple MPS (native PyTorch kernels on CPU/MPS; Triton kernels on CUDA). Added `timecopilot-tirex2>=0.1.0`, `mlstm-kernels>=2.0.4`, a [`tirex-family` example notebook](../examples/tirex-family.ipynb), and integration tests.
+
+    ```python
+    import pandas as pd
+    from timecopilot.models.foundation.tirex import TiRex
+
+    df = pd.read_csv(
+        "https://timecopilot.s3.amazonaws.com/public/data/air_passengers.csv",
+        parse_dates=["ds"],
+    )
+
+    model = TiRex(repo_id="NX-AI/TiRex-2", alias="TiRex-2")
+    fcst_df = model.forecast(df, h=12)
+    ```
+
+### Fixes
+
+* **T0 constructor config**: Pass `config` keyword arguments through to `T0` construction so custom model settings are applied correctly. See [#364](https://github.com/TimeCopilot/timecopilot/pull/364).
+
+* **Talk to Us booking link**: Updated the Talk to Us booking link. See [#363](https://github.com/TimeCopilot/timecopilot/pull/363).
+
+## New Contributors
+
+* @reneerosillo made their first contribution in [#363](https://github.com/TimeCopilot/timecopilot/pull/363)
+
+---
+
+**Full Changelog**: https://github.com/TimeCopilot/timecopilot/compare/v0.0.29...v0.0.30

docs/contributing.md

@@ -81,7 +81,7 @@ For further reference, see the [Google Developer Documentation Style Guide on fi
 
 In order for new models and adapters to show up in TimeCopilot's documentation properly there are a couple steps to take:
 
-- Add the model/adapter in its own file or the appropriate file in `docs/api/models` or the appropriate subdirectory. For adapters, it may also be good to create an example notebook in `docs/examples` and add it to the `Examples` section of [mkdocs.yml](https://github.com/TimeCopilot/timecopilot/blob/main/mkdocs.yml).
+- Add the model/adapter in its own file or the appropriate file in `docs/api/models` or the appropriate subdirectory. For adapters, it may also be good to create an example notebook in `docs/examples` and add it to the appropriate subsection under **Examples** in [mkdocs.yml](https://github.com/TimeCopilot/timecopilot/blob/main/mkdocs.yml) (e.g. Foundation Models → Family Guides, Integrations, Agent and LLMs).
     - when creating a new file in `docs/api/models` for this step, add that file in the API Reference section of [mkdocs.yml](https://github.com/TimeCopilot/timecopilot/blob/main/mkdocs.yml)
 - A docstring should be present for the new model/adapter with an example in the `forecast()` method in a similar style to the [TimeCopilot agent query method](https://timecopilot.dev/api/agent/#timecopilot.agent.TimeCopilot.query). Note: when adding an adapter with its own dependencies you may need to add those dependencies as dev dependencies with `uv add new_dep --group dev`.
 

docs/examples/index.md

@@ -0,0 +1,46 @@
+# Examples
+
+Interactive notebooks for common TimeCopilot workflows. Pick a section below based on your goal, or browse the sidebar.
+
+For model API details, see the [Model Hub](../model-hub.md).
+
+## Core Workflows
+
+| Notebook | What you'll learn | Prerequisites |
+|----------|-------------------|---------------|
+| [Forecaster Quickstart](forecaster-quickstart.ipynb) | Forecast and cross-validate with `TimeCopilotForecaster` and classical models | Python 3.10+ |
+| [Anomaly Detection](anomaly-detection-forecaster-quickstart.ipynb) | Detect anomalies with the forecaster API | Python 3.10+ |
+| [Cryptocurrency Walkthrough](cryptocurrency-quickstart.ipynb) | Handle missing data and multi-series forecasting on real-world data | Python 3.10+ |
+
+## Agent and LLMs
+
+| Notebook | What you'll learn | Prerequisites |
+|----------|-------------------|---------------|
+| [Agent Quickstart](agent-quickstart.ipynb) | Conversational forecasting with the TimeCopilot agent | OpenAI API key |
+| [LLM Providers](llm-providers.ipynb) | Configure OpenAI, Ollama, and other LLM backends | Provider API key or local runtime |
+| [AWS Bedrock](aws-bedrock.ipynb) | Run the agent with AWS-hosted models | AWS account with Bedrock access |
+| [Google LLMs](google-llms.ipynb) | Run the agent with Google AI Studio / Gemini | Google AI Studio API key |
+| [Explaining Models and Ensembles](explaining-foundation-models-and-ensembles.ipynb) | Ask the agent to explain and compare custom model setups | OpenAI API key |
+
+## Foundation Models
+
+| Notebook | What you'll learn | Prerequisites |
+|----------|-------------------|---------------|
+| [Compare Foundation Models](ts-foundation-models-comparison-quickstart.ipynb) | Benchmark multiple foundation models side by side | Python 3.10+; GPU optional |
+| [Chronos Family](chronos-family.ipynb) | Forecast with Chronos 1.x and 2.x checkpoints | Python 3.10+ |
+| [TiRex Family](tirex-family.ipynb) | Forecast with TiRex 1.0 and 2.0 | Python 3.11+ |
+| [Toto Family](toto-family.ipynb) | Forecast with Toto 1.0 and 2.0 | Python 3.10+ |
+| [Finetuning](finetuning.ipynb) | Adapt Chronos 2 and TimeGPT to your data | Python 3.10+; GPU recommended |
+
+## Benchmarks and Ensembles
+
+| Notebook | What you'll learn | Prerequisites |
+|----------|-------------------|---------------|
+| [GIFT-Eval](gift-eval.ipynb) | Evaluate a foundation model ensemble on GIFT-Eval | Python 3.10+; GPU recommended |
+| [Custom Ensembles](custom-ensembles.ipynb) | Build weighted and custom ensembles beyond `MedianEnsemble` | Python 3.10+ |
+
+## Integrations
+
+| Notebook | What you'll learn | Prerequisites |
+|----------|-------------------|---------------|
+| [sktime](sktime.ipynb) | Use sktime models through TimeCopilot | Python 3.10+; `sktime` extra |

docs/getting-started/quickstart.md

@@ -187,7 +187,7 @@ accuracy and reliability surpassing basic seasonal models.'
 TimeCopilot uses [Pydantic](https://docs.pydantic.dev/latest/) to make calls to LLM endpoints, so it should be compatible with all endpoints [Pydantic supports](https://ai.pydantic.dev/models/overview/). Instructions on using other models/endpoints with Pydantic can be found on the matching Pydantic docs page, such as this page for [Google's models](https://ai.pydantic.dev/models/google/#api-key-generative-language-api).
 
 
-For more details go to the [LLM Providers example](http://timecopilot.dev/examples/llm-providers/).
+For more details see the [Examples](../examples/index.md) section, starting with the [LLM Providers](http://timecopilot.dev/examples/llm-providers/) notebook.
 
 Note: models need support for tool use to function properly with TimeCopilot.
 

docs/model-hub.md

@@ -5,6 +5,9 @@ TimeCopilot provides a unified API for time series forecasting, integrating foun
 
 Here you'll find all the time series forecasting models available in TimeCopilot, organized by family. Click on any model name to jump to its detailed API documentation.
 
+!!! note "Family example notebooks"
+    Walkthrough notebooks for families of foundation models [here](examples/index.md#foundation-models).
+
 !!! tip "Forecast multiple models using a unified API"
 
     With the [TimeCopilotForecaster][timecopilot.forecaster.TimeCopilotForecaster] class, you can generate and cross-validate forecasts using a unified API. Here's an example:
@@ -46,7 +49,7 @@ TimeCopilot provides a unified interface to state-of-the-art foundation models f
 - [Sundial](api/models/foundation/models.md#timecopilot.models.foundation.sundial) ([arXiv:2502.00816](https://arxiv.org/pdf/2502.00816))
 - [T0](api/models/foundation/models.md#timecopilot.models.foundation.t0) ([model card](https://huggingface.co/theforecastingcompany/t0-alpha))
 - [TabPFN](api/models/foundation/models.md#timecopilot.models.foundation.tabpfn) ([arXiv:2501.02945](https://arxiv.org/abs/2501.02945))
-- [TiRex](api/models/foundation/models.md#timecopilot.models.foundation.tirex) ([arXiv:2505.23719](https://arxiv.org/abs/2505.23719))
+- [TiRex / TiRex-2](api/models/foundation/models.md#timecopilot.models.foundation.tirex) ([arXiv:2505.23719](https://arxiv.org/abs/2505.23719), [arXiv:2607.01204](https://arxiv.org/abs/2607.01204))
 - [TimeGPT](api/models/foundation/models.md#timecopilot.models.foundation.timegpt) ([arXiv:2310.03589](https://arxiv.org/abs/2310.03589))
 - [TimesFM](api/models/foundation/models.md#timecopilot.models.foundation.timesfm) ([arXiv:2310.10688](https://arxiv.org/abs/2310.10688))
 - [Toto](api/models/foundation/models.md#timecopilot.models.foundation.toto) ([arXiv:2505.14766](https://arxiv.org/abs/2505.14766))

mkdocs.yml

@@ -18,29 +18,32 @@ nav:
       - Introduction: getting-started/introduction.md
       - Quickstart: getting-started/quickstart.md
       - Installation: getting-started/installation.md
-      - Examples:
-          - examples/agent-quickstart.ipynb
-          - examples/llm-providers.ipynb
-          - examples/aws-bedrock.ipynb
-          - examples/google-llms.ipynb
-          - examples/forecaster-quickstart.ipynb
-          - examples/anomaly-detection-forecaster-quickstart.ipynb
-          - examples/ts-foundation-models-comparison-quickstart.ipynb
-          - examples/gift-eval.ipynb
-          - examples/chronos-family.ipynb
-          - examples/toto-family.ipynb
-          - examples/finetuning.ipynb
-          - examples/cryptocurrency-quickstart.ipynb
-          - examples/sktime.ipynb
-          - examples/custom-ensembles.ipynb
-          - examples/explaining-foundation-models-and-ensembles.ipynb
-
-  - Experiments:
-      - experiments/gift-eval.md
-      - experiments/fev.md
-  - Documentation:
       - Time Series Model Hub: model-hub.md
       - Forecasting Parameters: forecasting-parameters.md
+  - Examples:
+      - examples/index.md
+      - Core Workflows:
+          - Forecaster Quickstart: examples/forecaster-quickstart.ipynb
+          - Anomaly Detection: examples/anomaly-detection-forecaster-quickstart.ipynb
+          - Cryptocurrency Walkthrough: examples/cryptocurrency-quickstart.ipynb
+      - Agent and LLMs:
+          - Agent Quickstart: examples/agent-quickstart.ipynb
+          - LLM Providers: examples/llm-providers.ipynb
+          - AWS Bedrock: examples/aws-bedrock.ipynb
+          - Google LLMs: examples/google-llms.ipynb
+          - Explaining Models and Ensembles: examples/explaining-foundation-models-and-ensembles.ipynb
+      - Foundation Models:
+          - Compare Foundation Models: examples/ts-foundation-models-comparison-quickstart.ipynb
+          - Family Guides:
+              - Chronos: examples/chronos-family.ipynb
+              - TiRex: examples/tirex-family.ipynb
+              - Toto: examples/toto-family.ipynb
+          - Finetuning: examples/finetuning.ipynb
+      - Benchmarks and Ensembles:
+          - GIFT-Eval: examples/gift-eval.ipynb
+          - Custom Ensembles: examples/custom-ensembles.ipynb
+      - Integrations:
+          - sktime: examples/sktime.ipynb
   - API Reference:
       - api/agent.md
       - api/forecaster.md
@@ -53,8 +56,12 @@ nav:
       - api/models/adapters/adapters.md
       - api/models/utils/forecaster.md
       - api/gift-eval/gift-eval.md
+  - Experiments:
+      - experiments/gift-eval.md
+      - experiments/fev.md
   - Changelogs:
       - changelogs/index.md
+      - changelogs/v0.0.30.md
       - changelogs/v0.0.29.md
       - changelogs/v0.0.28.md
       - changelogs/v0.0.27.md

pyproject.toml

@@ -76,6 +76,7 @@ dependencies = [
   "logfire>=4.11.0",
   "matplotlib>=3.10.6",
   "mlforecast>=1.0.2",
+  "mlstm-kernels>=2.0.4 ; python_full_version >= '3.11'",
   "neuralforecast>=3.0.2",
   "nixtla>=0.7.0",
   "openai>=1.99.7",
@@ -97,10 +98,11 @@ dependencies = [
   "timecopilot-chronos-forecasting>=0.2.2",
   "timecopilot-granite-tsfm>=0.2.1 ; python_full_version >= '3.11' and python_full_version < '3.14'",
   "timecopilot-timesfm>=0.3.0",
+  "timecopilot-tirex2>=0.1.0 ; python_full_version >= '3.11'",
   "timecopilot-tirex>=0.1.1 ; python_full_version >= '3.11'",
   "timecopilot-toto-2>=0.1.1",
   "timecopilot-toto>=0.1.7",
-  "timecopilot-uni2ts>=0.1.2 ; python_full_version < '3.14'",
+  "timecopilot-uni2ts>=0.1.3 ; python_full_version < '3.14'",
   "torchmetrics>=1.8.2",
   "transformers>=4.41,<6 ; python_full_version < '3.13'",
   "transformers>=4.48,<6 ; python_full_version >= '3.13'",
@@ -113,7 +115,7 @@ license = "MIT"
 name = "timecopilot"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.0.29"
+version = "0.0.30"
 
 [project.optional-dependencies]
 distributed = [

tests/models/conftest.py

@@ -134,6 +134,13 @@ def disable_mps_session(monkeypatch):
     from timecopilot.models.foundation.tirex import TiRex
 
     models.append(TiRex())
+    models.append(
+        TiRex(
+            repo_id="NX-AI/TiRex-2",
+            alias="TiRex-2",
+            batch_size=2,
+        )
+    )
 
 if (3, 11) <= sys.version_info < (3, 14):
     from timecopilot.models.foundation.t0 import T0

tests/models/foundation/test_tirex.py

@@ -0,0 +1,28 @@
+import sys
+
+import pytest
+from utilsforecast.data import generate_series
+
+pytestmark = pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="TiRex requires Python >= 3.11",
+)
+
+
+def test_is_tirex2_dispatch():
+    from timecopilot.models.foundation.tirex import TiRex
+
+    assert not TiRex(repo_id="NX-AI/TiRex")._is_tirex2()
+    assert TiRex(repo_id="NX-AI/TiRex-2")._is_tirex2()
+    assert TiRex(repo_id="NX-AI/TiRex-2/")._is_tirex2()
+
+
+def test_tirex2_forecast():
+    from timecopilot.models.foundation.tirex import TiRex
+
+    df = generate_series(2, freq="D", min_length=50, max_length=50)
+    df["unique_id"] = df["unique_id"].astype(str)
+    model = TiRex(repo_id="NX-AI/TiRex-2", alias="TiRex-2", batch_size=2)
+    fcst = model.forecast(df, h=3, freq="D")
+    assert fcst.shape == (6, 3)
+    assert "TiRex-2" in fcst.columns