🌌 Aurora-X

A time series foundation model with native covariate support.

Aurora-X forecasts univariate and multivariate series zero-shot, and conditions on past-only and known-future covariates in a single forward pass.

• 🧠 Mixture-of-experts backbone — 16 routed + 1 shared expert, top-4, 12 layers
• 🔗 Group attention over variates — the channels of one series are modelled jointly
• 🎲 Probabilistic forecasts — IQN head: sample trajectories or query quantiles
• 📏 Long context — up to 8160 steps

---

📦 Installation

pip install aurorax-model

Weights (~4 GB) are downloaded from this repository on first load and cached under ~/.cache/huggingface/.

---

🚀 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)

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. Tensors come back on CPU.

📊 Point forecast and 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 requested quantile

mean is averaged over a dense interior quantile grid, which is deterministic and avoids the extreme tails a random sample can hit.

🔀 Multivariate

Wrap a (n_variates, context_length) array in a list so the variates are forecast jointly and group attention can use their 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 independent univariate series, not one multivariate case — pipe.predict(series, ...) would return 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; also list it in future_covariates when its 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 the target rows are returned

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

⚠️ A key in future_covariates without a matching entry in past_covariates raises an error rather than being silently dropped — the model needs the covariate's history to normalise its future values.

predict_quantiles takes the same covariate input:

quantiles, mean = pipe.predict_quantiles(cases, prediction_length=96,
                                         quantile_levels=[0.1, 0.5, 0.9])

📚 Batching

Pass a list to forecast many cases at once. Context lengths may differ; 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

---

🧩 Input forms

Form	Meaning
1D array (T,)	one univariate series
2D array (N, T)	N independent univariate series
3D array (N, V, T)	N multivariate cases, V variates each
list of 1D arrays	N univariate cases, context lengths may differ
list of 2D arrays (V, T)	N multivariate cases, forecast jointly
list of dicts	cases with covariates (see above)

✅ NumPy arrays and PyTorch tensors are interchangeable everywhere, on CPU or GPU. NaN marks missing values and is masked out.

⚙️ Arguments

Argument	Meaning	Default
prediction_length	forecast horizon	required
num_samples	sampled trajectories when mode="sample"	20
mode	"sample" or "quantile"	"sample"
quantile_levels	levels to return in quantile mode	—
inference_token_len	patch size; lower it (8/16/32) for short series	48
batch_size	max variate rows per forward pass	256
cross_learning	share attention across all cases in the batch	False

---

🔧 Lower-level access

load_model returns the raw model when you want to drive generate yourself. It comes back in eval mode — the config keeps dropout_rate=0.1, so a model left in train mode would produce stochastic forecasts.

import torch
from aurorax import load_model

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

Both loaders accept 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
)

---

📈 Results

Zero-shot, evaluated with the multivariate protocol (variates forecast jointly). Lower is better; 1.0 means "no better than Seasonal Naive".

Benchmark	Metric	Aurora-X
🥈 GIFT-Eval (97 configs)	MASE / WQL, rel. Seasonal Naive	0.659 / 0.467
📐 TIME (98 configs)	MASE / WQL, rel. Seasonal Naive	0.640 / 0.582
🥇 fev-bench (100 tasks)	MASE / SQL, rel. Seasonal Naive	0.617 / 0.517

ℹ️ GIFT-Eval does not mandate a variate protocol, and several leaderboard entries split multivariate series into univariate ones instead. Under that univariate protocol Aurora-X scores 0.682 / 0.479, so the protocol should be stated when comparing.

---

📄 License

Apache 2.0