Skip to content

fix(batch): validate Batch metadata dimensions in __post_init__ - #196

Open
Xerxes Chong (xerxeschongxian26) wants to merge 1 commit into
microsoft:mainfrom
xerxeschongxian26:fix/188-validate-metadata-time-and-atmos-level-length
Open

fix(batch): validate Batch metadata dimensions in __post_init__#196
Xerxes Chong (xerxeschongxian26) wants to merge 1 commit into
microsoft:mainfrom
xerxeschongxian26:fix/188-validate-metadata-time-and-atmos-level-length

Conversation

@xerxeschongxian26

@xerxeschongxian26 Xerxes Chong (xerxeschongxian26) commented Jul 25, 2026

Copy link
Copy Markdown

Fixes #188

Add __post_init__ to Batch that asserts len(metadata.time) equals the batch size and len(metadata.atmos_levels) equals the pressure-level dimension, raising ValueError on mismatch.

Use shape[-3] for the atmos level check so validation works for both input (b, t, c, h, w) and decoder output (b, c, h, w) batches.

Align test_aurora_small with the metadata.time contract for batch size two. Add unhappy-path tests in tests/test_batch.py.

Test plan

  • Ran full local suite except tests/test_headers.py (skipped locally due to .venv crawl) — 49 passed ...
  • Pre-commit hooks pass locally

Note: This PR was drafted with AI assistance.

Wessel (@wesselb)

- Add __post_init__ to Batch that asserts len(metadata.time) equals the
batch size and len(metadata.atmos_levels) equals the pressure-level
dimension, raising ValueError on mismatch.

- Use shape[-3] for the atmos level check so validation works for both
input (b, t, c, h, w) and decoder output (b, c, h, w) batches.

- Align test_aurora_small with the metadata.time contract for batch size
two. Add unhappy-path tests in tests/test_batch.py.
@xerxeschongxian26
Xerxes Chong (xerxeschongxian26) requested a review from a team July 25, 2026 12:03
@xerxeschongxian26

Copy link
Copy Markdown
Author

@microsoft-github-policy-service agree

@xerxeschongxian26

Copy link
Copy Markdown
Author

Hi Wessel (@wesselb) , just a gentle bump on this whenever you have a moment. No rush at all; I know reviews compete with plenty of other work!

Happy to make any changes if anything needs adjusting and fixing. Thanks again for the nudge to open the PR :)

@wesselb Wessel (wesselb) left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Xerxes Chong (@xerxeschongxian26) for putting this together! This looks good to me. :) I've left a few minor suggestions.

Comment thread aurora/batch.py
Comment thread aurora/batch.py
Comment thread tests/test_batch.py
Comment thread tests/test_batch.py
Comment thread tests/test_model.py
@xerxeschongxian26

Xerxes Chong (xerxeschongxian26) commented Aug 20, 2026

Copy link
Copy Markdown
Author

Thank you for the review Wessel (@wesselb) and your suggestions! I am glad to have been able to contribute!

I believe the PR is now ready for merging!

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.

Missing validation: Silent output-shape corruption when len(Metadata.time) != batch size

2 participants