Forecast a Multivariate Series using a Pre-trained Foundation Model

This guide shows you how to use a large, pre-trained foundation model (like Toto) within Chronax to generate zero-shot forecasts for high-dimensional, multivariate time series data. This approach is ideal when you need state-of-the-art performance without fine-tuning.

Prerequisites

Chronax installed (pip install chronax).
A pre-trained model checkpoint (e.g., weights from a model like Toto).
Input data y must be a 2-D jnp.ndarray of shape (T, C), where $T$ is the number of time steps and $C$ is the number of channels/variables, typically float32.

import jax.numpy as jnp
# TODO: Confirm the actual Chronax model name for foundation models (e.g., TotoFoundationModel)
from chronax.models import FoundationModel

Steps

1. Load the pre-trained model and weights

Foundation models require loading pre-trained weights to function effectively in a zero-shot manner. You must initialize the model and load the checkpoint state.

# Define the prediction horizon (h)
H = 336 

# 1a. Initialize the model structure
# TODO: Confirm the Chronax initialization arguments required for the FoundationModel.
# Foundation models often require context length and output dimension.
try:
    model = FoundationModel(
        context_length=4096, 
        prediction_length=H
    )
except NameError:
    print("Using placeholder model initialization.")

# 1b. Load the pre-trained weights
# Chronax models use Flax state management.
# TODO: Confirm the exact method for loading external weights (e.g., Hugging Face checkpoints) into the Chronax model state.
# For demonstration, we assume a standard Chronax initialization pattern:
key = jax.random.PRNGKey(0)
params = model.init(key, jnp.ones((1, 4096, 7), dtype=jnp.float32))['params'] 
print("Model initialized (weights not loaded).")

2. Prepare the multivariate input data

The input data y must be shaped as (T, C), where $T$ is the context length used during pre-training (e.g., 4096) and $C$ is the number of variables (e.g., 7).

# Example: 4096 time steps, 7 variables
T, C = 4096, 7
y_input = jnp.ones((T, C), dtype=jnp.float32) * jnp.arange(T)[:, None]

print(f"Input shape: {y_input.shape}")

3. Generate the probabilistic forecast

Foundation models typically generate probabilistic forecasts by sampling from the learned distribution (often a Student-T mixture). Specify the prediction horizon h and the number of samples (num_samples) to estimate the distribution.

h = H # 336 steps
num_samples = 256 # Number of samples for probabilistic estimation

# The forecast method uses the loaded parameters and the input data.
# TODO: Confirm the Chronax API for passing parameters and sampling arguments to FoundationModel.
out = model.forecast(
    params=params,
    y=y_input,
    h=h,
    num_samples=num_samples,
    # Foundation models often require a specific PRNG key for sampling
    rng_key=jax.random.PRNGKey(1), 
)

print(f"Forecast output keys: {out.keys()}")

4. Extract the point forecast and prediction intervals

The forecast method returns a dictionary. The point forecast is found under the "mean" key. Prediction intervals (e.g., 80% or 90%) are derived from the samples.

# Extract the point forecast (mean)
point_forecast = out["mean"]

# Extract prediction intervals (assuming 90% interval keys)
# TODO: Confirm the exact keys Chronax uses for prediction intervals 
# when generated via sampling (e.g., 'lo-90', 'hi-90', or 'samples').
try:
    lower_bound = out["lo-90"]
    upper_bound = out["hi-90"]
    print(f"Point forecast shape: {point_forecast.shape}")
    print(f"Lower bound shape: {lower_bound.shape}")
except KeyError:
    print("Warning: Could not find standard interval keys. Assuming 'mean' is available.")
    print(f"Point forecast shape: {point_forecast.shape}")

Full example

import jax
import jax.numpy as jnp
# TODO: Confirm the actual Chronax model name for foundation models
from chronax.models import FoundationModel 

# --- Configuration ---
T, C = 4096, 7  # Context length and number of channels
H = 336         # Prediction horizon
num_samples = 256

# --- 1. Load the pre-trained model (Placeholder) ---
# TODO: Confirm Chronax initialization and weight loading.
try:
    model = FoundationModel(
        context_length=T, 
        prediction_length=H
    )
except NameError:
    print("Using placeholder model initialization.")

key = jax.random.PRNGKey(0)
params = model.init(key, jnp.ones((1, T, C), dtype=jnp.float32))['params'] 

# --- 2. Prepare the multivariate input data ---
y_input = jnp.ones((T, C), dtype=jnp.float32) * jnp.arange(T)[:, None]

# --- 3. Generate the probabilistic forecast ---
out = model.forecast(
    params=params,
    y=y_input,
    h=H,
    num_samples=num_samples,
    rng_key=jax.random.PRNGKey(1), 
)

# --- 4. Extract results ---
point_forecast = out["mean"]

# TODO: Confirm Chronax keys for probabilistic output
try:
    lower_bound = out["lo-90"]
    upper_bound = out["hi-90"]
    print(f"Forecast generated for {H} steps across {C} channels.")
    print(f"Point forecast shape: {point_forecast.shape}")
    print(f"Interval shape: {lower_bound.shape}")
except KeyError:
    print("Warning: Could not extract interval keys. Check Chronax documentation for FoundationModel output structure.")
    print(f"Point forecast shape: {point_forecast.shape}")

Next steps

Review the FoundationModel API documentation to confirm required context_length and prediction_length parameters.
Learn how to load specific pre-trained weights using chronax.utils.load_checkpoint.
Explore the out["samples"] key if you need to calculate custom quantiles or visualize the full predictive distribution.
See the guide on "Evaluating Forecast Accuracy" to benchmark the zero-shot predictions.