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