diff --git a/src/chronos/chronos_bolt.py b/src/chronos/chronos_bolt.py index abe96c55..db44821f 100644 --- a/src/chronos/chronos_bolt.py +++ b/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): diff --git a/test/test_chronos2.py b/test/test_chronos2.py index c96c9a09..a98e2678 100644 --- a/test/test_chronos2.py +++ b/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") diff --git a/test/test_chronos_bolt.py b/test/test_chronos_bolt.py index cabe8c5a..95fb3fc6 100644 --- a/test/test_chronos_bolt.py +++ b/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(