Repository: Project-MONAI/MONAI
File: monai/transforms/inverse.py (Python)

Code before:
            lazy: whether to push the transform to pending_operations or applied_operations.

        Returns:

            For backward compatibility, if ``data`` is a dictionary, it returns the dictionary with
            updated ``data[key]``. Otherwise, this function returns a MetaObj with updated transform metadata.
        """
        data_t = data[key] if key is not None else data  # compatible with the dict data representation
        out_obj = MetaObj()
        # after deprecating metadict, we should always convert data_t to metatensor here
        if isinstance(data_t, MetaTensor):
            out_obj.copy_meta_from(data_t, keys=out_obj.__dict__.keys())

        if lazy and (not get_track_meta()):
            warnings.warn("metadata is not tracked, please call 'set_track_meta(True)' if doing lazy evaluation.")

        if not lazy and affine is not None and isinstance(data_t, MetaTensor):
            # not lazy evaluation, directly update the metatensor affine (don't push to the stack)
            orig_affine = data_t.peek_pending_affine()
            orig_affine = convert_to_dst_type(orig_affine, affine, dtype=torch.float64)[0]
            try:
                affine = orig_affine @ to_affine_nd(len(orig_affine) - 1, affine, dtype=torch.float64)
            except RuntimeError as e:
                raise RuntimeError(
                    "mismatch affine matrix, ensured that the batch dimension is not included in the calculation."
                ) from e
            out_obj.meta[MetaKeys.AFFINE] = convert_to_tensor(affine, device=torch.device("cpu"), dtype=torch.float64)

        if not (get_track_meta() and transform_info and transform_info.get(TraceKeys.TRACING)):
            if isinstance(data, Mapping):
                if not isinstance(data, dict):
                    data = dict(data)
                data[key] = data_t.copy_meta_from(out_obj) if isinstance(data_t, MetaTensor) else data_t
                return data
            return out_obj  # return with data_t as tensor if get_track_meta() is False

        info = transform_info.copy()
        # track the current spatial shape
        if orig_size is not None:
            info[TraceKeys.ORIG_SIZE] = orig_size
        elif isinstance(data_t, MetaTensor):
            info[TraceKeys.ORIG_SIZE] = data_t.peek_pending_shape()
        elif hasattr(data_t, "shape"):
            info[TraceKeys.ORIG_SIZE] = data_t.shape[1:]

        # add lazy status to the transform info
        info[TraceKeys.LAZY] = lazy

        # include extra_info
        if extra_info is not None:
            extra_info.pop(LazyAttr.SHAPE, None)

Review comment (performance):
```suggestion
            except RuntimeError as e:
                if orig_affine.ndim > 2:
                    if data_t.is_batch:
                        msg = "Transform applied to batched tensor, should be applied to instances only"
                    else:
                        msg = "Mismatch affine matrix, ensured that the batch dimension is not included in the calculation."
                    raise RuntimeError(msg) from e
                else:
                    raise
```

We should only say something about this if we know the affine matrix has an odd dimensionality (and maybe should check of it's square or not). We can check a MetaTensor to see if it's a batched tensor this way and raise the exception like this in this case. For vanilla tensors the error message may be a bit confusing but this is still better. I haven't tried this code works so please double check if you want to use it.

Code after:
            lazy: whether to push the transform to pending_operations or applied_operations.

        Returns:

            For backward compatibility, if ``data`` is a dictionary, it returns the dictionary with
            updated ``data[key]``. Otherwise, this function returns a MetaObj with updated transform metadata.
        """
        data_t = data[key] if key is not None else data  # compatible with the dict data representation
        out_obj = MetaObj()
        # after deprecating metadict, we should always convert data_t to metatensor here
        if isinstance(data_t, MetaTensor):
            out_obj.copy_meta_from(data_t, keys=out_obj.__dict__.keys())

        if lazy and (not get_track_meta()):
            warnings.warn("metadata is not tracked, please call 'set_track_meta(True)' if doing lazy evaluation.")

        if not lazy and affine is not None and isinstance(data_t, MetaTensor):
            # not lazy evaluation, directly update the metatensor affine (don't push to the stack)
            orig_affine = data_t.peek_pending_affine()
            orig_affine = convert_to_dst_type(orig_affine, affine, dtype=torch.float64)[0]
            try:
                affine = orig_affine @ to_affine_nd(len(orig_affine) - 1, affine, dtype=torch.float64)
            except RuntimeError as e:
                if orig_affine.ndim > 2:
                    if data_t.is_batch:
                        msg = "Transform applied to batched tensor, should be applied to instances only"
                    else:
                        msg = "Mismatch affine matrix, ensured that the batch dimension is not included in the calculation."
                    raise RuntimeError(msg) from e
                else:
                    raise
            out_obj.meta[MetaKeys.AFFINE] = convert_to_tensor(affine, device=torch.device("cpu"), dtype=torch.float64)

        if not (get_track_meta() and transform_info and transform_info.get(TraceKeys.TRACING)):
            if isinstance(data, Mapping):
                if not isinstance(data, dict):
                    data = dict(data)
                data[key] = data_t.copy_meta_from(out_obj) if isinstance(data_t, MetaTensor) else data_t
                return data
            return out_obj  # return with data_t as tensor if get_track_meta() is False

        info = transform_info.copy()
        # track the current spatial shape
        if orig_size is not None:
            info[TraceKeys.ORIG_SIZE] = orig_size
        elif isinstance(data_t, MetaTensor):
            info[TraceKeys.ORIG_SIZE] = data_t.peek_pending_shape()
        elif hasattr(data_t, "shape"):
            info[TraceKeys.ORIG_SIZE] = data_t.shape[1:]

        # add lazy status to the transform info