From 34a1a860c4f73582267c1575987af4a61e1ce0b8 Mon Sep 17 00:00:00 2001 From: Vikas Reddy Date: Mon, 2 Nov 2026 13:06:26 +0000 Subject: [PATCH 1/4] Fix compiled lazy resampling with align_corners=True Signed-off-by: Vikas Reddy --- monai/transforms/spatial/functional.py | 11 +++++++++++ tests/transforms/test_affine.py | 26 ++++++++++++++++++++++++-- tests/transforms/test_spacing.py | 2 +- 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/monai/transforms/spatial/functional.py b/monai/transforms/spatial/functional.py index 967020a3bf8..08d7a5cfc39 100644 --- a/monai/transforms/spatial/functional.py +++ b/monai/transforms/spatial/functional.py @@ -216,6 +216,17 @@ 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): + # Resample/AffineGrid re-apply the align_corners=True scalings; cancel them so `xform` + # (already in voxel space) reaches grid_pull unchanged. + _eye = torch.eye(spatial_rank + 1, dtype=xform.dtype, device=xform.device) + _s_inv, _d_inv, _tm = _eye.clone(), _eye.clone(), _eye.clone() + for _i in range(spatial_rank): + _din, _dout = max(2, int(in_sp_size[_i])), max(2, int(spatial_size[_i])) + _s_inv[_i, _i] = _din / (_din - 1) + _d_inv[_i, _i] = (_dout - 1) / _dout + _tm[_i, -1] = -0.5 + xform = _tm @ _s_inv @ xform @ _d_inv affine_xform = monai.transforms.Affine( affine=xform, spatial_size=spatial_size, diff --git a/tests/transforms/test_affine.py b/tests/transforms/test_affine.py index 8a29e31de87..cadce2383d6 100644 --- a/tests/transforms/test_affine.py +++ b/tests/transforms/test_affine.py @@ -18,11 +18,12 @@ import torch from parameterized import parameterized +from monai.config import USE_COMPILED 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 +292,27 @@ 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) + @unittest.skipUnless(USE_COMPILED, "requires the compiled extension (BUILD_MONAI=1)") + 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() diff --git a/tests/transforms/test_spacing.py b/tests/transforms/test_spacing.py index b8f838c4ac4..40da7ebd1e6 100644 --- a/tests/transforms/test_spacing.py +++ b/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]]]]) From 03ec6ed7174c372e37b090e8ad4967a397a265b9 Mon Sep 17 00:00:00 2001 From: Vikas Reddy Date: Mon, 2 Nov 2026 13:40:43 +0000 Subject: [PATCH 2/4] Fix lazy resampling backend mode handling Signed-off-by: Vikas Reddy --- monai/transforms/spatial/functional.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monai/transforms/spatial/functional.py b/monai/transforms/spatial/functional.py index 08d7a5cfc39..c8323eb03eb 100644 --- a/monai/transforms/spatial/functional.py +++ b/monai/transforms/spatial/functional.py @@ -45,6 +45,7 @@ fall_back_tuple, optional_import, ) +from monai.utils.enums import TransformBackends nib, has_nib = optional_import("nibabel") cupy, _ = optional_import("cupy") @@ -213,10 +214,11 @@ def spatial_resample( img = img.reshape(xform_shape) img = img.to(dtype_pt) _use_compiled = USE_COMPILED and not _compiled_unsupported(img.device) + backend, interp_mode, _, _ = resolves_modes(mode, padding_mode, backend=None, use_compiled=_use_compiled) 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): + if align_corners and backend == TransformBackends.NUMPY and interp_mode == 1: # Resample/AffineGrid re-apply the align_corners=True scalings; cancel them so `xform` # (already in voxel space) reaches grid_pull unchanged. _eye = torch.eye(spatial_rank + 1, dtype=xform.dtype, device=xform.device) From 921be4a85cd776145f65c0e3e94fed83f21094e2 Mon Sep 17 00:00:00 2001 From: Vikas Reddy Date: Tue, 3 Nov 2026 05:16:09 +0000 Subject: [PATCH 3/4] Revert "Fix lazy resampling backend mode handling" This reverts commit 03ec6ed7174c372e37b090e8ad4967a397a265b9. Signed-off-by: Vikas Reddy --- monai/transforms/spatial/functional.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/monai/transforms/spatial/functional.py b/monai/transforms/spatial/functional.py index c8323eb03eb..08d7a5cfc39 100644 --- a/monai/transforms/spatial/functional.py +++ b/monai/transforms/spatial/functional.py @@ -45,7 +45,6 @@ fall_back_tuple, optional_import, ) -from monai.utils.enums import TransformBackends nib, has_nib = optional_import("nibabel") cupy, _ = optional_import("cupy") @@ -214,11 +213,10 @@ def spatial_resample( img = img.reshape(xform_shape) img = img.to(dtype_pt) _use_compiled = USE_COMPILED and not _compiled_unsupported(img.device) - backend, interp_mode, _, _ = resolves_modes(mode, padding_mode, backend=None, use_compiled=_use_compiled) 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 backend == TransformBackends.NUMPY and interp_mode == 1: + if align_corners and not isinstance(mode, int): # Resample/AffineGrid re-apply the align_corners=True scalings; cancel them so `xform` # (already in voxel space) reaches grid_pull unchanged. _eye = torch.eye(spatial_rank + 1, dtype=xform.dtype, device=xform.device) From 46b69ea7e7c0925e6285ca89aba78044d27074e2 Mon Sep 17 00:00:00 2001 From: Vikas Reddy Date: Wed, 4 Nov 2026 13:49:27 +0000 Subject: [PATCH 4/4] Fix compiled lazy resampling with align_corners=True Signed-off-by: Vikas Reddy --- monai/transforms/spatial/functional.py | 47 ++++++++++++++++++++------ tests/transforms/test_affine.py | 2 -- 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/monai/transforms/spatial/functional.py b/monai/transforms/spatial/functional.py index 08d7a5cfc39..add5bda8bd9 100644 --- a/monai/transforms/spatial/functional.py +++ b/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: @@ -217,16 +253,7 @@ def spatial_resample( 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): - # Resample/AffineGrid re-apply the align_corners=True scalings; cancel them so `xform` - # (already in voxel space) reaches grid_pull unchanged. - _eye = torch.eye(spatial_rank + 1, dtype=xform.dtype, device=xform.device) - _s_inv, _d_inv, _tm = _eye.clone(), _eye.clone(), _eye.clone() - for _i in range(spatial_rank): - _din, _dout = max(2, int(in_sp_size[_i])), max(2, int(spatial_size[_i])) - _s_inv[_i, _i] = _din / (_din - 1) - _d_inv[_i, _i] = (_dout - 1) / _dout - _tm[_i, -1] = -0.5 - xform = _tm @ _s_inv @ xform @ _d_inv + xform = _cancel_align_corners_scaling(xform, in_sp_size, spatial_size) affine_xform = monai.transforms.Affine( affine=xform, spatial_size=spatial_size, diff --git a/tests/transforms/test_affine.py b/tests/transforms/test_affine.py index cadce2383d6..ec99c65dfdd 100644 --- a/tests/transforms/test_affine.py +++ b/tests/transforms/test_affine.py @@ -18,7 +18,6 @@ import torch from parameterized import parameterized -from monai.config import USE_COMPILED from monai.data import MetaTensor, set_track_meta from monai.transforms import Affine, RandRotate, Resize from monai.transforms.lazy.functional import apply_pending @@ -292,7 +291,6 @@ 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) - @unittest.skipUnless(USE_COMPILED, "requires the compiled extension (BUILD_MONAI=1)") def test_lazy_rotate_align_corners_compiled(self): """lazy resampling with the compiled backend and align_corners=True must match eager.""" init = dict(