Add generator argument to forward for reproducible noise in stochastic mode

aurora/model/aurora.py

@@ -26,7 +26,7 @@
 from aurora.model.encoder import Perceiver3DEncoder
 from aurora.model.lora import LoRAMode
 from aurora.model.perceiver import PerceiverAttention
-from aurora.model.swin3d import Swin3DTransformerBackbone, WindowAttention
+from aurora.model.swin3d import NoiseGenerator, Swin3DTransformerBackbone, WindowAttention
 from aurora.normalisation import log_transform, log_untransform
 
 __all__ = [
@@ -323,6 +323,9 @@ def __init__(
                 if isinstance(m, (WindowAttention, PerceiverAttention)):
                     m.use_fp16_safe_attention = True
 
+        # Warn only once when `generator` is passed to `forward` of a non-stochastic model.
+        self._generator_ignored_warned = False
+
     def reset_noise(self) -> None:
         """Flush the backbone noise cache.
 
@@ -339,18 +342,48 @@ 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,
+        *,
+        generator: NoiseGenerator = None,
+    ) -> Batch:
         """Forward pass.
 
         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.
+            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).
 
         Returns:
             :class:`Batch`: Prediction for the batch.
         """
+        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
+
         batch = self.batch_transform_hook(batch)
 
         # Get the first parameter. We'll derive the data type and device from this parameter.
@@ -433,6 +466,7 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba
                 lead_times=lead_times,
                 patch_res=patch_res,
                 rollout_step=batch.metadata.rollout_step,
+                generator=generator,
             )
         with context_decoder:
             pred = self.decoder(

aurora/model/swin3d.py

@@ -28,6 +28,11 @@
 
 __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."""
+
 
 class MLP(nn.Module):
     """A one-hidden-layer MLP with dropout after the hidden layer and at the end."""
@@ -922,6 +927,12 @@ def reset_noise(self) -> None:
 
         Call this to clear all cached noise tensors, e.g. at the beginning of a new forecast issue
         time. After reset, the next forward call starts building the cache afresh.
+
+        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.
         """
 
         if self.stochastic:
@@ -946,6 +957,57 @@ def set_noise_accumulation(self, n: int = 0) -> None:
             self._noise_cache_size = max(n, 0)
             self._accumulate_noise = self._noise_cache_size > 0
 
+    def _validate_generator(
+        self, generator: NoiseGenerator, batch_size: int, device: torch.device
+    ) -> None:
+        """Validate `generator` against the batch before any randomness is consumed.
+
+        Checking the tuple length and every device up front ensures that no generator has already
+        advanced when an error is raised for a later batch element.
+        """
+        if not isinstance(generator, tuple):
+            return
+        if len(generator) != batch_size:
+            raise ValueError(
+                f"Expected {batch_size} generators (one per batch element), got {len(generator)}."
+            )
+        for i, g in enumerate(generator):
+            if g is None:
+                continue
+            # Like PyTorch, treat an index-less device (e.g. `cuda`) as compatible with an
+            # indexed one (e.g. `cuda:0`); only compare indices when both are explicit.
+            if g.device.type != device.type or (
+                g.device.index is not None
+                and device.index is not None
+                and g.device.index != device.index
+            ):
+                raise ValueError(
+                    f"Generator for batch element {i} is on device `{g.device}`, but noise is "
+                    f"generated on device `{device}`."
+                )
+
+    def _sample_noise(
+        self,
+        shape: tuple[int, ...],
+        device: torch.device,
+        dtype: torch.dtype,
+        generator: NoiseGenerator = None,
+    ) -> torch.Tensor:
+        """Draw one noise sample of shape `(B, L, D)`.
+
+        A single generator produces the whole sample in one draw, so its stream depends on the
+        batch size. A tuple of generators instead produces one `(L, D)` draw per batch element, so
+        the sequence of draws for a given element does not depend on the batch composition. `None`,
+        or a `None` entry in a tuple, falls back to the global RNG. Because the two modes consume
+        a generator with differently shaped draws, a single generator and a tuple are not
+        interchangeable.
+        """
+        if isinstance(generator, tuple):
+            return torch.stack(
+                [torch.randn(shape[1:], device=device, dtype=dtype, generator=g) for g in generator]
+            )
+        return torch.randn(shape, device=device, dtype=dtype, generator=generator)
+
     def get_encoder_specs(
         self, patch_res: tuple[int, int, int]
     ) -> tuple[list[tuple[int, int, int]], list[tuple[int, int, int]]]:
@@ -968,6 +1030,7 @@ def forward(
         lead_times: torch.Tensor,
         rollout_step: int,
         patch_res: tuple[int, int, int],
+        generator: NoiseGenerator = None,
     ) -> torch.Tensor:
         """Run the backbone.
 
@@ -976,6 +1039,10 @@ def forward(
             lead_times (torch.Tensor): Lead times of shape `(batch,)` in hours.
             rollout_step (int): Roll-out step.
             patch_res (tuple[int, int, int]): Patch resolution of the form `(C, H, W)`.
+            generator (torch.Generator or tuple[torch.Generator | None, ...], optional): Source of
+                randomness for noise injection in stochastic mode. See :meth:`_sample_noise` and
+                :meth:`Aurora.forward` for the semantics. Only used when the model is stochastic.
+                Defaults to `None`, which draws from the global RNG.
 
         Returns:
             torch.Tensor: Output tokens of shape `(B, L, D)`.
@@ -997,13 +1064,21 @@ def forward(
 
         if self.stochastic:
             noise_shape = x.shape[:-1] + (self.embed_dim,)
-            noise = torch.randn(noise_shape, device=x.device, dtype=x.dtype)
+            self._validate_generator(generator, x.shape[0], x.device)
+            noise = self._sample_noise(noise_shape, x.device, x.dtype, generator)
             if self._accumulate_noise:
-                # Shape change (e.g. different batch size) invalidates the cache.
-                if self._noise_cache and self._noise_cache[0].shape != noise.shape:
+                # 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()
@@ -1014,7 +1089,7 @@ def forward(
                 # Fill any remaining slots so the cache is always exactly N entries.
                 while len(self._noise_cache) < self._noise_cache_size:
                     self._noise_cache.append(
-                        torch.randn(noise_shape, device=x.device, dtype=x.dtype)
+                        self._sample_noise(noise_shape, x.device, x.dtype, generator)
                     )
                 effective_noise = torch.stack(self._noise_cache).sum(dim=0) / (
                     self._noise_cache_size**0.5

aurora/rollout.py

@@ -8,6 +8,7 @@
 
 from aurora.batch import Batch
 from aurora.model.aurora import Aurora
+from aurora.model.swin3d import NoiseGenerator
 
 __all__ = ["rollout"]
 
@@ -48,6 +49,7 @@ def rollout(
     fine_lead_times: Optional[Sequence[float]] = None,
     use_noise_accumulation: bool = True,
     apply_rollout_input_clipping: bool = True,
+    generator: NoiseGenerator = None,
 ) -> Generator[Batch, None, None]:
     """Perform a roll-out to make long-term predictions.
 
@@ -90,6 +92,13 @@ def rollout(
             back into the model during roll-out, but may be undesirable if the model was not trained
             with clipping and the user wants to preserve the raw model predictions for analysis.
             Default: `True`.
+        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.
 
     Yields:
         :class:`aurora.Batch`: The prediction after every (sub-)step.
@@ -128,7 +137,7 @@ def rollout(
             # 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)
+                pred = model.forward(batch, lead_times=sub_lead_times, generator=generator)
 
                 yield pred
 
@@ -137,7 +146,7 @@ def rollout(
                 pred = model.apply_rollout_input_clipping(pred)
             batch = _advance_batch(batch, pred)
         else:
-            pred = model.forward(batch, lead_times=base_lead_times)
+            pred = model.forward(batch, lead_times=base_lead_times, generator=generator)
 
             yield pred
 

docs/models.md

@@ -450,3 +450,37 @@ 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.
+
+### 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.

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