Fix compiled lazy resampling with align_corners=True

monai/transforms/spatial/functional.py

@@ -16,6 +16,7 @@
 
 import math
 import warnings
+from collections.abc import Sequence
 from enum import Enum
 
 import numpy as np
@@ -106,6 +107,41 @@ def _maybe_new_metatensor(img, dtype=None, device=None):
     )
 
 
+def _cancel_align_corners_scaling(
+    xform: torch.Tensor, in_spatial_size: Sequence[int], out_spatial_size: Sequence[int]
+) -> torch.Tensor:
+    """
+    Pre-compensate a voxel-space affine for the scalings applied again downstream when
+    ``align_corners=True`` and the compiled ``Resample`` path is used.
+
+    ``xform`` is already expressed in voxel space. However, ``Affine`` passes it through ``AffineGrid``,
+    which applies a ``d_out / (d_out - 1)`` scaling per spatial dimension, and the compiled ``Resample``
+    applies a ``(d_in - 1) / d_in`` scaling before ``grid_pull``. This function returns
+    ``T @ S_inv @ xform @ D_inv``, where ``S_inv`` and ``D_inv`` cancel those two scalings and ``T`` is a
+    translation by -0.5 voxel that re-centres the result, so that ``grid_pull`` receives the original
+    ``xform`` unchanged.
+
+    Dimensions are clamped to a minimum of 2 to avoid division by zero for singleton axes.
+
+    Args:
+        xform: homogeneous affine of shape ``(rank + 1, rank + 1)`` in voxel space.
+        in_spatial_size: spatial shape of the input image.
+        out_spatial_size: spatial shape of the resampled output.
+
+    Returns:
+        the compensated homogeneous affine, same shape, dtype and device as ``xform``.
+    """
+    rank = len(xform) - 1
+    eye = torch.eye(rank + 1, dtype=xform.dtype, device=xform.device)
+    s_inv, d_inv, translate = eye.clone(), eye.clone(), eye.clone()
+    for i in range(rank):
+        d_in, d_out = max(2, int(in_spatial_size[i])), max(2, int(out_spatial_size[i]))
+        s_inv[i, i] = d_in / (d_in - 1)
+        d_inv[i, i] = (d_out - 1) / d_out
+        translate[i, -1] = -0.5
+    return translate @ s_inv @ xform @ d_inv
+
+
 def spatial_resample(
     img, dst_affine, spatial_size, mode, padding_mode, align_corners, dtype_pt, lazy, transform_info
 ) -> torch.Tensor:
@@ -216,6 +252,8 @@ def spatial_resample(
     if isinstance(mode, int) or _use_compiled:
         dst_xform = create_translate(spatial_rank, [float(d - 1) / 2 for d in spatial_size])
         xform = xform @ convert_to_dst_type(dst_xform, xform)[0]
+        if align_corners and not isinstance(mode, int):
+            xform = _cancel_align_corners_scaling(xform, in_sp_size, spatial_size)
         affine_xform = monai.transforms.Affine(
             affine=xform,
             spatial_size=spatial_size,

tests/transforms/test_affine.py

@@ -19,10 +19,10 @@
 from parameterized import parameterized
 
 from monai.data import MetaTensor, set_track_meta
-from monai.transforms import Affine, Resize
+from monai.transforms import Affine, RandRotate, Resize
 from monai.transforms.lazy.functional import apply_pending
 from monai.utils import optional_import
-from tests.lazy_transforms_utils import test_resampler_lazy
+from tests.lazy_transforms_utils import get_apply_param, test_resampler_lazy
 from tests.test_utils import TEST_NDARRAYS_ALL, assert_allclose, test_local_inversion
 
 TESTS = []
@@ -291,6 +291,26 @@ def method_3(im, ac):
                     ref = Resize(align_corners=ac, spatial_size=(sp_size, sp_size), mode="bilinear")(im)
                     assert_allclose(out, ref, rtol=1e-4, atol=1e-4, type_test=False)
 
+    def test_lazy_rotate_align_corners_compiled(self):
+        """lazy resampling with the compiled backend and align_corners=True must match eager."""
+        init = dict(
+            range_x=np.pi,
+            prob=1.0,
+            keep_size=False,
+            mode="nearest",
+            padding_mode="zeros",
+            dtype=np.float64,
+            align_corners=True,
+        )
+        img = MetaTensor(torch.arange(64 * 48, dtype=torch.float32).reshape(1, 64, 48))
+        xform = RandRotate(**init)
+        xform.set_random_state(seed=243)
+        eager = xform(img)
+        xform.set_random_state(seed=243)
+        xform.lazy = True
+        lazy = apply_pending(xform(img), overrides=get_apply_param(init, {"img": img}))[0]
+        assert_allclose(lazy, eager, rtol=1e-3, atol=1e-3)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/transforms/test_spacing.py

@@ -27,7 +27,7 @@
 from tests.lazy_transforms_utils import test_resampler_lazy
 from tests.test_utils import TEST_DEVICES, TEST_NDARRAYS_ALL, assert_allclose, dict_product, skip_if_quick
 
-_TEMPLATE_5_COMPILED = torch.tensor([[[[0.75, 0.75]], [[0.75, 0.75]], [[0.75, 0.75]]]])
+_TEMPLATE_5_COMPILED = torch.tensor([[[[1.0, 1.0]], [[1.0, 1.0]], [[1.0, 1.0]]]])
 _TEMPLATE_5_NATIVE = torch.tensor([[[[0.95527864, 0.95527864]], [[1.0, 1.0]], [[1.0, 1.0]]]])