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

The point forecast is the arithmetic mean of 19 fixed interior quantiles (0.05, 0.10, ..., 0.95), independently of the requested levels. It approximates the distribution mean; it is not an exact expectation. Quantile curves are per-time marginal predictions, not joint future sample paths.

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

inference_token_len applies to both historical and known-future covariate patches. For example, inference_token_len=96 uses 96-step patches and resamples each patch to the training resolution of 48 steps. This changes the time resolution, so speed gains do not imply unchanged forecast accuracy.

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, device=next(model.parameters()).device),
    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
)