Aurora-X

Aurora-X is a time series foundation model with native covariate support. It forecasts univariate and multivariate series zero-shot, and can condition on past-only and known-future covariates in a single forward pass.

Model weights: DecisionIntelligence/Aurora-X

Installation

pip install aurorax-model

The architecture config ships with the package; the weights (~4 GB) are pulled from the Hugging Face Hub the first time you load the model and cached locally.

Quick start

import numpy as np
from aurorax import load_pipeline

pipe = load_pipeline()          # downloads weights on first use

context = np.random.randn(512)                      # 512 historical steps
preds = pipe.predict(context, prediction_length=96) # list with one entry
print(preds[0].shape)                               # (1, 20, 96) = (n_targets, n_samples, horizon)

predict returns a list with one tensor per input case, shaped (n_targets, K, prediction_length), where K is num_samples in sampling mode or the number of quantile levels in quantile mode.

Point forecast + quantiles

quantiles, mean = pipe.predict_quantiles(
    context,
    prediction_length=96,
    quantile_levels=[0.1, 0.5, 0.9],
)
mean[0].shape         # (1, 96)     -> point forecast
quantiles[0].shape    # (1, 96, 3)  -> one column per quantile level

Multivariate

Wrap a (n_variates, context_length) array in a list to forecast its variates jointly, so the model can exploit cross-variate structure.

series = np.random.randn(3, 512)              # 3 variates, 512 steps

preds = pipe.predict([series], prediction_length=96)   # list of 2D -> one multivariate case
preds[0].shape        # (3, 20, 96)

Careful: a bare 2D array is read as a batch of univariate series, not as one multivariate case. pipe.predict(series, ...) returns three separate univariate forecasts. Use [series] (or the 3D form series[None]) whenever the variates belong together.

Covariates

Give each case a target plus optional covariates. Every covariate needs its history in past_covariates; add it to future_covariates as well when the future values are known ahead of time (calendar, weather forecast, promotions).

preds = pipe.predict(
    [{
        "target": np.random.randn(512),
        # past-only covariate: history is used, future is unknown
        "past_covariates":   {"sales": np.random.randn(512),
                              "temp":  np.random.randn(512)},
        # known-future covariate: future values are fed to the model
        "future_covariates": {"temp":  np.random.randn(96)},
    }],
    prediction_length=96,
)
preds[0].shape        # (1, 20, 96) -> only target rows are returned

Covariate rows condition the forecast but are never scored, so the output only covers the target rows.

Batching

Pass a list to forecast many cases in one go. Contexts may have different lengths; shorter ones are left-padded and masked automatically.

preds = pipe.predict(
    [np.random.randn(300), np.random.randn(512), np.random.randn(1024)],
    prediction_length=96,
)
len(preds)            # 3

Inputs

predict and predict_quantiles accept any of:

NumPy arrays and PyTorch tensors are interchangeable everywhere; NaN marks missing values and is masked out. Returned tensors are always on CPU.

Key arguments

Lower-level access

load_model returns the raw model when you want to drive generate yourself:

import torch
from aurorax import load_model

model = load_model()            # eval mode, on cuda if available
preds = model.generate(
    inputs=torch.randn(4, 512),
    max_output_length=96,
    num_samples=20,
)                               # (4, 20, 96)

Both loaders take the same arguments:

load_pipeline(
    repo_id="DecisionIntelligence/Aurora-X",  # or a local checkpoint directory
    cache_dir=None,                           # where to cache the weights
    force_download=False,
    device=None,                              # default: cuda if available
)