Skip to content

fix: support custom output quantiles via ForecastConfig (#286)#365

Open
aydinomer00 wants to merge 1 commit intogoogle-research:masterfrom
aydinomer00:fix/custom-quantiles-286
Open

fix: support custom output quantiles via ForecastConfig (#286)#365
aydinomer00 wants to merge 1 commit intogoogle-research:masterfrom
aydinomer00:fix/custom-quantiles-286

Conversation

@aydinomer00
Copy link
Copy Markdown

Problem

When passing custom quantiles (e.g. quantiles=[0.3, 0.4]) to ForecastConfig,
the model raised a shape mismatch error because the output architecture was assumed
to always return all 10 default quantile values.

Closes #286

Solution

The model always outputs 10 values internally (1 point + 9 quantiles: 0.1–0.9).
This PR adds an output_quantiles field to ForecastConfig that slices the output
tensor after decoding — no model weights or checkpoint formats changed.

Fully backwards-compatible: output_quantiles=None (default) returns all 9
quantiles as before.

Changes

  • src/timesfm/configs.py: Added DEFAULT_QUANTILES constant and
    output_quantiles: tuple[float, ...] | None = None to ForecastConfig
  • src/timesfm/timesfm_2p5/timesfm_2p5_torch.py: Validate at compile() and slice output
  • src/timesfm/timesfm_2p5/timesfm_2p5_flax.py: Same fix for Flax backend

Usage

model.compile(
    timesfm.ForecastConfig(
        max_context=1024,
        max_horizon=256,
        output_quantiles=(0.1, 0.5, 0.9),
    )
)
point, quantiles = model.forecast(horizon=12, inputs=[...])
# quantiles.shape → (1, 12, 3)  ✓

…rch#286)

Add output_quantiles field to ForecastConfig that allows users to
select a subset of the 9 default quantiles to return at inference
time. Model weights and checkpoint format are unchanged; the fix
slices the output tensor after decoding.

Changes:
- Add DEFAULT_QUANTILES constant to configs.py
- Add output_quantiles: tuple[float, ...] | None = None to ForecastConfig
- Validate selected quantiles at compile() time in both Torch and Flax
  backends, raising ValueError for values outside DEFAULT_QUANTILES
- Slice full_forecast output tensor to return only requested quantiles

Closes google-research#286
@google-cla
Copy link
Copy Markdown

google-cla bot commented Feb 20, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: model cannot be loaded with custom quantiles

1 participant