Chronos-2 ONNX (fp16 weights, fp32 compute)

An ONNX export of amazon/chronos-2 at revision 29ec3766d36d6f73f0696f85560a422f50e8498c, for CPU inference with ONNX Runtime, including onnxruntime-node. It keeps Chronos-2's covariate interface, and every dimension is dynamic: the number of series, the context length, and the forecast length.

Files

Keep chronos2.onnx.data in the same folder as chronos2.onnx.

Load the model with graph optimizations disabled. The weights are stored as fp16 and cast to fp32 in the graph; with optimizations on, ONNX Runtime folds the casts into an fp32 copy of every weight, which doubles the memory for no speed gain.

Interface

This is the network call only (Chronos2Model.forward). Prepare the tensors as chronos-forecasting 2.3.2 does in Chronos2Dataset (test mode) and Chronos2Pipeline._predict_step.

Inputs:

• context: float32[rows, time]. Each series is its target row followed by its covariate rows, left-padded with NaN to the longest context in the batch. NaN means missing.
• group_ids: int64[rows]. The series each row belongs to. Rows of different series do not attend to each other.
• future_covariates: float32[rows, 16 × patches]. The known future values of the covariates. Use NaN on target rows and on past-only covariates. Pad with NaN to a whole number of 16-step output patches. The model reads the number of output patches (1 to 64) from this width.

Output:

• quantiles: float32[rows, 21, 16 × patches], at the levels 0.01, 0.05, 0.1, 0.15, …, 0.9, 0.95, 0.99. Read each series' target row.

Changes from the original model

The export makes three changes. The first two change no value.

1. In _prepare_patched_future, the check that raises on NaN future covariates is removed. NaNs are already replaced by 0 under the mask before the check, and the check depends on data, so it cannot be exported.
2. In _construct_and_invert_group_time_mask, the boolean group mask is cast to the attention mask's float type before the einsum. ONNX Runtime's CPU Einsum has no boolean kernel.
3. Every weight tensor of 1,024 values or more is stored as fp16 and cast back to fp32 in the graph. All computation stays fp32. Full fp16 computation overflows on high-volume series, so it is not used.

Checks

The model was compared with the PyTorch pipeline (chronos-forecasting 2.3.2, torch 2.14.1) on synthetic weekly series of 20 to 300 weeks, with four seasonal covariates, at forecast lengths of 13, 52, 80 and 140 weeks:

• The fp32 graph matches PyTorch to a relative difference of 1.8e-6 or less.
• This fp16-weights graph matches PyTorch to a relative difference of 4.2e-4 or less.

Exported with torch.onnx.export (dynamo) and onnx 1.23.2. Tested with onnxruntime 1.31.0 (Python) and onnxruntime-node 1.23.2.

License

Apache-2.0, the license of the original model. The original model, Chronos-2, is by Amazon. This export is provided as is, without warranty of any kind.