Diff Coverage

    from .reentrant_checkpoint import checkpoint as _reentrant  # pylint: disable=C0415

    preserve_rng_state = kwargs.pop("preserve_rng_state", True)
    if not isinstance(preserve_rng_state, bool):
        raise ValueError("preserve_rng_state must be bool")
    if context_fn is not None and not callable(context_fn):
        raise ValueError("context_fn must be callable")
    state_context = partial(create_recompute_contexts, preserve_rng_state=preserve_rng_state)
    reentrant_context = (
        _compose_context_fns((state_context, context_fn)) if context_fn is not None else state_context
    )

) -> None:
    """Validate options unsupported by eager reentrant checkpointing."""
    unsupported = []
    if swap_inputs:
        unsupported.append("swap_inputs")
    if policy_fn is not None:
        unsupported.append("policy_fn")
    if group_swap:
        unsupported.append("group_swap")
    if cpu_pool is not None:
        unsupported.append("cpu_pool")
    if early_stop is not True:
        unsupported.append("early_stop")
    for name in ("debug", "determinism_check"):
        if name in kwargs:
            unsupported.append(name)
    if is_compiling():
        unsupported.append("use_reentrant=True under torch.compile")
    if unsupported:
        raise ValueError("Reentrant checkpoint does not support: " + ", ".join(unsupported))

    """
    if not isinstance(early_stop, bool):
        raise ValueError(f"early_stop must be bool, but got {type(early_stop).__name__}.")
    if not isinstance(use_reentrant, bool):
        raise ValueError("use_reentrant must be bool")
    if use_reentrant:
        _validate_reentrant_options(
            swap_inputs, policy_fn, group_swap, cpu_pool, early_stop, kwargs
        )

            return self.function(*args, **kwargs)

    def _gradients(self, values: tuple, grads: tuple) -> tuple:
        """Compute replay-input gradients without returning captured parameter edges."""
        if self.split and self.phase == "dx":
            targets = [item for item, is_param in zip(self.replay_inputs, self.param_input_mask)
                       if item.requires_grad and not is_param]
            if not targets:
                return (None,) * self.input_count
            calculated = iter(torch.autograd.grad(values, targets, grads, retain_graph=True, allow_unused=True))
            return tuple(next(calculated) if item.requires_grad and not is_param else None
                         for item, is_param in zip(self.replay_inputs, self.param_input_mask))

        for item in self.replay_inputs:
            if item.requires_grad:
                item.grad = None
        torch.autograd.backward(values, grads, retain_graph=self.retain_graph)
        return tuple(item.grad if item.requires_grad else None for item in self.replay_inputs)

    def recompute(self) -> None:
        """Replay this frame once; retain its graph until DW or explicit cleanup."""
        if self.released:
            raise RuntimeError("Cannot replay a released reentrant checkpoint")
        if self.replay_outputs is not None:
            return
        context_ref = self.context_ref
        if context_ref is None:
            raise RuntimeError("Reentrant checkpoint output graph was released before replay")
        ctx = _read_context(context_ref)
        if ctx is None:
            raise RuntimeError("Reentrant checkpoint output graph was released before replay")
        inputs = ctx.saved_tensors
        self.replay_inputs = tuple(
            tensor.detach().requires_grad_(tensor.requires_grad) if tensor.is_floating_point() or tensor.is_complex()
            else tensor.detach() for tensor in inputs

        self.replaying = True
        try:
            def replay(*unused: Any) -> None:
                """Store the differentiable replay output for the owning frame."""
                del unused
                self.replay_outputs = self._call(self.replay_inputs)

            with torch.enable_grad():
                self.state.recompute_fn(replay, {})
        finally:

    def set_phase(self, phase: str) -> None:
        """Select the phase used by this frame's engine callbacks (dx or dw)."""
        if not self.split or phase not in {"dx", "dw"}:
            raise ValueError("Split reentrant handles require phase='dx' or phase='dw'")
        if self.released:
            raise RuntimeError("Cannot set phase on a released reentrant checkpoint")
        self.phase = phase

    def release(self) -> None:
        """Release all retained replay and excluded local graphs."""

    def release(self) -> None:
        """Release all retained replay and excluded local graphs."""
        self.exclude_cache.clear()
        for entry in self.entries:
            entry.tape.release()
            entry.input_bindings.clear()
            entry.output = None
        self.entries.clear()
        self.excluded_output_mask.clear()
        self.replay_outputs = None
        self.replay_inputs = ()

    def bind_session(self, session_id: Any) -> None:
        """Register retained state for cleanup through the unified session API."""
        activation = _RECOMPUTE_SESSION.get()
        if activation is not None and activation.session_id != session_id:
            raise ValueError("Active session does not match the reentrant handle session")
        with _SESSIONS_LOCK:
            if self.active_session is not None and self.active_session.session_id != session_id:
                raise RuntimeError("Concurrent reentrant recompute sessions on one frame are unsupported")
            _SESSIONS.setdefault(session_id, weakref.WeakSet()).add(self)
            if activation is not None:
                self.active_session = activation
                self.retain_graph = activation.retain_on_unpack

def activate_session(activation: Any) -> None:
    """Bind a session's retention policy before autograd enters worker threads."""
    with _SESSIONS_LOCK:
        for handle in _SESSIONS.get(activation.session_id, ()):
            if handle.active_session is not None and handle.active_session is not activation:
                raise RuntimeError("Concurrent reentrant recompute sessions on one frame are unsupported")
            handle.active_session = activation
            handle.retain_graph = activation.retain_on_unpack


def deactivate_session(activation: Any) -> None:
    """Unbind the policy without discarding registered frames or retained graph."""

            for index, value in enumerate(leaves)
        )
        outputs = tuple(value.detach() for value in leaves if isinstance(value, torch.Tensor))
        if not outputs:
            raise ValueError("Reentrant checkpoint requires at least one Tensor output")
        originals = (value for value in leaves if isinstance(value, torch.Tensor))
        def differentiable(original: torch.Tensor) -> bool:
            """Match only the same live detached object, never a reused Python id."""
            matched = frame.excluded_output_mask.get(id(original))

    def _run_backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple:
        """Perform the replay and route gradients according to the explicit phase."""
        frame = ctx.frame
        if frame.split and frame.phase == "ordinary":
            raise RuntimeError("Set dx or dw phase on reentrant handles before split backward")
        frame.recompute()
        outputs = tuple(value for value in tree_flatten(frame.replay_outputs)[0] if isinstance(value, torch.Tensor))
        selected = [
            (value, grad) for value, grad in zip(outputs, grad_outputs)
            if value.requires_grad and grad is not None
        ]
        input_grads = (frame._gradients(*zip(*selected))  # pylint: disable=protected-access
                       if selected else (None,) * frame.input_count)
        if (not frame.split and not frame.retain_graph) or (frame.split and frame.phase == "dw"):
            frame.release()
        return (None, None, *input_grads, None)


@contextlib.contextmanager
def reentrant_checkpoint_scope(

        include every ``backward_trigger`` in DW's autograd ``inputs`` tuple.
        Release handles on cancellation; DW releases them after consumption.
    """
    if _SCOPE.get() is not None:
        raise ValueError("Nested reentrant checkpoint scopes are not supported")
    scope = [parameters, [], None]
    token = _SCOPE.set(scope)
    try:
        yield scope[1]

        if isinstance(item, torch.Tensor):
            template.append(_InputSlot(len(inputs)))
            inputs.append(item)
        else:
            template.append(item)
    return tuple(inputs), template, spec


def _scope_parameters(scope: Optional[list]) -> Optional[tuple]:

        return None
    if scope[2] is None:
        supplied = tuple(scope[0]() if callable(scope[0]) else scope[0])
        if any(not isinstance(param, torch.Tensor) for param in supplied):
            raise ValueError("Scope parameters must be Tensors")
        scope[2] = tuple(dict((id(param), param) for param in supplied if param.requires_grad).values())
    return scope[2]

    @staticmethod
    def backward(ctx: Any, *grads: torch.Tensor) -> tuple:
        """Record local input sensitivities for the replay bridge."""
        ctx.slots[:] = grads
        return None, None


class _LocalTape:
    """Keep the original excluded graph and only lightweight input gradient edges."""

        self.param_aliases = ()

    def gradients(self, grad_outputs: tuple) -> tuple:
        """Compute local gradients according to the owning frame's explicit phase."""
        frame = self.frame_ref()
        if frame is None:
            raise RuntimeError("Reentrant checkpoint frame was released before exclude backward")
        pairs = [(out, grad) for out, grad in zip(self.outputs, grad_outputs) if out.requires_grad and grad is not None]
        if not pairs:
            return (None,) * len(self.input_edges)
        outputs, grads = zip(*pairs)
        if frame.phase in {"ordinary", "dw"}:
            return self._backward(outputs, grads, frame.retain_graph)
        return self._dx_gradients(outputs, grads)

    def _backward(self, outputs: tuple, grads: tuple, retain_graph: bool) -> tuple:
        """Accumulate captured parameters once and return explicit input sensitivities."""
        self.slots[:] = [None] * len(self.slots)
        for alias in self.param_aliases:
            if alias is not None:
                alias.grad = None
        torch.autograd.backward(outputs, grads, retain_graph=retain_graph)
        slot_grads = iter(self.slots)
        seen_aliases = set()

        def alias_grad(alias: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
            """Return a shared explicit proxy gradient only at its first input edge."""
            if alias is None or id(alias) in seen_aliases:
                return None
            seen_aliases.add(id(alias))
            return alias.grad

        return tuple(
            alias_grad(alias) if is_param else next(slot_grads)
            for is_param, alias in zip(self.param_input_mask, self.param_aliases)
        )

        )

    def _dx_gradients(self, outputs: tuple, grads: tuple) -> tuple:
        """Target only activation boundaries, leaving parameter hooks untouched."""
        targets = [edge for edge, is_param in zip(self.input_edges, self.param_input_mask)
                   if edge is not None and not is_param]
        if not targets:
            return (None,) * len(self.input_edges)
        result = iter(torch.autograd.grad(outputs, targets, grads, retain_graph=True, allow_unused=True))
        return tuple(
            next(result) if edge is not None and not is_param else None
            for edge, is_param in zip(self.input_edges, self.param_input_mask)
        )

    @staticmethod
    def forward(ctx: Any, entry: Any, input_count: int, *tensors: torch.Tensor) -> tuple:
        """Return detached outputs with the recorded differentiability mask."""
        del tensors
        ctx.entry = entry
        ctx.input_count = input_count
        ctx.set_materialize_grads(False)
        outputs = tuple(tensor.detach() for tensor in entry.tape.outputs)
        ctx.mark_non_differentiable(*(
            output for output, differentiable in zip(outputs, entry.tape.output_mask)
            if not differentiable
        ))
        return outputs

    @staticmethod
    def backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple:
        """Propagate through the original local graph using the frame phase."""
        input_grads = ctx.entry.tape.gradients(grad_outputs)
        return (None, None, *input_grads, None)


def _replace_tensor_leaves(value, tensors):
    leaves, spec = tree_flatten(value)

def _replay_exclude(wrapper: Any, args: tuple, kwargs: dict, frame: Any) -> Any:
    """Bind recipes and bridge cached outputs to the replay graph."""
    entry = frame.exclude_cache.pop(id(wrapper))
    frame.entries.append(entry)
    recipes._restore_rng_state(entry.rng_state)  # pylint: disable=protected-access
    recipes._materialize_recompute_inputs(entry, args, kwargs)  # pylint: disable=protected-access
    current = _tensor_leaves((args, kwargs))
    if len(current) != len(entry.tape.input_edges):
        raise RuntimeError("Checkpoint exclusion Tensor input count changed during replay")
    result = _ReplayBridge.apply(entry, len(current), *current, torch.empty(0, requires_grad=True))
    tensors = (result,) if isinstance(result, torch.Tensor) else result
    return _replace_tensor_leaves(entry.output, tensors)


def _local_inputs(original: tuple, param_ids: set) -> tuple:
    """Build a local activation boundary and explicit parameter proxies."""

        finally:
            frame.recording_exclude = False
            recipes._ACTIVE_PACK_STATE.reset(token)  # pylint: disable=protected-access
    if swap_hooks is not None:
        wrapper._register_exclude_outputs(output)  # pylint: disable=protected-access
    outputs = _tensor_leaves(output)
    if not outputs:
        raise ValueError("Reentrant checkpoint exclusion requires at least one Tensor output")
    tape = _LocalTape(outputs, _input_edges(original, boundary, param_ids, param_aliases), slots, frame)
    tape.param_input_mask = tuple(id(item) in param_ids for item in original)
    tape.param_aliases = tuple(param_aliases.get(id(item)) for item in original)
    entry = recipes._ExcludeCacheEntry(  # pylint: disable=protected-access

def run_exclude(wrapper: Any, module: Any, args: tuple, kwargs: dict, frame: Any) -> Any:
    """Record an excluded body once, or bridge its output during replay."""
    if frame.replaying:
        return _replay_exclude(wrapper, args, kwargs, frame)
    return _record_exclude(wrapper, module, args, kwargs, frame)


def current_frame() -> Optional[Any]:

            scope = reentrant_checkpoint_scope(self._get_trainable_params)
        with scope as frames:
            try:
                out = self.submodule(*composite_args, **composite_kwargs)
            except BaseException:
                for frame in frames:
                    if isinstance(frame, ReentrantCheckpointHandle):
                        frame.release()
                raise
        reentrant_frames = tuple(frame for frame in frames if isinstance(frame, ReentrantCheckpointHandle))
        if reentrant_frames:
            self._reentrant_frames[micro_index] = reentrant_frames
        out_tuple = out if isinstance(out, tuple) else (out,)

        if phase == "full":
            if not self._has_sharded_params:
                self._bind_reentrant(micro_index, "dw")
            torch.autograd.backward(outputs, grad_tensors=grad_tensors)
            return
        inputs = []
        if phase != "dw":
            inputs = [info.buffer for info in self.args_recv_info[micro_index] if info.requires_grad]
        if self._has_sharded_params:

        local_output = self._filter_grad_outputs(fwd_output)

        if not local_output:
            # Nothing to backprop through (e.g. all forward outputs detached).
            self._release_reentrant(micro_index)
            self._clear_recv_buffer(self.grad_recv_info, micro_index)
            self._clear_recv_buffer(self.args_recv_info, micro_index)
            return

        if micro_index in self._reentrant_frames:
            try:
                self._dw_cache[micro_index] = (local_output, grad_tensors)
                self._reentrant_backward(micro_index, "full")
                if not self.is_first_stage:
                    self._populate_bwd_cache(micro_index)
            finally:
                self._release_reentrant(micro_index)
                self._clear_recv_buffer(self.grad_recv_info, micro_index)
                self._clear_recv_buffer(self.args_recv_info, micro_index)
            return
        # MPipe owner-backward shares the tower's all-gather node across micro
        # graphs, so freeing it on the first backward breaks the later ones.
        retain_graph = getattr(self, "retain_backward_graph", False)
        torch.autograd.backward(local_output, grad_tensors=grad_tensors,

        if micro_index in self._reentrant_frames:
            try:
                self._dw_cache[micro_index] = (local_output, grad_tensors)
                if not local_output:
                    self._dw_cache[micro_index] = ([], [])
                    self._populate_bwd_cache(micro_index)
                    return
                self._reentrant_backward(micro_index, "dx")
                self._populate_bwd_cache(micro_index)
            except BaseException:
                self._release_reentrant(micro_index)

                self._release_reentrant(micro_index)
                self._clear_recv_buffer(self.grad_recv_info, micro_index)
                self._clear_recv_buffer(self.args_recv_info, micro_index)
                if self.is_last_stage:
                    self.fwd_outputs_cache.pop(micro_index, None)
            return

        param_groups = self._dw_cache.pop(micro_index)
        stage_backward_weight(iter(self._get_trainable_params()), param_groups)
        self._clear_recv_buffer(self.grad_recv_info, micro_index)