TSFM-ai/chronos-2-onnx

Chronos-2 ONNX With Future Covariates

This repository contains a real FP32 ONNX export of amazon/chronos-2 with the Chronos-2 future-covariate tensor interface enabled.

This is a tensor-level ONNX model. It does not package the high-level Chronos2Pipeline.predict preprocessing API, and it does not define a serving API. Prepare tensors before calling ONNX Runtime.

Files

• model.onnx: FP32 ONNX model, 456 MB.
• config.json: Chronos-2 model config copied from the source model.
• generation_config.json: generation config copied from the source model.
• reports/full_parity_report.json: PyTorch vs ONNX Runtime parity report.
• reports/PROOF.md: concise local runtime proof.
• reports/WEB_AND_PARITY_SUMMARY.md: optional ONNX Runtime Web proof notes.
• reports/*: supporting parity and runtime reports.

Model SHA256:

ce8cbf5cd70fc33a347d4943796dd9366d5adc91279ccd1df72450eb30c3d41a

ONNX Tensor Interface

Inputs:

• context: float32[batch_size, 512]
• group_ids: int64[batch_size]
• attention_mask: float32[batch_size, 512]
• future_covariates: float32[batch_size, 64]
• num_output_patches: int64[]

Output:

• quantile_preds: float32[batch_size, 21, 64]

Batch size is dynamic. Context length, future covariate length, and prediction length are fixed in this export.

Python Usage

import numpy as np
import onnxruntime as ort

session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

batch_size = 2
inputs = {
    "context": np.random.randn(batch_size, 512).astype(np.float32),
    "group_ids": np.arange(batch_size, dtype=np.int64),
    "attention_mask": np.ones((batch_size, 512), dtype=np.float32),
    "future_covariates": np.random.randn(batch_size, 64).astype(np.float32),
    "num_output_patches": np.array(4, dtype=np.int64),
}

quantile_preds = session.run(None, inputs)[0]
print(quantile_preds.shape)  # (2, 21, 64)

Parity Proof

The ONNX model was compared against PyTorch using onnxruntime==1.27.0 on CPU.

The validation covered:

• batch sizes 1 through 5
• shared and distinct group_ids
• zero, sinusoidal, cosine, ramp, and random future_covariates
• missing context values
• missing future covariate values
• missing context and future covariates together

All parity cases passed at:

• rtol=1e-4
• atol=1e-4
• max_abs_limit=5e-4
• mean_abs_limit=2e-5

See reports/full_parity_report.json for details.

Runtime Notes

This model has also been loaded with onnxruntime-web WASM in Node and in a browser using a local fixture. Those reports are included for convenience, but the primary supported artifact here is the ONNX Runtime tensor model.

The web runtime proof is functional, not a browser performance recommendation. The FP32 model is 456 MB; quantization and runtime-specific packaging should be evaluated before production browser use.

Limitations

• Fixed context length: 512
• Fixed prediction length: 64
• Fixed future covariate length: 64
• Dynamic batch only
• No high-level list-of-dicts or DataFrame preprocessing wrapper
• No INT8 quantized artifact in this repository
• No arbitrary-horizon autoregressive wrapper in this repository

Source

This artifact was produced from amazon/chronos-2 while testing ONNX export fixes for amazon-science/chronos-forecasting.