Pr/richer batch

models/docs/usage/create_model.rst

@@ -188,7 +188,8 @@ actual model (see :ref:`overview`).
    -  ``y_norm = model_interface.forward(x_norm)`` with ``x_in`` and
       ``y_pred`` are normalized.
    -  ``y = model_interface.predict_step(x)`` with ``x`` and ``y`` are
-      absolute values.
+      absolute values. Transport models additionally require a
+      ``target_template`` batch so sampling knows the output geometry.
 
 *******************
  The PyTorch Model

models/src/anemoi/models/data/batch.py

@@ -76,7 +76,7 @@ class Batch:
         Per-dataset input. For gridded datasets a single stacked tensor of
         shape ``(batch, time, ensemble, grid, vars)``; for sparse
         observation datasets a ``list[torch.Tensor]`` of length ``batch``,
-        one entry per sample with shape ``(ensemble=1, grid_i, vars)``.
+        one entry per sample with shape ``(grid_i, vars)``.
     coordinates : dict[str, torch.Tensor | list[torch.Tensor]]
         Per-dataset ``(N, 2)`` coordinate tensor stacking
         ``(latitudes, longitudes)`` in **radians**. For static-grid
@@ -215,8 +215,8 @@ def __getitem__(self, dataset_name: str) -> "SourceView":
         return create_source_view(
             name=dataset_name,
             data=self.data[dataset_name],
-            variables=self.variables.get(dataset_name),
-            statistics=self.statistics[dataset_name],
+            variables=self.variables.get(dataset_name, []),
+            statistics=self.statistics.get(dataset_name, {}),
             coordinates=self.coordinates.get(dataset_name),
             is_static=self.is_static_coords(dataset_name),
             timedeltas=self.timedeltas.get(dataset_name),
@@ -343,19 +343,24 @@ def with_data(self, new_data: dict[str, torch.Tensor | list[torch.Tensor]]) -> "
             )
 
         new_data_keys = set(new_data.keys())
+        unknown_keys = new_data_keys - set(self.data.keys())
+        if unknown_keys:
+            msg = f"Replacement data contains unknown dataset names: {sorted(unknown_keys)}."
+            raise ValueError(msg)
+
         metadata_static_coords = self.metadata.get(STATIC_COORDS_META_KEY, frozenset())
         metadata_static_coords &= new_data_keys
         return Batch(
             new_data,
-            coordinates={name: self.coordinates[name] for name in new_data_keys},
+            coordinates={name: self.coordinates[name] for name in new_data_keys if name in self.coordinates},
             metadata={STATIC_COORDS_META_KEY: metadata_static_coords}
             | {name: self.metadata[name] for name in new_data_keys if name in self.metadata},
-            grid_sizes={name: self.grid_sizes[name] for name in new_data_keys},
+            grid_sizes={name: self.grid_sizes[name] for name in new_data_keys if name in self.grid_sizes},
             timedeltas={name: self.timedeltas[name] for name in new_data_keys if name in self.timedeltas},
             shard_sizes={name: self.shard_sizes[name] for name in new_data_keys if name in self.shard_sizes},
-            layouts={name: self.layouts[name] for name in new_data_keys},
-            variables={name: self.variables[name] for name in new_data_keys},
-            statistics=self.statistics,
+            layouts={name: self.layouts[name] for name in new_data_keys if name in self.layouts},
+            variables={name: self.variables[name] for name in new_data_keys if name in self.variables},
+            statistics={name: self.statistics[name] for name in new_data_keys if name in self.statistics},
         )
 
     def update_source(self, source_name: str, source_view: SourceView) -> "Batch":
@@ -448,7 +453,7 @@ def collate(
         * **Sparse** — ``payload["metadata"]["boundaries"]`` is present (set
           by :meth:`anemoi.training.data.data_reader.ObservationDataReader._unpack_sample`).
           ``payload["data"]`` is a per-sample :class:`torch.Tensor` of
-          shape ``(E=1, N_i, V)`` whose ``N_i`` varies between samples.
+          shape ``(N_i, V)`` whose ``N_i`` varies between samples.
           ``data[name]``, ``coordinates[name]`` and ``timedeltas[name]``
           each become a ``list[torch.Tensor]`` of length ``B``;
           per-sample ``payload["metadata"]`` is collected into

models/src/anemoi/models/data/tensor_layout.py

@@ -40,9 +40,9 @@ class TensorLayout:
     Notes
     -----
     For sparse observation datasets (``time_in_grid=True``) the per-sample
-    inner tensor has shape ``(ensemble, grid, variables)`` and the batch
-    dimension is represented by the outer Python list (one tensor per
-    sample). ``batch`` therefore stays ``None`` even after collation.
+    inner tensor has shape ``(grid, variables)`` and the batch dimension is
+    represented by the outer Python list (one tensor per sample). ``batch``
+    therefore stays ``None`` even after collation.
     """
 
     batch: int | None = None

models/src/anemoi/models/data/views.py

@@ -32,10 +32,11 @@
 
 def create_source_view(**kwargs) -> "SourceView":
     """Factory function to create a SourceView for a source dataset."""
-    if kwargs.pop("is_static"):
-        return GriddedSourceView(**kwargs)
+    kwargs.pop("is_static")
+    if kwargs["layout"].time_in_grid:
+        return TabularSourceView(**kwargs)
 
-    return TabularSourceView(**kwargs)
+    return GriddedSourceView(**kwargs)
 
 
 def _fancy_variable_index(
@@ -232,8 +233,34 @@ def flatten(self) -> FlatView:
         flattened_data = einops.rearrange(self.data, f"{current_pattern} -> {self.pattern_for_2d}")
         device = self.data.device
 
-        batch_size = self.data.shape[self.layout.batch]
-        coordinates = einops.repeat(self.coordinates, "grid latlon -> (batch grid) latlon", batch=batch_size)
+        if self.coordinates is None:
+            msg = f"{self.__class__.__name__} requires coordinates for flattening."
+            raise ValueError(msg)
+        if isinstance(self.coordinates, list):
+            msg = f"{self.__class__.__name__} coordinates must be a tensor, not a list."
+            raise TypeError(msg)
+
+        batch_size = self.data.shape[self.layout.axis("batch", ndim=self.data.ndim)]
+        ensemble_size = self.data.shape[self.layout.axis("ensemble", ndim=self.data.ndim)]
+        if self.coordinates.ndim == 2:
+            coordinates = einops.repeat(
+                self.coordinates,
+                "grid latlon -> (batch ensemble grid) latlon",
+                batch=batch_size,
+                ensemble=ensemble_size,
+            )
+        elif self.coordinates.ndim == 3:
+            coordinates = einops.repeat(
+                self.coordinates,
+                "batch grid latlon -> (batch ensemble grid) latlon",
+                ensemble=ensemble_size,
+            )
+        else:
+            msg = (
+                f"{self.__class__.__name__} coordinates must have shape (grid, 2) "
+                f"or (batch, grid, 2), got {tuple(self.coordinates.shape)}."
+            )
+            raise ValueError(msg)
 
         return FlatView(
             data=flattened_data,
@@ -482,6 +509,18 @@ def apply_loss(self, other: "TabularSourceView", loss_func: Callable, **kwargs)
             assert torch.all(
                 self.coordinates[i] == other.coordinates[i]
             ), f"Sample {i} of both views must have the same coordinates; got {self.coordinates[i]} and {other.coordinates[i]}."
+            sample_kwargs = {
+                key: (
+                    value[i]
+                    if (
+                        isinstance(value, list)
+                        and len(value) == len(self.data)
+                        and all(isinstance(item, torch.Tensor) for item in value)
+                    )
+                    else value
+                )
+                for key, value in kwargs.items()
+            }
 
             losses.append(
                 loss_func(
@@ -490,18 +529,22 @@ def apply_loss(self, other: "TabularSourceView", loss_func: Callable, **kwargs)
                     layout=self.layout,
                     statistics=self.statistics,
                     name_to_index=self.name_to_index,
-                    **kwargs,
+                    **sample_kwargs,
                 )
             )
             # Handle empty batches: a fully-empty worker returns a graph-connected 0
             non_empty.append(pred.shape[self.layout.grid] > 0)
 
+        if not losses:
+            msg = "Cannot apply a loss to an empty sparse source view."
+            raise ValueError(msg)
+
         stacked = torch.stack(losses)
         num_non_empty = sum(non_empty)
         # Divide by the number of non-empty samples (>= 1) rather than the batch size.
         # When every sample is empty, the stacked tensor is all-zero and graph-connected,
         # so summing and dividing by 1 preserves the zero gradient path.
-        return stacked.sum() / max(num_non_empty, 1)
+        return stacked.sum(dim=0) / max(num_non_empty, 1)
 
     def allgather(self, group: ProcessGroup | None) -> "TabularSourceView":
         """Allgather this view across the given process group.

models/src/anemoi/models/interface/__init__.py

@@ -224,7 +224,8 @@ def predict_step(
         gather_out : bool, optional
             Whether to gather the output, by default True.
         **kwargs
-            Additional prediction keyword arguments.
+            Additional prediction keyword arguments. Transport models require
+            ``target_template`` here so sampling knows the output geometry.
 
         Returns
         -------

models/src/anemoi/models/layers/mapper.py

@@ -359,7 +359,7 @@ def mapper_forward_with_edge_sharding(
         )
 
         out_channels = self.out_channels_dst if self.out_channels_dst is not None else self.hidden_dim
-        out_type = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else x_dst.dtype
+        out_type = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else x_dst.dtype
         out_dst = torch.empty((*x_dst.shape[:-1], out_channels), device=x_dst.device, dtype=out_type)
 
         for chunk_id in range(chunk_partition.num_parts):

models/src/anemoi/models/models/encoder_processor_decoder.py

@@ -276,13 +276,17 @@ def forward(
             tensors used by dynamic graph providers / node attributes. Per-dataset
             grid sharding is carried by the batch and read through the source
             views (``view.flatten().shard_sizes``).
+        target : Optional[Batch], optional
+            Output-time batch used to refresh decoder-side conditioning, by default None.
         model_comm_group : Optional[ProcessGroup], optional
-            Model communication group, by default None
+            Model communication group, by default None.
+        **kwargs
+            Additional keyword arguments forwarded to the mappers and processor.
 
         Returns
         -------
         dict[str, Tensor]
-            Output of the model, with the same shape as the input (sharded if input is sharded)
+            Output of the model, with the same shape as the input (sharded if input is sharded).
         """
         dataset_names = list(batch.keys())
 
@@ -340,7 +344,7 @@ def forward(
                 model_comm_group=model_comm_group,
                 **graph_batch_kwargs,
             )
-            encoder_edge_attr = encoder_edge_attr.to(x_data_latent.device)
+            encoder_edge_attr = encoder_edge_attr.to(x_data_latent.device)  # todo SL: remove device movement
             encoder_edge_index = encoder_edge_index.to(x_data_latent.device)
 
             enc_shard_info = BipartiteGraphShardInfo(
@@ -372,7 +376,7 @@ def forward(
             batch_size=batch_size,
             model_comm_group=model_comm_group,
         )
-        processor_edge_attr = processor_edge_attr.to(x_latent.device)
+        processor_edge_attr = processor_edge_attr.to(x_latent.device)  # todo SL: remove device movement
         processor_edge_index = processor_edge_index.to(x_latent.device)
 
         x_latent_proc = self.processor(
@@ -422,7 +426,7 @@ def forward(
                 model_comm_group=model_comm_group,
                 **graph_batch_kwargs,
             )
-            decoder_edge_attr = decoder_edge_attr.to(x_latent.device)
+            decoder_edge_attr = decoder_edge_attr.to(x_latent.device)  # todo SL: remove device movement
             decoder_edge_index = decoder_edge_index.to(x_latent.device)
 
             dec_shard_info = BipartiteGraphShardInfo(