Repository navigation
Make the number of ensemble members internal to the Aurora model. #192
Description
Aurora 1.5 Ensemble has the option of being based on an ensemble with N members. This is currently achieved by running the model N times in the context in which the ensemble approach is required.
There should be an option to make the ensemble members internal to Aurora, so that getting all of the ensemble results is done using a single call to e.g. forward, as opposed to running N of them in a loop.
This would expose more opportunities for making the full use of GPUs. Currently, each ensemble member is a completely separate model with separate inputs and outputs. This means that a high-performance GPU may be under-utilised compared to representing all ensemble members using the same tensor.
Activity
wesselb commentedon Aug 19, 2026
Hey Agnieszka Słowik (@Slowika)! Sorry also for the very delayed reply here.
I think you might be able to generate multiple ensemble members simultaneously by copying the input batch along the batch dimension. The model will then generate a different noise for each batch elements. How would that sound? Would you be able to test whether that works as expected?
Slowika commentedon Aug 20, 2026
Hi Wessel (@wesselb) ! Thank you for the reply.
Are you suggesting that we should redefine what an Aurora batch is? This has some implications for backward compatibility and would affect all users of the
Batchclass: code usingBatchhas to be aware of how this per-ensemble-member information is embedded when handling the raw data.In the PR, the approach I took was that copying of the input batch data across a number of dimensions was handled during forward passes on stochastic models with fused data representations for each ensemble member (see
_tile_batch,_split_batch: they change the size ofBatchbut are helpers mainly forAurora). The alternative is to expose the fuse/unfuse functionality if we want to leave the implementation of Aurora alone, but this means that the user is responsible for this logic outside of the forward pass. Any other approach will mean the code will not type-check.Changing the definition of
Batchwould probably work just fine for the majority of uses.If this is the approach you are suggesting, I can make these changes.