fix: address CodeRabbit review comments on align_corners changes

- scale_affine: handle dst dim == 1 special case (avoid spurious non-zero
  scale via max(n-1, 1)) and document that `centered` is ignored when
  align_corners=True.
- kwargs_from_pending: docstring now notes align_corners extraction from
  extra_info.
- test_affine_transform: use detach().cpu().numpy() consistently and pin
  identity_affine dtype.
- test_affine: switch the known method_0/align_corners=True mismatch to
  subTest + SkipTest so it stays visible in CI; deepcopy lazy params.

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

Author: aymuos15

monai/transforms/lazy/utils.py

@@ -90,7 +90,11 @@ def affine_from_pending(pending_item):
 
 
 def kwargs_from_pending(pending_item):
-    """Extract kwargs from a pending transform item."""
+    """Extract kwargs from a pending transform item.
+
+    When ``pending_item`` is a dict, ``align_corners`` is also extracted from its ``extra_info`` entry
+    (if present and boolean) so the lazy pipeline preserves the original transform's alignment.
+    """
     if not isinstance(pending_item, dict):
         return {}
     ret = {

monai/transforms/utils.py

@@ -2105,6 +2105,7 @@ def scale_affine(spatial_size, new_spatial_size, centered: bool = True, align_co
         spatial_size: original spatial size.
         new_spatial_size: new spatial size.
         centered: whether the scaling is with respect to the image center (True, default) or corner (False).
+            Ignored when ``align_corners=True``, since corner-aligned scaling is inherently centered.
         align_corners: if True, use (size-1) based scaling to match torch.nn.functional.interpolate behavior.
 
     Returns:
@@ -2115,9 +2116,10 @@ def scale_affine(spatial_size, new_spatial_size, centered: bool = True, align_co
     if spatial_size == new_spatial_size:
         return np.eye(r + 1)
     if align_corners:
-        # Match interpolate behavior: (src-1)/(dst-1)
+        # Match interpolate behavior: (src-1)/(dst-1); when dst == 1 the scale collapses to 0
         s = np.array(
-            [(float(o) - 1) / max(float(n) - 1, 1) for o, n in zip(spatial_size, new_spatial_size)], dtype=float
+            [0.0 if float(n) == 1 else (float(o) - 1) / (float(n) - 1) for o, n in zip(spatial_size, new_spatial_size)],
+            dtype=float,
         )
     else:
         # Standard scaling: src/dst

tests/networks/layers/test_affine_transform.py

@@ -392,17 +392,19 @@ def test_align_corners_consistency(self):
         image = torch.arange(1.0, 13.0).view(1, 1, 3, 4)
 
         # Identity affine in pixel space (i, j, k convention with reverse_indexing=True)
-        identity_affine = torch.eye(3).unsqueeze(0)
+        identity_affine = torch.eye(3, dtype=torch.float32).unsqueeze(0)
 
         # Test with align_corners=True (the default)
         xform_true = AffineTransform(align_corners=True)
         out_true = xform_true(image, identity_affine)
-        np.testing.assert_allclose(out_true.numpy(), image.numpy(), atol=1e-5, rtol=_rtol)
+        np.testing.assert_allclose(out_true.detach().cpu().numpy(), image.detach().cpu().numpy(), atol=1e-5, rtol=_rtol)
 
         # Test with align_corners=False
         xform_false = AffineTransform(align_corners=False)
         out_false = xform_false(image, identity_affine)
-        np.testing.assert_allclose(out_false.numpy(), image.numpy(), atol=1e-5, rtol=_rtol)
+        np.testing.assert_allclose(
+            out_false.detach().cpu().numpy(), image.detach().cpu().numpy(), atol=1e-5, rtol=_rtol
+        )
 
     def test_align_corners_true_translation(self):
         """
@@ -425,7 +427,7 @@ def test_align_corners_true_translation(self):
 
         # Expected: shift columns left by 1, rightmost column becomes 0
         expected = torch.tensor([[[[2, 3, 4, 0], [6, 7, 8, 0], [10, 11, 12, 0], [14, 15, 16, 0]]]], dtype=torch.float32)
-        np.testing.assert_allclose(out.numpy(), expected.numpy(), atol=1e-4, rtol=_rtol)
+        np.testing.assert_allclose(out.detach().cpu().numpy(), expected.detach().cpu().numpy(), atol=1e-4, rtol=_rtol)
 
 
 if __name__ == "__main__":

tests/transforms/test_affine.py

@@ -189,7 +189,7 @@ def test_affine(self, input_param, input_data, expected_val):
         set_track_meta(True)
 
         # test lazy
-        lazy_input_param = input_param.copy()
+        lazy_input_param = deepcopy(input_param)
         for align_corners in [True, False]:
             lazy_input_param["align_corners"] = align_corners
             resampler = Affine(**lazy_input_param)
@@ -238,13 +238,16 @@ def method_3(im, ac):
 
         for call in (method_0, method_1, method_2, method_3):
             for ac in (False, True):
-                # Skip method_0 with align_corners=True due to known issue with lazy pipeline
-                # padding_mode override when using align_corners=True in optimized path
-                if call == method_0 and ac:
-                    continue
-                out = call(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)
+                with self.subTest(method=call.__name__, align_corners=ac):
+                    if call is method_0 and ac:
+                        # Known issue: lazy pipeline padding_mode override mismatches
+                        # when using align_corners=True in the optimized path.
+                        raise unittest.SkipTest(
+                            "method_0 with align_corners=True is a known mismatch in the lazy pipeline."
+                        )
+                    out = call(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)
 
 
 if __name__ == "__main__":