Internal ensembling in Aurora models

[Agnieszka Słowik (Slowika)]

Addressed Issue #192.
Developed with the aid of AI, in line with the Code of Conduct.
Jonathan Weyn (@jweyn) Wessel (@wesselb)

Problem

Running an N-member ensemble currently requires a loop that calls Aurora.forward()/rollout() once per member, and then manually combining the results. This under-utilises the GPU (N separate launches) when the GPU is capable of storing all of the ensemble state.

Change

Add a num_ensemble_members constructor argument to Aurora (default 1, fully backwards compatible). When set to N > 1, forward()/rollout() run all N members as a single, fused batched computation internally, rather than N separate calls. This is useful when combined with stochastic=True: since the backbone's existing per-batch-element noise injection means every instance receives independent noise.

This is purely an additional option: looping over forward()/rollout() to implement ensembling remains fully supported and unaffected.

Design notes

• Batch's public shape contract is untouched: no new dimension, no new methods. The batch-dimension tiling used to fuse the computation is a private implementation detail (_tile_batch/_split_batch in aurora/batch.py), never exposed on Batch itself.
• forward()'s return type is now Batch | list[Batch]: a plain Batch when num_ensemble_members == 1 (no change from today), or a list[Batch] of N standard-shaped batches: pred[m] is member m's ordinary, individually inspectable Batch.
• rollout() follows the same contract per yielded step, keeping the tiled representation internal across autoregressive steps for efficiency, and temporarily forcing model.num_ensemble_members = 1 during its loop (restored via try/finally, even on early generator closure) so nested forward() calls don't re-tile.
• Warns if num_ensemble_members > 1 is requested on a non-stochastic model, since all members would then be identical.

Tests

Added tests/v1p5/test_ensemble.py covering:

• _tile_batch/_split_batch round-tripping
• constructor validation/warnings
• forward()'s single-Batch vs. list[Batch] return contract
• member divergence under stochastic=True vs. identity under stochastic=False
• rollout()'s per-step output shape plus num_ensemble_members restoration (including on early .close()).

[Wessel (wesselb)]

Thanks Agnieszka Słowik (@Slowika) for opening a PR! I replied on the issue and mentioned the idea of using the batch size to produce multiple ensemble members simultaneously. Do you think that approach would suffice, or does this capability need to be added to the model explicitly?

[Agnieszka Słowik (Slowika)]

Thanks Agnieszka Słowik (Agnieszka Słowik (@Slowika)) for opening a PR! I replied on the issue and mentioned the idea of using the batch size to produce multiple ensemble members simultaneously. Do you think that approach would suffice, or does this capability need to be added to the model explicitly?

I've addressed your feedback. Should be ready for review now! Jonathan Weyn (@jweyn)

[Wessel (wesselb)]

Thanks for the changes, Agnieszka Słowik (@Slowika)! This is looking much simpler now. Really great! :) I've left some Claude-assisted comments. After those, I think this is ready to be merged!

[Agnieszka Słowik (Slowika)]

Hi Wessel (Wessel (@wesselb))! Thank you for the review. I addressed all comments and made some of the suggested changes. For the remaining Claude suggestions, they contradict your suggestions in the issue, especially regarding tiling. Could you please have another look?

[Wessel (wesselb)]

Thanks, Agnieszka Słowik (@Slowika)! I've replied to all outstanding comments.

[Agnieszka Słowik (Slowika)]

Hi Wessel (Wessel (@wesselb))! Thank you for the review :). I have addressed all comments and made the suggested changes.

[Wessel (wesselb)]

Thanks, Agnieszka Słowik (@Slowika)! This is looking great now. A few more comments and suggestions. I'd be happy to merge this in after those. :)

snip: deleted Claude comment about the PR description

[Agnieszka Słowik (Slowika)]

Thanks for the review, Wessel (@wesselb)! Your comments have been addressed and I agree with all of your suggestions.

I just need one final re-approval before this PR can be merged.