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.

Author: shchur

src/chronos/chronos_bolt.py

@@ -121,8 +121,12 @@ def forward(
 
         return scaled_x.to(orig_dtype), (loc, scale)
 
-    def inverse(self, x: torch.Tensor, loc_scale: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
-        orig_dtype = x.dtype
+    def inverse(
+        self,
+        x: torch.Tensor,
+        loc_scale: tuple[torch.Tensor, torch.Tensor],
+        output_dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
         x = x.to(torch.float32)
         loc, scale = loc_scale
 
@@ -131,7 +135,7 @@ def inverse(self, x: torch.Tensor, loc_scale: tuple[torch.Tensor, torch.Tensor])
 
         x = x * scale + loc
 
-        return x.to(orig_dtype)
+        return x if output_dtype is None else x.to(output_dtype)
 
 
 class ResidualBlock(nn.Module):

test/test_chronos2.py

@@ -44,6 +44,14 @@ def test_chronos2_encoder_accepts_config_without_is_decoder():
     assert len(encoder.block) == config.num_layers
 
 
+def test_when_chronos2_model_uses_bfloat16_then_unscaled_predictions_use_float32(pipeline):
+    pipeline.model.to(torch.bfloat16)
+
+    output = pipeline.model(context=torch.rand(1, 16), num_output_patches=1)
+
+    assert output.quantile_preds.dtype == torch.float32
+
+
 def test_base_chronos2_pipeline_loads_from_s3():
     BaseChronosPipeline.from_pretrained("s3://autogluon/chronos-2", device_map="cpu")
 

test/test_chronos_bolt.py

@@ -22,6 +22,14 @@ def pipeline() -> ChronosBoltPipeline:
     return BaseChronosPipeline.from_pretrained(DUMMY_MODEL_PATH, device_map="cpu")
 
 
+def test_when_chronos_bolt_model_uses_bfloat16_then_unscaled_predictions_use_float32():
+    pipeline = ChronosBoltPipeline.from_pretrained(DUMMY_MODEL_PATH, device_map="cpu", torch_dtype=torch.bfloat16)
+
+    output = pipeline.model(context=torch.rand(1, 16))
+
+    assert output.quantile_preds.dtype == torch.float32
+
+
 def test_base_chronos_pipeline_loads_from_huggingface():
     BaseChronosPipeline.from_pretrained("amazon/chronos-bolt-tiny", device_map="cpu")
 
@@ -355,6 +363,21 @@ def test_when_instancenorm_applied_and_reversed_then_output_correct():
     assert torch.allclose(output, input_)
 
 
+def test_when_instancenorm_reversed_to_float32_then_precision_is_preserved():
+    inorm = InstanceNorm()
+    normalized = torch.tensor([[0.125]], dtype=torch.bfloat16)
+    loc = torch.tensor([[1_000_000.0]], dtype=torch.float32)
+    scale = torch.tensor([[100.0]], dtype=torch.float32)
+
+    output = inorm.inverse(normalized, (loc, scale))
+
+    assert output.dtype == torch.float32
+    torch.testing.assert_close(
+        output,
+        torch.tensor([[1_000_012.5]], dtype=torch.float32),
+    )
+
+
 @pytest.mark.parametrize("task_kwargs", [{}, {"eval_metric": "WQL", "quantile_levels": [0.1, 0.2]}])
 def test_pipeline_can_evaluate_on_dummy_fev_task(task_kwargs):
     pipeline = BaseChronosPipeline.from_pretrained(