Separate prognostic and boundary tensors end-to-end

Thread prognostic and boundary inputs as separate tensors through the
data pipeline, stepper, base model, and all model implementations.
Previously they were concatenated along the channel dim before reaching
the model; now each model's forward_once receives (prog, boundary, ctx).

FOMO concatenates before its single-stream encoder (the dual-perceiver
encoder that enables cross-resolution fusion lands in a follow-up).
Samudra and FOMini concatenate inside their forward_once as before.

Also includes unrelated changes from previously-merged PRs that were
on this branch: DataLoadingConfig (#668), Nsight profiling (#674),
samudra highres configs (#679), torch profiling improvements (#677).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

# Conflicts:
#	src/ocean_emulators/stepper.py
#	tests/test_stepper.py

Author: alxmrs

src/ocean_emulators/config.py

@@ -665,7 +665,8 @@ class BaseModelConfig(BaseConfig, abc.ABC):
     @abc.abstractmethod
     def build(
         self,
-        in_channels: int,
+        prog_channels: int,
+        boundary_channels: int,
         out_channels: int,
         hist: int,
         static_data_for_corrector: xr.Dataset | None,
@@ -688,7 +689,8 @@ class SamudraConfig(BaseModelConfig):
 
     def build(
         self,
-        in_channels: int,
+        prog_channels: int,
+        boundary_channels: int,
         out_channels: int,
         hist: int,
         static_data_for_corrector: xr.Dataset | None,
@@ -704,6 +706,7 @@ def build(
             corrector = self.corrector.build(
                 hist, src.spherical_area_weights, static_data_for_corrector
             )
+        in_channels = prog_channels + boundary_channels
         total_in_channels = (
             in_channels + self.pos_channels + (3 if self.add_3d_coordinates else 0)
         )
@@ -751,7 +754,8 @@ class FOMOConfig(BaseModelConfig):
 
     def build(
         self,
-        in_channels: int,
+        prog_channels: int,
+        boundary_channels: int,
         out_channels: int,
         hist: int,
         static_data_for_corrector: xr.Dataset | None,
@@ -773,8 +777,16 @@ def build(
                 "Please set `use_bfloat16=True` or `perceiver_implementation='naive'`."
             )
 
+        in_channels = prog_channels + boundary_channels
+        total_in_channels = in_channels + (3 if self.add_3d_coordinates else 0)
+
         encoder = self.encoder.build(
-            in_channels, self.embedding_dim, extent, max_lat_size, max_lon_size, impl
+            total_in_channels,
+            self.embedding_dim,
+            extent,
+            max_lat_size,
+            max_lon_size,
+            impl,
         )
         processor = self.processor.build(
             self.embedding_dim,
@@ -788,7 +800,6 @@ def build(
             impl,
         )
 
-        total_in_channels = in_channels + (3 if self.add_3d_coordinates else 0)
         add_3d_coordinates = Concat3dCoordinates() if self.add_3d_coordinates else None
         return FOMO(
             in_channels=total_in_channels,
@@ -838,7 +849,8 @@ class FOMiniConfig(BaseModelConfig):
 
     def build(
         self,
-        in_channels: int,
+        prog_channels: int,
+        boundary_channels: int,
         out_channels: int,
         hist: int,
         static_data_for_corrector: xr.Dataset | None,
@@ -857,6 +869,7 @@ def build(
                 "Please set `use_bfloat16=True` or `perceiver_implementation='naive'`."
             )
 
+        in_channels = prog_channels + boundary_channels
         perceiver_io = self.perceiver.build_io(
             self.embedding_dim,
             self.queries_dim,

src/ocean_emulators/constants.py

@@ -29,7 +29,9 @@
 # So, we'll leave this default and use symbolic axes locally.
 type Input = Float[Grid, "*batch total_vars"]
 
-Example = tuple[Input, Prognostic]
+Example = tuple[
+    Prognostic, Boundary, Prognostic
+]  # (prognostic_input, boundary_input, label)
 
 GridMask = Bool[Tensor, "lat lon"]
 PrognosticMask = Bool[GridMask, "prognostic_vars"]

src/ocean_emulators/datasets.py

@@ -13,6 +13,7 @@
 from xarray_einstats.einops import rearrange as xr_rearrange  # noqa: F401
 
 from ocean_emulators.constants import (
+    Boundary,
     BoundaryVarNames,
     Example,
     GridMask,
@@ -138,9 +139,9 @@ def inference_target(self, step: int | slice):
         label = self._get_label(x_index)
         return label
 
-    def get_initial_input(self):
-        data = self.__getitem__(0)[0]
-        return data
+    def get_initial_input(self) -> tuple[Prognostic, Boundary]:
+        prog, boundary, _ = self.__getitem__(0)
+        return prog, boundary
 
     def get_target_time(self, start_step: int, num_steps: int):
         x_index = self._get_x_index(start_step)
@@ -154,20 +155,28 @@ def get_target_time(self, start_step: int, num_steps: int):
             )
         )
 
-    def merge_prognostic_and_boundary(self, prognostic: torch.Tensor, step: int):
+    def get_boundary_for_prognostic(
+        self, prognostic: Prognostic, step: int
+    ) -> tuple[Prognostic, Boundary]:
+        """Return ``(prognostic, boundary)`` at the requested step.
+
+        Prognostic is passed through verbatim; boundary is fetched from the
+        boundary source and returned on the same device as ``prognostic``. The
+        two tensors may live on different spatial grids once the cross-
+        resolution boundary source lands — the encoder is responsible for
+        fusing them.
+        """
         x_index = self._get_x_index(step)
         boundary = self._get_boundary(x_index).to(prognostic.device)
-        data = torch.cat((prognostic, boundary), dim=1)
-        return data
+        return prognostic, boundary
 
     @elapsed(level=logging.DEBUG)
     def __getitem__(self, idx):
         x_index = self._get_x_index(idx)
-        data_in = self._get_prognostic(x_index)
+        data_in_prog = self._get_prognostic(x_index)
         data_in_boundary = self._get_boundary(x_index)
-        data_in = torch.cat((data_in, data_in_boundary), dim=1)
         label = self._get_label(x_index)
-        return (data_in, label)
+        return (data_in_prog, data_in_boundary, label)
 
     def _get_x_index(self, idx):
         if isinstance(idx, slice):
@@ -357,44 +366,40 @@ def pin_memory(self):
 class TrainData:
     """A single batch of training data.
 
-    A single batch contains multiple steps worth of `Example`s (i.e., input/output pairs). These steps are used during
-    autoregressive rollout in the training and inference process.
-
-    Constraint: The `Input` tensor is a combination of (flattened) prognostic variables (at all depth levels) and
-    boundary forcings. The top `num_prognostic_channels` number of channels must be prognostic variables whereas the
-    remaining bottom channels are boundary forcings.
+    A single batch contains multiple steps worth of ``Example`` entries, each
+    of which is a ``(prognostic_input, boundary_input, label)`` triple. The
+    prognostic and boundary tensors are carried *separately* (they are no
+    longer concatenated along the channel dim) so that the encoder can fuse
+    them at the token level, potentially at different spatial resolutions.
     """
 
-    def __init__(self, num_prognostic_channels: int, ctx: GridContext):
+    def __init__(
+        self, num_prognostic_channels: int, num_boundary_channels: int, ctx: GridContext
+    ):
         self.num_prognostic_channels = num_prognostic_channels
+        self.num_boundary_channels = num_boundary_channels
         self.ctx = ctx
         self.example_by_step: list[Example] = []
         self.load_stats: LoadStats | None = None
 
-    def append(self, input_: Input, label: Prognostic):
+    def append(
+        self, prognostic_input: Prognostic, boundary_input: Boundary, label: Prognostic
+    ) -> None:
         """Add another Example as a new step."""
-        self.example_by_step.append((input_, label))
+        self.example_by_step.append((prognostic_input, boundary_input, label))
 
-    def get_initial_input(self) -> Input:
+    def get_initial_input(self) -> tuple[Prognostic, Boundary]:
         return self.get_input(0)
 
-    def get_input(self, step: int) -> Input:
-        return self[step][0]
+    def get_input(self, step: int) -> tuple[Prognostic, Boundary]:
+        prog, boundary, _ = self.example_by_step[step]
+        return prog, boundary
 
     def get_label(self, step: int) -> Prognostic:
-        return self[step][1]
-
-    def merge_prognostic_and_boundary(self, prognostic: torch.Tensor, step: int):
-        input_ = self.get_input(step)
-        merged = input_.clone()
-        merged[:, : self.num_prognostic_channels] = prognostic
-        return merged
-
-    def values(self):
-        return self.example_by_step
+        return self.example_by_step[step][2]
 
     def __getitem__(self, step: int) -> Example:
-        """Converts index (step) into (data, label) tuple."""
+        """Converts index (step) into (prognostic, boundary, label) triple."""
         return self.example_by_step[step]
 
     def __len__(self) -> int:
@@ -405,16 +410,20 @@ def __iter__(self):
 
     def to(self, device: torch.device) -> None:
         for step in self:
+            prog, boundary, label = self.example_by_step[step]
             self.example_by_step[step] = (
-                self[step][0].to(device, non_blocking=True),
-                self[step][1].to(device, non_blocking=True),
+                prog.to(device, non_blocking=True),
+                boundary.to(device, non_blocking=True),
+                label.to(device, non_blocking=True),
             )
 
     def pin_memory(self):
         for step in self:
+            prog, boundary, label = self.example_by_step[step]
             self.example_by_step[step] = (
-                self[step][0].pin_memory(),
-                self[step][1].pin_memory(),
+                prog.pin_memory(),
+                boundary.pin_memory(),
+                label.pin_memory(),
             )
         return self
 
@@ -484,6 +493,7 @@ def __init__(
         self._concurrent_compute = concurrent_compute_
 
         self.num_prognostic_channels: int = (hist + 1) * len(prognostic_var_names)
+        self.num_boundary_channels: int = (hist + 1) * len(boundary_var_names)
         assert np.array_equal(srcs[0].data.time, srcs[-1].data.time), (
             "src and dst DataSource have different time slices!"
         )
@@ -605,9 +615,13 @@ def to_train_data(
         Returns:
             TrainData with tensors on the target device
         """
-        train_data = TrainData(self.num_prognostic_channels, self.ctx.to(device))
+        train_data = TrainData(
+            self.num_prognostic_channels,
+            self.num_boundary_channels,
+            self.ctx.to(device),
+        )
         for input_, boundary, label in raw_train_data.raw_data:
-            input_, label = self._to_example(
+            prog_input, boundary_input, label_tensor = self._to_example(
                 OceanData.from_data_source(
                     input_,
                     self.wet_prognostic[0],
@@ -622,43 +636,27 @@ def to_train_data(
                     label, self.wet_prognostic[-1], self.prognostic_srcs[-1]
                 ).to(device=device, non_blocking=True),
             )
-            train_data.append(input_, label)
+            train_data.append(prog_input, boundary_input, label_tensor)
         train_data.load_stats = raw_train_data.load_stats
         return train_data
 
     def _to_example(
         self, input_: OceanData, boundary: OceanData, label: OceanData
-    ) -> tuple[Input, Prognostic]:
+    ) -> Example:
         # Input/boundary only include current steps; label only includes forecasted steps.
-        total_input = self._prep_tensor_steps(input_, boundary)
+        prog_input = self._prep_tensor_steps(input_)
+        boundary_input = self._prep_tensor_steps(boundary)
         label_tensor = self._prep_tensor_steps(label)
-        return total_input, label_tensor
+        return prog_input, boundary_input, label_tensor
 
-    def _prep_tensor_steps(
-        self,
-        prognostic: OceanData,
-        boundary: OceanData | None = None,
-    ) -> Input:
-        """Prepare tensor steps by normalizing, masking and flattening dimensions."""
-        prognostic_steps = prognostic.normalize_and_mask(
+    def _prep_tensor_steps(self, ocean_data: OceanData) -> Input:
+        """Normalize, mask, and flatten (time, variable) dims into a channel dim."""
+        steps = ocean_data.normalize_and_mask(
             self.normalize_before_mask, self.masked_fill_value
         )
-
-        # Flatten time and variable dimensions
-        def flatten_dims(tensor: torch.Tensor) -> torch.Tensor:
-            return rearrange(
-                tensor, "batch time variable lat lon -> batch (time variable) lat lon"
-            )
-
-        prognostic_steps = flatten_dims(prognostic_steps)
-        if boundary is not None:
-            boundary_steps = boundary.normalize_and_mask(
-                self.normalize_before_mask, self.masked_fill_value
-            )
-            boundary_steps = flatten_dims(boundary_steps)
-            return torch.cat((prognostic_steps, boundary_steps), dim=1)
-
-        return prognostic_steps
+        return rearrange(
+            steps, "batch time variable lat lon -> batch (time variable) lat lon"
+        )
 
     def _get_x_index(self, idx: int, step: int) -> xr.DataArray:
         assert isinstance(idx, int)

src/ocean_emulators/eval.py

@@ -18,7 +18,7 @@
     construct_metadata,
 )
 from ocean_emulators.datasets import InferenceDataset
-from ocean_emulators.stepper import Stepper
+from ocean_emulators.stepper import inference
 from ocean_emulators.utils.data import (
     Normalize,
     get_inference_steps,
@@ -76,8 +76,10 @@ def __init__(self, cfg: EvalConfig) -> None:
         self.N_bound = len(self.boundary_var_names)
         self.N_prog = len(self.prognostic_var_names)
 
-        self.num_in = int((cfg.data.hist + 1) * (self.N_prog + self.N_bound))
-        self.num_out = int((cfg.data.hist + 1) * self.N_prog)
+        self.num_prog_in = int((cfg.data.hist + 1) * self.N_prog)
+        self.num_boundary_in = int((cfg.data.hist + 1) * self.N_bound)
+        self.num_in = self.num_prog_in + self.num_boundary_in
+        self.num_out = self.num_prog_in
 
         self.tensor_map = TensorMap.init_instance(
             cfg.experiment.prognostic_vars_key, cfg.experiment.boundary_vars_key
@@ -110,7 +112,8 @@ def __init__(self, cfg: EvalConfig) -> None:
 
         # Model
         self.model = cfg.model.build(
-            in_channels=self.num_in,
+            prog_channels=self.num_prog_in,
+            boundary_channels=self.num_boundary_in,
             out_channels=self.num_out,
             hist=cfg.data.hist,
             static_data_for_corrector=self.static_data,
@@ -207,7 +210,7 @@ def standalone_inference(self):
             self.prognostic_var_names,
         )
 
-        Stepper.inference(
+        inference(
             model=self.model,
             dataset=self.inference_dataset,
             inf_aggregator=inf_aggregator,