Diff Coverage

        self._stack.__enter__()
        try:
            for ctx in self._ctxs:
                self._stack.enter_context(ctx)
        except BaseException as exc:
            self._stack.__exit__(type(exc), exc, exc.__traceback__)
            raise
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return self._stack.__exit__(exc_type, exc_val, exc_tb)

    """Store excluded-region outputs for one checkpoint invocation."""

    def __init__(self) -> None:
        """Initialize an empty per-checkpoint output cache."""
        self._outputs: Dict[int, Deque[Any]] = defaultdict(deque)

    def save(self, wrapper_id: int, output: Any) -> None:
        """Save one output produced by a checkpoint-excluded region."""
        self._outputs[wrapper_id].append(output)

    def pop(self, wrapper_id: int) -> Any:
        """Return the matching forward output during recomputation."""
        outputs = self._outputs.get(wrapper_id)
        if not outputs:
            raise RuntimeError("No cached forward output is available for this checkpoint exclusion wrapper")
        output = outputs.popleft()
        if not outputs:
            self._outputs.pop(wrapper_id)
        return output

    def clear(self) -> None:
        """Release outputs not consumed because recomputation stopped early."""
        self._outputs.clear()


def _identity(tensor: Any) -> Any:
    """Return a saved tensor unchanged."""
    return tensor


def _saved_tensors_context() -> Any:
    """Create an inner hook that stores real tensors instead of placeholders."""
    import mindspore as ms  # pylint: disable=C0415
    return ms.saved_tensors_hooks(_identity, _identity)


_CHECKPOINT_EXCLUDE_CACHE_KEY = object()

    """Exclude a callable region from checkpoint recomputation."""

    def __init__(self, module: Callable[..., Any]) -> None:
        """Initialize a checkpoint exclusion wrapper for a MindSpore Cell or function."""
        if not callable(module):
            raise ValueError("module must be a MindSpore Cell or callable")
        super().__init__(module, track_overlaps=False)

    def construct(self, *args: Any, **kwargs: Any) -> Any:
        """Execute normally outside recompute and return the cached output in recompute."""
        state = get_recompute_state()
        if state is None:
            return self._ckpt_wrapped_module(*args, **kwargs)
        cache = state.get_resource(_CHECKPOINT_EXCLUDE_CACHE_KEY, _CheckpointExcludeCache)
        if state.is_recomputing:
            return cache.pop(id(self))

        with _saved_tensors_context():
            output = self._ckpt_wrapped_module(*args, **kwargs)
        cache.save(id(self), output)
        return output


def checkpoint_exclude_wrapper(module: Callable[..., Any]) -> CheckpointExcludeWrapper:
    """Wrap a MindSpore Cell or function so its region is not recomputed.

    Note:
        This feature requires MindSpore PyNative mode and a surrounding
        HyperParallel checkpoint configured with ``use_reentrant=False``.
    """
    return CheckpointExcludeWrapper(module)

        Returns:
            The platform-specific checkpoint exclusion wrapper.
        """
        # pylint: disable=C0415
        from hyper_parallel.platform.mindspore.activation_checkpoint.checkpoint_exclude_wrapper import checkpoint_exclude_wrapper
        return checkpoint_exclude_wrapper(module)

    @staticmethod
    def swap_wrapper(module, policy_fn=None, group_swap=False):
        # pylint: disable=C0415