5991 add utility enhancements for lazy resampling (#6017)

Part of #5991 .

### Description

This PR adds the minor utility enhancements from #5860 relate to lazy
resampling, and it belongs to the first part mentioned in
#5991

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: Yiheng Wang <vennw@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: yiheng-wang-nv

monai/data/meta_tensor.py

@@ -25,7 +25,7 @@
 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_numpy, convert_to_tensor
+from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_numpy, convert_to_tensor
 
 __all__ = ["MetaTensor"]
 
@@ -461,7 +461,7 @@ 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"))
+        self.meta[MetaKeys.AFFINE] = torch.as_tensor(d, device=torch.device("cpu"), dtype=torch.double)
 
     @property
     def pixdim(self):
@@ -479,10 +479,18 @@ def peek_pending_shape(self):
         return tuple(convert_to_numpy(self.shape, wrap_sequence=True).tolist()[1:]) if res is None else res
 
     def peek_pending_affine(self):
-        res = None
-        if self.pending_operations:
-            res = self.pending_operations[-1].get(LazyAttr.AFFINE, None)
-        return self.affine if res is None else res
+        res = self.affine
+        for p in self.pending_operations:
+            next_matrix = convert_to_tensor(p.get(LazyAttr.AFFINE))
+            if next_matrix is None:
+                continue
+            res = convert_to_dst_type(res, next_matrix)[0]
+            res = monai.transforms.lazy.utils.combine_transforms(res, next_matrix)
+        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))
 
     def new_empty(self, size, dtype=None, device=None, requires_grad=False):
         """

monai/transforms/lazy/utils.py

@@ -38,7 +38,7 @@ def is_affine_shaped(data):
             return False
         if not hasattr(data, "shape") or len(data.shape) < 2:
             return False
-        return data.shape[-1] in (3, 4) and data.shape[-2] in (3, 4) and data.shape[-1] == data.shape[-2]
+        return data.shape[-1] in (3, 4) and data.shape[-1] == data.shape[-2]
 
 
 class DisplacementField:
@@ -129,6 +129,6 @@ def resample(data: torch.Tensor, matrix: NdarrayOrTensor, spatial_size, kwargs:
         "padding_mode": kwargs.pop(LazyAttr.PADDING_MODE, None),
     }
     resampler = monai.transforms.SpatialResample(**init_kwargs)
-    # resampler.lazy_evaluation = False
-    with resampler.trace_transform(False):  # don't track this transform in `data`
+    # resampler.lazy_evaluation = False  # resampler is a lazytransform
+    with resampler.trace_transform(False):  # don't track this transform in `img`
         return resampler(img=img, **call_kwargs)

monai/transforms/transform.py

@@ -254,18 +254,17 @@ class LazyTransform(Transform, LazyTrait):
     dictionary transforms to simplify implementation of new lazy transforms.
     """
 
-    def __init__(self, lazy_evaluation: bool | None = True):
-        self.lazy_evaluation = lazy_evaluation
+    _lazy_evaluation: bool = False
 
     @property
     def lazy_evaluation(self):
-        return self.lazy_evaluation
+        return self._lazy_evaluation
 
     @lazy_evaluation.setter
     def lazy_evaluation(self, lazy_evaluation: bool):
         if not isinstance(lazy_evaluation, bool):
             raise TypeError(f"lazy_evaluation must be a bool but is of type {type(lazy_evaluation)}")
-        self.lazy_evaluation = lazy_evaluation
+        self._lazy_evaluation = lazy_evaluation
 
 
 class RandomizableTransform(Randomizable, Transform):

monai/transforms/utils.py

@@ -873,6 +873,7 @@ def create_translate(
         backend: APIs to use, ``numpy`` or ``torch``.
     """
     _backend = look_up_option(backend, TransformBackends)
+    spatial_dims = int(spatial_dims)
     if _backend == TransformBackends.NUMPY:
         return _create_translate(spatial_dims=spatial_dims, shift=shift, eye_func=np.eye, array_func=np.asarray)
     if _backend == TransformBackends.TORCH:

tests/test_meta_tensor.py

@@ -515,9 +515,11 @@ def test_pending_ops(self):
         self.assertEqual(m.pending_operations, [])
         self.assertEqual(m.peek_pending_shape(), (10, 8))
         self.assertIsInstance(m.peek_pending_affine(), torch.Tensor)
+        self.assertTrue(m.peek_pending_rank() >= 1)
         m.push_pending_operation({})
         self.assertEqual(m.peek_pending_shape(), (10, 8))
         self.assertIsInstance(m.peek_pending_affine(), torch.Tensor)
+        self.assertTrue(m.peek_pending_rank() >= 1)
 
     @parameterized.expand(TESTS)
     def test_multiprocessing(self, device=None, dtype=None):