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()).