Add explicit spatial_ndim tracking to MetaTensor (Fixes #6397)

Fixes dimension-mismatch crashes when einops.rearrange() or other reshape
operations change tensor ndim by decoupling spatial rank from tensor shape.

- Add _spatial_ndim attribute to MetaObj, derived from affine in MetaTensor
- Expose spatial_ndim property with getter/setter and validation
- Sync spatial_ndim on affine assignment and propagate through collation
- Update transforms to use spatial_ndim instead of ndim-1 heuristic
- Add 18 new tests for spatial_ndim behavior

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>

Author: aymuos15

monai/data/__init__.py

@@ -71,7 +71,7 @@
     monai_to_itk_ddf,
 )
 from .meta_obj import MetaObj, get_track_meta, set_track_meta
-from .meta_tensor import MetaTensor
+from .meta_tensor import MetaTensor, get_spatial_ndim
 from .samplers import DistributedSampler, DistributedWeightedRandomSampler
 from .synthetic import create_test_image_2d, create_test_image_3d
 from .test_time_augmentation import TestTimeAugmentation

monai/data/meta_obj.py

@@ -84,6 +84,7 @@ def __init__(self) -> None:
         self._applied_operations: list = MetaObj.get_default_applied_operations()
         self._pending_operations: list = MetaObj.get_default_applied_operations()  # the same default as applied_ops
         self._is_batch: bool = False
+        self._spatial_ndim: int = 3  # default: 3 spatial dimensions
 
     @staticmethod
     def flatten_meta_objs(*args: Iterable):

monai/data/meta_tensor.py

@@ -21,14 +21,25 @@
 import torch
 
 import monai
-from monai.config.type_definitions import NdarrayTensor
+from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor
 from monai.data.meta_obj import MetaObj, get_track_meta
 from monai.data.utils import affine_to_spacing, decollate_batch, list_data_collate, remove_extra_metadata
 from monai.utils import look_up_option
 from monai.utils.enums import LazyAttr, MetaKeys, PostFix, SpaceKeys
 from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_numpy, convert_to_tensor
 
-__all__ = ["MetaTensor"]
+__all__ = ["MetaTensor", "get_spatial_ndim"]
+
+
+def get_spatial_ndim(img: NdarrayOrTensor) -> int:
+    """Return the number of spatial dimensions assuming channel-first layout.
+
+    Uses ``MetaTensor.spatial_ndim`` when available, otherwise falls back to
+    ``img.ndim - 1``.
+    """
+    if isinstance(img, MetaTensor):
+        return img.spatial_ndim
+    return img.ndim - 1
 
 
 @functools.lru_cache(None)
@@ -111,6 +122,7 @@ def __new__(
         meta: dict | None = None,
         applied_operations: list | None = None,
         *args,
+        spatial_ndim: int | None = None,
         **kwargs,
     ) -> MetaTensor:
         _kwargs = {"device": kwargs.pop("device", None), "dtype": kwargs.pop("dtype", None)} if kwargs else {}
@@ -123,6 +135,7 @@ def __init__(
         meta: dict | None = None,
         applied_operations: list | None = None,
         *_args,
+        spatial_ndim: int | None = None,
         **_kwargs,
     ) -> None:
         """
@@ -134,6 +147,8 @@ def __init__(
                 the list is typically maintained by `monai.transforms.TraceableTransform`.
                 See also: :py:class:`monai.transforms.TraceableTransform`
             _args: additional args (currently not in use in this constructor).
+            spatial_ndim: optional number of spatial dimensions. If ``None``, derived
+                from the affine matrix clamped by the tensor shape.
             _kwargs: additional kwargs (currently not in use in this constructor).
 
         Note:
@@ -158,6 +173,12 @@ def __init__(
             self.affine = self.meta[MetaKeys.AFFINE]
         else:
             self.affine = self.get_default_affine()
+        # derive spatial_ndim from affine, clamped by tensor shape
+        if spatial_ndim is not None:
+            self.spatial_ndim = spatial_ndim
+        elif self.affine.ndim == 2:
+            self.spatial_ndim = min(self.affine.shape[-1] - 1, max(self.ndim - 1, 1))
+
         # applied_operations
         if applied_operations is not None:
             self.applied_operations = applied_operations
@@ -468,14 +489,29 @@ def affine(self) -> torch.Tensor:
     @affine.setter
     def affine(self, d: NdarrayTensor) -> None:
         """Set the affine."""
-        self.meta[MetaKeys.AFFINE] = torch.as_tensor(d, device=torch.device("cpu"), dtype=torch.float64)
+        a = torch.as_tensor(d, device=torch.device("cpu"), dtype=torch.float64)
+        self.meta[MetaKeys.AFFINE] = a
+        if a.ndim == 2:  # non-batched: sync spatial_ndim
+            self.spatial_ndim = a.shape[-1] - 1
+
+    @property
+    def spatial_ndim(self) -> int:
+        """Get the number of spatial dimensions."""
+        return getattr(self, "_spatial_ndim", 3)
+
+    @spatial_ndim.setter
+    def spatial_ndim(self, val: int) -> None:
+        """Set the number of spatial dimensions."""
+        if val < 1:
+            raise ValueError(f"spatial_ndim must be >= 1, got {val}")
+        self._spatial_ndim = val
 
     @property
     def pixdim(self):
         """Get the spacing"""
         if self.is_batch:
-            return [affine_to_spacing(a) for a in self.affine]
-        return affine_to_spacing(self.affine)
+            return [affine_to_spacing(a, r=self.spatial_ndim) for a in self.affine]
+        return affine_to_spacing(self.affine, r=self.spatial_ndim)
 
     def peek_pending_shape(self):
         """
@@ -490,7 +526,7 @@ def peek_pending_shape(self):
 
     def peek_pending_affine(self):
         res = self.affine
-        r = len(res) - 1
+        r = res.shape[-1] - 1 if res.ndim >= 2 else self.spatial_ndim
         if r not in (2, 3):
             warnings.warn(f"Only 2d and 3d affine are supported, got {r}d input.")
         for p in self.pending_operations:
@@ -503,8 +539,10 @@ def peek_pending_affine(self):
         return res
 
     def peek_pending_rank(self):
-        a = self.pending_operations[-1].get(LazyAttr.AFFINE, None) if self.pending_operations else self.affine
-        return 1 if a is None else int(max(1, len(a) - 1))
+        if self.pending_operations:
+            a = self.pending_operations[-1].get(LazyAttr.AFFINE, None)
+            return 1 if a is None else int(max(1, len(a) - 1))
+        return self.spatial_ndim
 
     def new_empty(self, size, dtype=None, device=None, requires_grad=False):  # type: ignore[override]
         """

monai/data/utils.py

@@ -432,6 +432,7 @@ def collate_meta_tensor_fn(batch, *, collate_fn_map=None):
     collated.meta = default_collate(meta_dicts)
     collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]
     collated.is_batch = True
+    collated.spatial_ndim = getattr(batch[0], "spatial_ndim", 3)  # assumes uniform spatial_ndim
     return collated
 
 

monai/transforms/croppad/functional.py

@@ -22,7 +22,7 @@
 
 from monai.config.type_definitions import NdarrayTensor
 from monai.data.meta_obj import get_track_meta
-from monai.data.meta_tensor import MetaTensor
+from monai.data.meta_tensor import MetaTensor, get_spatial_ndim
 from monai.data.utils import to_affine_nd
 from monai.transforms.inverse import TraceableTransform
 from monai.transforms.utils import convert_pad_mode, create_translate
@@ -132,7 +132,7 @@ def crop_or_pad_nd(img: torch.Tensor, translation_mat, spatial_size: tuple[int,
         mode: the padding mode.
         kwargs: other arguments for the `np.pad` or `torch.pad` function.
     """
-    ndim = len(img.shape) - 1
+    ndim = get_spatial_ndim(img)
     matrix_np = np.round(to_affine_nd(ndim, convert_to_numpy(translation_mat, wrap_sequence=True).copy()))
     matrix_np = to_affine_nd(len(spatial_size), matrix_np)
     cc = np.asarray(np.meshgrid(*[[0.5, x - 0.5] for x in spatial_size], indexing="ij"))

monai/transforms/intensity/array.py

@@ -26,6 +26,7 @@
 from monai.config import DtypeLike
 from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor
 from monai.data.meta_obj import get_track_meta
+from monai.data.meta_tensor import get_spatial_ndim
 from monai.data.ultrasound_confidence_map import UltrasoundConfidenceMap
 from monai.data.utils import get_random_patch, get_valid_patch_size
 from monai.networks.layers import GaussianFilter, HilbertTransform, MedianFilter, SavitzkyGolayFilter
@@ -1580,7 +1581,7 @@ def __init__(self, radius: Sequence[int] | int = 1) -> None:
     def __call__(self, img: NdarrayTensor) -> NdarrayTensor:
         img = convert_to_tensor(img, track_meta=get_track_meta())
         img_t, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float)
-        spatial_dims = img_t.ndim - 1
+        spatial_dims = get_spatial_ndim(img)
         r = ensure_tuple_rep(self.radius, spatial_dims)
         median_filter_instance = MedianFilter(r, spatial_dims=spatial_dims)
         out_t: torch.Tensor = median_filter_instance(img_t)
@@ -1616,7 +1617,7 @@ def __call__(self, img: NdarrayTensor) -> NdarrayTensor:
             sigma = [torch.as_tensor(s, device=img_t.device) for s in self.sigma]
         else:
             sigma = torch.as_tensor(self.sigma, device=img_t.device)
-        gaussian_filter = GaussianFilter(img_t.ndim - 1, sigma, approx=self.approx)
+        gaussian_filter = GaussianFilter(get_spatial_ndim(img), sigma, approx=self.approx)
         out_t: torch.Tensor = gaussian_filter(img_t.unsqueeze(0)).squeeze(0)
         out, *_ = convert_to_dst_type(out_t, dst=img, dtype=out_t.dtype)
 
@@ -1673,7 +1674,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen
         if not self._do_transform:
             return img
 
-        sigma = ensure_tuple_size(vals=(self.x, self.y, self.z), dim=img.ndim - 1)
+        sigma = ensure_tuple_size(vals=(self.x, self.y, self.z), dim=get_spatial_ndim(img))
         return GaussianSmooth(sigma=sigma, approx=self.approx)(img)
 
 
@@ -1723,7 +1724,7 @@ def __call__(self, img: NdarrayTensor) -> NdarrayTensor:
         img_t, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float32)
 
         gf1, gf2 = (
-            GaussianFilter(img_t.ndim - 1, sigma, approx=self.approx).to(img_t.device)
+            GaussianFilter(get_spatial_ndim(img), sigma, approx=self.approx).to(img_t.device)
             for sigma in (self.sigma1, self.sigma2)
         )
         blurred_f = gf1(img_t.unsqueeze(0))
@@ -1811,8 +1812,9 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen
 
         if self.x2 is None or self.y2 is None or self.z2 is None or self.a is None:
             raise RuntimeError("please call the `randomize()` function first.")
-        sigma1 = ensure_tuple_size(vals=(self.x1, self.y1, self.z1), dim=img.ndim - 1)
-        sigma2 = ensure_tuple_size(vals=(self.x2, self.y2, self.z2), dim=img.ndim - 1)
+        _sp = get_spatial_ndim(img)
+        sigma1 = ensure_tuple_size(vals=(self.x1, self.y1, self.z1), dim=_sp)
+        sigma2 = ensure_tuple_size(vals=(self.x2, self.y2, self.z2), dim=_sp)
         return GaussianSharpen(sigma1=sigma1, sigma2=sigma2, alpha=self.a, approx=self.approx)(img)
 
 

monai/transforms/inverse.py

@@ -213,7 +213,7 @@ def track_transform_meta(
             orig_affine = data_t.peek_pending_affine()
             orig_affine = convert_to_dst_type(orig_affine, affine, dtype=torch.float64)[0]
             try:
-                affine = orig_affine @ to_affine_nd(len(orig_affine) - 1, affine, dtype=torch.float64)
+                affine = orig_affine @ to_affine_nd(orig_affine.shape[-1] - 1, affine, dtype=torch.float64)
             except RuntimeError as e:
                 if orig_affine.ndim > 2:
                     if data_t.is_batch:

monai/transforms/lazy/functional.py

@@ -256,9 +256,11 @@ def apply_pending(data: torch.Tensor | MetaTensor, pending: list | None = None,
     if not pending:
         return data, []
 
+    _rank = data.spatial_ndim if isinstance(data, MetaTensor) else 3
+
     cumulative_xform = affine_from_pending(pending[0])
-    if cumulative_xform.shape[0] == 3:
-        cumulative_xform = to_affine_nd(3, cumulative_xform)
+    if cumulative_xform.shape[0] < _rank + 1:
+        cumulative_xform = to_affine_nd(_rank, cumulative_xform)
 
     cur_kwargs = kwargs_from_pending(pending[0])
     override_kwargs: dict[str, Any] = {}
@@ -283,8 +285,8 @@ def apply_pending(data: torch.Tensor | MetaTensor, pending: list | None = None,
             data = resample(data.to(device), cumulative_xform, _cur_kwargs)
 
         next_matrix = affine_from_pending(p)
-        if next_matrix.shape[0] == 3:
-            next_matrix = to_affine_nd(3, next_matrix)
+        if next_matrix.shape[0] < _rank + 1:
+            next_matrix = to_affine_nd(_rank, next_matrix)
 
         cumulative_xform = combine_transforms(cumulative_xform, next_matrix)
         cur_kwargs.update(new_kwargs)

monai/transforms/post/array.py

@@ -23,7 +23,7 @@
 
 from monai.config.type_definitions import NdarrayOrTensor
 from monai.data.meta_obj import get_track_meta
-from monai.data.meta_tensor import MetaTensor
+from monai.data.meta_tensor import MetaTensor, get_spatial_ndim
 from monai.networks import one_hot
 from monai.networks.layers import GaussianFilter, apply_filter, separable_filtering
 from monai.transforms.inverse import InvertibleTransform
@@ -624,7 +624,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
         img = convert_to_tensor(img, track_meta=get_track_meta())
         img_: torch.Tensor = convert_to_tensor(img, track_meta=False)
-        spatial_dims = len(img_.shape) - 1
+        spatial_dims = get_spatial_ndim(img)
         img_ = img_.unsqueeze(0)  # adds a batch dim
         if spatial_dims == 2:
             kernel = torch.tensor([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype=torch.float32)
@@ -1104,7 +1104,7 @@ def __call__(self, image: NdarrayOrTensor) -> torch.Tensor:
         image_tensor = convert_to_tensor(image, track_meta=get_track_meta())
 
         # Check/set spatial axes
-        n_spatial_dims = image_tensor.ndim - 1  # excluding the channel dimension
+        n_spatial_dims = get_spatial_ndim(image_tensor)
         valid_spatial_axes = list(range(n_spatial_dims)) + list(range(-n_spatial_dims, 0))
 
         # Check gradient axes to be valid

monai/transforms/spatial/array.py

@@ -27,7 +27,7 @@
 from monai.config.type_definitions import NdarrayOrTensor
 from monai.data.box_utils import BoxMode, StandardMode
 from monai.data.meta_obj import get_track_meta, set_track_meta
-from monai.data.meta_tensor import MetaTensor
+from monai.data.meta_tensor import MetaTensor, get_spatial_ndim
 from monai.data.utils import AFFINE_TOL, affine_to_spacing, compute_shape_offset, iter_patch, to_affine_nd, zoom_affine
 from monai.networks.layers import AffineTransform, GaussianFilter, grid_pull
 from monai.networks.utils import meshgrid_ij
@@ -848,12 +848,14 @@ def __call__(
         anti_aliasing = self.anti_aliasing if anti_aliasing is None else anti_aliasing
         anti_aliasing_sigma = self.anti_aliasing_sigma if anti_aliasing_sigma is None else anti_aliasing_sigma
 
-        input_ndim = img.ndim - 1  # spatial ndim
+        input_ndim = get_spatial_ndim(img)
         if self.size_mode == "all":
             output_ndim = len(ensure_tuple(self.spatial_size))
             if output_ndim > input_ndim:
                 input_shape = ensure_tuple_size(img.shape, output_ndim + 1, 1)
                 img = img.reshape(input_shape)
+                if isinstance(img, MetaTensor):
+                    img.spatial_ndim = output_ndim
             elif output_ndim < input_ndim:
                 raise ValueError(
                     "len(spatial_size) must be greater or equal to img spatial dimensions, "
@@ -1034,7 +1036,7 @@ def inverse_transform(self, data: torch.Tensor, transform) -> torch.Tensor:
         out = convert_to_dst_type(out, dst=data, dtype=out.dtype)[0]
         if isinstance(out, MetaTensor):
             affine = convert_to_tensor(out.peek_pending_affine(), track_meta=False)
-            mat = to_affine_nd(len(affine) - 1, transform_t)
+            mat = to_affine_nd(out.spatial_ndim, transform_t)
             out.affine @= convert_to_dst_type(mat, affine)[0]
         return out
 
@@ -1131,7 +1133,7 @@ def __call__(
                 during initialization for this call. Defaults to None.
         """
         img = convert_to_tensor(img, track_meta=get_track_meta())
-        _zoom = ensure_tuple_rep(self.zoom, img.ndim - 1)  # match the spatial image dim
+        _zoom = ensure_tuple_rep(self.zoom, get_spatial_ndim(img))
         _mode = self.mode if mode is None else mode
         _padding_mode = padding_mode or self.padding_mode
         _align_corners = self.align_corners if align_corners is None else align_corners
@@ -1519,7 +1521,7 @@ def randomize(self, data: NdarrayOrTensor) -> None:
         super().randomize(None)
         if not self._do_transform:
             return None
-        self._axis = self.R.randint(data.ndim - 1)
+        self._axis = self.R.randint(get_spatial_ndim(data))
 
     def __call__(self, img: torch.Tensor, randomize: bool = True, lazy: bool | None = None) -> torch.Tensor:
         """
@@ -1629,13 +1631,14 @@ def randomize(self, img: NdarrayOrTensor) -> None:
         super().randomize(None)
         if not self._do_transform:
             return None
+        _sp = get_spatial_ndim(img)
         self._zoom = [self.R.uniform(l, h) for l, h in zip(self.min_zoom, self.max_zoom)]
         if len(self._zoom) == 1:
             # to keep the spatial shape ratio, use same random zoom factor for all dims
-            self._zoom = ensure_tuple_rep(self._zoom[0], img.ndim - 1)
-        elif len(self._zoom) == 2 and img.ndim > 3:
+            self._zoom = ensure_tuple_rep(self._zoom[0], _sp)
+        elif len(self._zoom) == 2 and _sp > 2:
             # if 2 zoom factors provided for 3D data, use the first factor for H and W dims, second factor for D dim
-            self._zoom = ensure_tuple_rep(self._zoom[0], img.ndim - 2) + ensure_tuple(self._zoom[-1])
+            self._zoom = ensure_tuple_rep(self._zoom[0], _sp - 1) + ensure_tuple(self._zoom[-1])
 
     def __call__(
         self,
@@ -2350,7 +2353,7 @@ def inverse(self, data: torch.Tensor) -> torch.Tensor:
         out.meta = data.meta  # type: ignore
         affine = convert_data_type(out.peek_pending_affine(), torch.Tensor)[0]
         xform, *_ = convert_to_dst_type(
-            Affine.compute_w_affine(len(affine) - 1, inv_affine, data.shape[1:], orig_size), affine
+            Affine.compute_w_affine(out.spatial_ndim, inv_affine, data.shape[1:], orig_size), affine
         )
         out.affine @= xform
         return out
@@ -2619,7 +2622,7 @@ def inverse(self, data: torch.Tensor) -> torch.Tensor:
         out.meta = data.meta  # type: ignore
         affine = convert_data_type(out.peek_pending_affine(), torch.Tensor)[0]
         xform, *_ = convert_to_dst_type(
-            Affine.compute_w_affine(len(affine) - 1, inv_affine, data.shape[1:], orig_size), affine
+            Affine.compute_w_affine(out.spatial_ndim, inv_affine, data.shape[1:], orig_size), affine
         )
         out.affine @= xform
         return out
@@ -3032,7 +3035,7 @@ def __call__(
             raise ValueError("the spatial size of `img` does not match with the length of `distort_steps`")
 
         all_ranges = []
-        num_cells = ensure_tuple_rep(self.num_cells, len(img.shape) - 1)
+        num_cells = ensure_tuple_rep(self.num_cells, get_spatial_ndim(img))
         if isinstance(img, MetaTensor) and img.pending_operations:
             warnings.warn("MetaTensor img has pending operations, transform may return incorrect results.")
         for dim_idx, dim_size in enumerate(img.shape[1:]):