Internal ensembling in Aurora models

[Agnieszka Słowik (Slowika)]

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

Problem

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

Change

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

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

Design notes

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

Tests

Added tests/v1p5/test_ensemble.py covering:

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

[Wessel (wesselb)]

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

[Agnieszka Słowik (Slowika)]

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

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

[Wessel (wesselb)]

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

[Wessel (wesselb)]

Comment generated by Claude.

Document rollout_ensemble.

rollout_ensemble is a new public function, so it should be listed in docs/api.rst under "Roll-Outs" and mentioned in the "Recommended Use" section for Aurora 1.5 Ensemble in docs/models.md, which currently says that ensemble members are generated by running the model multiple times. For example,

.. autofunction:: aurora.rollout_ensemble

and

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)
    ]

[Agnieszka Słowik (Slowika)]

We don't document any functions using autofunction, and rollout_ensemble() is part of the rollout class (and hence gets document automatically anyway, just like rollout()). I think it shouldn't be listed under docs/api.rst. Please let me know if my assumptions and understanding are wrong.

I am adding some documentation to models.md, including references to tile_batch and split_batch.

[Wessel (wesselb)]

Comment generated by Claude.

Module placement and stale docstring.

Move these two helpers into aurora/rollout.py as _tile_batch and _split_batch, next to _advance_batch, and shorten the docstrings:

Suggested change

	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 _tile_batch(batch: Batch, n: int) -> Batch:
    """Repeat `batch` `n` times along the batch dimension."""
    return dataclasses.replace(
        batch,
        surf_vars={k: torch.cat(n * [v]) for k, v in batch.surf_vars.items()},
        atmos_vars={k: torch.cat(n * [v]) for k, v in batch.atmos_vars.items()},
        metadata=dataclasses.replace(batch.metadata, time=n * batch.metadata.time),
    )


def _split_batch(batch: Batch, n: int) -> list[Batch]:
    """Split `batch` into `n` equally sized batches along the batch dimension."""
    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)
    ]

rollout_ensemble is their only consumer, and rollout.py already keeps its private batch helper _advance_batch next to rollout. Moving them also resolves their current half-public state: public names, but excluded from __all__ and documented as not public. The docstring of tile_batch is out of date, since Aurora.forward no longer uses it, and the rest of it describes how the function is used rather than what it does. Finally, torch.cat(n * [v]) does the same as v.repeat(n, *([1] * (v.dim() - 1))) without the dimension arithmetic.

[Agnieszka Słowik (Slowika)]

Jonathan Weyn (@jweyn) Wessel (@wesselb) I disagree. split_batch and tile_batch are now meant to be used publicly! rollout_ensemble is the only user within the repository, but they otherwise need to be used in all internal ensembling cases (as discussed in the relevant issue). The unit tests show how they are used.

[Wessel (wesselb)]

Comment generated by Claude.

Compress the tests.

The new functionality is small, so two tests suffice: a round trip for the helpers and one parametrised test of rollout_ensemble that checks the yielded structure and that members actually differ.

Suggested change

	"""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).
	"""
	import warnings
	from datetime import datetime
	import torch
	from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5
	from aurora import Aurora, Batch, Metadata, rollout_ensemble
	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:
	"""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_forward_ensemble_members_differ_when_stochastic():
	n = 3
	torch.manual_seed(0)
	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
	# (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)
	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]
	with torch.inference_mode():
	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
	)
	def test_rollout_ensemble_yields_list_of_standard_shaped_batches_across_steps():
	n = 2
	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, num_ensemble_members=n))
	assert len(preds) == 3
	for step_pred in preds:
	assert len(step_pred) == n
	for member in step_pred:
	for v in member.surf_vars.values():
	assert v.shape[0] == b
	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)
	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
	"""Copyright (c) Microsoft Corporation. Licensed under the MIT license.
	Tests for `rollout_ensemble`.
	"""
	import itertools
	import pytest
	import torch
	from ._helpers import BATCH, _OUTPUT_ONLY_SURF, _SURF_VARS, _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
	def _unzero_adaptive_layer_norms(model: Aurora, std: float = 0.1) -> None:
	"""Nudge every `AdaptiveLayerNorm`'s modulation away from its zero initialisation.
	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):
	with torch.no_grad():
	m.ln_modulation[-1].weight.normal_(std=std)
	m.ln_modulation[-1].bias.normal_(std=std)
	def test_tile_and_split_batch_roundtrip():
	b, n = 2, 3
	batch = _make_batch(batch_size=b)
	tiled = tile_batch(batch, n)
	assert len(tiled.metadata.time) == n * b
	for v in (*tiled.surf_vars.values(), *tiled.atmos_vars.values()):
	assert v.shape[0] == n * b
	members = split_batch(tiled, n)
	assert len(members) == n
	for member in members:
	assert member.metadata.time == batch.metadata.time
	for k, v in member.surf_vars.items():
	assert torch.equal(v, batch.surf_vars[k])
	for k, v in member.atmos_vars.items():
	assert torch.equal(v, batch.atmos_vars[k])
	@pytest.mark.parametrize("num_ensemble_members", [1, 3])
	def test_rollout_ensemble(num_ensemble_members):
	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()
	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, num_ensemble_members))
	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] == 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().max() > 1e-2

This needs a batch_size parameter on the shared _make_batch in tests/v1p5/_helpers.py, in place of the parallel _make_ensemble_test_batch:

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_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_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=batch_size * (datetime(2023, 6, 15, 12, 0),),
            atmos_levels=(100, 250, 500, 850),
        ),
    )

In detail:

• test_forward_ensemble_members_identical_without_stochastic checks that a deterministic model treats identical batch elements identically. That is a property of the model and of PyTorch rather than of this PR. Its warnings.catch_warnings is also dead: nothing warns for stochastic=False any more.
• test_forward_ensemble_members_differ_when_stochastic checks the property rollout_ensemble relies on, but by hand-tiling and calling forward. Checking it through rollout_ensemble tests the public function directly, and the two rollout_ensemble shape tests then fold into the same parametrised test.
• In the round trip, the three loops over slices of tiled pin down the tiling layout, which only needs to be consistent between tile_batch and split_batch; the round trip already checks that. torch.equal is appropriate since tiling copies exactly.
• The test names still refer to forward_ensemble, which no longer exists, and the helper docstring used a double hyphen where the code base uses a single one.

[Agnieszka Słowik (Slowika)]

I don't agree with the main point here. Hand-tiling is exposed to the user. I think it makes sense to have tests for these functions in detail and check that they do what they are supposed to.

Parameterised tests won't work now that we have input validation logic for rollout_ensemble.

On naming: I test_forward_ensemble_* doesn't mean that we are testing a function named forward_ensemble; we are testing forward on ensemble input/batches. Naming could be better; do you have a suggestion?