thread global mean removal state through forward/inverse explicitly

GlobalMeanRemoval used to cache per-step state (offsets / shifts and
the normalized extras dict) on the instance between forward_transform
and inverse_transform / extras_normalized. Correctness then depended on
call order — interleaved or skipped calls silently produced wrong
outputs.

Refactor so forward_transform returns an opaque GlobalMeanRemovalState
that callers thread through to inverse_transform and
extras_normalized. The instance carries no per-step state; the
RuntimeError("called before forward_transform") guard is no longer
needed because the contract is enforced by the type signature.

Changes:
- fme.core.step.global_mean_removal: add GlobalMeanRemovalState
  dataclass; forward_transform returns (TensorDict, state);
  inverse_transform and extras_normalized take state. Remove the
  per-instance cached attributes.
- fme.core.step.single_module.step_with_adjustments: capture state
  from forward_transform and pass it back to extras_normalized and
  inverse_transform.
- Tests cover the new state-threading contract with interleaved
  forward/forward/inverse/inverse round-trips for both shared and
  per-channel variants.

Author: mcgibbon

fme/core/step/global_mean_removal.py

@@ -45,37 +45,53 @@ def extra_channel_source_field(name: str) -> str | None:
     return None
 
 
+@dataclasses.dataclass
+class GlobalMeanRemovalState:
+    """Opaque token produced by ``forward_transform`` and consumed by the
+    same ``GlobalMeanRemoval`` instance via ``inverse_transform`` and
+    ``extras_normalized``.
+
+    Callers should thread this value from ``forward_transform`` through
+    to those methods, but not read its fields directly — layout is an
+    implementation detail.
+    """
+
+    shifts: dict[str, torch.Tensor]
+    extras: TensorDict
+
+
 class GlobalMeanRemoval(abc.ABC):
     """Removes global means from fields before normalization and restores
     them after denormalization.
 
     ``forward_transform`` shifts input fields toward their climatological
-    means.  ``inverse_transform`` reverses that shift on output fields.
-    Fields listed in ``field_names`` that appear only in the output (e.g.
-    diagnostic temperatures) are not shifted by ``forward_transform``
-    (they are not inputs), but *are* un-shifted by ``inverse_transform``.
-    The network learns to compensate for this through end-to-end training,
-    so all listed fields — whether input, output, or both — share the
-    same global-mean offset.
+    means and returns a ``GlobalMeanRemovalState`` that must be passed
+    back to ``inverse_transform`` to reverse the shift on output fields.
+    Threading the state explicitly (rather than caching it on the
+    instance) makes the transform stateless and order-independent.
 
     Optional synthetic input channels (when ``append_as_input=True``) are
     produced by ``forward_transform`` already in *normalized* space and
-    exposed via ``extras_normalized()`` as a name -> tensor mapping
+    exposed via ``extras_normalized(state)`` as a name -> tensor mapping
     keyed by ``extra_channel_names``.  Callers append the sentinel names
     to their input packer's name list and merge the dict into the
     normalized-input dict, so the extras flow through packing and the
     channel-mask machinery uniformly with real input channels.
 
+    Fields listed in ``field_names`` that appear only in the output (e.g.
+    diagnostic temperatures) are not shifted by ``forward_transform``
+    (they are not inputs), but *are* un-shifted by ``inverse_transform``.
+    The network learns to compensate for this through end-to-end training,
+    so all listed fields — whether input, output, or both — share the
+    same global-mean offset.
+
     Note:
         "Global mean" here refers to a *cellwise* (unweighted) spatial
         mean of each sample, not the area-weighted global mean used
         elsewhere in ACE for stats and metrics.  The two differ slightly
         on a non-uniform grid; this transform uses the cellwise mean
         for simplicity, since the network learns to compensate during
         end-to-end training.
-
-    Call sequence per step: ``forward_transform`` -> ``extras_normalized``
-    -> ``inverse_transform``.
     """
 
     @property
@@ -98,26 +114,38 @@ def forward_transform(
         self,
         input: TensorMapping,
         data_mask: TensorMapping | None,
-    ) -> TensorDict:
+    ) -> tuple[TensorDict, GlobalMeanRemovalState]:
         """Remove global means from denormalized input fields.
 
-        Caches internal state needed by ``inverse_transform`` and
-        ``extras_normalized``.
+        Returns:
+            ``(shifted_input, state)``.  ``state`` is an opaque value
+            that must be passed back to ``inverse_transform`` to reverse
+            the shift, and to ``extras_normalized`` to obtain the
+            synthetic input channels.
         """
 
     @abc.abstractmethod
-    def inverse_transform(self, output: TensorDict) -> TensorDict:
-        """Restore global means on denormalized output fields."""
+    def inverse_transform(
+        self,
+        output: TensorDict,
+        state: GlobalMeanRemovalState,
+    ) -> TensorDict:
+        """Restore global means on denormalized output fields using
+        ``state`` produced by ``forward_transform``.
+        """
 
-    @abc.abstractmethod
-    def extras_normalized(self) -> TensorDict:
+    def extras_normalized(
+        self,
+        state: GlobalMeanRemovalState,
+    ) -> TensorDict:
         """Return the synthetic input channels in *normalized* space,
         keyed by the names in ``extra_channel_names``.
 
-        Each value has shape ``[batch, *spatial]`` (no channel dim).
-        Returns an empty dict when no extras are configured.  Must be
-        called after ``forward_transform``.
+        Each value has shape ``[batch, *spatial]`` (no channel dim) so
+        it can be fed to the same packer that handles real inputs.
+        Returns an empty dict if no extra channels are configured.
         """
+        return dict(state.extras)
 
 
 class NoGlobalMeanRemoval(GlobalMeanRemoval):
@@ -131,15 +159,16 @@ def forward_transform(
         self,
         input: TensorMapping,
         data_mask: TensorMapping | None,
-    ) -> TensorDict:
-        return dict(input)
+    ) -> tuple[TensorDict, GlobalMeanRemovalState]:
+        return dict(input), GlobalMeanRemovalState(shifts={}, extras={})
 
-    def inverse_transform(self, output: TensorDict) -> TensorDict:
+    def inverse_transform(
+        self,
+        output: TensorDict,
+        state: GlobalMeanRemovalState,
+    ) -> TensorDict:
         return output
 
-    def extras_normalized(self) -> TensorDict:
-        return {}
-
 
 def _broadcast_to_spatial(
     scalar_per_sample: torch.Tensor, spatial_shape: tuple[int, ...]
@@ -174,8 +203,6 @@ def __init__(
         self._append_as_input = append_as_input
         self._reference_mean = reference_mean
         self._reference_std = reference_std
-        self._cached_offset: torch.Tensor | None = None
-        self._cached_extras: TensorDict = {}
 
     @property
     def extra_channel_names(self) -> list[str]:
@@ -187,7 +214,7 @@ def forward_transform(
         self,
         input: TensorMapping,
         data_mask: TensorMapping | None,
-    ) -> TensorDict:
+    ) -> tuple[TensorDict, GlobalMeanRemovalState]:
         ref_name = self._reference_field
         if ref_name not in input:
             raise ValueError(
@@ -202,18 +229,14 @@ def forward_transform(
         ref = input[ref_name]
         sample_mean = ref.mean(dim=tuple(range(1, ref.ndim)))
         offset = self._reference_mean - sample_mean
-        self._cached_offset = offset
         spatial_shape = tuple(ref.shape[1:])
 
+        extras: TensorDict = {}
         if self._append_as_input:
             normalized_mean = -offset / self._reference_std
-            self._cached_extras = {
-                _extra_channel_name(ref_name): _broadcast_to_spatial(
-                    normalized_mean, spatial_shape
-                )
-            }
-        else:
-            self._cached_extras = {}
+            extras[_extra_channel_name(ref_name)] = _broadcast_to_spatial(
+                normalized_mean, spatial_shape
+            )
 
         result = dict(input)
         for name in self._field_names:
@@ -222,24 +245,26 @@ def forward_transform(
             t = result[name]
             broadcast = offset.view(offset.shape[0], *(1,) * (t.ndim - 1))
             result[name] = t + broadcast
-        return result
 
-    def inverse_transform(self, output: TensorDict) -> TensorDict:
-        if self._cached_offset is None:
-            raise RuntimeError("inverse_transform() called before forward_transform().")
-        offset = self._cached_offset
+        # All listed fields share the same offset; inverse will un-shift
+        # those that appear in the output (including output-only fields).
+        shifts = {name: offset for name in self._field_names}
+        return result, GlobalMeanRemovalState(shifts=shifts, extras=extras)
+
+    def inverse_transform(
+        self,
+        output: TensorDict,
+        state: GlobalMeanRemovalState,
+    ) -> TensorDict:
         result = dict(output)
-        for name in self._field_names:
+        for name, shift in state.shifts.items():
             if name not in result:
                 continue
             t = result[name]
-            broadcast = offset.view(offset.shape[0], *(1,) * (t.ndim - 1))
+            broadcast = shift.view(shift.shape[0], *(1,) * (t.ndim - 1))
             result[name] = t - broadcast
         return result
 
-    def extras_normalized(self) -> TensorDict:
-        return dict(self._cached_extras)
-
 
 class PerChannelGlobalMeanRemoval(GlobalMeanRemoval):
     """Shift each field's per-sample cellwise spatial mean to its climatology.
@@ -265,8 +290,6 @@ def __init__(
         self._append_as_input = append_as_input
         self._means = means
         self._stds = stds
-        self._cached_shifts: dict[str, torch.Tensor] | None = None
-        self._cached_extras: TensorDict = {}
 
     @property
     def extra_channel_names(self) -> list[str]:
@@ -278,7 +301,7 @@ def forward_transform(
         self,
         input: TensorMapping,
         data_mask: TensorMapping | None,
-    ) -> TensorDict:
+    ) -> tuple[TensorDict, GlobalMeanRemovalState]:
         result = dict(input)
         shifts: dict[str, torch.Tensor] = {}
         spatial_shape: tuple[int, ...] | None = None
@@ -298,8 +321,6 @@ def forward_transform(
             broadcast = shift.view(shift.shape[0], *(1,) * (t.ndim - 1))
             result[name] = t + broadcast
 
-        self._cached_shifts = shifts
-
         extras: TensorDict = {}
         if self._append_as_input and spatial_shape is not None:
             for name in self._field_names:
@@ -309,26 +330,23 @@ def forward_transform(
                 extras[_extra_channel_name(name)] = _broadcast_to_spatial(
                     normalized, spatial_shape
                 )
-        self._cached_extras = extras
 
-        return result
+        return result, GlobalMeanRemovalState(shifts=shifts, extras=extras)
 
-    def inverse_transform(self, output: TensorDict) -> TensorDict:
-        if self._cached_shifts is None:
-            raise RuntimeError("inverse_transform() called before forward_transform().")
+    def inverse_transform(
+        self,
+        output: TensorDict,
+        state: GlobalMeanRemovalState,
+    ) -> TensorDict:
         result = dict(output)
-        for name in self._field_names:
-            if name not in result or name not in self._cached_shifts:
+        for name, shift in state.shifts.items():
+            if name not in result:
                 continue
             t = result[name]
-            sh = self._cached_shifts[name]
-            broadcast = sh.view(sh.shape[0], *(1,) * (t.ndim - 1))
+            broadcast = shift.view(shift.shape[0], *(1,) * (t.ndim - 1))
             result[name] = t - broadcast
         return result
 
-    def extras_normalized(self) -> TensorDict:
-        return dict(self._cached_extras)
-
 
 @dataclasses.dataclass
 class SharedGlobalMeanRemovalConfig:

fme/core/step/single_module.py

@@ -25,6 +25,7 @@
 from fme.core.step.global_mean_removal import (
     GlobalMeanRemoval,
     GlobalMeanRemovalConfigUnion,
+    GlobalMeanRemovalState,
     NoGlobalMeanRemoval,
     extra_channel_source_field,
 )
@@ -550,24 +551,25 @@ def step_with_adjustments(
     """
     if prescribed_prognostic_names is None:
         prescribed_prognostic_names = []
+    gmr_state: GlobalMeanRemovalState | None = None
     if global_mean_removal is not None:
-        network_input: TensorMapping = global_mean_removal.forward_transform(
+        network_input, gmr_state = global_mean_removal.forward_transform(
             input, data_mask
         )
     else:
-        network_input = input
+        network_input = dict(input)
     input_norm = normalizer.normalize(network_input)
-    if global_mean_removal is not None:
+    if global_mean_removal is not None and gmr_state is not None:
         # Synthetic GMR channels are produced in normalized space; merge
         # them in after normalization so the network sees a single uniform
         # input dict.
-        input_norm = {**input_norm, **global_mean_removal.extras_normalized()}
+        input_norm = {**input_norm, **global_mean_removal.extras_normalized(gmr_state)}
     output_norm = network_calls(input_norm)
     if residual_prediction:
         output_norm = add_names(input_norm, output_norm, prognostic_names)
     output = normalizer.denormalize(output_norm)
-    if global_mean_removal is not None:
-        output = global_mean_removal.inverse_transform(output)
+    if global_mean_removal is not None and gmr_state is not None:
+        output = global_mean_removal.inverse_transform(output, gmr_state)
     if corrector is not None:
         corrector_state: CorrectorState | None = (
             stepper_state.corrector_state if stepper_state is not None else None