Diff Coverage

                f"ac_region_source must be one of {AC_REGION_SOURCES}, "
                f"got {self.ac_region_source!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

    """Build a scope from the adapter-declared recompute regions."""
    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"
        )
    excluded = _select_excluded_targets(module_by_fqn, roots, recompute_spec)
    return _CheckpointScope(

        if pass_config.ac_mode == "full":
            recompute_spec = _adapter_recompute_spec(model, required=explicit_regions)
            if recompute_spec is not None:
                if not recompute_spec.region_patterns:
                    raise ValueError(
                        "model adapter recompute requires "
                        "RecomputePolicy.region_patterns"
                    )
                return _adapter_scope(model, block_fqns, block_fqns, recompute_spec)

                "activation_checkpoint.selection.source=model_adapter_safe_regions "
                "requires ModelAdapterSpec.recompute"
            )
        return _CheckpointScope(region_fqns=selected_layers, unit_fqns=block_fqns)
    if not recompute_spec.region_patterns:
        if explicit_regions:
            raise ValueError(
                "model adapter recompute requires RecomputePolicy.region_patterns"
            )
        return _CheckpointScope(region_fqns=selected_layers, unit_fqns=block_fqns)
    return _adapter_scope(model, selected_layers, block_fqns, recompute_spec)


def prepare_activation_checkpoint(
    model: Optional[torch.nn.Module], pass_config: PassConfig

    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:

    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 block unit, 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:

def _mutated_arg_nodes(node: fx.Node) -> List[fx.Node]:
    """Return the input nodes whose tensors ``node`` overwrites in place."""
    schema = getattr(node.target, "_schema", None)
    if schema is None:
        return []
    mutated: List[fx.Node] = []
    for position, schema_arg in enumerate(schema.arguments):
        alias_info = getattr(schema_arg, "alias_info", None)
        if alias_info is None or not alias_info.is_write:

    """
    if node.op != "call_function":
        return []
    if node.target is operator.getitem:
        producer = node.args[0] if node.args else None
        return [producer] if isinstance(producer, fx.Node) else []
    schema = getattr(node.target, "_schema", None)
    if schema is None:
        return []
    result_names: Set[str] = set()
    for result in schema.returns:
        alias_info = getattr(result, "alias_info", None)
        if alias_info is not None:

            or wildcard_result
            or "*" in arg_names
            or bool(arg_names & result_names)
        ):
            continue
        arg = (
            node.args[position]
            if position < len(node.args)
            else node.kwargs.get(schema_arg.name)

        if node.meta.get("autograd_backward", False):
            continue
        for arg in _aliased_input_nodes(node):
            if arg.meta.get("autograd_backward", False):
                continue
            neighbors.setdefault(node, set()).add(arg)
            neighbors.setdefault(arg, set()).add(node)
    return neighbors

    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

    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

        Returns:
            The transformed graph module.
        """
        if pass_config.ac_mode == "off":
            return graph_module
        scope = kwargs.get("ac_scope")
        if scope is None:
            scope = prepare_activation_checkpoint(kwargs.get("model"), pass_config)
        already_rematted = any(