fix(lazy): keep the image's affine rank on the lazy crop/pad path

monai/transforms/lazy/utils.py

@@ -193,6 +193,11 @@ def resample(data: torch.Tensor, matrix: NdarrayOrTensor, kwargs: dict | None =
     ndim = len(matrix) - 1
     img = convert_to_tensor(data=data, track_meta=monai.data.get_track_meta())
     init_affine = monai.data.to_affine_nd(ndim, img.affine)
+    # The crop/pad shortcut below writes the affine directly, so compose at the image's own affine rank: a 2D image
+    # with a 4x4 affine must keep its out-of-plane spacing and origin, as the eager transforms do.
+    full_rank = len(img.affine) - 1
+    full_xform = monai.data.to_affine_nd(full_rank, matrix)
+    full_dst_affine = img.affine @ monai.utils.convert_to_dst_type(full_xform, img.affine)[0]
     spatial_size = kwargs.get(LazyAttr.SHAPE, None)
     out_spatial_size = img.peek_pending_shape() if spatial_size is None else spatial_size
     out_spatial_size = convert_to_numpy(out_spatial_size, wrap_sequence=True)
@@ -225,13 +230,13 @@ def resample(data: torch.Tensor, matrix: NdarrayOrTensor, kwargs: dict | None =
             and len(in_shape) == len(out_spatial_size)
             and allclose(convert_to_numpy(in_shape, wrap_sequence=True), out_spatial_size)
         ):
-            img.affine = call_kwargs["dst_affine"]
+            img.affine = full_dst_affine
             img = img.to(torch.float32)  # consistent with monai.transforms.spatial.functional.spatial_resample
             return img
         # pyrefly: ignore [bad-argument-type, implicit-import]
         img = monai.transforms.crop_or_pad_nd(img, matrix_np, out_spatial_size, mode=call_kwargs["padding_mode"])
         img = img.to(torch.float32)  # consistent with monai.transforms.spatial.functional.spatial_resample
-        img.affine = call_kwargs["dst_affine"]
+        img.affine = full_dst_affine
         return img
 
     # pyrefly: ignore [bad-argument-type, implicit-import]

tests/transforms/functional/test_resample.py

@@ -16,6 +16,8 @@
 import torch
 from parameterized import parameterized
 
+from monai.data import MetaTensor
+from monai.transforms import CenterSpatialCropd, Compose, Flipd, SpatialPadd
 from monai.transforms.lazy.functional import resample
 from monai.utils import convert_to_tensor
 from tests.test_utils import assert_allclose, get_arange_img
@@ -45,6 +47,26 @@ def test_resample_function_impl(self, img, matrix, expected):
         out_1 = resample(img, matrix, {"lazy_resample_mode": "other value", "lazy_dtype": torch.float})
         self.assertIs(out.dtype, out_1.dtype)  # testing dtype in different lazy_resample_mode
 
+    @parameterized.expand(
+        [
+            ["flip", lambda: Flipd("img", spatial_axis=0)],
+            ["crop", lambda: CenterSpatialCropd("img", (5, 4))],
+            ["pad", lambda: SpatialPadd("img", (12, 12))],
+        ]
+    )
+    def test_lazy_crop_pad_keeps_affine_rank(self, _, make_transform):
+        """A 2D image with a 4x4 affine keeps its out-of-plane spacing and origin on the lazy crop/pad path."""
+        affine = torch.diag(torch.tensor([1.5, 0.8, 2.0, 1.0], dtype=torch.float64))
+        affine[:3, 3] = torch.tensor([3.0, -2.0, 1.0], dtype=torch.float64)
+        img = MetaTensor(torch.rand(1, 10, 8), affine=affine)
+
+        eager = Compose([make_transform()], lazy=False)({"img": img.clone()})["img"]
+        lazy = Compose([make_transform()], lazy=True)({"img": img.clone()})["img"]
+
+        self.assertEqual(lazy.affine.shape, (4, 4))
+        assert_allclose(lazy.affine, eager.affine, type_test=False)
+        assert_allclose(lazy, eager, type_test=False)
+
 
 if __name__ == "__main__":
     unittest.main()