Fix inverse matching with mixed postprocessing history

eclipse0922 · eclipse0922 · commit c42261bf7592 · 2026-03-06T22:38:01.000+09:00
Signed-off-by: sewon.jeon &lt;irocks0922@gmail.com&gt;
diff --git a/monai/transforms/compose.py b/monai/transforms/compose.py
@@ -255,90 +255,6 @@ def __init__(
         self.set_random_state(seed=get_seed())
         self.overrides = overrides
 
-        # Automatically assign group ID to child transforms for inversion tracking
-        self._set_transform_groups()
-
-    def _set_transform_groups(self):
-        """
-        Automatically set group IDs on child transforms for inversion tracking.
-
-        This allows Invertd to identify which transforms belong to this
-        ``Compose`` instance, including wrapped transforms (for example,
-        array transforms inside dictionary transforms).
-
-        Args:
-            None.
-
-        Returns:
-            None.
-        """
-        from monai.transforms.inverse import TraceableTransform
-
-        group_id = str(id(self))
-        visited = set()  # Track visited objects to avoid infinite recursion
-
-        def set_group_recursive(obj, gid, allow_compose: bool = False):
-            """
-            Recursively set a group ID on a transform and its wrapped transforms.
-
-            Args:
-                obj: Transform instance to process.
-                gid: Group identifier to assign.
-                allow_compose: Whether to set group on ``Compose`` instances.
-                    ``Compose`` internals are not traversed to preserve nested
-                    pipeline boundaries.
-
-            Returns:
-                None.
-            """
-            if obj is None or isinstance(obj, (bool, int, float, str, bytes)):
-                return
-
-            # Avoid infinite recursion
-            obj_id = id(obj)
-            if obj_id in visited:
-                return
-            visited.add(obj_id)
-
-            if isinstance(obj, Compose):
-                if allow_compose:
-                    obj._group = gid
-                return
-
-            if isinstance(obj, TraceableTransform):
-                obj._group = gid
-
-            if isinstance(obj, Mapping):
-                for attr in obj.values():
-                    set_group_recursive(attr, gid)
-                return
-
-            if isinstance(obj, (list, tuple, set)):
-                for attr in obj:
-                    set_group_recursive(attr, gid)
-                return
-
-            attrs: list[Any] = []
-            if hasattr(obj, "__dict__"):
-                attrs.extend(vars(obj).values())
-
-            slots = getattr(type(obj), "__slots__", ())
-            if isinstance(slots, str):
-                slots = (slots,)
-            for slot in slots:
-                if slot.startswith("__"):
-                    continue
-                try:
-                    attrs.append(getattr(obj, slot))
-                except AttributeError:
-                    continue
-
-            for attr in attrs:
-                set_group_recursive(attr, gid)
-
-        for transform in self.transforms:
-            set_group_recursive(transform, group_id, allow_compose=True)
-
     @LazyTransform.lazy.setter  # type: ignore
     def lazy(self, val: bool):
         self._lazy = val
diff --git a/monai/transforms/inverse.py b/monai/transforms/inverse.py
@@ -82,10 +82,6 @@ def _init_trace_threadlocal(self):
         if not hasattr(self._tracing, "value"):
             self._tracing.value = MONAIEnvVars.trace_transform() != "0"
 
-        # Initialize group identifier (set by Compose for automatic group tracking)
-        if not hasattr(self, "_group"):
-            self._group: str | None = None
-
     def __getstate__(self):
         """When pickling, remove the `_tracing` member from the output, if present, since it's not picklable."""
         _dict = dict(getattr(self, "__dict__", {}))  # this makes __dict__ always present in the unpickled object
@@ -123,7 +119,6 @@ def get_transform_info(self) -> dict:
         """
         Return a dictionary with the relevant information pertaining to an applied transform.
         """
-        # Ensure _group is initialized
         self._init_trace_threadlocal()
 
         vals = (
@@ -132,13 +127,7 @@ def get_transform_info(self) -> dict:
             self.tracing,
             self._do_transform if hasattr(self, "_do_transform") else True,
         )
-        info = dict(zip(self.transform_info_keys(), vals))
-
-        # Add group if set (automatically set by Compose)
-        if self._group is not None:
-            info[TraceKeys.GROUP] = self._group
-
-        return info
+        return dict(zip(self.transform_info_keys(), vals))
 
     def push_transform(self, data, *args, **kwargs):
         """
@@ -314,24 +303,33 @@ def track_transform_meta(
 
     def check_transforms_match(self, transform: Mapping) -> None:
         """Check transforms are of same instance."""
-        xform_id = transform.get(TraceKeys.ID, "")
-        if xform_id == id(self):
-            return
-        # TraceKeys.NONE to skip the id check
-        if xform_id == TraceKeys.NONE:
+        if self._transforms_match(transform):
             return
+
+        xform_id = transform.get(TraceKeys.ID, "")
         xform_name = transform.get(TraceKeys.CLASS_NAME, "")
         warning_msg = transform.get(TraceKeys.EXTRA_INFO, {}).get("warn")
         if warning_msg:
             warnings.warn(warning_msg)
-        # basic check if multiprocessing uses 'spawn' (objects get recreated so don't have same ID)
-        if torch.multiprocessing.get_start_method() in ("spawn", None) and xform_name == self.__class__.__name__:
-            return
         raise RuntimeError(
             f"Error {self.__class__.__name__} getting the most recently "
             f"applied invertible transform {xform_name} {xform_id} != {id(self)}."
         )
 
+    def _transforms_match(self, transform: Mapping) -> bool:
+        """Return whether a traced transform entry matches this transform instance."""
+        xform_id = transform.get(TraceKeys.ID, "")
+        if xform_id == id(self):
+            return True
+        # TraceKeys.NONE to skip the id check
+        if xform_id == TraceKeys.NONE:
+            return True
+        xform_name = transform.get(TraceKeys.CLASS_NAME, "")
+        # basic check if multiprocessing uses 'spawn' (objects get recreated so don't have same ID)
+        if torch.multiprocessing.get_start_method() in ("spawn", None) and xform_name == self.__class__.__name__:
+            return True
+        return False
+
     def get_most_recent_transform(self, data, key: Hashable = None, check: bool = True, pop: bool = False):
         """
         Get most recent matching transform for the current class from the sequence of applied operations.
@@ -363,10 +361,16 @@ def get_most_recent_transform(self, data, key: Hashable = None, check: bool = Tr
         if not all_transforms:
             raise ValueError(f"Item of type {type(data)} (key: {key}, pop: {pop}) has empty 'applied_operations'")
 
+        match_idx = len(all_transforms) - 1
         if check:
-            self.check_transforms_match(all_transforms[-1])
+            for idx in range(len(all_transforms) - 1, -1, -1):
+                if self._transforms_match(all_transforms[idx]):
+                    match_idx = idx
+                    break
+            else:
+                self.check_transforms_match(all_transforms[-1])
 
-        return all_transforms.pop(-1) if pop else all_transforms[-1]
+        return all_transforms.pop(match_idx) if pop else all_transforms[match_idx]
 
     def pop_transform(self, data, key: Hashable = None, check: bool = True):
         """
diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py
@@ -48,7 +48,7 @@
 from monai.transforms.transform import MapTransform
 from monai.transforms.utility.array import ToTensor
 from monai.transforms.utils import allow_missing_keys_mode, convert_applied_interp_mode
-from monai.utils import PostFix, TraceKeys, convert_to_tensor, ensure_tuple, ensure_tuple_rep
+from monai.utils import PostFix, convert_to_tensor, ensure_tuple, ensure_tuple_rep
 from monai.utils.type_conversion import convert_to_dst_type
 
 __all__ = [
@@ -859,27 +859,6 @@ def __init__(
         self.post_func = ensure_tuple_rep(post_func, len(self.keys))
         self._totensor = ToTensor()
 
-    def _filter_transforms_by_group(self, all_transforms: list[dict]) -> list[dict]:
-        """Filter applied operations to only include transforms from the target pipeline.
-
-        Uses automatic group tracking where ``Compose`` assigns its ID to child transforms.
-
-        Args:
-            all_transforms: Full list of applied transform metadata dictionaries.
-
-        Returns:
-            Subset whose ``TraceKeys.GROUP`` matches ``str(id(self.transform))``, or the original
-            list when no match is found for backward compatibility.
-        """
-        # Get the group ID of the transform (Compose instance)
-        target_group = str(id(self.transform))
-
-        # Filter transforms that match the target group
-        filtered = [xform for xform in all_transforms if xform.get(TraceKeys.GROUP) == target_group]
-
-        # If no transforms match (backward compatibility), return all transforms
-        return filtered if filtered else all_transforms
-
     def __call__(self, data: Mapping[Hashable, Any]) -> dict[Hashable, Any]:
         d = dict(data)
         for (
@@ -915,13 +894,10 @@ def __call__(self, data: Mapping[Hashable, Any]) -> dict[Hashable, Any]:
 
             orig_meta_key = orig_meta_key or f"{orig_key}_{meta_key_postfix}"
             if orig_key in d and isinstance(d[orig_key], MetaTensor):
-                all_transforms = d[orig_key].applied_operations
+                transform_info = d[orig_key].applied_operations
                 meta_info = d[orig_key].meta
-
-                # Automatically filter by Compose instance group ID
-                transform_info = self._filter_transforms_by_group(all_transforms)
             else:
-                transform_info = self._filter_transforms_by_group(d[InvertibleTransform.trace_key(orig_key)])
+                transform_info = d[InvertibleTransform.trace_key(orig_key)]
                 meta_info = d.get(orig_meta_key, {})
             if nearest_interp:
                 transform_info = convert_applied_interp_mode(
diff --git a/monai/utils/enums.py b/monai/utils/enums.py
@@ -334,7 +334,6 @@ class TraceKeys(StrEnum):
     TRACING: str = "tracing"
     STATUSES: str = "statuses"
     LAZY: str = "lazy"
-    GROUP: str = "group"
 
 
 class TraceStatusKeys(StrEnum):
diff --git a/tests/transforms/compose/test_compose.py b/tests/transforms/compose/test_compose.py
@@ -268,43 +268,6 @@ def test_data_loader_2(self):
             self.assertAlmostEqual(out_1.cpu().item(), 0.28602141572)
         set_determinism(None)
 
-    def test_set_transform_groups_on_wrapped_transform_attributes(self):
-        class _IdentityInvertible(mt.InvertibleTransform):
-            def __call__(self, data):
-                return data
-
-            def inverse(self, data):
-                return data
-
-        class _WrapperWithTransform:
-            def __init__(self):
-                self.transform = _IdentityInvertible()
-
-            def __call__(self, data):
-                return self.transform(data)
-
-        class _WrapperWithTransforms:
-            def __init__(self):
-                self.transforms = [_IdentityInvertible(), {"inner": _IdentityInvertible()}]
-
-            def __call__(self, data):
-                for transform in self.transforms:
-                    if isinstance(transform, dict):
-                        for nested_transform in transform.values():
-                            data = nested_transform(data)
-                    else:
-                        data = transform(data)
-                return data
-
-        wrapped_transform = _WrapperWithTransform()
-        wrapped_transforms = _WrapperWithTransforms()
-        composed = mt.Compose([wrapped_transform, wrapped_transforms])
-        expected_group = str(id(composed))
-
-        self.assertEqual(getattr(wrapped_transform.transform, "_group", None), expected_group)
-        self.assertEqual(getattr(wrapped_transforms.transforms[0], "_group", None), expected_group)
-        self.assertEqual(getattr(wrapped_transforms.transforms[1]["inner"], "_group", None), expected_group)
-
     def test_flatten_and_len(self):
         x = mt.EnsureChannelFirst(channel_dim="no_channel")
         t1 = mt.Compose([x, x, x, x, mt.Compose([mt.Compose([x, x]), x, x])])
diff --git a/tests/transforms/inverse/test_invertd.py b/tests/transforms/inverse/test_invertd.py
