Diff Coverage

        """Validate the activation checkpointing mode and layer selectors."""
        if not isinstance(self.ac_mode, str) or self.ac_mode not in AC_MODES:
            raise ValueError(f"ac_mode must be one of {AC_MODES}, got {self.ac_mode!r}")
        if not isinstance(self.ac_save_forward_all_gather, bool):
            raise ValueError("ac_save_forward_all_gather must be a bool")
        self._validate_ac_layer_count()
        self._validate_ac_layer_indices()
        if self.ac_layer_count is not None and self.ac_layer_indices is not None:
            raise ValueError(

                raise ValueError(f"{name} must be a bool")
        if self.fsdp_bucket_modules is not None and not isinstance(
            self.fsdp_bucket_modules, (list, tuple)
        ):
            raise ValueError(
                "fsdp_bucket_modules must be None or a list/tuple of module groups"
            )
        value = self.max_in_flight_gb
        if (

                "activation_checkpoint.swap_inputs is not supported in graph "
                "mode; input swapping will be disabled."
            )
        if getattr(ac_config, "allow_cache_entry_mutation", False):
            _LOG.info(
                "activation_checkpoint.allow_cache_entry_mutation has no "
                "effect in graph mode; ignoring it."
            )

    """Return the c10d ``wait_tensor`` overload when this torch ships it."""
    functional = getattr(torch.ops, "_c10d_functional", None)
    wait = getattr(functional, "wait_tensor", None)
    if wait is None:
        return None
    return getattr(wait, "default", None)


# getitem / wait_tensor extract or await their producer's value: their

        return _CheckpointScope(region_fqns=frozenset(), unit_fqns=block_fqns)
    recompute_spec = _adapter_recompute_spec(model, required=False)
    if recompute_spec is None or not recompute_spec.region_patterns:
        return _CheckpointScope(region_fqns=selected_layers, unit_fqns=block_fqns)
    module_by_fqn = dict(model.named_modules())
    roots = _select_layer_region_roots(module_by_fqn, selected_layers, recompute_spec)
    if not roots:
        raise ValueError(
            "model adapter recompute patterns matched no region in the selected layers"
        )
    region_fqns = frozenset(fqn for fqn, _ in roots)
    excluded = _select_excluded_targets(module_by_fqn, roots, recompute_spec)
    exclude_fqns = frozenset(fqn for fqn, _ in excluded)
    return _CheckpointScope(
        region_fqns=region_fqns,
        exclude_fqns=exclude_fqns,
        unit_fqns=block_fqns,
    )

    checkpointing targets transformer-block forwards, mirroring the eager
    path that wraps the discovered blocks only.
    """
    if not scope.region_fqns:
        return False
    fqns = _node_stack_fqns(node)
    if not fqns:
        return False
    if _first_container(fqns, scope.exclude_fqns) is not None:
        return False
    return _first_container(fqns, scope.region_fqns) is not None


def _make_policy_fn(pass_config: PassConfig) -> Callable[[fx.Node], CheckpointPolicy]:

    def selective_policy(node: fx.Node) -> CheckpointPolicy:
        """Apply the shared selective-AC operator policy."""
        if node.target in force_recompute:
            return CheckpointPolicy.MUST_RECOMPUTE
        if node.target in must_save:
            return CheckpointPolicy.MUST_SAVE
        return CheckpointPolicy.PREFER_RECOMPUTE

def _unit_of(node: fx.Node, scope: _CheckpointScope) -> Optional[str]:
    """Return ``node``'s boundary unit (block or adapter region root), if any."""
    if not scope.unit_fqns:
        return None
    return _first_container(_node_stack_fqns(node), scope.unit_fqns)


def _tag_policy_pass(

    for node in graph_module.graph.nodes:
        if node.op != "call_function" or node.meta.get("autograd_backward", False):
            continue
        if node.target in _TRANSPARENT_OPS:
            parent = node.args[0]
            if isinstance(parent, fx.Node) and "recompute" in parent.meta:
                node.meta["recompute"] = parent.meta["recompute"]
            continue
        if not _is_eligible(node, scope):
            continue
        if _is_rng_op(node) or node.target in comm_ops:
            node.meta["recompute"] = CheckpointPolicy.MUST_SAVE

            node.meta["recompute"] = CheckpointPolicy.MUST_SAVE
            continue
        value = node.meta.get("val")
        if isinstance(value, torch.SymInt):
            node.meta["recompute"] = CheckpointPolicy.MUST_SAVE
            continue
        node.meta["recompute"] = policy_fn(node)


def _flatten_getitems(graph_module: fx.GraphModule) -> None:

def _is_mutable_op(target: Any) -> bool:
    """Return whether ``target`` is an in-place ``OpOverload``."""
    if not isinstance(target, OpOverload):
        return False
    return bool(getattr(getattr(target, "_schema", None), "is_mutable", False))


def _is_rng_op(node: fx.Node) -> bool:

    while stack:
        current = stack.pop()
        for user in current.users:
            if user in seen:
                continue
            seen.add(user)
            if user.op == "output" or user.meta.get("autograd_backward", False):
                continue
            if (

                user.op == "call_function"
                and user.target in _TRANSPARENT_OPS
                and user.meta.get("recompute") == policy
            ):
                stack.append(user)
                continue
            if _unit_of(user, scope) != node_unit:
                return True
    return False

    duplicates: Dict[fx.Node, fx.Node] = {}
    for arg in fwd_node.all_input_nodes:
        if arg.op != "get_attr":
            continue
        if not isinstance(getattr(graph_module, arg.target, None), fx.GraphModule):
            continue
        with graph_module.graph.inserting_before(bwd_target):
            dup_attr = graph_module.graph.get_attr(arg.target)
            dup_attr.meta.update(arg.meta)
        duplicates[arg] = dup_attr
    return duplicates


def _remat(graph_module: fx.GraphModule) -> int:

    def remat_arg(arg: Any, get_attr_dups: Dict[fx.Node, fx.Node]) -> Any:
        """Redirect duplicate args to their recompute copies."""
        if isinstance(arg, fx.Node):
            if arg in get_attr_dups:
                return get_attr_dups[arg]
            if arg in recomputed:
                return recomputed[arg]
        return arg

        Returns:
            The transformed graph module.
        """
        if pass_config.ac_mode == "off":
            return graph_module
        scope = _resolve_scope(kwargs.get("model"), pass_config)
        _tag_sac_policy(graph_module, pass_config, scope)
        duplicated = _remat(graph_module)
        _LOG.info(