Skip to content

Fix endpoint handling for batched quantile grids - #526

Open
tandede wants to merge 2 commits into
amazon-science:mainfrom
tandede:fix/batched-quantile-endpoints
Open

Fix endpoint handling for batched quantile grids#526
tandede wants to merge 2 commits into
amazon-science:mainfrom
tandede:fix/batched-quantile-endpoints

Conversation

@tandede

@tandede tandede commented Aug 20, 2026

Copy link
Copy Markdown

Summary

interpolate_quantiles accepts a tensor of quantile levels matching original_values, allowing each batch element to use a different quantile grid. Endpoint padding was previously decided from the global minimum and maximum across the entire batch.

If one row contained 0 or 1 while another did not, the latter row was left unpadded. Queries below that row's minimum produced a negative gather index, while queries above its maximum were extrapolated instead of using the endpoint value.

This change pads every flattened batch row with 0 and 1 and repeats that row's first and last values. Duplicate endpoints are safe for rows that already contain them.

Test plan

  • .venv/bin/python -m pytest -q test/test_utils.py — 7 passed
  • .venv/bin/mypy src test — no issues in 24 source files
  • .venv/bin/python -m pytest -q — 554 passed, 20 xfailed
  • uvx ruff check --select E,F,I,B src/chronos/utils.py test/test_utils.py
  • uvx ruff format --check src/chronos/utils.py test/test_utils.py
  • git diff --check

@tandede
tandede marked this pull request as ready for review August 20, 2026 08:35
@tandede

tandede commented Aug 30, 2026

Copy link
Copy Markdown
Author

Hi maintainers, I wanted to follow up on this PR in case it was overlooked. It fixes endpoint handling for batched quantile grids and is ready for review. Please let me know if any changes or additional validation would be helpful. Thank you!

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.

1 participant