From b5d8773c7efba2b0637c7c9668b1714c2c4edec8 Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Thu, 13 Aug 2026 12:46:50 +0100 Subject: [PATCH 01/13] 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 02/13] 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 03/13] 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 04/13] 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 From aa60f79e107bd9b3ae8bb82ee4a54035ec882c7e Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Thu, 27 Aug 2026 14:40:10 +0100 Subject: [PATCH 05/13] Address wesselb feedback. --- aurora/batch.py | 8 +-- aurora/model/aurora.py | 73 +-------------------- aurora/rollout.py | 8 +-- tests/v1p5/test_ensemble.py | 123 ++++++------------------------------ 4 files changed, 29 insertions(+), 183 deletions(-) diff --git a/aurora/batch.py b/aurora/batch.py index dca332df..adf7ef32 100644 --- a/aurora/batch.py +++ b/aurora/batch.py @@ -311,13 +311,13 @@ def from_netcdf(cls, path: str | Path) -> "Batch": ) -def _tile_batch(batch: Batch, n: int) -> 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. + `split_batch` before any result derived from it is returned to a caller. """ return dataclasses.replace( batch, @@ -327,8 +327,8 @@ def _tile_batch(batch: Batch, n: int) -> Batch: ) -def _split_batch(batch: Batch, n: int) -> list[Batch]: - """Undo `_tile_batch`, splitting a tiled batch back into `n` standard-shaped batches.""" +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 [ diff --git a/aurora/model/aurora.py b/aurora/model/aurora.py index 7886955a..059dad24 100644 --- a/aurora/model/aurora.py +++ b/aurora/model/aurora.py @@ -14,7 +14,7 @@ apply_activation_checkpointing, ) -from aurora.batch import Batch, _split_batch, _tile_batch +from aurora.batch import Batch from aurora.insolation import insolation from aurora.model.compat import ( _adapt_checkpoint_air_pollution, @@ -103,7 +103,6 @@ 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, @@ -201,17 +200,6 @@ 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_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 @@ -248,10 +236,6 @@ 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 " @@ -300,14 +284,6 @@ 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, @@ -373,47 +349,7 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba `variable_lead_time=True`. Ignored otherwise. Returns: - :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. + :class:`Batch`: Prediction for the batch. """ batch = self.batch_transform_hook(batch) @@ -425,11 +361,6 @@ def _forward_impl(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) 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) - batch = _tile_batch(batch, self.num_ensemble_members) - H, W = batch.spatial_shape patch_res = ( self.encoder.latent_levels, diff --git a/aurora/rollout.py b/aurora/rollout.py index af166a03..b41d5f14 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -6,7 +6,7 @@ import torch -from aurora.batch import Batch, _split_batch, _tile_batch +from aurora.batch import Batch, split_batch, tile_batch from aurora.model.aurora import Aurora __all__ = ["rollout", "rollout_ensemble"] @@ -154,6 +154,7 @@ def rollout_ensemble( model: Aurora, batch: Batch, steps: int, + num_ensemble_members: int, fine_lead_times: Optional[Sequence[float]] = None, use_noise_accumulation: bool = True, apply_rollout_input_clipping: bool = True, @@ -168,12 +169,11 @@ def rollout_ensemble( `Batch`\\ s after every (sub-)step, one per ensemble member; see :meth:`aurora.Aurora.forward_ensemble`. """ - num_ensemble_members = model.num_ensemble_members # 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) + batch = tile_batch(batch, num_ensemble_members) model.num_ensemble_members = 1 try: for pred in rollout( @@ -184,7 +184,7 @@ def rollout_ensemble( use_noise_accumulation=use_noise_accumulation, apply_rollout_input_clipping=apply_rollout_input_clipping, ): - yield _split_batch(pred, num_ensemble_members) + 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 db7a3755..a571f514 100644 --- a/tests/v1p5/test_ensemble.py +++ b/tests/v1p5/test_ensemble.py @@ -1,17 +1,17 @@ """Copyright (c) Microsoft Corporation. Licensed under the MIT license. -Tests for internal ensemble members (`num_ensemble_members`). +Tests for internal ensemble members (`tile_batch` and `split_batch` functions +and the `rollout_ensemble` utility). """ 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 Aurora, Batch, Metadata, rollout, rollout_ensemble -from aurora.batch import _split_batch, _tile_batch +from aurora import Aurora, Batch, Metadata, rollout_ensemble +from aurora.batch import split_batch, tile_batch from aurora.model.film import AdaptiveLayerNorm @@ -51,7 +51,7 @@ def test_tile_and_split_batch_roundtrip(): b, n = 2, 3 batch = _make_ensemble_test_batch(b) - tiled = _tile_batch(batch, n) + tiled = tile_batch(batch, n) v = tiled.surf_vars["2t"] assert v.shape[0] == n * b @@ -72,7 +72,7 @@ def test_tile_and_split_batch_roundtrip(): # Splitting undoes the tiling: every member is identical to the original, standard-shaped # batch (tiling itself introduces no randomness). - members = _split_batch(tiled, n) + 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"]) @@ -80,70 +80,10 @@ def test_tile_and_split_batch_roundtrip(): 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_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_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() - 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_ensemble(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 torch.manual_seed(0) - model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + model = _make_small_v1p5(stochastic=True) # 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 @@ -152,17 +92,19 @@ def test_forward_ensemble_members_differ_when_stochastic(): 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) + batch = tile_batch(batch, n) b = next(iter(batch.surf_vars.values())).shape[0] with torch.inference_mode(): - pred = model.forward_ensemble(batch, lead_times=torch.full((b,), 6.0)) + pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + members = split_batch(pred, n) # 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): - diff = (pred[i].surf_vars["2t"] - pred[j].surf_vars["2t"]).abs().max() + diff = (members[i].surf_vars["2t"] - members[j].surf_vars["2t"]).abs().max() assert diff > 1e-2 @@ -170,34 +112,36 @@ 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 = _make_small_v1p5(stochastic=False) 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) + batch = tile_batch(batch, n) b = next(iter(batch.surf_vars.values())).shape[0] with torch.inference_mode(): - pred = model.forward_ensemble(batch, lead_times=torch.full((b,), 6.0)) + pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + members = split_batch(pred, n) 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 + members[0].surf_vars["2t"], members[m].surf_vars["2t"], atol=1e-3, rtol=1e-3 ) 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 = _make_small_v1p5(stochastic=True) 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_ensemble(model, batch, steps=3)) + preds = list(rollout_ensemble(model, batch, steps=3, num_ensemble_members=n)) assert len(preds) == 3 for step_pred in preds: @@ -206,34 +150,6 @@ def test_rollout_ensemble_yields_list_of_standard_shaped_batches_across_steps(): 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_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() - 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_ensemble(model, batch, steps=5) - next(gen) - gen.close() - - assert model.num_ensemble_members == n - def test_rollout_ensemble_num_ensemble_members_one_still_works(): model = _make_small_v1p5() @@ -243,10 +159,9 @@ def test_rollout_ensemble_num_ensemble_members_one_still_works(): b = next(iter(batch.surf_vars.values())).shape[0] with torch.inference_mode(): - preds = list(rollout_ensemble(model, batch, steps=2)) + preds = list(rollout_ensemble(model, batch, steps=2, num_ensemble_members=1)) 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 From 225ef3a884df88a6bef5a7cd845c78ccd3a782d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Agnieszka=20S=C5=82owik?= Date: Sun, 6 Sep 2026 11:51:15 +0100 Subject: [PATCH 06/13] Update `rollout_ensemble` docstring. Co-authored-by: Wessel --- aurora/rollout.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/aurora/rollout.py b/aurora/rollout.py index b41d5f14..ac023623 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -159,15 +159,17 @@ def rollout_ensemble( 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. + """Perform a roll-out for `num_ensemble_members` ensemble members simultaneously. - All arguments are identical to `rollout`; see there for details. + The members are computed in one pass by repeating `batch` along the batch dimension, so every + member receives independent noise. All other arguments are as for :func:`rollout`. + + Args: + num_ensemble_members (int): Number of ensemble members. 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`. + list[:class:`aurora.Batch`]: After every (sub-)step, one prediction per ensemble member, + each with the batch size of `batch`. """ # Tile the batch once up front, then temporarily disable further expansion so it isn't From ecce6df8bb1591ffe4a48af36b6f83b1a4bd5d32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Agnieszka=20S=C5=82owik?= Date: Sun, 6 Sep 2026 11:52:14 +0100 Subject: [PATCH 07/13] Remove stale `try/except` block. Co-authored-by: Wessel --- aurora/rollout.py | 28 +++++++++------------------- 1 file changed, 9 insertions(+), 19 deletions(-) diff --git a/aurora/rollout.py b/aurora/rollout.py index ac023623..5a418a88 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -171,23 +171,13 @@ def rollout_ensemble( list[:class:`aurora.Batch`]: After every (sub-)step, one prediction per ensemble member, each with the batch size of `batch`. """ - - # 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: - 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. - model.num_ensemble_members = num_ensemble_members + 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) From ff7db215374741061e8dc94aa5520a429ac58c84 Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Sun, 6 Sep 2026 12:12:37 +0100 Subject: [PATCH 08/13] Add input validation to `rollout_ensemble`. --- aurora/rollout.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/aurora/rollout.py b/aurora/rollout.py index 5a418a88..002ab4fd 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -171,6 +171,13 @@ def rollout_ensemble( list[:class:`aurora.Batch`]: After every (sub-)step, one prediction per ensemble member, each with the batch size of `batch`. """ + if num_ensemble_members < 1: + raise ValueError( + f"`num_ensemble_members` must be at least `1`, but is `{num_ensemble_members}`." + ) + if not model.backbone.stochastic: + raise ValueError("`rollout_ensemble` requires a stochastic model.") + batch = tile_batch(batch, num_ensemble_members) for pred in rollout( model, From 43126872cc2d71814b3d402d693aa67d5ce59d45 Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Sun, 6 Sep 2026 13:03:29 +0100 Subject: [PATCH 09/13] Update `models.md` doc. --- docs/models.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/models.md b/docs/models.md index 874edddb..6657fea6 100644 --- a/docs/models.md +++ b/docs/models.md @@ -449,3 +449,26 @@ When using `rollout` with `fine_lead_times`, noise accumulation is enabled by de smoother intra-step transitions while using independent effective noise between main steps, matching the training regimen. Set `use_noise_accumulation=False` to draw independent noise at each sub-step instead, though this is not recommended. + +### Internal Ensembling + +You can avoid running the model multiple times to generate different ensemble outputs. +In Aurora 1.5 you can run the model only once, but the forward pass must operate on +batch in which the data is tiled. +This is entirely optional and consumes more memory, but may result in faster inference +on some hardware. +Functions `aurora.batch.tile_batch` and `aurora.batch.split_batch` expose this logic: +the former modifies a `aurora.batch.Batch` to support a single forward pass, and the +latter unpacks a tiled batch into a list of batches. + +In the rollout scenario, you can avoid using these functions and just use `rollout_ensemble`: + +```python +from aurora import rollout_ensemble + +with torch.inference_mode(): + preds = [ + [member.to("cpu") for member in members] + for members in rollout_ensemble(model, batch, steps=10, num_ensemble_members=4) + ] +``` From f7a578368b8d4dac61bd497858e5be681b338e73 Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Sun, 6 Sep 2026 13:03:49 +0100 Subject: [PATCH 10/13] Update test to demonstrate an exception is raised in `rollout_ensemble`. --- tests/v1p5/test_ensemble.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py index a571f514..5b5025db 100644 --- a/tests/v1p5/test_ensemble.py +++ b/tests/v1p5/test_ensemble.py @@ -7,6 +7,7 @@ import warnings from datetime import datetime +import pytest import torch from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5 @@ -151,17 +152,13 @@ def test_rollout_ensemble_yields_list_of_standard_shaped_batches_across_steps(): assert v.shape[0] == b -def test_rollout_ensemble_num_ensemble_members_one_still_works(): +def test_rollout_ensemble_num_ensemble_members_one_raises(): 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_ensemble(model, batch, steps=2, num_ensemble_members=1)) - - for step_pred in preds: - assert len(step_pred) == 1 - for v in step_pred[0].surf_vars.values(): - assert v.shape[0] == b + with pytest.raises(ValueError): + with torch.inference_mode(): + preds = list(rollout_ensemble(model, batch, steps=2, num_ensemble_members=1)) From 4d449f0453f2706c9a562351b070f6aca562767e Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Sat, 12 Sep 2026 11:10:45 +0100 Subject: [PATCH 11/13] Address CI failures in test_rollout_ensemble_num_ensemble_members_one_raises. --- tests/v1p5/test_ensemble.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py index 5b5025db..62b51b90 100644 --- a/tests/v1p5/test_ensemble.py +++ b/tests/v1p5/test_ensemble.py @@ -157,8 +157,7 @@ def test_rollout_ensemble_num_ensemble_members_one_raises(): 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] + _ = next(iter(batch.surf_vars.values())).shape[0] - with pytest.raises(ValueError): - with torch.inference_mode(): - preds = list(rollout_ensemble(model, batch, steps=2, num_ensemble_members=1)) + with pytest.raises(ValueError), torch.inference_mode(): + _ = list(rollout_ensemble(model, batch, steps=2, num_ensemble_members=1)) From 6297886edf35350bafd5f06327a429b0cf072387 Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Tue, 15 Sep 2026 11:53:43 +0100 Subject: [PATCH 12/13] Docs and API. --- aurora/__init__.py | 4 +++- aurora/batch.py | 25 +++++++++++++++++++------ docs/api.rst | 9 +++++++-- docs/models.md | 19 +++++++++---------- 4 files changed, 38 insertions(+), 19 deletions(-) diff --git a/aurora/__init__.py b/aurora/__init__.py index 8f414aac..5ca4f317 100644 --- a/aurora/__init__.py +++ b/aurora/__init__.py @@ -1,6 +1,6 @@ """Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" -from aurora.batch import Batch, Metadata +from aurora.batch import Batch, Metadata, split_batch, tile_batch from aurora.insolation import insolation from aurora.model.aurora import ( Aurora, @@ -30,6 +30,8 @@ "AuroraV1p5Ensemble", "Batch", "Metadata", + "tile_batch", + "split_batch", "insolation", "rollout", "rollout_ensemble", diff --git a/aurora/batch.py b/aurora/batch.py index adf7ef32..00ed6b80 100644 --- a/aurora/batch.py +++ b/aurora/batch.py @@ -17,7 +17,7 @@ unnormalise_surf_var, ) -__all__ = ["Metadata", "Batch"] +__all__ = ["Metadata", "Batch", "tile_batch", "split_batch"] @dataclasses.dataclass @@ -314,10 +314,15 @@ 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. + Used to run `n` ensemble members as a single fused computation. + Results derived from the tiling should be undone with `split_batch`. + + Args: + batch (:class:`aurora.Batch`): The batch to tile. + n (int): Number of times to tile. + + Returns: + :class:`aurora.Batch`: `batch` tiled `n` times along the batch dimension. """ return dataclasses.replace( batch, @@ -328,7 +333,15 @@ def tile_batch(batch: Batch, n: int) -> Batch: def split_batch(batch: Batch, n: int) -> list[Batch]: - """Undo `tile_batch`, splitting a tiled batch back into `n` standard-shaped batches.""" + """Undo `tile_batch`, splitting a tiled batch back into `n` standard-shaped batches. + + Args: + batch (:class:`aurora.Batch`): The tiled batch to split. + n (int): Number of batches `batch` was tiled into. + + Returns: + list[:class:`aurora.Batch`]: `batch` split into `n` standard-shaped batches. + """ b = next(iter(batch.surf_vars.values())).shape[0] // n time = batch.metadata.time return [ diff --git a/docs/api.rst b/docs/api.rst index 7540ae0d..236fabcf 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -9,10 +9,15 @@ Batch .. autoclass:: aurora.Metadata :members: +.. autofunction:: aurora.tile_batch + +.. autofunction:: aurora.split_batch + Roll-Outs --------- -.. autoclass:: aurora.rollout - :members: +.. autofunction:: aurora.rollout + +.. autofunction:: aurora.rollout_ensemble Tropical Cyclone Tracking ------------------------- diff --git a/docs/models.md b/docs/models.md index 6657fea6..3d605ca4 100644 --- a/docs/models.md +++ b/docs/models.md @@ -452,16 +452,8 @@ noise at each sub-step instead, though this is not recommended. ### Internal Ensembling -You can avoid running the model multiple times to generate different ensemble outputs. -In Aurora 1.5 you can run the model only once, but the forward pass must operate on -batch in which the data is tiled. -This is entirely optional and consumes more memory, but may result in faster inference -on some hardware. -Functions `aurora.batch.tile_batch` and `aurora.batch.split_batch` expose this logic: -the former modifies a `aurora.batch.Batch` to support a single forward pass, and the -latter unpacks a tiled batch into a list of batches. - -In the rollout scenario, you can avoid using these functions and just use `rollout_ensemble`: +For a stochastic model, such as `AuroraV1p5Ensemble`, you can use `rollout_ensemble` to +generate an autoregressive forecast with multiple ensemble members at once: ```python from aurora import rollout_ensemble @@ -472,3 +464,10 @@ with torch.inference_mode(): for members in rollout_ensemble(model, batch, steps=10, num_ensemble_members=4) ] ``` + +Alternatively, you can produce an ensemble forecast by calling `rollout` multiple times, +once for each ensemble member. +`rollout_ensemble` can be faster, at the cost of requiring more memory. + +Internally, `rollout_ensemble` use `aurora.tile_batch` and `aurora.split_batch`. +These functions can also be used for fused forward passes. From a584479214e0d78dbdb58f68bbee8ddc57dcf88f Mon Sep 17 00:00:00 2001 From: Aga Slowik Date: Tue, 15 Sep 2026 12:49:55 +0100 Subject: [PATCH 13/13] Simplify tests. --- tests/v1p5/_helpers.py | 7 +- tests/v1p5/test_ensemble.py | 136 ++++++++++++------------------------ 2 files changed, 47 insertions(+), 96 deletions(-) diff --git a/tests/v1p5/_helpers.py b/tests/v1p5/_helpers.py index 4ddd7a65..575f357b 100644 --- a/tests/v1p5/_helpers.py +++ b/tests/v1p5/_helpers.py @@ -26,16 +26,17 @@ def _make_batch( surf_vars: tuple[str, ...] = _SURF_VARS, static_vars: tuple[str, ...] = _STATIC_VARS, atmos_vars: tuple[str, ...] = _ATMOS_VARS, + batch_size: int = BATCH, ) -> Batch: """Create a minimal synthetic batch.""" return Batch( - surf_vars={k: torch.randn(BATCH, HISTORY, H, W) for k in surf_vars}, + surf_vars={k: torch.randn(batch_size, HISTORY, H, W) for k in surf_vars}, static_vars={k: torch.randn(H, W) for k in static_vars}, - atmos_vars={k: torch.randn(BATCH, HISTORY, N_LEVELS, H, W) for k in atmos_vars}, + atmos_vars={k: torch.randn(batch_size, HISTORY, N_LEVELS, H, W) for k in atmos_vars}, metadata=Metadata( lat=torch.linspace(90, -90, H), lon=torch.linspace(0, 360, W + 1)[:-1], - time=(datetime(2023, 6, 15, 12, 0),), + time=(datetime(2023, 6, 15, 12, 0),) * batch_size, atmos_levels=(100, 250, 500, 850), ), ) diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py index 62b51b90..9a940736 100644 --- a/tests/v1p5/test_ensemble.py +++ b/tests/v1p5/test_ensemble.py @@ -1,17 +1,15 @@ """Copyright (c) Microsoft Corporation. Licensed under the MIT license. -Tests for internal ensemble members (`tile_batch` and `split_batch` functions -and the `rollout_ensemble` utility). +Tests for ensemble rollout and tiling utility functions. """ -import warnings -from datetime import datetime +import itertools import pytest import torch -from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5 -from aurora import Aurora, Batch, Metadata, rollout_ensemble +from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, BATCH, _make_batch, _make_small_v1p5 +from aurora import Aurora, rollout_ensemble from aurora.batch import split_batch, tile_batch from aurora.model.film import AdaptiveLayerNorm @@ -19,11 +17,8 @@ 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. + A freshly constructed model is exactly insensitive to its conditioning signal, and hence to + the injected noise, because `AdaptiveLayerNorm.ln_modulation` is zero-initialised. """ for m in model.modules(): if isinstance(m, AdaptiveLayerNorm): @@ -32,59 +27,29 @@ def _unzero_adaptive_layer_norms(model: Aurora, std: float = 0.1) -> None: m.ln_modulation[-1].bias.normal_(std=std) -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) + batch = _make_batch(batch_size=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"]) + for v in (*tiled.surf_vars.values(), *tiled.atmos_vars.values()): + assert v.shape[0] == n * b - # 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 + for k, v in member.surf_vars.items(): + torch.testing.assert_close(v, batch.surf_vars[k]) + for k, v in member.atmos_vars.items(): + torch.testing.assert_close(v, batch.atmos_vars[k]) -def test_forward_ensemble_members_differ_when_stochastic(): - n = 3 +@pytest.mark.parametrize("stochastic", [True, False]) +def test_forward_ensemble(stochastic: bool): torch.manual_seed(0) - model = _make_small_v1p5(stochastic=True) + model = _make_small_v1p5(stochastic=stochastic) # 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 @@ -93,30 +58,7 @@ def test_forward_ensemble_members_differ_when_stochastic(): 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) - batch = tile_batch(batch, n) - b = next(iter(batch.surf_vars.values())).shape[0] - - with torch.inference_mode(): - pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) - members = split_batch(pred, n) - - # 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): - diff = (members[i].surf_vars["2t"] - members[j].surf_vars["2t"]).abs().max() - assert diff > 1e-2 - - -def test_forward_ensemble_members_identical_without_stochastic(): n = 3 - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - model = _make_small_v1p5(stochastic=False) - 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) batch = tile_batch(batch, n) b = next(iter(batch.surf_vars.values())).shape[0] @@ -124,35 +66,43 @@ def test_forward_ensemble_members_identical_without_stochastic(): pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) members = split_batch(pred, n) - 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( - members[0].surf_vars["2t"], members[m].surf_vars["2t"], atol=1e-3, rtol=1e-3 - ) + if stochastic: + # Every member receives independent noise, so members must differ. + for member1, member2 in itertools.combinations(members, 2): + assert (member1.surf_vars["2t"] - member2.surf_vars["2t"]).abs().mean() > 1e-2 + else: + # Check that all are equal. + for m in range(1, n): + torch.testing.assert_close( + members[0].surf_vars["2t"], members[m].surf_vars["2t"], atol=1e-3, rtol=1e-3 + ) -def test_rollout_ensemble_yields_list_of_standard_shaped_batches_across_steps(): - n = 2 +def test_rollout_ensemble(): + num_ensemble_members = 3 + torch.manual_seed(0) model = _make_small_v1p5(stochastic=True) + _unzero_adaptive_layer_norms(model) # Otherwise, the noise has no effect on the output. 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] + batch = _make_batch(surf_vars=tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF)) + steps = 2 with torch.inference_mode(): - preds = list(rollout_ensemble(model, batch, steps=3, num_ensemble_members=n)) + preds = list(rollout_ensemble(model, batch, steps, num_ensemble_members)) - assert len(preds) == 3 - for step_pred in preds: - assert len(step_pred) == n - for member in step_pred: + assert len(preds) == steps + for members in preds: + assert len(members) == num_ensemble_members + for member in members: for v in member.surf_vars.values(): - assert v.shape[0] == b + assert v.shape[0] == BATCH + # Every member receives independent noise, so members must differ. + for member1, member2 in itertools.combinations(members, 2): + assert (member1.surf_vars["2t"] - member2.surf_vars["2t"]).abs().mean() > 1e-2 def test_rollout_ensemble_num_ensemble_members_one_raises(): + num_ensemble_members = 1 model = _make_small_v1p5() model.eval() surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) @@ -160,4 +110,4 @@ def test_rollout_ensemble_num_ensemble_members_one_raises(): _ = next(iter(batch.surf_vars.values())).shape[0] with pytest.raises(ValueError), torch.inference_mode(): - _ = list(rollout_ensemble(model, batch, steps=2, num_ensemble_members=1)) + _ = list(rollout_ensemble(model, batch, steps=2, num_ensemble_members=num_ensemble_members))