Swap

def forward(self, data: Any) -> Any:
    """Apply the transform.

    The output type always matches the input type.

    Args:
        data: Input data to transform.
    """
    if self.copy:
        data = _copy.deepcopy(data)
    batch, unwrap = self._wrap(data)
    if torch.rand(1).item() > self.p:
        return unwrap(batch)
    params = self.make_params(batch)
    batch = self.apply_transform(batch, params)
    # Record history on the batch
    trace = AppliedTransform(name=type(self).__name__, params=params)
    if not hasattr(batch, "applied_transforms"):
        batch.applied_transforms = []
    batch.applied_transforms.append(trace)
    result = unwrap(batch)
    # Propagate history to outputs that can carry it
    if (
        hasattr(batch, "applied_transforms")
        and not isinstance(result, (SubjectsBatch, Tensor, np.ndarray))
        and not isinstance(result, dict)
    ):
        with contextlib.suppress(AttributeError):
            result.applied_transforms = list(batch.applied_transforms)
    return result

def inverse(self, params: dict[str, Any]) -> Transform:
    """Return a transform that undoes this one.

    Override in invertible subclasses. The returned transform,
    when applied, reverses the effect of the forward pass with
    the given parameters.

    Args:
        params: The parameters recorded in the forward pass.

    Returns:
        A new `Transform` instance that inverts this one.
    """
    msg = f"{type(self).__name__} is not invertible"
    raise NotImplementedError(msg)

def to_hydra(self) -> dict[str, Any]:
    """Export as a Hydra-compatible config dict.

    Returns a dict with `_target_` set to the fully qualified
    class name and only non-default field values included.

    Returns:
        Dict suitable for `hydra.utils.instantiate()`.
    """
    from .parameter_range import _ParameterRange

    cls = type(self)
    target = f"torchio.{cls.__qualname__}"
    cfg: dict[str, Any] = {"_target_": target}

    for name, default in _collect_init_params(cls).items():
        value = getattr(self, name, default)
        if isinstance(value, _ParameterRange):
            if value._original == default:
                continue
            value = _hydra_value(value._original)
        elif value == default:
            continue
        else:
            value = _hydra_value(value)
        cfg[name] = value
    return cfg

def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
    """Sample swap locations for each image."""
    n = max(1, round(self.num_iterations.sample_1d()))

    # Warn if label maps are present.
    for _name, img_batch in batch.images.items():
        if issubclass(img_batch._image_class, LabelMap):
            warnings.warn(
                "Swap is applied to a subject containing LabelMap "
                "images. The spatial rearrangement will make labels "
                "inconsistent with the swapped image. This transform "
                "is intended for self-supervised learning.",
                stacklevel=2,
            )
            break

    # Sample one set of locations per image (shared across batch).
    any_img = next(iter(batch.images.values()))
    spatial_shape = any_img.data.shape[2:]  # (I, J, K)
    locations = _sample_swap_locations(
        spatial_shape,
        self.patch_size,
        n,
    )
    return {"locations": locations}

def apply_transform(
    self,
    batch: SubjectsBatch,
    params: dict[str, Any],
) -> SubjectsBatch:
    """Swap patches in each selected image."""
    locations = params["locations"]
    for _name, img_batch in self._get_images(batch).items():
        img_batch.data = _apply_swaps(
            img_batch.data,
            locations,
            self.patch_size,
        )
    return batch