Preserve precision when unscaling Chronos-2 and Chronos-Bolt forecasts (#532)
## Problem
Inverse normalization is computed in float32, but its result was cast
back to the model dtype. For bfloat16 models, this discards low-order
bits when the time-series level is large relative to its variation.
Converting predictions to float32 later in the pipeline cannot recover
that information.
## Reproduction
```python
import torch
from chronos.chronos_bolt import InstanceNorm
scaler, base = InstanceNorm(), torch.linspace(-2, 2, 9)[None]
for level, variation in [(1e4, 100), (1e5, 100), (1e6, 100), (1e6, 1e3), (1e6, 1e4), (1e7, 1e4)]:
x = level + variation * base
normalized, stats = scaler(x)
bf16 = normalized.bfloat16()
reference = scaler.inverse(normalized, stats, torch.float32)
outputs = (scaler.inverse(bf16, stats).float(), scaler.inverse(bf16, stats, torch.float32))
print(level, variation, *((output - reference).abs().max().item() for output in outputs))
```
The table shows the maximum absolute delta relative to performing both
normalization and inverse normalization in float32:
| Series level | Series variation | Max delta before | Max delta after |
|---:|---:|---:|---:|
| 10,000 | 100 | 30.000 | 0.299 |
| 100,000 | 100 | 252.000 | 0.297 |
| 1,000,000 | 100 | 776.000 | 0.312 |
| 1,000,000 | 1,000 | 2,020.000 | 3.000 |
| 1,000,000 | 10,000 | 1,960.000 | 29.938 |
| 10,000,000 | 10,000 | 32,008.000 | 30.000 |
## Changes
- Add an optional output dtype to InstanceNorm inverse scaling while
preserving its existing default behavior.
- Keep inverse-scaled Chronos-2 and Chronos-Bolt predictions in float32.
- Add bfloat16 model-level regression tests and a numerical precision
test.
## Testing
- Focused Chronos-2, Chronos-Bolt, and InstanceNorm suite: 16 passed.
- Verified that the Chronos-2 regression test fails on the previous code
because model predictions are returned as bfloat16. O
Oleksandr Shchur committed
ff27460c184635a991c4fa1ec1f3edc114a3d008
Parent: 8589d19
Committed by GitHub <noreply@github.com>
on 9/8/2026, 1:58:38 PM