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

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(