Preserve precision when unscaling Chronos-2 and Chronos-Bolt forecasts

[shchur]

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

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.

[abdulfatir]

Is there ever a usecase where we actually want the inverse to be cast back to orig_dtype? I am wondering if we should just always return fp32. Simply return x.