From b5d8773c7efba2b0637c7c9668b1714c2c4edec8 Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Thu, 13 Aug 2026 12:46:50 +0100 Subject: [PATCH 1/4] Add support for internal ensembling in Aurora models. Add tests. --- aurora/batch.py | 31 ++++++ aurora/model/aurora.py | 48 ++++++++- aurora/rollout.py | 101 ++++++++++------- tests/v1p5/test_ensemble.py | 208 ++++++++++++++++++++++++++++++++++++ 4 files changed, 348 insertions(+), 40 deletions(-) create mode 100644 tests/v1p5/test_ensemble.py diff --git a/aurora/batch.py b/aurora/batch.py index d0616d7b..dca332df 100644 --- a/aurora/batch.py +++ b/aurora/batch.py @@ -311,6 +311,37 @@ def from_netcdf(cls, path: str | Path) -> "Batch": ) +def _tile_batch(batch: Batch, n: int) -> Batch: + """Tile `batch` along the batch dimension `n` times. + + Not part of the public `Batch` API. Used only by `aurora.Aurora.forward` and + `aurora.rollout.rollout` to run `n` ensemble members as a single fused computation. The + tiled batch dimension is an internal implementation detail and must be undone with + `_split_batch` before any result derived from it is returned to a caller. + """ + return dataclasses.replace( + batch, + surf_vars={k: v.repeat(n, *([1] * (v.dim() - 1))) for k, v in batch.surf_vars.items()}, + atmos_vars={k: v.repeat(n, *([1] * (v.dim() - 1))) for k, v in batch.atmos_vars.items()}, + metadata=dataclasses.replace(batch.metadata, time=batch.metadata.time * n), + ) + + +def _split_batch(batch: Batch, n: int) -> list[Batch]: + """Undo `_tile_batch`, splitting a tiled batch back into `n` standard-shaped batches.""" + b = next(iter(batch.surf_vars.values())).shape[0] // n + time = batch.metadata.time + return [ + dataclasses.replace( + batch, + surf_vars={k: v[m * b : (m + 1) * b] for k, v in batch.surf_vars.items()}, + atmos_vars={k: v[m * b : (m + 1) * b] for k, v in batch.atmos_vars.items()}, + metadata=dataclasses.replace(batch.metadata, time=time[m * b : (m + 1) * b]), + ) + for m in range(n) + ] + + def _np(x: torch.Tensor) -> np.ndarray: return x.detach().cpu().numpy() diff --git a/aurora/model/aurora.py b/aurora/model/aurora.py index 059dad24..71c7c50a 100644 --- a/aurora/model/aurora.py +++ b/aurora/model/aurora.py @@ -14,7 +14,7 @@ apply_activation_checkpointing, ) -from aurora.batch import Batch +from aurora.batch import Batch, _split_batch, _tile_batch from aurora.insolation import insolation from aurora.model.compat import ( _adapt_checkpoint_air_pollution, @@ -103,6 +103,7 @@ def __init__( clamp_at_first_step: bool = False, simulate_indexing_bug: bool = False, stochastic: bool = False, + num_ensemble_members: int = 1, use_updated_lead_time_embedding: bool = False, variable_lead_time: bool = False, rollout_input_clipping: Optional[dict[str, dict[str, Optional[float]]]] = None, @@ -200,6 +201,15 @@ def __init__( to the original implementation. Defaults to `False`. stochastic (bool, optional): If `True`, enable stochastic mode with noise injection. Defaults to `False`. + num_ensemble_members (int, optional): Number of ensemble members to produce + *internally* on every call to :meth:`forward`, as an alternative to looping over + separate calls and combining the results externally yourself (which remains + perfectly valid, e.g. if you need more control over how members are seeded or + combined). When set to a value greater than `1`, the batch is tiled + `num_ensemble_members` times internally and run through the model in a single, + fully-batched pass, which is far more efficient on a GPU than looping. This is + most useful in combination with `stochastic=True`, since every tiled copy then + receives independent noise. Defaults to `1`, i.e. no internal ensembling. use_updated_lead_time_embedding (bool, optional): Whether to use the updated lead time embedding with a minimum wavelength of 2 hours. Defaults to `False`. variable_lead_time (bool, optional): If `True`, use per-sample lead times passed @@ -236,6 +246,10 @@ def __init__( self.output_only_surf_vars = output_only_surf_vars self.output_only_atmos_vars = output_only_atmos_vars + if num_ensemble_members < 1: + raise ValueError("`num_ensemble_members` must be at least `1`.") + self.num_ensemble_members = num_ensemble_members + if self.surf_stats: warnings.warn( f"The normalisation statics for the following surface-level variables are manually " @@ -284,6 +298,14 @@ def __init__( use_updated_lead_time_embedding=use_updated_lead_time_embedding, ) + if num_ensemble_members > 1 and not self.backbone.stochastic: + warnings.warn( + f"`num_ensemble_members={num_ensemble_members}` was requested, but `stochastic=" + f"False`, so the model has no source of randomness. All ensemble members will be " + f"identical.", + stacklevel=2, + ) + self.decoder = Perceiver3DDecoder( surf_vars=surf_vars, atmos_vars=atmos_vars, @@ -339,7 +361,9 @@ def set_noise_accumulation(self, n: int = 0) -> None: """ self.backbone.set_noise_accumulation(n) - def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Batch: + def forward( + self, batch: Batch, lead_times: Optional[torch.Tensor] = None + ) -> Batch | list[Batch]: """Forward pass. Args: @@ -349,7 +373,12 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba `variable_lead_time=True`. Ignored otherwise. Returns: - :class:`Batch`: Prediction for the batch. + :class:`Batch` | list[:class:`Batch`]: Prediction for `batch`. If + `self.num_ensemble_members == 1` (the default, i.e. no internal ensembling), this + is a single `Batch`, exactly as `batch`. If `self.num_ensemble_members > 1`, all + members are computed internally as a single fused pass, but the result is a list + of `num_ensemble_members` standard-shaped `Batch`\\ s, one per ensemble member, + each with the same batch dimension as `batch`. """ batch = self.batch_transform_hook(batch) @@ -361,6 +390,17 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba batch = batch.crop(patch_size=self.patch_size) batch = batch.to(p.device) + if self.num_ensemble_members > 1: + if lead_times is not None: + lead_times = lead_times.repeat(self.num_ensemble_members) + # This tiling implements *internal* ensembling only, as a private implementation + # detail: it lets every ensemble member run through the encoder/backbone/decoder as a + # single fused batch instead of `num_ensemble_members` separate calls. Ensembling by + # externally looping over `forward` yourself remains equally valid and is unaffected. + # The tiled batch is split back apart into standard-shaped batches below, right before + # `forward` returns. + batch = _tile_batch(batch, self.num_ensemble_members) + H, W = batch.spatial_shape patch_res = ( self.encoder.latent_levels, @@ -486,6 +526,8 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba pred = self._post_unnorm_hook(batch, pred) + if self.num_ensemble_members > 1: + return _split_batch(pred, self.num_ensemble_members) return pred def batch_transform_hook(self, batch: Batch) -> Batch: diff --git a/aurora/rollout.py b/aurora/rollout.py index 3a84ce6b..462f72c6 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -2,11 +2,11 @@ import dataclasses import math -from typing import Generator, Optional, Sequence +from typing import Generator, Optional, Sequence, cast import torch -from aurora.batch import Batch +from aurora.batch import Batch, _split_batch, _tile_batch from aurora.model.aurora import Aurora __all__ = ["rollout"] @@ -48,7 +48,7 @@ def rollout( fine_lead_times: Optional[Sequence[float]] = None, use_noise_accumulation: bool = True, apply_rollout_input_clipping: bool = True, -) -> Generator[Batch, None, None]: +) -> Generator[Batch | list[Batch], None, None]: """Perform a roll-out to make long-term predictions. For Aurora models prior to Aurora 1.5, the rollout is straightforward: iteratively make a @@ -92,7 +92,12 @@ def rollout( Default: `True`. Yields: - :class:`aurora.Batch`: The prediction after every (sub-)step. + :class:`aurora.Batch` | list[:class:`aurora.Batch`]: The prediction after every (sub-)step. + If `model.num_ensemble_members == 1` (the default, i.e. no internal ensembling), this + is a single `Batch`. If `model.num_ensemble_members > 1`, all members are computed + internally as a single fused pass, but each yielded value is a list of + `num_ensemble_members` standard-shaped `Batch`\\ s, one per ensemble member; see + :meth:`aurora.Aurora.forward`. """ # We will need to concatenate data, so ensure that everything is already of the right form. batch = model.batch_transform_hook(batch) # This might modify the available variables. @@ -114,37 +119,59 @@ def rollout( f"of {base_timestep_hours} hours. Found {fine_lead_times[-1]} hours." ) - # Enable noise accumulation when the model is stochastic and sub-stepping. - if use_noise_accumulation and fine_lead_times is not None: - model.set_noise_accumulation(n=len(fine_lead_times)) - - # Pre-compute the base lead-time tensor for models with variable lead time support. - base_lead_times: Optional[torch.Tensor] = None - if model.variable_lead_time: - base_lead_times = _make_lead_time_tensor(batch, model.timestep.total_seconds() / 3600.0) - - for _ in range(steps): - if fine_lead_times is not None: - # Inner loop: iterate over sub-step lead times. - for lt_hours in fine_lead_times: - sub_lead_times = _make_lead_time_tensor(batch, lt_hours) - pred = model.forward(batch, lead_times=sub_lead_times) - - yield pred - - # If desired, apply clipping before feeding predictions back as inputs. - if apply_rollout_input_clipping: - pred = model.apply_rollout_input_clipping(pred) - batch = _advance_batch(batch, pred) - else: - pred = model.forward(batch, lead_times=base_lead_times) - - yield pred - - if apply_rollout_input_clipping: - pred = model.apply_rollout_input_clipping(pred) - batch = _advance_batch(batch, pred) + # If the model produces ensemble members internally, tile the batch once up front, then + # temporarily disable further expansion so it isn't repeated on every step's `forward` call. + # The tiled representation is kept purely internal to this loop: every `pred` yielded to the + # caller is split back into standard-shaped batches first. + num_ensemble_members = model.num_ensemble_members + if num_ensemble_members > 1: + batch = _tile_batch(batch, num_ensemble_members) + model.num_ensemble_members = 1 + + try: + # Enable noise accumulation when the model is stochastic and sub-stepping. + if use_noise_accumulation and fine_lead_times is not None: + model.set_noise_accumulation(n=len(fine_lead_times)) + + # Pre-compute the base lead-time tensor for models with variable lead time support. + base_lead_times: Optional[torch.Tensor] = None + if model.variable_lead_time: + base_lead_times = _make_lead_time_tensor( + batch, model.timestep.total_seconds() / 3600.0 + ) - # Disable noise accumulation after roll-out is complete, in case the model will be used for - # normal inference or training afterwards. - model.set_noise_accumulation(n=0) + for _ in range(steps): + if fine_lead_times is not None: + # Inner loop: iterate over sub-step lead times. + for lt_hours in fine_lead_times: + sub_lead_times = _make_lead_time_tensor(batch, lt_hours) + # `num_ensemble_members` is forced to `1` for the duration of this loop, so + # `forward` always returns a plain `Batch` here, never a `list[Batch]`. + pred = cast(Batch, model.forward(batch, lead_times=sub_lead_times)) + + yield _split_batch(pred, num_ensemble_members) if ( + num_ensemble_members > 1 + ) else pred + + # If desired, apply clipping before feeding predictions back as inputs. + if apply_rollout_input_clipping: + pred = model.apply_rollout_input_clipping(pred) + batch = _advance_batch(batch, pred) + else: + pred = cast(Batch, model.forward(batch, lead_times=base_lead_times)) + + yield _split_batch(pred, num_ensemble_members) if ( + num_ensemble_members > 1 + ) else pred + + if apply_rollout_input_clipping: + pred = model.apply_rollout_input_clipping(pred) + batch = _advance_batch(batch, pred) + + # Disable noise accumulation after roll-out is complete, in case the model will be used for + # normal inference or training afterwards. + model.set_noise_accumulation(n=0) + finally: + # Restore the model's ensemble configuration, whether the roll-out ran to completion or + # was abandoned early. + model.num_ensemble_members = num_ensemble_members diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py new file mode 100644 index 00000000..41c548a2 --- /dev/null +++ b/tests/v1p5/test_ensemble.py @@ -0,0 +1,208 @@ +"""Copyright (c) Microsoft Corporation. Licensed under the MIT license. + +Tests for internal ensemble members (`num_ensemble_members`). +""" + +import warnings +from datetime import datetime + +import pytest +import torch + +from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5 +from aurora import Batch, Metadata, rollout +from aurora.batch import _split_batch, _tile_batch + + +def _make_ensemble_test_batch(b: int = 2) -> Batch: + """A small batch with a configurable batch size `b`, used to test tiling/splitting.""" + h, w = 8, 8 + return Batch( + surf_vars={"2t": torch.randn(b, 2, h, w)}, + static_vars={"lsm": torch.randn(h, w)}, + atmos_vars={"z": torch.randn(b, 2, 2, h, w)}, + metadata=Metadata( + lat=torch.linspace(90, -90, h), + lon=torch.linspace(0, 360, w + 1)[:-1], + time=tuple(datetime(2023, 6, 15, i, 0) for i in range(b)), + atmos_levels=(500, 850), + ), + ) + + +def test_tile_and_split_batch_roundtrip(): + b, n = 2, 3 + batch = _make_ensemble_test_batch(b) + + tiled = _tile_batch(batch, n) + + v = tiled.surf_vars["2t"] + assert v.shape[0] == n * b + for m in range(n): + torch.testing.assert_close(v[m * b : (m + 1) * b], batch.surf_vars["2t"]) + + v = tiled.atmos_vars["z"] + assert v.shape[0] == n * b + for m in range(n): + torch.testing.assert_close(v[m * b : (m + 1) * b], batch.atmos_vars["z"]) + + assert len(tiled.metadata.time) == n * b + for m in range(n): + assert tiled.metadata.time[m * b : (m + 1) * b] == batch.metadata.time + + # Static variables have no batch dimension and are untouched. + torch.testing.assert_close(tiled.static_vars["lsm"], batch.static_vars["lsm"]) + + # Splitting undoes the tiling: every member is identical to the original, standard-shaped + # batch (tiling itself introduces no randomness). + members = _split_batch(tiled, n) + assert len(members) == n + for member in members: + torch.testing.assert_close(member.surf_vars["2t"], batch.surf_vars["2t"]) + torch.testing.assert_close(member.atmos_vars["z"], batch.atmos_vars["z"]) + assert member.metadata.time == batch.metadata.time + + +def test_num_ensemble_members_must_be_positive(): + with pytest.raises(ValueError, match="num_ensemble_members"): + _make_small_v1p5(num_ensemble_members=0) + + +def test_num_ensemble_members_warns_without_stochastic(): + with pytest.warns(UserWarning, match="stochastic"): + _make_small_v1p5(num_ensemble_members=2, stochastic=False) + + +def test_num_ensemble_members_no_warning_with_stochastic(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + _make_small_v1p5(num_ensemble_members=2, stochastic=True) + + +def test_forward_returns_single_batch_when_num_ensemble_members_one(): + model = _make_small_v1p5() + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + + with torch.inference_mode(): + pred = model.forward(batch, lead_times=torch.full((1,), 6.0)) + + assert isinstance(pred, Batch) + + +def test_forward_returns_list_of_standard_shaped_batches(): + n = 3 + model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + + assert isinstance(pred, list) + assert len(pred) == n + for member in pred: + assert isinstance(member, Batch) + for v in member.surf_vars.values(): + assert v.shape[0] == b + for v in member.static_vars.values(): + # Static variables have no batch dimension. + assert v.dim() == 2 + + +def test_forward_ensemble_members_differ_when_stochastic(): + n = 3 + model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + + for i in range(n): + for j in range(i + 1, n): + # Use exact equality: with a small, randomly-initialised model, the noise's effect on + # the output can be numerically tiny, so `allclose` may not reliably distinguish them. + assert not torch.equal(pred[i].surf_vars["2t"], pred[j].surf_vars["2t"]) + + +def test_forward_ensemble_members_identical_without_stochastic(): + n = 3 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + model = _make_small_v1p5(stochastic=False, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + + for m in range(1, n): + # Loose tolerance: floating-point ops (e.g. batched matmul/softmax reductions) are not + # strictly invariant to how many other (tiled) rows share the batch, so bitwise equality + # isn't guaranteed even though the members are mathematically identical computations. + torch.testing.assert_close( + pred[0].surf_vars["2t"], pred[m].surf_vars["2t"], atol=1e-3, rtol=1e-3 + ) + + +def test_rollout_yields_list_of_standard_shaped_batches_across_steps(): + n = 2 + model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + preds = list(rollout(model, batch, steps=3)) + + assert len(preds) == 3 + for step_pred in preds: + assert isinstance(step_pred, list) + assert len(step_pred) == n + for member in step_pred: + for v in member.surf_vars.values(): + assert v.shape[0] == b + + # The model's ensemble configuration is restored after the roll-out completes. + assert model.num_ensemble_members == n + + +def test_rollout_restores_num_ensemble_members_on_early_close(): + n = 2 + model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + + with torch.inference_mode(): + gen = rollout(model, batch, steps=5) + next(gen) + gen.close() + + assert model.num_ensemble_members == n + + +def test_rollout_num_ensemble_members_one_is_unaffected(): + model = _make_small_v1p5() + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + preds = list(rollout(model, batch, steps=2)) + + for pred in preds: + assert isinstance(pred, Batch) + for v in pred.surf_vars.values(): + assert v.shape[0] == b + assert model.num_ensemble_members == 1 From c1a9cad3c4189bda52d50df0403638c07e4c0ba4 Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Thu, 13 Aug 2026 14:43:45 +0100 Subject: [PATCH 2/4] Address formatting complaint. --- aurora/rollout.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/aurora/rollout.py b/aurora/rollout.py index 462f72c6..b619b0b3 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -136,9 +136,7 @@ def rollout( # Pre-compute the base lead-time tensor for models with variable lead time support. base_lead_times: Optional[torch.Tensor] = None if model.variable_lead_time: - base_lead_times = _make_lead_time_tensor( - batch, model.timestep.total_seconds() / 3600.0 - ) + base_lead_times = _make_lead_time_tensor(batch, model.timestep.total_seconds() / 3600.0) for _ in range(steps): if fine_lead_times is not None: @@ -149,9 +147,11 @@ def rollout( # `forward` always returns a plain `Batch` here, never a `list[Batch]`. pred = cast(Batch, model.forward(batch, lead_times=sub_lead_times)) - yield _split_batch(pred, num_ensemble_members) if ( - num_ensemble_members > 1 - ) else pred + yield ( + _split_batch(pred, num_ensemble_members) + if num_ensemble_members > 1 + else pred + ) # If desired, apply clipping before feeding predictions back as inputs. if apply_rollout_input_clipping: @@ -160,9 +160,9 @@ def rollout( else: pred = cast(Batch, model.forward(batch, lead_times=base_lead_times)) - yield _split_batch(pred, num_ensemble_members) if ( - num_ensemble_members > 1 - ) else pred + yield ( + _split_batch(pred, num_ensemble_members) if num_ensemble_members > 1 else pred + ) if apply_rollout_input_clipping: pred = model.apply_rollout_input_clipping(pred) From a4c12def38670e5bbec89cfce0287c5c790b6761 Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Thu, 13 Aug 2026 15:12:42 +0100 Subject: [PATCH 3/4] Fix test. --- tests/v1p5/test_ensemble.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py index 41c548a2..15c6b98f 100644 --- a/tests/v1p5/test_ensemble.py +++ b/tests/v1p5/test_ensemble.py @@ -10,8 +10,25 @@ import torch from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5 -from aurora import Batch, Metadata, rollout +from aurora import Aurora, Batch, Metadata, rollout from aurora.batch import _split_batch, _tile_batch +from aurora.model.film import AdaptiveLayerNorm + + +def _unzero_adaptive_layer_norms(model: Aurora, std: float = 0.1) -> None: + """Nudge every `AdaptiveLayerNorm`'s modulation away from its zero initialisation. + + At construction, `AdaptiveLayerNorm.ln_modulation` is exactly zero-initialised (the + `adaLN-Zero` trick), which makes a freshly-built, untrained model exactly insensitive to its + conditioning signal `c` -- which is what carries the ensemble noise. Without this, no output + difference a test observes between ensemble members can be attributed to noise, since noise + provably has zero effect on such a model. + """ + for m in model.modules(): + if isinstance(m, AdaptiveLayerNorm): + with torch.no_grad(): + m.ln_modulation[-1].weight.normal_(std=std) + m.ln_modulation[-1].bias.normal_(std=std) def _make_ensemble_test_batch(b: int = 2) -> Batch: @@ -115,7 +132,12 @@ def test_forward_returns_list_of_standard_shaped_batches(): def test_forward_ensemble_members_differ_when_stochastic(): n = 3 + torch.manual_seed(0) model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + # Un-zero the modulation so noise has a real, appreciable effect (see helper docstring); + # otherwise this test cannot distinguish genuine noise sensitivity from incidental + # floating-point batching noise (see `test_forward_ensemble_members_identical_without_stochastic`). + _unzero_adaptive_layer_norms(model) model.eval() surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) batch = _make_batch(surf_vars=surf_vars) @@ -124,11 +146,13 @@ def test_forward_ensemble_members_differ_when_stochastic(): with torch.inference_mode(): pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + # Threshold well above the ~1e-3 floating-point batching floor established in + # `test_forward_ensemble_members_identical_without_stochastic`, so a pass here can only be + # explained by the injected noise actually differing per member, not incidental rounding. for i in range(n): for j in range(i + 1, n): - # Use exact equality: with a small, randomly-initialised model, the noise's effect on - # the output can be numerically tiny, so `allclose` may not reliably distinguish them. - assert not torch.equal(pred[i].surf_vars["2t"], pred[j].surf_vars["2t"]) + diff = (pred[i].surf_vars["2t"] - pred[j].surf_vars["2t"]).abs().max() + assert diff > 1e-2 def test_forward_ensemble_members_identical_without_stochastic(): From f0fa3683d8cf605fface49b5073c0d882948a0ea Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Thu, 13 Aug 2026 15:40:53 +0100 Subject: [PATCH 4/4] Split forward()/rollout() from forward_ensemble()/rollout_ensemble() for backward compatibility. --- aurora/__init__.py | 3 +- aurora/model/aurora.py | 77 ++++++++++++++------- aurora/rollout.py | 130 ++++++++++++++++++++---------------- tests/v1p5/test_ensemble.py | 54 ++++++++++----- 4 files changed, 163 insertions(+), 101 deletions(-) diff --git a/aurora/__init__.py b/aurora/__init__.py index 9403fb7f..8f414aac 100644 --- a/aurora/__init__.py +++ b/aurora/__init__.py @@ -14,7 +14,7 @@ AuroraV1p5Ensemble, AuroraWave, ) -from aurora.rollout import rollout +from aurora.rollout import rollout, rollout_ensemble from aurora.tracker import Tracker __all__ = [ @@ -32,5 +32,6 @@ "Metadata", "insolation", "rollout", + "rollout_ensemble", "Tracker", ] diff --git a/aurora/model/aurora.py b/aurora/model/aurora.py index 71c7c50a..7886955a 100644 --- a/aurora/model/aurora.py +++ b/aurora/model/aurora.py @@ -202,14 +202,16 @@ def __init__( stochastic (bool, optional): If `True`, enable stochastic mode with noise injection. Defaults to `False`. num_ensemble_members (int, optional): Number of ensemble members to produce - *internally* on every call to :meth:`forward`, as an alternative to looping over - separate calls and combining the results externally yourself (which remains - perfectly valid, e.g. if you need more control over how members are seeded or - combined). When set to a value greater than `1`, the batch is tiled - `num_ensemble_members` times internally and run through the model in a single, - fully-batched pass, which is far more efficient on a GPU than looping. This is - most useful in combination with `stochastic=True`, since every tiled copy then - receives independent noise. Defaults to `1`, i.e. no internal ensembling. + *internally* on every call to :meth:`forward_ensemble`, as an alternative to + looping over separate :meth:`forward` calls and combining the results externally + yourself (which remains perfectly valid, e.g. if you need more control over how + members are seeded or combined). When set to a value greater than `1`, the batch + is tiled `num_ensemble_members` times internally and run through the model in a + single, fully-batched pass, which is far more efficient on a GPU than looping. + This is most useful in combination with `stochastic=True`, since every tiled copy + then receives independent noise. When greater than `1`, plain :meth:`forward` + raises, since it can only ever return a single `Batch`; use + :meth:`forward_ensemble` instead. Defaults to `1`, i.e. no internal ensembling. use_updated_lead_time_embedding (bool, optional): Whether to use the updated lead time embedding with a minimum wavelength of 2 hours. Defaults to `False`. variable_lead_time (bool, optional): If `True`, use per-sample lead times passed @@ -361,9 +363,7 @@ def set_noise_accumulation(self, n: int = 0) -> None: """ self.backbone.set_noise_accumulation(n) - def forward( - self, batch: Batch, lead_times: Optional[torch.Tensor] = None - ) -> Batch | list[Batch]: + def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Batch: """Forward pass. Args: @@ -373,12 +373,47 @@ def forward( `variable_lead_time=True`. Ignored otherwise. Returns: - :class:`Batch` | list[:class:`Batch`]: Prediction for `batch`. If - `self.num_ensemble_members == 1` (the default, i.e. no internal ensembling), this - is a single `Batch`, exactly as `batch`. If `self.num_ensemble_members > 1`, all - members are computed internally as a single fused pass, but the result is a list - of `num_ensemble_members` standard-shaped `Batch`\\ s, one per ensemble member, - each with the same batch dimension as `batch`. + :class:`Batch`: Prediction for `batch`. + """ + if self.num_ensemble_members > 1: + raise RuntimeError( + f"This model was constructed with `num_ensemble_members=" + f"{self.num_ensemble_members}`. Use `forward_ensemble` instead of `forward` to " + f"obtain all ensemble members." + ) + return self._forward_impl(batch, lead_times) + + def forward_ensemble( + self, batch: Batch, lead_times: Optional[torch.Tensor] = None + ) -> list[Batch]: + """Forward pass producing all `self.num_ensemble_members` ensemble members internally. + + All members are computed internally as a single fused pass through the + encoder/backbone/decoder, rather than looping over separate `forward` calls and combining + the results externally yourself (which remains equally valid and unaffected). This is most + useful in combination with `stochastic=True`, since every internally-tiled copy then + receives independent noise. + + Args: + batch (:class:`aurora.Batch`): Batch to run the model on. + lead_times (:class:`torch.Tensor`, optional): Per-sample lead times of shape + `(batch,)` in hours. Required when the model was configured with + `variable_lead_time=True`. Ignored otherwise. + + Returns: + list[:class:`Batch`]: A list of `self.num_ensemble_members` standard-shaped `Batch`\\ + s, one per ensemble member, each with the same batch dimension as `batch`. + """ + pred = self._forward_impl(batch, lead_times) + return _split_batch(pred, self.num_ensemble_members) + + def _forward_impl(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Batch: + """Shared implementation for `forward` and `forward_ensemble`. + + Internally tiles `batch` by `self.num_ensemble_members` before running it through the + encoder/backbone/decoder as a single fused batch, when greater than `1`. The tiled batch + dimension is a private implementation detail: `forward` forbids it (see above) and + `forward_ensemble` splits it back apart before returning. """ batch = self.batch_transform_hook(batch) @@ -393,12 +428,6 @@ def forward( if self.num_ensemble_members > 1: if lead_times is not None: lead_times = lead_times.repeat(self.num_ensemble_members) - # This tiling implements *internal* ensembling only, as a private implementation - # detail: it lets every ensemble member run through the encoder/backbone/decoder as a - # single fused batch instead of `num_ensemble_members` separate calls. Ensembling by - # externally looping over `forward` yourself remains equally valid and is unaffected. - # The tiled batch is split back apart into standard-shaped batches below, right before - # `forward` returns. batch = _tile_batch(batch, self.num_ensemble_members) H, W = batch.spatial_shape @@ -526,8 +555,6 @@ def forward( pred = self._post_unnorm_hook(batch, pred) - if self.num_ensemble_members > 1: - return _split_batch(pred, self.num_ensemble_members) return pred def batch_transform_hook(self, batch: Batch) -> Batch: diff --git a/aurora/rollout.py b/aurora/rollout.py index b619b0b3..af166a03 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -2,14 +2,14 @@ import dataclasses import math -from typing import Generator, Optional, Sequence, cast +from typing import Generator, Optional, Sequence import torch from aurora.batch import Batch, _split_batch, _tile_batch from aurora.model.aurora import Aurora -__all__ = ["rollout"] +__all__ = ["rollout", "rollout_ensemble"] def _make_lead_time_tensor(batch: Batch, lead_time_hours: float) -> torch.Tensor: @@ -48,7 +48,7 @@ def rollout( fine_lead_times: Optional[Sequence[float]] = None, use_noise_accumulation: bool = True, apply_rollout_input_clipping: bool = True, -) -> Generator[Batch | list[Batch], None, None]: +) -> Generator[Batch, None, None]: """Perform a roll-out to make long-term predictions. For Aurora models prior to Aurora 1.5, the rollout is straightforward: iteratively make a @@ -92,12 +92,7 @@ def rollout( Default: `True`. Yields: - :class:`aurora.Batch` | list[:class:`aurora.Batch`]: The prediction after every (sub-)step. - If `model.num_ensemble_members == 1` (the default, i.e. no internal ensembling), this - is a single `Batch`. If `model.num_ensemble_members > 1`, all members are computed - internally as a single fused pass, but each yielded value is a list of - `num_ensemble_members` standard-shaped `Batch`\\ s, one per ensemble member; see - :meth:`aurora.Aurora.forward`. + :class:`aurora.Batch`: The prediction after every (sub-)step. """ # We will need to concatenate data, so ensure that everything is already of the right form. batch = model.batch_transform_hook(batch) # This might modify the available variables. @@ -119,58 +114,77 @@ def rollout( f"of {base_timestep_hours} hours. Found {fine_lead_times[-1]} hours." ) - # If the model produces ensemble members internally, tile the batch once up front, then - # temporarily disable further expansion so it isn't repeated on every step's `forward` call. - # The tiled representation is kept purely internal to this loop: every `pred` yielded to the - # caller is split back into standard-shaped batches first. + # Enable noise accumulation when the model is stochastic and sub-stepping. + if use_noise_accumulation and fine_lead_times is not None: + model.set_noise_accumulation(n=len(fine_lead_times)) + + # Pre-compute the base lead-time tensor for models with variable lead time support. + base_lead_times: Optional[torch.Tensor] = None + if model.variable_lead_time: + base_lead_times = _make_lead_time_tensor(batch, model.timestep.total_seconds() / 3600.0) + + for _ in range(steps): + if fine_lead_times is not None: + # Inner loop: iterate over sub-step lead times. + for lt_hours in fine_lead_times: + sub_lead_times = _make_lead_time_tensor(batch, lt_hours) + pred = model.forward(batch, lead_times=sub_lead_times) + + yield pred + + # If desired, apply clipping before feeding predictions back as inputs. + if apply_rollout_input_clipping: + pred = model.apply_rollout_input_clipping(pred) + batch = _advance_batch(batch, pred) + else: + pred = model.forward(batch, lead_times=base_lead_times) + + yield pred + + if apply_rollout_input_clipping: + pred = model.apply_rollout_input_clipping(pred) + batch = _advance_batch(batch, pred) + + # Disable noise accumulation after roll-out is complete, in case the model will be used for + # normal inference or training afterwards. + model.set_noise_accumulation(n=0) + + +def rollout_ensemble( + model: Aurora, + batch: Batch, + steps: int, + fine_lead_times: Optional[Sequence[float]] = None, + use_noise_accumulation: bool = True, + apply_rollout_input_clipping: bool = True, +) -> Generator[list[Batch], None, None]: + """Like `rollout`, but produces `model.num_ensemble_members` ensemble members internally on + every step, as a single fused pass, instead of `rollout` yielding one `Batch` per step. + + All arguments are identical to `rollout`; see there for details. + + Yields: + list[:class:`aurora.Batch`]: A list of `model.num_ensemble_members` standard-shaped + `Batch`\\ s after every (sub-)step, one per ensemble member; see + :meth:`aurora.Aurora.forward_ensemble`. + """ num_ensemble_members = model.num_ensemble_members - if num_ensemble_members > 1: - batch = _tile_batch(batch, num_ensemble_members) - model.num_ensemble_members = 1 + # Tile the batch once up front, then temporarily disable further expansion so it isn't + # repeated by `_forward_impl` on every step. The tiled representation is kept purely internal + # to this loop: every `pred` yielded to the caller is split back into standard-shaped batches. + batch = _tile_batch(batch, num_ensemble_members) + model.num_ensemble_members = 1 try: - # Enable noise accumulation when the model is stochastic and sub-stepping. - if use_noise_accumulation and fine_lead_times is not None: - model.set_noise_accumulation(n=len(fine_lead_times)) - - # Pre-compute the base lead-time tensor for models with variable lead time support. - base_lead_times: Optional[torch.Tensor] = None - if model.variable_lead_time: - base_lead_times = _make_lead_time_tensor(batch, model.timestep.total_seconds() / 3600.0) - - for _ in range(steps): - if fine_lead_times is not None: - # Inner loop: iterate over sub-step lead times. - for lt_hours in fine_lead_times: - sub_lead_times = _make_lead_time_tensor(batch, lt_hours) - # `num_ensemble_members` is forced to `1` for the duration of this loop, so - # `forward` always returns a plain `Batch` here, never a `list[Batch]`. - pred = cast(Batch, model.forward(batch, lead_times=sub_lead_times)) - - yield ( - _split_batch(pred, num_ensemble_members) - if num_ensemble_members > 1 - else pred - ) - - # If desired, apply clipping before feeding predictions back as inputs. - if apply_rollout_input_clipping: - pred = model.apply_rollout_input_clipping(pred) - batch = _advance_batch(batch, pred) - else: - pred = cast(Batch, model.forward(batch, lead_times=base_lead_times)) - - yield ( - _split_batch(pred, num_ensemble_members) if num_ensemble_members > 1 else pred - ) - - if apply_rollout_input_clipping: - pred = model.apply_rollout_input_clipping(pred) - batch = _advance_batch(batch, pred) - - # Disable noise accumulation after roll-out is complete, in case the model will be used for - # normal inference or training afterwards. - model.set_noise_accumulation(n=0) + for pred in rollout( + model, + batch, + steps, + fine_lead_times=fine_lead_times, + use_noise_accumulation=use_noise_accumulation, + apply_rollout_input_clipping=apply_rollout_input_clipping, + ): + yield _split_batch(pred, num_ensemble_members) finally: # Restore the model's ensemble configuration, whether the roll-out ran to completion or # was abandoned early. diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py index 15c6b98f..db7a3755 100644 --- a/tests/v1p5/test_ensemble.py +++ b/tests/v1p5/test_ensemble.py @@ -10,7 +10,7 @@ import torch from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5 -from aurora import Aurora, Batch, Metadata, rollout +from aurora import Aurora, Batch, Metadata, rollout, rollout_ensemble from aurora.batch import _split_batch, _tile_batch from aurora.model.film import AdaptiveLayerNorm @@ -96,7 +96,7 @@ def test_num_ensemble_members_no_warning_with_stochastic(): _make_small_v1p5(num_ensemble_members=2, stochastic=True) -def test_forward_returns_single_batch_when_num_ensemble_members_one(): +def test_forward_returns_batch_when_num_ensemble_members_one(): model = _make_small_v1p5() model.eval() surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) @@ -108,7 +108,17 @@ def test_forward_returns_single_batch_when_num_ensemble_members_one(): assert isinstance(pred, Batch) -def test_forward_returns_list_of_standard_shaped_batches(): +def test_forward_raises_when_num_ensemble_members_greater_than_one(): + model = _make_small_v1p5(stochastic=True, num_ensemble_members=3) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + + with pytest.raises(RuntimeError, match="forward_ensemble"): + model.forward(batch, lead_times=torch.full((1,), 6.0)) + + +def test_forward_ensemble_returns_list_of_standard_shaped_batches(): n = 3 model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) model.eval() @@ -117,7 +127,7 @@ def test_forward_returns_list_of_standard_shaped_batches(): b = next(iter(batch.surf_vars.values())).shape[0] with torch.inference_mode(): - pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + pred = model.forward_ensemble(batch, lead_times=torch.full((b,), 6.0)) assert isinstance(pred, list) assert len(pred) == n @@ -136,7 +146,8 @@ def test_forward_ensemble_members_differ_when_stochastic(): model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) # Un-zero the modulation so noise has a real, appreciable effect (see helper docstring); # otherwise this test cannot distinguish genuine noise sensitivity from incidental - # floating-point batching noise (see `test_forward_ensemble_members_identical_without_stochastic`). + # floating-point batching noise + # (see `test_forward_ensemble_members_identical_without_stochastic`). _unzero_adaptive_layer_norms(model) model.eval() surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) @@ -144,7 +155,7 @@ def test_forward_ensemble_members_differ_when_stochastic(): b = next(iter(batch.surf_vars.values())).shape[0] with torch.inference_mode(): - pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + pred = model.forward_ensemble(batch, lead_times=torch.full((b,), 6.0)) # Threshold well above the ~1e-3 floating-point batching floor established in # `test_forward_ensemble_members_identical_without_stochastic`, so a pass here can only be @@ -166,7 +177,7 @@ def test_forward_ensemble_members_identical_without_stochastic(): b = next(iter(batch.surf_vars.values())).shape[0] with torch.inference_mode(): - pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + pred = model.forward_ensemble(batch, lead_times=torch.full((b,), 6.0)) for m in range(1, n): # Loose tolerance: floating-point ops (e.g. batched matmul/softmax reductions) are not @@ -177,7 +188,7 @@ def test_forward_ensemble_members_identical_without_stochastic(): ) -def test_rollout_yields_list_of_standard_shaped_batches_across_steps(): +def test_rollout_ensemble_yields_list_of_standard_shaped_batches_across_steps(): n = 2 model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) model.eval() @@ -186,11 +197,10 @@ def test_rollout_yields_list_of_standard_shaped_batches_across_steps(): b = next(iter(batch.surf_vars.values())).shape[0] with torch.inference_mode(): - preds = list(rollout(model, batch, steps=3)) + preds = list(rollout_ensemble(model, batch, steps=3)) assert len(preds) == 3 for step_pred in preds: - assert isinstance(step_pred, list) assert len(step_pred) == n for member in step_pred: for v in member.surf_vars.values(): @@ -200,7 +210,17 @@ def test_rollout_yields_list_of_standard_shaped_batches_across_steps(): assert model.num_ensemble_members == n -def test_rollout_restores_num_ensemble_members_on_early_close(): +def test_rollout_raises_when_num_ensemble_members_greater_than_one(): + model = _make_small_v1p5(stochastic=True, num_ensemble_members=2) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + + with torch.inference_mode(), pytest.raises(RuntimeError, match="forward_ensemble"): + next(rollout(model, batch, steps=1)) + + +def test_rollout_ensemble_restores_num_ensemble_members_on_early_close(): n = 2 model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) model.eval() @@ -208,14 +228,14 @@ def test_rollout_restores_num_ensemble_members_on_early_close(): batch = _make_batch(surf_vars=surf_vars) with torch.inference_mode(): - gen = rollout(model, batch, steps=5) + gen = rollout_ensemble(model, batch, steps=5) next(gen) gen.close() assert model.num_ensemble_members == n -def test_rollout_num_ensemble_members_one_is_unaffected(): +def test_rollout_ensemble_num_ensemble_members_one_still_works(): model = _make_small_v1p5() model.eval() surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) @@ -223,10 +243,10 @@ def test_rollout_num_ensemble_members_one_is_unaffected(): b = next(iter(batch.surf_vars.values())).shape[0] with torch.inference_mode(): - preds = list(rollout(model, batch, steps=2)) + preds = list(rollout_ensemble(model, batch, steps=2)) - for pred in preds: - assert isinstance(pred, Batch) - for v in pred.surf_vars.values(): + for step_pred in preds: + assert len(step_pred) == 1 + for v in step_pred[0].surf_vars.values(): assert v.shape[0] == b assert model.num_ensemble_members == 1