Diff Coverage

def get_device_handle(device_type: str = "npu"):
    """Return the torch device module (e.g. ``torch.npu`` or ``torch.cuda``)."""
    try:
        return getattr(torch, device_type)
    except AttributeError as e:
        raise RuntimeError(f"expect got device handle: 'torch.{device_type}' failed.") from e


def new_stream():
    """Create a new device stream on the current accelerator."""
    return get_device_handle().Stream()


def get_stream_context():
    """Return the stream context manager (``torch.npu.stream`` / ``torch.cuda.stream``)."""
    return get_device_handle().stream


def get_current_stream():
    """Return the current device stream."""
    return get_device_handle().current_stream()


def new_event():
    """Create a new device event on the current accelerator."""
    return get_device_handle().Event()


def no_grad():
    """Return ``torch.no_grad()``."""
    return torch.no_grad()


def preserve_version_counter(tensor):
    """Temporarily keep the tensor version counter unchanged across an in-place write."""
    return torch.autograd._unsafe_preserve_version_counter(tensor)  # pylint: disable=W0212


def cat(tensors, dim=0):
    """Concatenate tensors along *dim*."""
    return torch.cat(tensors, dim=dim)


def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
    """Allocate an uninitialized tensor shaped like *tensor*."""
    return torch.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)


def tree_map(fn, tree):
    """Apply *fn* to every leaf of *tree* and rebuild the same structure."""

def get_element_size(tensor) -> int:
    """Return the size in bytes of one element of *tensor*."""
    return tensor.element_size()


def alloc_tensor_buffer(numel: int, dtype, device="cpu", pin_memory: bool = False):
    """Allocate an uninitialized 1-D tensor buffer."""
    if pin_memory:
        return torch.empty(numel, dtype=dtype, device="cpu", pin_memory=True)
    return torch.empty(numel, dtype=dtype, device=device)


def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
    """Register a forward pre-hook on *module*, ignoring *prepend*."""

        self.active_session: Optional[_SessionActivation] = None

    def check_recomputed_tensors_match(self, session_id: Any) -> None:
        """Validate saved-tensor count and metadata after recomputation."""
        if self.ignore_saved_mismatch:
            return
        if len(self.weak_holders) != self.recomp_counter[session_id]:
            raise CheckpointError(
                "Hyper checkpoint saved a different number of tensors during forward and recomputation. "
                f"Forward saved {len(self.weak_holders)} tensors, but recomputation saved "
                f"{self.recomp_counter[session_id]} tensors."
            )

                f"Forward saved {len(self.weak_holders)} tensors, but recomputation saved "
                f"{self.recomp_counter[session_id]} tensors."
            )

        mismatches = []
        for index, weak_holder in enumerate(self.weak_holders):
            holder = weak_holder()
            if holder is None:
                continue
            handle = holder.handles.get(session_id)
            _internal_assert(handle is not None, "Missing recomputed tensor handle during metadata validation.")
            _internal_assert(
                handle in self.recomputed[session_id],
                "Missing recomputed tensor during metadata validation.",
            )
            recomputed_tensor = self.recomputed[session_id][handle]
            recomputed_metadata = self.metadata_fn(recomputed_tensor)
            if self.x_metadatas[index] != recomputed_metadata:
                mismatches.append((index, self.x_metadatas[index], recomputed_metadata))

        if mismatches:
            details = "\n".join(
                f"tensor {index}: forward={forward_metadata}, recompute={recomputed_metadata}"
                for index, forward_metadata, recomputed_metadata in mismatches
            )
            raise CheckpointError(
                "Hyper checkpoint detected different tensor metadata during recomputation:\n" + details
            )

    def clear_session(self, session_id: Any) -> None:

def _bind_session_activation(frame: _CheckpointFrame, activation: _SessionActivation) -> None:
    """Bind one activation to a frame outside the unpack hot path."""
    if frame.active_session is activation:
        return
    if frame.active_session is not None:
        raise CheckpointError("Concurrent recompute sessions on the same checkpoint frame are not supported.")
    frame.active_session = activation
    activation.frames.add(frame)

    if not device_types:
        return DefaultDeviceType.get_device_type()
    if "cuda" in device_types_set:
        return "cuda"
    return device_types[0]


def _get_device_module(device_type: str) -> Any:
    if device_type == "meta":

    frame_ref = weakref.ref(frame)

    def pack_hook(tensor: Any) -> Any:
        """Store recomputed tensors in their forward holders."""
        tensor = tensor.detach() if tensor.requires_grad else tensor
        target_frame = frame_ref()
        _internal_assert(target_frame is not None, "Checkpoint frame was released during recomputation.")
        recompute_index = target_frame.recomp_counter[session_id]
        target_frame.recomp_counter[session_id] += 1

        if recompute_index >= len(target_frame.weak_holders):
            if not target_frame.early_stop and not target_frame.forward_completed:
                target_frame.ignore_saved_mismatch = True
                return tensor
            raise CheckpointError(
                "Hyper checkpoint tried to save more tensors during recomputation than during forward."
            )

        holder = target_frame.weak_holders[recompute_index]()
        if holder is not None:
            _internal_assert(
                holder.handles.get(session_id) is None,
                "A recomputed tensor handle already exists for this session.",
            )
            handle = _Handle()
            holder.handles[session_id] = handle
            target_frame.recomputed[session_id][handle] = tensor

        if target_frame.early_stop and target_frame.recomp_counter[session_id] == len(target_frame.weak_holders):
            raise _StopRecomputationError
        return tensor

    def unpack_hook(tensor: Any) -> Any:
        """Return tensors saved by operations inside the recomputation."""
        return tensor

    return torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

def _run_recomputation(frame: _CheckpointFrame, session_id: Any) -> None:
    """Run and validate a frame recomputation for the given session."""
    if frame.is_recomputed[session_id]:
        return

    activation = frame.active_session
    if activation is not None:
        _internal_assert(activation.session_id == session_id, "Active session key does not match recomputation key.")

        _internal_assert(activation.session_id == session_id, "Active session key does not match recomputation key.")
    previous_activation = _RECOMPUTE_SESSION.get()
    token = None
    if activation is not None and previous_activation is not activation:
        token = _RECOMPUTE_SESSION.set(activation)
    try:
        input_context = frame.input_saver.grad_fn
        args = input_context.get_args(input_context.saved_tensors)
        try:

        try:
            with _create_recomputation_hooks(frame, session_id), torch.autograd.enable_grad():
                _run_fn_with_dynamo_disabled(frame.recompute_fn, *args)
        except _StopRecomputationError:
            pass
    finally:
        if token is not None:
            _RECOMPUTE_SESSION.reset(token)
    frame.is_recomputed[session_id] = True
    frame.check_recomputed_tensors_match(session_id)


def _create_checkpoint_hooks(frame: _CheckpointFrame) -> Any:
    """Create hooks that lazily recompute tensors saved during forward."""

    def unpack_hook(holder: _Holder) -> Any:
        """Return the corresponding tensor from lazy or prefired recomputation."""
        activation = frame.active_session
        if activation is not None:
            session_id = activation.session_id
            retain_on_unpack = activation.retain_on_unpack
        else:
            session_id = torch._C._current_graph_task_id()  # pylint: disable=W0212
            if session_id == -1:
                session_id = int(uuid.uuid4())
            retain_on_unpack = False

        _run_recomputation(frame, session_id)
        _internal_assert(session_id in holder.handles, "No recomputed tensor was saved for this checkpoint value.")
        handle = holder.handles[session_id]
        if handle is None:
            raise CheckpointError("A checkpoint tensor was unpacked more than once in the same recompute session.")
        _internal_assert(handle in frame.recomputed[session_id], "The recomputed tensor has already been released.")
        tensor = frame.recomputed[session_id][handle]
        if not retain_on_unpack:
            holder.handles[session_id] = None
        return tensor

    return torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

    forward_cpu_state = None
    if preserve_rng_state:
        forward_cpu_state = torch.get_rng_state()
        if getattr(device_module, "_initialized", False):
            had_device_in_forward = True
            forward_devices, forward_device_states = _get_device_states(device_type, *args)

    def recompute_fn(*inputs: Any) -> None:
        """Restore execution state and rerun the checkpointed function."""
        function_kwargs, *function_args = inputs

        ):
            if preserve_rng_state:
                torch.set_rng_state(forward_cpu_state)
                if had_device_in_forward:
                    _set_device_states(device_type, forward_devices, forward_device_states)

            device_autocast_context = contextlib.nullcontext()
            if device_autocast_kwargs is not None:
                device_autocast_context = torch.amp.autocast(device_type=device_type, **device_autocast_kwargs)
            with device_autocast_context, torch.amp.autocast("cpu", **cpu_autocast_kwargs), recompute_context:
                function(*function_args, **function_kwargs)

    frame = _CheckpointFrame(recompute_fn, early_stop, metadata_fn)
    dummy = torch.empty((0,), requires_grad=True)
    frame.input_saver = _NoopSaveInputs.apply(dummy, kwargs, *args)

        return

    activation = _RECOMPUTE_SESSION.get()
    if activation is not None:
        raise CheckpointError("Nested checkpoint is not supported during scheduled recomputation.")

    collector = _RECOMPUTE_COLLECTOR.get()
    if collector is not None:
        collector.append(frame)

        and activation.retain_on_unpack
    ):
        _register_session_frame(handle, session_id, activation)
        _run_recomputation(handle, session_id)
        return
    if activation is not None:
        raise CheckpointError("recompute_handle cannot enter another active recompute session.")

    _register_session_frame(handle, session_id)

def _to_torch_checkpoint_policy(policy: Any) -> Any:
    """Convert supported HyperParallel policies to the native Torch enum."""
    from torch.utils import checkpoint as torch_checkpoint  # pylint: disable=C0415

    torch_policy_cls = torch_checkpoint.CheckpointPolicy
    supported_native = {
        getattr(torch_policy_cls, name) for name in _SUPPORTED_POLICY_NAMES
    }
    if isinstance(policy, torch_policy_cls):
        if policy in supported_native:
            return policy
        raise ValueError(
            f"Torch checkpoint policy {policy.name} is not supported by "
            "HyperParallel compile mode. Only SAVE and RECOMPUTE policies are supported."
        )
    if isinstance(policy, CheckpointPolicy):
        if policy.name in _SUPPORTED_POLICY_NAMES:
            return getattr(torch_policy_cls, policy.name)
        raise ValueError(
            f"HyperParallel checkpoint policy {policy.name} is not supported in compile mode. "
            "Only SAVE and RECOMPUTE policies are supported."
        )
    raise TypeError(
        "Selective checkpoint policy_fn must return a HyperParallel or Torch "
        f"CheckpointPolicy, but got {type(policy).__name__}."
    )

def _torch_policy_adapter(
    policy_fn: Callable, torch_context: Any, op: Any, *args: Any, **kwargs: Any
) -> Any:
    """Pass native Torch inputs through and adapt only the policy result."""
    return _to_torch_checkpoint_policy(policy_fn(torch_context, op, *args, **kwargs))


def create_native_selective_checkpoint_contexts(policy_fn: Callable) -> Tuple[Any, Any]:
    """Create Torch-native selective-checkpoint contexts for compile capture."""
    if not callable(policy_fn):
        raise TypeError("policy_fn must be callable in HyperParallel compile mode.")
    from torch.utils import checkpoint as torch_checkpoint  # pylint: disable=C0415

    return torch_checkpoint.create_selective_checkpoint_contexts(
        partial(_torch_policy_adapter, policy_fn)
    )

        # event objects so each runtime event is queried only once per scan.
        event_results: Dict[int, Tuple[Any, bool]] = {}
        for capacity in list(self._pending):
            if capacity < minimum_capacity:
                continue
            still_pending = []
            for block in self._pending[capacity]:
                event = block.event
                event_key = id(event)

        """Associate a checked-out view with its owning allocation."""
        key = _storage_key(view)
        owner = self._blocks.get(key)
        if owner is not None and owner is not block:
            raise RuntimeError("Two host buffers exposed the same storage identity.")
        old_key = block.storage_key
        if old_key is not None and old_key != key and self._blocks.get(old_key) is block:
            del self._blocks[old_key]
        self._blocks[key] = block

            block = self._find_available_locked(aligned_size)
            if block is not None:
                try:
                    return self._checkout_locked(block, size)
                except Exception:
                    block.state = _AVAILABLE
                    self._available[block.capacity].append(block)
                    raise

            if self._total_allocated + aligned_size <= self._max_host_bytes:
                self._total_allocated += aligned_size
                reserved = True

            block = _Block(buffer, aligned_size, _IN_USE)
            try:
                with self._lock:
                    view = self._checkout_locked(block, size)
            except Exception:
                with self._lock:
                    self._total_allocated -= aligned_size
                raise
            return view

        event = block.event
        try:

        event = block.event
        try:
            event.synchronize()
        except Exception:
            with self._lock:
                block.state = _PENDING
                block.event = event
                self._pending[block.capacity].insert(0, block)
            raise
        block.event = None
        try:
            with self._lock:
                return self._checkout_locked(block, size)
        except Exception:
            with self._lock:
                block.state = _AVAILABLE
                self._available[block.capacity].append(block)
            raise

    def release(self, tensor: Any, event: Optional[Any] = None) -> None:
        """Return an acquired view to the pool, optionally after an async event."""
        key = _storage_key(tensor)

    # There should probably be a better way to do this...
    # NOTE: unify _is_compiling across all compile stacks
    for arg in args:
        if isinstance(arg, torch.Tensor) and is_fun(arg):
            return True
    return False


class _VersionWrapper:

        is_compiling = _is_compiling(func, args, kwargs)

        if is_compiling:
            # Overwrite each node's "recompute" tag to add in the user annotation.
            fx_traceback.current_meta["recompute"] = policy

        out = func(*args, **kwargs)

        has_alias = any(ret.alias_info is not None for ret in func._schema.returns)

        self._swap_cleared = False

    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        if func in SAC_IGNORED_OPS:
            return func(*args, **kwargs)

        kwargs = {} if kwargs is None else kwargs
        policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=True),
                                func, *args, **kwargs)

                    "on any region computed under selective activation checkpoint."
                )
            out = tree_map(lambda x: x.get_val(self.allow_cache_entry_mutation), storage.pop(0))
        else:
            out = func(*args, **kwargs)
        return out


def create_selective_checkpoint_contexts(

    def async_group_load(self, source):
        """Copy a packed device-buffer slice back into the original storage."""
        if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap:
            return
        if not self._group_managed:
            return
        if self._state != self.STATE_HOST:
            warnings.warn(
                f"[SwapTensor.async_group_load] Invalid state: current={self._state}, "
                f"expected 'host'. Operation skipped."
            )
            return
        if self.val.untyped_storage().size() != self.storage_size:
            raise RuntimeError(
                f"Cannot load grouped tensor from {self.funcname}: device storage was not restored. "
                f"expected size:{self.storage_size}, current size:{self.val.untyped_storage().size()}"
            )
        with _backend.preserve_version_counter(self.val):

    def release_cpu_buffer(self, event=None):
        """Release an explicitly pooled host buffer exactly once."""
        if self.cpu_pool is None or self._cpu_pool_buffer is None:
            return
        release_tensor = self.val_cpu if self.val_cpu is not None else self._cpu_pool_buffer
        self.cpu_pool.release(release_tensor, event=event)
        self._cpu_pool_buffer = None
        self.val_cpu = None

    def wait_load(self, release_event=None):
        """change state to device after async load is done"""
        if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap:

                self.val_cpu = _backend.empty_like(
                    self.val, device="cpu", pin_memory=True
                )
            else:
                logical_bytes = self.val.numel() * _backend.get_element_size(self.val)
                self._cpu_pool_buffer = self.cpu_pool.acquire(logical_bytes)
                try:
                    self.val_cpu = self._cpu_pool_buffer.view(self.val.dtype).reshape(self.val.shape)
                except Exception:
                    self.release_cpu_buffer()
                    raise
        try:
            if self.cpu_pool is not None or self.is_slice_tensor:
                self.val_cpu.copy_(self.val, non_blocking=True)
            else:
                self.val_cpu.untyped_storage().copy_(self.val.untyped_storage(), non_blocking=True)
        except Exception:
            if self.cpu_pool is not None and self._cpu_pool_buffer is not None:
                release_event = _backend.new_event()
                release_event.record(_backend.get_current_stream())
                self.release_cpu_buffer(release_event)
            self.val_cpu = None
            raise
        self._state = self.STATE_D2H

    def wait_offload(self):
        """wait offload to host and free device memory"""

        self.clear()

    def release_cpu_buffers(self, event=None):
        """Release all explicitly pooled host buffers held by this storage."""
        def _release(x):
            if isinstance(x, SwapTensor):
                x.release_cpu_buffer(event=event)
            return x

        for storage_list in self.values():
            for item in storage_list:
                _backend.tree_map(_release, item)

    def wait_offload(self):
        """wait offload for all tensors in swap storage"""
        def _wait_offload(x):

                       or not x.val.is_contiguous())
            if no_pack:
                return x
            if x.storage_size != x.val.untyped_storage().size():
                raise RuntimeError(
                    f"There is a tensor from {x.funcname} cannot be SWAPPED! Its storage has been resized "
                    f"presize:{x.storage_size}, current size:{x.val.untyped_storage().size()}"
                )
            if x.ver != x.val._version:

                        cpu_pool = bucket["cpu_pool"]
                        if cpu_pool is None:
                            cpu_buf = _get_cpu_pinned_buf(dtype_key, numel, bucket["dtype"])
                        else:
                            raw_buf = cpu_pool.acquire(bucket["total_bytes"])
                            try:
                                cpu_buf = raw_buf.view(bucket["dtype"])
                            except Exception:
                                cpu_pool.release(raw_buf)
                                raise
                        group_cpu_bufs[bucket_key] = cpu_buf
                        cpu_buf[:numel].copy_(group_device_bufs[bucket_key], non_blocking=True)
                except Exception:
                    release_event = _backend.new_event()
                    release_event.record(copy_stream)
                    for bucket_key, cpu_buf in group_cpu_bufs.items():
                        bucket = self._packed_buckets[bucket_key]
                        if bucket["cpu_pool"] is not None:
                            bucket["cpu_pool"].release(cpu_buf, event=release_event)
                        else:
                            _return_cpu_pinned_buf(cpu_buf)
                    raise
                self._group_device_buf = group_device_bufs
                self._group_cpu_buf = group_cpu_bufs

            # Slice tensors use the existing per-tensor path.

        with _backend.no_grad(), stream_context(copy_stream):
            compute_event.wait(copy_stream)

            if self._packed_tensor_info and self._group_cpu_buf is not None:
                group_device_bufs = {}
                for bucket_key, bucket in self._packed_buckets.items():
                    cpu_buf = self._group_cpu_buf.get(bucket_key)
                    if cpu_buf is None:
                        continue
                    numel = bucket["total_numel"]
                    group_device_bufs[bucket_key] = _backend.alloc_tensor_buffer(
                        numel, bucket["dtype"], bucket["device"]
                    )
                    # One-shot H2D per packed bucket.
                    group_device_bufs[bucket_key].copy_(cpu_buf[:numel], non_blocking=True)
                self._group_device_buf = group_device_bufs

                # Unpack with D2D copies into the original storages. Rebinding
                # st.val with set_() would leave existing aliases on freed storage.
                for st, bucket_key, element_offset in self._packed_tensor_info:
                    group_device_buf = group_device_bufs.get(bucket_key)
                    if group_device_buf is None:
                        continue
                    source = group_device_buf[element_offset:element_offset + st.val.numel()]
                    st.async_group_load(source)

            # Slice tensors use the existing per-tensor path.
            # Group-managed tensors skip async_load via _group_managed flag.
            for storage in self._storages:

    def release_cpu_buffers(self, event=None):
        """Release staging buffers immediately or defer until ``event`` completes."""
        if self._group_cpu_buf is None:
            return
        for bucket_key, buf in self._group_cpu_buf.items():
            bucket = self._packed_buckets.get(bucket_key)
            if bucket is not None and bucket["cpu_pool"] is not None:
                bucket["cpu_pool"].release(buf, event=event)
            else:
                _return_cpu_pinned_buf(buf)
        self._group_cpu_buf = None

    def wait_load(self):
        """Wait for grouped H2D and D2D loads to complete."""
        if self._load_event is None:

        for event in (group._offload_event, group._load_event):
            if event is not None:
                event.synchronize()
        for storage in group._storages:
            storage.release_cpu_buffers()
        group.release_cpu_buffers()
        group._storages.clear()

    def get_current_group_name(self) -> str:

        Ensures idempotency: safe to call multiple times on the same layer pair.
        """
        if first_layer is second_layer:
            warnings.warn(
                "set_forward_prefetch_layer: "
                "Prefetching between identical layers has no effect.",
                UserWarning,
                stacklevel=2,

        if submodule is module:
            continue
        wrapped_callable = _get_wrapped_callable(submodule)
        if wrapped_callable is not None and _is_callable_exempt_from_overlap_check(wrapped_callable):
            continue
        if getattr(submodule, '_is_wrapped', False):
            if wrapped_callable is not None:
                _raise_callable_already_wrapped(wrapped_callable)
            # A param-free module shared across siblings (one rotary embedding
            # per decoder layer) is never its own region; skip, do not flag.
            if next(submodule.parameters(recurse=True), None) is None:
                continue

        swap_manager = SwapManager()

        def pack_to_cpu(tensor: torch.Tensor):
            if not base_check_fn(tensor):
                return tensor.detach()
            if policy_fn is not None:
                if policy_fn(tensor) == CheckpointPolicy.MUST_SAVE:
                    return tensor.detach()
                if policy_fn(tensor) != CheckpointPolicy.MUST_SWAP:
                    raise RuntimeError(f"Swap :set an invalid policy {policy_fn(tensor)}")
            group_name = swap_manager.get_current_group_name()
            if not group_name:
                return tensor.detach()
            if not self.add_to_storage:
                swap_manager.add_storage(group_name, self.storage)
                self.add_to_storage = True
            funcname = f"{group_name}::{tensor.shape}"

    def _register_tensor(tensor):
        nonlocal count_idx
        if not base_check_fn(tensor):
            return tensor

        tensor_tag = tag or f"{group_name}_swap_tensor"
        funcname = f"{tensor_tag}::{tuple(tensor.shape)}"
        storage[count_idx].append(

        self.checkpoint_kwargs = checkpoint_kwargs

    def _do_checkpoint(self, wrapped_module: Any, *args: Any, **kwargs: Any) -> Any:
        # Checkpoint may save inputs before the wrapped module's pre-hook runs.
        group_name = getattr(self, "_swap_group_name", None)
        if group_name is not None:
            SwapManager().set_current_group_name(group_name)

        from hyper_parallel.core.activation_memory.api import checkpoint  # pylint: disable=C0415
        return checkpoint(
            wrapped_module,
            *args,
            **self.checkpoint_kwargs,
            **kwargs,

            **kwargs,
        )

    def forward(self, *args: Any, **kwargs: Any) -> Any:
        return self._do_checkpoint(self._wrapped_module, *args, **kwargs)


def ckpt_wrapper(module: Union[nn.Module, Callable], **checkpoint_kwargs: Any) -> CheckpointWrapper:
    """Wrap *module* with activation checkpointing."""

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

    def __init__(self, module: Callable[..., Any], *, save_output: bool = True) -> None:
        """Initialize a checkpoint exclusion wrapper for a PyTorch module or function."""
        if not callable(module):
            raise ValueError("module must be a PyTorch Module or callable")
        if not isinstance(save_output, bool):
            raise ValueError(f"save_output must be a bool, got {type(save_output).__name__}")
        super().__init__(module, track_overlaps=False)
        self.save_output = save_output

    def forward(self, *args: Any, **kwargs: Any) -> Any:
        """Execute normally outside recompute and return the cached output in recompute."""
        wrapped_module = cast(Callable[..., Any], self._wrapped_module)
        state = get_recompute_state()
        if state is None:
            return wrapped_module(*args, **kwargs)
        cache = state.get_resource(_checkpoint_exclude._EXCLUDE_CACHE_KEY, _checkpoint_exclude._ExcludeCache)
        if state.is_recomputing:
            entry = cache.pop(id(self))
            _checkpoint_exclude._materialize_recompute_inputs(entry, args, kwargs)
            output = (
                entry.output
                if self.save_output
                else _checkpoint_exclude._make_replay_placeholder_output(entry.output_tensor_count)
            )
            return _checkpoint_exclude._finalize_save_outputs(
                output,
                _checkpoint_exclude._has_used_input(entry.input_bindings),
                None,
            )

                _checkpoint_exclude._has_used_input(entry.input_bindings),
                None,
            )

        input_bindings, previous_handles = _checkpoint_exclude._mark_recompute_inputs(
            state.invocation_id,
            args,
            kwargs,
        )
        try:
            with _checkpoint_exclude._saved_tensors_context():
                output = wrapped_module(*args, **kwargs)
        finally:
            _checkpoint_exclude._restore_recompute_inputs(previous_handles)
        needs_recompute_boundary = _checkpoint_exclude._has_used_input(input_bindings)
        tensor_leaf_count = None if self.save_output else [0]
        finalized_output = _checkpoint_exclude._finalize_save_outputs(
            output,
            needs_recompute_boundary,
            state.invocation_id,
            tensor_leaf_count,

            needs_recompute_boundary,
            state.invocation_id,
            tensor_leaf_count,
        )
        replay_output = output if self.save_output else None
        output_tensor_count = 1 if tensor_leaf_count is None else tensor_leaf_count[0]
        cache.save(
            id(self),
            _checkpoint_exclude._ExcludeCacheEntry(replay_output, input_bindings, output_tensor_count),
        )
        return finalized_output


def checkpoint_exclude_wrapper(
    module: Callable[..., Any],

        This feature requires eager mode and a surrounding HyperParallel
        checkpoint configured with ``use_reentrant=False``. Nested checkpoint
        exclusion wrappers are not supported.
    """
    return CheckpointExcludeWrapper(module, save_output=save_output)

    Args:
        ops: Backend operators to ignore. ``None`` entries represent optional
            operators that are unavailable in the installed PyTorch version.
    """
    _ignore_sac_ops(ops)


def ensure_profiler_ops_sac_ignored() -> None:
    """Keep profiler record-function operators out of selective-AC replay.

        Returns:
            The ``(forward_context, recompute_context)`` pair expected by the
            non-reentrant checkpointing ``context_fn`` contract.
        """
        return create_selective_checkpoint_contexts(
            _make_selective_checkpoint_policy_fn()
        )

    return selective_checkpoint_context_fn