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)

    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

        Returns:
            The wrapped module or callable.
        """
        raise NotImplementedError("checkpoint_exclude_wrapper is currently only supported on MindSpore")

    @staticmethod
    def swap_wrapper(module, policy_fn=None, group_swap=False):
        """Wrap a module with activation swap functionality.