fix: change Chronos default dtype from bfloat16 to float32

[rebot-eng]

Summary

Changes the default torch_dtype for Chronos models from bfloat16 to float32 to improve numerical precision.

Problem

As reported in #307 by @abdulfatir, using bfloat16 can cause precision issues that significantly impact forecast accuracy. In time-bench benchmarks, switching to float32 improved Chronos model rankings from 5th to 1st place in terms of MASE.

Changes

• Updated torch_dtype from torch.bfloat16 to torch.float32 in _get_model()
• Updated docstring to reflect the change

Trade-offs

• Memory: float32 uses 2x memory compared to bfloat16
• Speed: Slightly slower inference on hardware with native bfloat16 support
• Accuracy: Significantly improved forecast precision (the main benefit)

The accuracy gains outweigh the performance costs for most use cases.

Closes #307

cc @AzulGarza

[Copilot]

Pull request overview

Updates Chronos model loading defaults to prioritize numerical precision by switching the default torch_dtype from bfloat16 to float32, addressing reported accuracy regressions tied to reduced-precision weights.

Changes:

• Update Chronos _get_model() to load model weights with torch_dtype=torch.float32.
• Update the Chronos class docstring to reflect the new default dtype and rationale.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

This PR changes a default that affects runtime behavior (memory/precision). There doesn’t appear to be a unit test asserting the torch_dtype passed to BaseChronosPipeline.from_pretrained; adding a small mocked test would help prevent regressions (e.g., accidentally reverting to bfloat16).

[rebot-eng]

Addressed review comments in ce3eb38:

1. Dataset dtype mismatch: Updated TimeSeriesDataset.from_df default dtype from bfloat16 to float32 in utils.py. This ensures inputs match model precision throughout the inference pipeline.

2. Unit tests: Added tests/models/foundation/test_chronos.py with two tests:

   • test_timeseries_dataset_default_dtype_is_float32: Verifies dataset tensor dtype
   • test_chronos_model_uses_float32: Mocked test ensuring from_pretrained is called with torch.float32

[AzulGarza]

@rebot-eng thanks! i think it's best if we add the attribute self.dtype as in flowstate defaulted to float32.

[rebot-eng]

Updated in 9f323c1 — added dtype parameter to Chronos following FlowState's pattern:

class Chronos(Forecaster):
    def __init__(
        self,
        repo_id: str = "amazon/chronos-t5-large",
        batch_size: int = 16,
        alias: str = "Chronos",
        dtype: torch.dtype = torch.float32,  # new
    ):
        ...
        self.dtype = dtype

Now uses self.dtype for both model loading and dataset creation. Users can override if needed:

# Default float32 (recommended)
model = Chronos()

# Explicit bfloat16 for memory-constrained environments
model = Chronos(dtype=torch.bfloat16)

[AzulGarza]

lgtm

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

This test is located in test_chronos.py but tests TimeSeriesDataset, which is a shared utility class used by multiple models (Chronos, FlowState, TimesFM, etc.). Consider moving this test to a dedicated test_utils.py file in the tests/models/foundation directory to better reflect its scope and make it more discoverable.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.