Add generator argument to forward for reproducible noise in stochastic mode

[Syota-Sasaki (s-sasaki-earthsea-wizard)]

Closes #191.

This implements the design suggested by Wessel (@wesselb) in #191 (comment) (see the discussion there): instead of a constructor-level seed, Aurora.forward accepts a keyword-only generator: torch.Generator | tuple[torch.Generator | None, ...] | None = None, which is passed through to Swin3DTransformerBackbone.forward and used in the torch.randn calls. RNG state is owned entirely by the caller, so no reset method is needed: re-seeding the generator(s) restores the noise sequence.

Semantics

• Single generator: one stream for the whole batch, drawn in a single (B, L, D) call. The stream therefore depends on the batch size, matching the semantics of a plain seeded torch.randn.
• Tuple of generators: one entry per batch element (ensemble member), in the current order of the batch dimension. Each member is drawn separately with shape (L, D) from its own generator, so a given member's noise sequence is independent of the batch composition: member i produces the same sequence whether it runs in a batch of 1 or a batch of N. Entries may be None to fall back to the global RNG for that member, and passing the same generator object in several slots deliberately shares one stream. Because the two modes draw with different shapes, a single generator and a tuple are not interchangeable.
• generator=None (default) preserves the current behaviour exactly (global RNG).

Design decisions

• The tuple length and the device of every tuple entry are validated before any randomness is consumed, so a ValueError cannot leave some generators already advanced. For a single generator, a device mismatch surfaces as PyTorch's usual RuntimeError.
• To reproduce a run, the caller must re-seed the generator(s) and call Aurora.reset_noise(): noise cached by noise accumulation in a previous run would otherwise contaminate the reproduced sequence. This is documented on forward, reset_noise, and rollout.
• The noise-accumulation cache is now invalidated on device or dtype changes as well, not only on shape changes; a stale same-shape cache would otherwise silently mix into a reproduced run.
• Passing generator to a non-stochastic model warns once (UserWarning) and ignores the argument, so the mistake is surfaced without spamming roll-outs.
• rollout() also accepts and passes through generator. This goes slightly beyond the design suggested in the issue, but reproducing an inference run in practice means reproducing a roll-out; the generators advance across (sub-)steps and a tuple stays bound to the batch-dim order throughout. Happy to drop this part if you prefer a smaller PR.

Verification

tests/v1p5/test_forward_generator.py covers: re-seeding a generator (or a fresh same-seed generator on a fresh model instance) reproduces the exact noise sequence, also with noise accumulation enabled; without re-seeding the stream keeps advancing; per-member reproducibility with a tuple, including independence from the batch composition (batch of 1 vs batch of 3) and of None entries from their neighbours' generators; a length mismatch raises before consuming any randomness; a non-stochastic model warns once and does not consume the generator; generator=None preserves the current global-RNG semantics; and the roll-out pass-through. Two CUDA-only tests additionally check that an index-less torch.Generator(device="cuda") is accepted on a cuda:0 model (matching PyTorch's own device semantics) and that a device mismatch raises before consuming any randomness.

All 49 tests under tests/v1p5/ pass (34 existing + 15 new; the two CUDA tests are skipped without a GPU), as do the existing roll-out, batch, and header tests.

docs/models.md gains a short "Reproducible Noise" subsection in the Aurora 1.5 Ensemble section with usage examples for both modes.

As in the prototype discussed in the issue, the tests record the sampled noise directly instead of comparing model outputs, because with randomly initialised weights the adaptive-LN modulation is zero-initialised and the noise context does not affect the output at initialization.

Minimal reproduction

With the pretrained checkpoint, the effect is visible directly in the forecasts. The following script rolls out two ensemble members with per-member generators, then reproduces the exact same forecasts by re-seeding the generators:

from datetime import datetime

import torch

from aurora import AuroraV1p5Ensemble, Batch, Metadata, rollout

# fp16 autocast overflows on the far-out-of-distribution random inputs below, so run in fp32.
model = AuroraV1p5Ensemble(autocast=False)
model.load_checkpoint()
model = model.cuda().eval()

# Random input data at a small resolution, with two ensemble members in the batch.
torch.manual_seed(0)
b, h, w = 2, 17, 32
batch = Batch(
    surf_vars={
        k: torch.rand(b, 2, h, w) for k in model.surf_vars if k not in model.output_only_surf_vars
    },
    static_vars={k: torch.rand(h, w) for k in model.static_vars},
    atmos_vars={k: torch.rand(b, 2, 4, h, w) for k in model.atmos_vars},
    metadata=Metadata(
        lat=torch.linspace(90, -90, h),
        lon=torch.linspace(0, 360, w + 1)[:-1],
        time=(datetime(2023, 6, 15, 12),) * b,
        atmos_levels=(100, 250, 500, 850),
    ),
)

device = next(model.parameters()).device
generators = tuple(torch.Generator(device=device).manual_seed(seed) for seed in (1, 2))


def run():
    with torch.inference_mode():
        return [p.surf_vars["2t"].cpu() for p in rollout(model, batch, steps=2, generator=generators)]


first = run()

# Re-seeding the generators (and flushing the noise cache) reproduces the exact forecasts.
generators[0].manual_seed(1)
generators[1].manual_seed(2)
model.reset_noise()
second = run()

# Without re-seeding, the generators keep advancing: fresh noise, different forecasts.
third = run()

print(all(torch.equal(a, b) for a, b in zip(first, second)))  # True
print(all(not torch.equal(a, b) for a, b in zip(first, third)))  # True

Both checks print True with the released checkpoint (verified on an RTX 5080): the reproduced forecasts are exactly equal, and a run without re-seeding is not. The same works on CPU by dropping .cuda() and creating CPU generators.

Disclosure

This PR was developed with AI assistance. I have reviewed, tested, and take responsibility for all of the changes.

[Syota-Sasaki (s-sasaki-earthsea-wizard)]

@microsoft-github-policy-service agree

[Wessel (wesselb)]

Thanks for the PR, Syota-Sasaki (@s-sasaki-earthsea-wizard)! I like the design. :) I've done a first pass and have left some Claude-assisted feedback.

[Wessel (wesselb)]

Comment generated by Claude.

Unnecessary state.

Delete. See the comment on the warning in forward below: without the warning, this flag has no
use.

Suggested change

	# Warn only once when `generator` is passed to `forward` of a non-stochastic model.
	self._generator_ignored_warned = False

[Wessel (wesselb)]

Comment generated by Claude.

Too much detail.

This is the one place that should document the semantics; the other docstrings can cross-reference
it. Shorten to the interface:

Suggested change

	generator (:class:`torch.Generator` or tuple of :class:`torch.Generator` or `None`,
	optional): Source of randomness for the noise injection in stochastic mode. A
	single generator drives one stream for the whole batch. A tuple must contain one
	entry per batch element (ensemble member), in the current order of the batch
	dimension; every element then draws from its own stream, so the noise sequence of
	a given member does not depend on the batch composition. Tuple entries may be
	`None` to fall back to the global RNG for that member, and passing the same
	generator object in several slots makes those members share one stream. Because
	the two modes draw with different shapes, a single generator and a tuple are not
	interchangeable. Generators must live on the same device as the model, and they
	advance on every forward pass. To reproduce a run, re-seed the generators (e.g.
	with `manual_seed`) *and* call :meth:`reset_noise`, so that noise cached by noise
	accumulation in a previous run cannot contaminate the reproduced sequence. When
	the model is not stochastic, this argument is ignored with a warning. Defaults to
	`None`, which draws from the global RNG (the previous behaviour).
	generator (:class:`torch.Generator` or tuple of :class:`torch.Generator` or `None`,
	optional): Generator for the noise in stochastic mode. A single generator is used
	for the whole batch. A tuple gives one generator per batch element, so that the
	noise of an element does not depend on the other elements in the batch; entries can
	be `None` to use the global RNG. Generators must be on the device of the model. To
	reproduce a run, re-seed the generators and call :meth:`reset_noise`. Ignored when
	the model is not stochastic. Defaults to `None`, which uses the global RNG.

The shapes of the draws are an implementation detail, that generators advance is how generators
work, and that the same object in several slots shares a stream is self-evident.

[Wessel (wesselb)]

Comment generated by Claude.

Unnecessary check.

Delete and simply pass generator through: the backbone only touches it when stochastic is
True, so a non-stochastic model ignores it either way.

Suggested change

	if generator is not None and not self.backbone.stochastic:
	if not self._generator_ignored_warned:
	warnings.warn(
	"`generator` is ignored because stochastic noise is disabled.",
	stacklevel=2,
	)
	self._generator_ignored_warned = True
	generator = None

[Wessel (wesselb)]

Comment generated by Claude.

Public API.

NoiseGenerator is imported by aurora.py and rollout.py, so add it to __all__, like LoRAMode
in lora.py. The docstring repeats what Aurora.forward documents and can go.

Suggested change

	__all__ = ["Swin3DTransformerBackbone"]
	NoiseGenerator = torch.Generator | tuple[torch.Generator | None, ...] | None
	"""Source of randomness for noise injection in stochastic mode: a single generator driving one
	stream for the whole batch, a tuple with one generator per batch element, or `None` for the
	global RNG."""
	__all__ = ["Swin3DTransformerBackbone", "NoiseGenerator"]
	NoiseGenerator = torch.Generator | tuple[torch.Generator | None, ...] | None

[Wessel (wesselb)]

Comment generated by Claude.

Duplicated documentation.

Delete. Aurora.forward already says that reproducing a run requires re-seeding and calling
reset_noise.

Suggested change

	To reproduce a run that controls the noise with `generator` (see :meth:`forward`), re-seed
	the generators *and* call this method: cached noise left over from a previous run would
	otherwise contaminate the reproduced sequence. The same applies after changing the order or
	composition of the ensemble members that a tuple of generators corresponds to, which cannot
	be detected from the noise tensors themselves.

[Wessel (wesselb)]

Comment generated by Claude.

Simplify validation.

With _validate_generator gone (see above), keep the length check here:

Suggested change

	self._validate_generator(generator, x.shape[0], x.device)
	if isinstance(generator, tuple) and len(generator) != x.shape[0]:
	raise ValueError(
	f"Expected one generator per batch element, but got `{len(generator)}` "
	f"generators for a batch of size `{x.shape[0]}`."
	)

[Wessel (wesselb)]

Comment generated by Claude.

Out of scope.

Revert to the shape-only check, to keep this PR to the generator argument. Happy to take the
device and dtype invalidation as a separate PR if you think it is needed.

Suggested change

	# A shape (e.g. different batch size), device, or dtype change invalidates the
	# cache.
	cached = self._noise_cache[0] if self._noise_cache else None
	if cached is not None and (
	cached.shape != noise.shape
	or cached.device != noise.device
	or cached.dtype != noise.dtype
	):
	warnings.warn(
	f"Noise shape changed from {self._noise_cache[0].shape} to "
	f"{noise.shape}; clearing noise cache.",
	f"Cached noise of shape {cached.shape} ({cached.dtype} on "
	f"{cached.device}) is incompatible with new noise of shape {noise.shape} "
	f"({noise.dtype} on {noise.device}); clearing noise cache.",
	stacklevel=2,
	)
	self._noise_cache.clear()
	# Shape change (e.g. different batch size) invalidates the cache.
	if self._noise_cache and self._noise_cache[0].shape != noise.shape:
	warnings.warn(
	f"Noise shape changed from {self._noise_cache[0].shape} to "
	f"{noise.shape}; clearing noise cache.",
	stacklevel=2,
	)
	self._noise_cache.clear()

[Wessel (wesselb)]

Comment generated by Claude.

Duplicated documentation.

Cross-reference only:

Suggested change

	generator (:class:`torch.Generator` or tuple of :class:`torch.Generator` or `None`,
	optional): Source of randomness for the noise injection of stochastic models, passed
	to every `forward` call of the roll-out. See :meth:`aurora.Aurora.forward` for the
	semantics. The generators advance on every (sub-)step, and a tuple corresponds to the
	order of the batch dimension throughout the roll-out. To reproduce a roll-out, re-seed
	the generators and call `model.reset_noise()` before starting. Default: `None`, which
	draws from the global RNG.
	generator (:class:`torch.Generator` or tuple of :class:`torch.Generator` or `None`,
	optional): Generator for the noise in stochastic mode, passed to every forward pass.
	See :meth:`aurora.Aurora.forward`. Default: `None`.

The reset_noise advice is not needed here: rollout flushes the cache itself, through
set_noise_accumulation at the start when sub-stepping and at the end of every roll-out.

[Wessel (wesselb)]

Comment generated by Claude.

Too much detail.

Condense, and match the surrounding examples, which run under torch.inference_mode() and move
predictions to the CPU. [pred for pred in ...] is a no-op comprehension.

Suggested change

	### Reproducible Noise
	By default, the injected noise is drawn from the global RNG, so the noise of an individual
	ensemble member cannot easily be reproduced. To control the noise, pass a `torch.Generator`
	to `Aurora.forward` or `rollout`. The generator must live on the same device as the model:
	```python
	import torch
	from aurora import rollout
	device = next(model.parameters()).device
	generator = torch.Generator(device=device).manual_seed(42)
	preds = [pred for pred in rollout(model, batch, steps=10, generator=generator)]
	```
	When generating multiple ensemble members simultaneously by using a batch size, pass a tuple
	with one generator per batch element (ensemble member) to control the noise of every member
	separately. Every member then draws from its own stream, so a member's noise sequence does not
	depend on the batch composition. Tuple entries may be `None` to fall back to the global RNG for
	that member. Note that a single generator and a tuple of generators draw with different shapes
	and are therefore not interchangeable.
	```python
	# `batch` contains three ensemble members.
	generators = tuple(torch.Generator(device=device).manual_seed(seed) for seed in (1, 2, 3))
	preds = [pred for pred in rollout(model, batch, steps=10, generator=generators)]
	```
	Generators advance on every forward pass.
	To reproduce a run, re-seed the generators with `manual_seed` *and* call `model.reset_noise()`:
	the latter clears noise cached by noise accumulation, which would otherwise contaminate the
	reproduced sequence.
	### Reproducible Noise
	By default, the noise is drawn from the global RNG. To control the noise, pass a `torch.Generator`
	on the device of the model to `Aurora.forward` or `rollout`:
	```python
	device = next(model.parameters()).device
	generator = torch.Generator(device=device).manual_seed(42)
	with torch.inference_mode():
	preds = [pred.to("cpu") for pred in rollout(model, batch, steps=4, generator=generator)]
	```
	When generating multiple ensemble members simultaneously by using a batch size, pass a tuple with
	one generator per batch element to control the noise of every member separately. The noise of a
	member then does not depend on the other members in the batch. Entries can be `None` to use the
	global RNG for that member.
	```python
	# `batch` contains three ensemble members.
	generators = tuple(torch.Generator(device=device).manual_seed(seed) for seed in (1, 2, 3))
	with torch.inference_mode():
	preds = [pred.to("cpu") for pred in rollout(model, batch, steps=4, generator=generators)]
	```
	To reproduce a run, re-seed the generators and call `model.reset_noise()`, which clears noise
	cached by noise accumulation.

[Wessel (wesselb)]

Comment generated by Claude.

Simplify the tests.

Fifteen tests can become four. Concretely:

• Use a fixture for the stochastic model instead of repeating _make_small_v1p5(stochastic=True)
  and model.eval() in every test, and drop the trailing set_noise_accumulation(n=0) calls, since
  every test gets a fresh model.
• Let _record_noise take a callable, so that the roll-out test reuses it instead of duplicating
  the recording logic. The clone() and the try/finally are not needed.
• Merge the single-generator tests (re-seed reproduces, no re-seed advances, fresh instance,
  independent of the global RNG, accumulation) into one test parametrised over the noise
  accumulation size, and the three tuple tests into one.
• Drop test_generator_none_preserves_global_rng_semantics (it tests torch.randn),
  test_same_seed_fresh_instance_reproduces_noise (subsumed), the non-stochastic warning test and
  the two CUDA tests (the code they test is removed above), and the generator-state assertions in
  the length-mismatch test.

A sketch of the resulting module (untested):

"""Copyright (c) Microsoft Corporation. Licensed under the MIT license.

Tests for the `generator` argument of `Aurora.forward`.
"""

import pytest
import torch

from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5
from aurora import rollout

_INPUT_SURF_VARS = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF)


@pytest.fixture
def model():
    model = _make_small_v1p5(stochastic=True)
    model.eval()
    return model


def _record_noise(model, run):
    """Call `run` and return the noise samples drawn."""
    recorded = []
    sample_noise = model.backbone._sample_noise

    def record(*args):
        recorded.append(sample_noise(*args))
        return recorded[-1]

    model.backbone._sample_noise = record
    with torch.inference_mode():
        run()
    model.backbone._sample_noise = sample_noise
    return recorded


def _forward_noise(model, generator, batch_size=1, n_forwards=3):
    """Return the noise samples drawn by `n_forwards` forward passes."""
    batch = _make_batch(surf_vars=_INPUT_SURF_VARS, batch_size=batch_size)
    lead_times = torch.full((batch_size,), 6.0)
    return _record_noise(
        model,
        lambda: [
            model.forward(batch, lead_times=lead_times, generator=generator)
            for _ in range(n_forwards)
        ],
    )


def _equal(a, b):
    return len(a) == len(b) and all(torch.equal(x, y) for x, y in zip(a, b))


def _member(recorded, i):
    return [noise[i] for noise in recorded]


@pytest.mark.parametrize("n", [0, 2])
def test_single_generator(model, n):
    model.set_noise_accumulation(n)
    generator = torch.Generator().manual_seed(42)
    torch.manual_seed(0)
    first = _forward_noise(model, generator)
    # Re-seeding the generator and flushing the cache reproduces the noise, regardless of the
    # global RNG.
    generator.manual_seed(42)
    model.reset_noise()
    torch.manual_seed(1)
    assert _equal(first, _forward_noise(model, generator))
    # Without re-seeding, the generator keeps advancing.
    assert not _equal(first, _forward_noise(model, generator))


def test_tuple_of_generators(model):
    def forward_noise(seeds, batch_size):
        torch.manual_seed(0)
        generators = tuple(None if s is None else torch.Generator().manual_seed(s) for s in seeds)
        return _forward_noise(model, generators, batch_size=batch_size)

    first = forward_noise((1, None, 3), batch_size=3)
    second = forward_noise((2, None, 3), batch_size=3)
    alone = forward_noise((3,), batch_size=1)
    # The noise of a member depends only on its own generator, ...
    assert not _equal(_member(first, 0), _member(second, 0))
    assert _equal(_member(first, 2), _member(second, 2))
    assert _equal(_member(first, 2), _member(alone, 0))
    # ... and a `None` entry uses the global RNG.
    assert _equal(_member(first, 1), _member(second, 1))


def test_tuple_length_mismatch(model):
    batch = _make_batch(surf_vars=_INPUT_SURF_VARS, batch_size=2)
    with torch.inference_mode(), pytest.raises(ValueError, match="one generator per batch element"):
        model.forward(batch, lead_times=torch.full((2,), 6.0), generator=(torch.Generator(),))


def test_rollout(model):
    generator = torch.Generator().manual_seed(42)
    batch = _make_batch(surf_vars=_INPUT_SURF_VARS)

    def run():
        list(rollout(model, batch, steps=2, generator=generator))

    first = _record_noise(model, run)
    generator.manual_seed(42)
    assert len(first) == 2
    assert _equal(first, _record_noise(model, run))