Diff Coverage

    from torchdata.stateful_dataloader import StatefulDataLoader as _TorchStatefulDataLoader
except ModuleNotFoundError as error:
    missing_module = str(error.name or "")
    if missing_module != "torchdata" and not missing_module.startswith("torchdata."):
        raise
    _TORCHDATA_IMPORT_ERROR = error
    _StatefulDataLoaderBase = _UnavailableStatefulDataLoader
else:
    _StatefulDataLoaderBase = _TorchStatefulDataLoader

logger = get_dataset_logger(__name__)

        Raises:
            RuntimeError: If the process was not launched by ``torchrun``.
            ValueError: If a launcher variable is outside its valid range.
        """
        required_names = ("RANK", "WORLD_SIZE", "LOCAL_RANK")
        missing_names = [name for name in required_names if name not in os.environ]
        if missing_names:
            raise RuntimeError(
                "Multi-rank HyperDryRun must be launched with torchrun; "
                f"missing environment variables: {missing_names}"
            )
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ["LOCAL_RANK"])
        if world_size < 1:
            raise ValueError(f"WORLD_SIZE must be >= 1, got {world_size}")
        if rank < 0 or rank >= world_size:
            raise ValueError(
                f"RANK must be in [0, {world_size}), got {rank}"
            )
        if local_rank < 0:
            raise ValueError(f"LOCAL_RANK must be >= 0, got {local_rank}")
        return cls(rank=rank, world_size=world_size, local_rank=local_rank)


class DryRunBatchMocker:
    """Erase one CPU training batch while preserving its complete structure."""

        if isinstance(value, Mapping):
            return self._mock_mapping(value)
        if isinstance(value, (list, tuple)):
            return self._mock_sequence(value)
        if is_dataclass(value) and not isinstance(value, type):
            return self._mock_dataclass(value)
        if hasattr(value, "__dict__") and self._contains_tensor(vars(value)):
            raise ValueError(
                "Dry-run cannot mock tensor-bearing batch carrier "
                f"{type(value).__module__}.{type(value).__qualname__}"
            )
        return value

    def _mock_mapping(self, value: Mapping[Any, Any]) -> Mapping[Any, Any]:
        """Mock mapping values while preserving the mapping carrier type."""
        mocked_items = [(key, self.mock(item)) for key, item in value.items()]

        """Mock mapping values while preserving the mapping carrier type."""
        mocked_items = [(key, self.mock(item)) for key, item in value.items()]
        if type(value) is dict:  # pylint: disable=unidiomatic-typecheck
            return dict(mocked_items)
        try:
            mocked_mapping = copy.copy(value)
            mocked_mapping.clear()
            mocked_mapping.update(mocked_items)
            return mocked_mapping
        except (AttributeError, TypeError):
            try:
                return type(value)(mocked_items)
            except TypeError as exc:
                raise ValueError(
                    "Dry-run cannot preserve batch mapping carrier "
                    f"{type(value).__module__}.{type(value).__qualname__}"
                ) from exc

        """Mock list and tuple values while preserving named tuples."""
        mocked = tuple(self.mock(item) for item in value)
        if isinstance(value, list):
            return list(mocked)
        if hasattr(value, "_fields"):
            return type(value)(*mocked)
        return mocked

    def _mock_dataclass(self, value: Any) -> Any:
        """Mock initialized and deferred dataclass fields."""
        updates = {
            field.name: self.mock(getattr(value, field.name))
            for field in fields(value)
            if field.init
        }
        mocked_dataclass = replace(value, **updates)
        for field in fields(value):
            if not field.init:
                object.__setattr__(
                    mocked_dataclass,
                    field.name,
                    self.mock(getattr(value, field.name)),
                )
        return mocked_dataclass

    def _contains_tensor(self, value: Any) -> bool:
        """Return whether an unsupported object recursively owns a tensor."""
        import torch  # pylint: disable=C0415

        if isinstance(value, torch.Tensor):
            return True
        if isinstance(value, Mapping):
            return any(self._contains_tensor(item) for item in value.values())
        if isinstance(value, (list, tuple)):
            return any(self._contains_tensor(item) for item in value)
        return False


def derive_tp_target_counts(
        loss_inputs: Mapping[str, Any],

    labels = loss_inputs.get("shift_labels")
    if labels is None:
        labels = loss_inputs.get("labels")
        if labels is None:
            raise ValueError("Dry-run loss_inputs must contain labels")
        labels = labels[..., 1:]
    if not isinstance(labels, torch.Tensor):
        raise ValueError("Dry-run labels must be a torch.Tensor")
    valid_mask = labels.ne(IGNORE_INDEX)
    loss_mask = loss_inputs.get("loss_mask")
    if isinstance(loss_mask, torch.Tensor):
        if loss_mask.shape != labels.shape:
            loss_mask = loss_mask[..., -labels.shape[-1]:]
        valid_mask = valid_mask & loss_mask.reshape_as(labels).bool()
    valid_labels = labels[valid_mask]
    chunk_size = (vocab_size + tp_size - 1) // tp_size
    owned = tuple(

        for rank in range(tp_size)
    )
    valid_count = int(valid_labels.numel())
    if sum(owned) != valid_count:
        raise ValueError("Dry-run labels contain token IDs outside the configured vocabulary")
    return valid_count, owned


@dataclass(frozen=True)

    paths: tuple[str, ...] = ()

    def discover(self, model: Any, context: Any) -> list[str]:
        """Return logical targets supported by this handler."""
        del model, context
        return []

    def recognizes(self, target: str, module: Any) -> bool:
        """Return whether this handler can explain an unconfigured target."""
        del target, module
        return False

    def validate(self, target: str, path: str, inputs: Dict[str, Any], context: Any) -> None:
        """Validate one expanded rule."""
        del target, inputs, context

    def validate(self, target: str, path: str, inputs: Dict[str, Any], context: Any) -> None:
        """Validate one expanded rule."""
        del target, inputs, context
        if self.paths and path not in self.paths:
            raise ValueError(
                f"value dependency handler {self.name!r} path must be one of "
                f"{self.paths}, got {path!r}"
            )

            self, stage: ValueDependencyStage, target: str, path: str,
            inputs: Dict[str, Any], context: Any,
    ) -> ContextManager[Any]:
        """Return scoped runtime state for one expanded rule."""
        del stage, target, path, inputs, context
        return nullcontext()

    def prepare_batch(
            self, target: str, path: str, inputs: Dict[str, Any],
            batch: Dict[str, Any], context: Any,

            self, target: str, path: str, inputs: Dict[str, Any],
            batch: Dict[str, Any], context: Any,
    ) -> None:
        """Apply shape-only batch changes for one expanded rule."""
        del target, path, inputs, batch, context

    def resolved_metadata(self) -> Dict[str, Any]:
        """Return handler-specific resolved report metadata."""
        return {}


_VALUE_DEPENDENCY_HANDLER_FACTORIES: Dict[str, Callable[[], ValueDependencyHandler]] = {}

    Raises:
        ValueError: If the name is empty, already registered, or inconsistent
            with the constructed handler.
    """
    if not isinstance(name, str) or not name.strip():
        raise ValueError("value dependency handler name must be a non-empty string")
    if name in _VALUE_DEPENDENCY_HANDLER_FACTORIES:
        raise ValueError(f"value dependency handler {name!r} is already registered")
    handler = factory()
    if not isinstance(handler, ValueDependencyHandler) or handler.name != name:
        raise ValueError(
            f"value dependency handler factory for {name!r} must return "
            f"ValueDependencyHandler(name={name!r})"
        )
    _VALUE_DEPENDENCY_HANDLER_FACTORIES[name] = factory


def _is_supported_moe_module(module: Any) -> bool:
    """Return whether a module exposes a supported routed-expert layout."""

    if experts is None or (
            getattr(module, "gate", None) is None
            and getattr(module, "router", None) is None
    ):
        return False
    layouts = (
        ("w1", "w2", "w3"),
        ("gate_up_proj", "down_proj"),
        ("gate_proj", "up_proj", "down_proj"),

        return [name for name, module in model.named_modules() if name and _is_supported_moe_module(module)]

    def recognizes(self, target: str, module: Any) -> bool:
        """Recognize supported routed-expert modules."""
        del target
        return _is_supported_moe_module(module)

    def validate(self, target: str, path: str, inputs: Dict[str, Any], context: Any) -> None:
        """Validate path-specific MoE rule inputs."""
        super().validate(target, path, inputs, context)

            "explicit": {"source_expert_loads"},
        }[path]
        unknown = sorted(set(inputs) - allowed)
        if unknown:
            raise ValueError(f"MoE rule for {target!r} contains unsupported inputs: {unknown}")
        if path == "hotspot" and (
                not isinstance(inputs.get("hotspot_expert"), int)
                or isinstance(inputs.get("hotspot_expert"), bool)
                or inputs["hotspot_expert"] < 0

                not isinstance(inputs.get("hotspot_expert"), int)
                or isinstance(inputs.get("hotspot_expert"), bool)
                or inputs["hotspot_expert"] < 0
        ):
            raise ValueError(f"MoE hotspot rule for {target!r} requires non-negative hotspot_expert")
        if path == "explicit" and not isinstance(inputs.get("source_expert_loads"), list):
            raise ValueError(f"MoE explicit rule for {target!r} requires source_expert_loads")

    def install(
            self, stage: ValueDependencyStage, target: str, path: str,
            inputs: Dict[str, Any], context: Any,

            self, stage: ValueDependencyStage, target: str, path: str,
            inputs: Dict[str, Any], context: Any,
    ) -> ContextManager[Any]:
        """Install routed-expert replay once for the configured fake step."""
        del target, path, inputs
        if stage is ValueDependencyStage.FAKE_STEP:
            return _MoERoutingValueDependencyRuntime(context.profile, context.base)
        return nullcontext()


class BranchValueDependencyHandler(ValueDependencyHandler):
    """Provide explicit values to instrumented Python branch decisions."""

    name = "branch"

    def discover(self, model: Any, context: Any) -> list[str]:
        """Expose module FQNs and logical execution regions."""
        del context
        return [name for name, _ in model.named_modules() if name] + ["loss", "backward", "optimizer"]

    def validate(self, target: str, path: str, inputs: Dict[str, Any], context: Any) -> None:
        """Validate named scalar branch decisions."""
        del path, context
        decisions = inputs.get("decisions")
        if set(inputs) != {"decisions"} or not isinstance(decisions, dict) or not decisions:
            raise ValueError(f"branch rule for {target!r} requires a non-empty decisions mapping")
        for name, value in decisions.items():
            if not isinstance(name, str) or not name.strip():
                raise ValueError(f"branch rule for {target!r} contains an invalid decision name")
            if not isinstance(value, (bool, int, float, str)) and value is not None:
                raise ValueError(
                    f"branch decision {name!r} for {target!r} must be a YAML scalar"
                )

    name = "operator_debug"

    def discover(self, model: Any, context: Any) -> list[str]:
        """Expose module FQNs and logical execution regions."""
        del context
        return [name for name, _ in model.named_modules() if name] + ["loss", "backward", "optimizer"]

    def validate(self, target: str, path: str, inputs: Dict[str, Any], context: Any) -> None:
        """Validate source anchors and mock return envelopes."""
        del path, context
        mocks = inputs.get("mocks")
        if set(inputs) != {"mocks"} or not isinstance(mocks, list) or not mocks:
            raise ValueError(f"operator_debug rule for {target!r} requires a non-empty mocks list")
        selectors = set()
        for index, mock in enumerate(mocks):
            location = f"operator_debug mock {index} for {target!r}"
            selector = _operator_debug_selector(mock, location)
            if selector in selectors:
                raise ValueError(f"{location} duplicates an earlier selector")
            selectors.add(selector)
            _validate_operator_return(mock["return"], location)


def _operator_debug_selector(mock: Any, location: str) -> tuple[Any, ...]:
    """Validate and return the stable selector for one operator mock."""
    if not isinstance(mock, dict) or set(mock) != {"source", "op", "occurrence", "return"}:
        raise ValueError(f"{location} must contain source, op, occurrence, and return")
    source = mock["source"]
    if not isinstance(source, dict) or set(source) != {"file", "function", "line"}:
        raise ValueError(f"{location}.source must contain file, function, and line")
    if (
            not isinstance(source["file"], str)
            or not isinstance(source["function"], str)
            or not isinstance(source["line"], int)
            or isinstance(source["line"], bool)

            or not isinstance(source["line"], int)
            or isinstance(source["line"], bool)
            or source["line"] < 1
    ):
        raise ValueError(f"{location} has an invalid source anchor")
    if not isinstance(mock["op"], str) or not mock["op"].startswith("aten."):
        raise ValueError(f"{location}.op must be an ATen overload name")
    occurrence = mock["occurrence"]
    if not isinstance(occurrence, int) or isinstance(occurrence, bool) or occurrence < 0:
        raise ValueError(f"{location}.occurrence must be a non-negative integer")
    return source["file"], source["function"], source["line"], mock["op"], occurrence


def _validate_operator_return(spec: Any, location: str) -> None:
    """Validate one recursive operator-debug return specification."""
    if not isinstance(spec, dict) or len(spec) != 1:
        raise ValueError(f"{location}.return must contain exactly one return kind")
    kind, value = next(iter(spec.items()))
    if kind == "scalar":
        if not isinstance(value, (bool, int, float)):
            raise ValueError(f"{location}.return.scalar must be bool, int, or float")
        return
    if kind == "tensor":
        _validate_operator_tensor_return(value, location)
        return
    if kind in ("tuple", "list"):
        _validate_operator_sequence_return(value, kind, location)
        return
    if kind == "by_global_rank":
        _validate_operator_rank_return(value, location)
        return
    raise ValueError(f"{location}.return uses unsupported kind {kind!r}")


def _validate_operator_tensor_return(value: Any, location: str) -> None:
    """Validate a tensor-shaped operator-debug return."""
    if not isinstance(value, dict) or set(value) != {"shape", "dtype"}:
        raise ValueError(f"{location}.return.tensor must contain shape and dtype")
    shape = value["shape"]
    if not isinstance(shape, list) or any(
            not isinstance(size, int) or isinstance(size, bool) or size < 0 for size in shape
    ):
        raise ValueError(f"{location}.return.tensor.shape must contain non-negative integers")
    if not isinstance(value["dtype"], str) or not value["dtype"]:
        raise ValueError(f"{location}.return.tensor.dtype must be a non-empty string")


def _validate_operator_sequence_return(value: Any, kind: str, location: str) -> None:
    """Validate a tuple- or list-shaped operator-debug return."""
    if not isinstance(value, list):
        raise ValueError(f"{location}.return.{kind} must be a list")
    for index, item in enumerate(value):
        _validate_operator_return(item, f"{location}.return.{kind}[{index}]")


def _validate_operator_rank_return(value: Any, location: str) -> None:
    """Validate a rank-indexed operator-debug return."""
    if not isinstance(value, dict) or not value:
        raise ValueError(f"{location}.return.by_global_rank must be a non-empty mapping")
    for rank, item in value.items():
        if not isinstance(rank, (int, str)) or not str(rank).isdigit():
            raise ValueError(f"{location}.return.by_global_rank keys must be non-negative ranks")
        _validate_operator_return(item, f"{location}.return.by_global_rank[{rank}]")


for _builtin_handler in (
        MoERoutingValueDependencyHandler,

    @staticmethod
    def _parse_rules(value_dependencies: Any) -> tuple[_ValueDependencyRule, ...]:
        """Validate the common rule envelope before model discovery."""
        if not isinstance(value_dependencies, dict):
            raise ValueError("dry_run.value_dependencies must be a mapping")
        unknown = sorted(set(value_dependencies) - {"rules"})
        if unknown:
            raise ValueError(f"dry_run.value_dependencies contains unknown fields: {unknown}")
        raw_rules = value_dependencies.get("rules", [])
        if not isinstance(raw_rules, list):
            raise ValueError("dry_run.value_dependencies.rules must be a list")
        rules = []
        for index, raw_rule in enumerate(raw_rules):
            location = f"dry_run.value_dependencies.rules[{index}]"
            if not isinstance(raw_rule, dict):
                raise ValueError(f"{location} must be a mapping")
            unknown_rule_fields = sorted(
                set(raw_rule) - {"match", "handler", "path", "inputs", "optional"}
            )
            if unknown_rule_fields:
                raise ValueError(f"{location} contains unknown fields: {unknown_rule_fields}")
            match = raw_rule.get("match")
            handler = raw_rule.get("handler")
            path = raw_rule.get("path")
            inputs = raw_rule.get("inputs", {})

            path = raw_rule.get("path")
            inputs = raw_rule.get("inputs", {})
            optional = raw_rule.get("optional", False)
            if not isinstance(match, str) or not match.strip():
                raise ValueError(f"{location}.match must be a non-empty string")
            if not isinstance(handler, str) or not handler.strip():
                raise ValueError(f"{location}.handler must be a non-empty string")
            if not isinstance(path, str) or not path.strip():
                raise ValueError(f"{location}.path must be a non-empty string")
            if not isinstance(inputs, dict):
                raise ValueError(f"{location}.inputs must be a mapping")
            if not isinstance(optional, bool):
                raise ValueError(f"{location}.optional must be a boolean")
            rules.append(_ValueDependencyRule(index, match, handler, path, inputs, optional))
        return tuple(rules)

    def bind_model(self, model: Any, context: Any = None) -> None:

        candidates: Dict[tuple[str, str], list[tuple[int, _ValueDependencyRule]]] = {}
        for rule in self._rules:
            handler = self._handlers.get(rule.handler)
            if handler is None:
                raise ValueError(
                    f"value dependency rule {rule.index} uses unknown handler {rule.handler!r}; "
                    f"registered handlers: {sorted(self._handlers)}"
                )
            supported_targets = handler.discover(model, context)

                )
            supported_targets = handler.discover(model, context)
            matches = [target for target in supported_targets if fnmatchcase(target, rule.match)]
            if not matches and not rule.optional:
                raise ValueError(
                    f"value dependency rule {rule.index} match {rule.match!r} did not match "
                    f"any target supported by handler {rule.handler!r}"
                )
            exact = int(rule.match in matches)

        for (handler_name, target), target_candidates in candidates.items():
            best_priority = max(priority for priority, _ in target_candidates)
            best_rules = [rule for priority, rule in target_candidates if priority == best_priority]
            if len(best_rules) > 1:
                indexes = [rule.index for rule in best_rules]
                raise ValueError(
                    f"value dependency rules {indexes} conflict for target {target!r} "
                    f"and handler {handler_name!r}"
                )
            rule = best_rules[0]

        ``bind_model`` must run on the complete model before this projection.
        Rules are intentionally not glob-expanded again: a valid global target
        may be owned by a different pipeline rank.
        """
        self._target_modules = dict(model.named_modules())
        projected = {}
        for handler_name, targets in self._resolved.items():
            handler = self._handlers[handler_name]
            supported_targets = set(handler.discover(model, context))
            local_targets = {}
            for target, rule in targets.items():
                for local_target in supported_targets:
                    if local_target == target or local_target.endswith(f".{target}"):
                        local_targets[local_target] = rule
            if local_targets:
                projected[handler_name] = local_targets
        self._resolved = projected

    def resolved_rules(self, handler: str) -> Dict[str, _ValueDependencyRule]:
        """Return expanded rules for one handler."""
        return dict(self._resolved.get(handler, {}))

        return dict(self._resolved.get(handler, {}))

    def handler(self, name: str) -> ValueDependencyHandler:
        """Return one fresh profile-owned handler instance."""
        return self._handlers[name]

    @property
    def moe_layers(self) -> dict[str, Any]:
        """Return per-layer routing overrides."""

    @property
    def moe_enabled(self) -> bool:
        """Return whether the MoE adapter is enabled."""
        return bool(self._resolved.get("moe_routing"))

    def metadata(self) -> Dict[str, Any]:
        """Return a JSON-safe profile summary for the memory report."""
        return {

            )

        hotspot = int(layer.get("hotspot_expert", 0))
        if policy == "hotspot" and not 0 <= hotspot < num_experts:
            raise ValueError(f"MoE hotspot_expert for {module_name!r} must be in [0, {num_experts})")
        if policy == "hotspot":
            row = tuple(expected if index == hotspot else 0 for index in range(num_experts))
        else:
            base, remainder = divmod(expected, num_experts)
            row = tuple(base + int(index < remainder) for index in range(num_experts))
        return tuple(row for _ in range(ep_size))

            local_expert_count: int,
    ) -> int:
        """Validate the expert layout and return this rank's first expert."""
        if local_expert_count < 1:
            raise ValueError(f"MoE {module_name!r} has no local experts")
        if local_expert_count == num_experts:
            return 0
        if num_experts % ep_size or local_expert_count != num_experts // ep_size:
            raise ValueError(f"MoE {module_name!r} local expert layout is incompatible with EP={ep_size}")
        return ep_rank * local_expert_count

    @staticmethod
    def _expert_ranges(

            local_expert_count: int,
    ) -> tuple[tuple[int, int], ...]:
        """Return the global expert range owned by each destination rank."""
        if local_expert_count == num_experts:
            return ((0, num_experts),)
        return tuple(
            (rank * local_expert_count, (rank + 1) * local_expert_count)
            for rank in range(ep_size)
        )

        num_experts = int(module.num_experts)
        expected = local_tokens * int(module.top_k)
        layer = self.moe_layers.get(module_name)
        if layer is None:
            raise ValueError(f"MoE module {module_name!r} has no value-dependency rule")
        matrix = self._moe_routing_matrix(layer, ep_size, num_experts, expected, module_name)
        start = self._local_expert_start(
            module_name,
            ep_size,

            expected: int, module_name: str,
    ) -> tuple[tuple[int, ...], ...]:
        """Validate the configured aggregate EP traffic matrix."""
        if not isinstance(source_loads, list) or len(source_loads) != ep_size:
            raise ValueError(f"MoE explicit source_expert_loads for {module_name!r} must have {ep_size} rows")
        matrix = []
        for source_rank, row in enumerate(source_loads):
            if not isinstance(row, list) or len(row) != num_experts:
                raise ValueError(
                    f"MoE explicit source_expert_loads row {source_rank} for {module_name!r} "
                    f"must have {num_experts} entries"
                )
            if any(

            if any(
                    not isinstance(count, int) or isinstance(count, bool) or count < 0
                    for count in row
            ):
                raise ValueError("MoE explicit source_expert_loads entries must be non-negative integers")
            if sum(row) != expected:
                raise ValueError(
                    f"MoE explicit source_expert_loads row {source_rank} for {module_name!r} must sum to "
                    f"local_tokens * top_k ({expected}), got {sum(row)}"
                )
            matrix.append(tuple(row))

        return tuple(matrix)

    def tp_cross_entropy_target_count(self, target_count: int, tp_size: int, tp_rank: int) -> tuple[int, int]:
        """Return CPU-derived valid and local-owned target counts for TP CE."""
        counts = self._tp_cross_entropy_counts
        if counts is None:
            raise ValueError("TP cross-entropy counts were not derived from the training batch")
        valid, owned = counts
        if len(owned) != tp_size or not 0 <= tp_rank < tp_size:
            raise ValueError(
                f"CPU-derived TP target counts contain {len(owned)} ranks, expected {tp_size}"
            )
        if valid > target_count:
            raise ValueError(
                f"CPU-derived valid token count {valid} exceeds FakeTensor target size {target_count}"
            )
        return valid, owned[tp_rank]

    def configure_tp_cross_entropy_counts(
            self,
            valid_token_count: int,

            valid_token_count: int,
            target_tokens_per_rank: tuple[int, ...],
    ) -> None:
        """Install token ownership computed from the real CPU training batch."""
        if valid_token_count < 0 or sum(target_tokens_per_rank) != valid_token_count:
            raise ValueError("CPU-derived TP target counts must be non-negative and sum to the valid count")
        self._tp_cross_entropy_counts = (valid_token_count, target_tokens_per_rank)


_T = TypeVar("_T")
_ACTIVE_VALUE_DEPENDENCY_MANAGER: ContextVar[Optional["ValueDependencyManager"]] = ContextVar(

    Raises:
        ValueError: If ``name`` is empty or ``default_factory`` is not callable.
    """
    if not isinstance(name, str) or not name.strip():
        raise ValueError("value dependency decision name must be a non-empty string")
    if not callable(default_factory):
        raise ValueError("value dependency decision default_factory must be callable")
    manager = _ACTIVE_VALUE_DEPENDENCY_MANAGER.get()
    if manager is not None:
        configured, value = manager.resolve_decision(name)
        if configured:
            return value
    return default_factory()


class UnconfiguredValueDependencyError(RuntimeError):
    """Actionable error for an unconfigured FakeTensor value dependency."""

    @contextmanager
    def parallelize_context(self) -> Iterator[None]:
        """Install value dependencies required while parallelization is applied."""
        with ExitStack() as stack:
            self._install_handler_contexts(
                stack,
                ValueDependencyStage.PARALLELIZE,
                SimpleNamespace(profile=self._profile, model=self._model, runtime=self._runtime),
            )
            yield

    def fake_step_context(self, base: Any) -> "ValueDependencyManager":
        """Prepare this manager to enter the complete fake-step scope."""
        self._base = base

    def __enter__(self) -> "ValueDependencyManager":
        """Install semantic overrides, module scopes, and dispatch interception."""
        if self._base is None:
            raise RuntimeError("ValueDependencyManager.fake_step_context must be configured before entry")
        self._install_handler_contexts(
            self._stack,
            ValueDependencyStage.FAKE_STEP,
            SimpleNamespace(profile=self._profile, base=self._base, runtime=self._runtime),

            SimpleNamespace(profile=self._profile, base=self._base, runtime=self._runtime),
        )
        mesh = getattr(self._base, "mesh", None)
        if bool(getattr(mesh, "loss_parallel", False)):
            self._stack.enter_context(_TPCrossEntropyValueDependencyRuntime(self._profile))
        self._install_module_hooks()
        self._active_token = _ACTIVE_VALUE_DEPENDENCY_MANAGER.set(self)
        self._stack.enter_context(self._build_operator_debug_mode())
        return self

            self, stack: ExitStack, stage: ValueDependencyStage, context: Any,
    ) -> None:
        """Install semantic handlers once and external handlers per expanded target."""
        for handler_name, rules in self._profile._resolved.items():
            if handler_name in ("branch", "operator_debug") or not rules:
                continue
            handler = self._profile.handler(handler_name)
            selected_rules = [next(iter(rules.items()))] if handler_name in self._SEMANTIC_HANDLERS else rules.items()
            for target, rule in selected_rules:
                stack.enter_context(handler.install(stage, target, rule.path, rule.inputs, context))

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
        """Restore every scope and reject unused configured decisions/mocks."""
        del traceback

            _VALUE_DEPENDENCY_TARGET_STACK.reset(token)

    def prepare_batch(self, batch: Dict[str, Any], context: Any) -> None:
        """Run configured handler batch preparation hooks."""
        for handler_name, rules in self._profile._resolved.items():
            handler = self._profile.handler(handler_name)
            for target, rule in rules.items():
                handler.prepare_batch(target, rule.path, rule.inputs, batch, context)

    def resolve_decision(self, name: str) -> tuple[bool, Any]:
        """Resolve a configured branch decision for the active innermost target."""
        target = self._current_target()
        if target is None:
            return False, None
        rule = self._profile.resolved_rules("branch").get(target)
        if rule is None or name not in rule.inputs["decisions"]:
            return False, None
        self._consumed_decisions.add((target, name))
        return True, rule.inputs["decisions"][name]

    def _install_module_hooks(self) -> None:
        """Track the innermost executing module with exception-safe hooks."""
        for module_name, module in self._model.named_modules():

            self._hook_handles.append(module.register_forward_pre_hook(_pre_hook))
            try:
                handle = module.register_forward_hook(_post_hook, always_call=True)
            except TypeError:
                handle = module.register_forward_hook(_post_hook)
            self._hook_handles.append(handle)

    def _clear_target_operator_counts(self, target: str) -> None:
        """Reset occurrence counters for one new target invocation."""

    @staticmethod
    def _source_anchor() -> Dict[str, Any]:
        """Resolve the first user-code frame for the current ATen dispatch."""
        this_file = os.path.abspath(__file__)
        for frame_info in inspect.stack()[2:]:
            filename = os.path.abspath(frame_info.filename)
            if filename == this_file or f"{os.sep}site-packages{os.sep}torch{os.sep}" in filename:
                continue
            owner = frame_info.frame.f_locals.get("self")
            function = frame_info.function
            if owner is not None:
                function = f"{type(owner).__qualname__}.{function}"
            return {"file": frame_info.filename, "function": function, "line": frame_info.lineno}
        return {"file": "<unknown>", "function": "<unknown>", "line": 0}

    @staticmethod
    def _source_matches(configured: Dict[str, Any], actual: Dict[str, Any]) -> bool:
        """Return whether an actual frame matches a strict configured anchor."""
        return (
            actual["file"].endswith(configured["file"])
            and actual["function"] == configured["function"]
            and actual["line"] == configured["line"]
        )

        target = self._current_target()
        op_name = str(func)
        rule = self._profile.resolved_rules("operator_debug").get(target or "")
        if rule is not None:
            matching_op_mocks = [mock for mock in rule.inputs["mocks"] if mock["op"] == op_name]
            if matching_op_mocks:
                source = self._source_anchor()
                anchored = [
                    mock for mock in matching_op_mocks
                    if self._source_matches(mock["source"], source)
                ]
                if anchored:
                    source_key = (source["file"], source["function"], source["line"])
                    count_key = (target, *source_key, op_name)
                    occurrence = self._operator_counts.get(count_key, 0)
                    self._operator_counts[count_key] = occurrence + 1
                    selected = [mock for mock in anchored if mock["occurrence"] == occurrence]
                    if len(selected) != 1:
                        raise ValueError(
                            f"operator_debug target {target!r} op {op_name} source {source} "
                            f"has no unique mock for occurrence {occurrence}"
                        )
                    mock = selected[0]
                    self._consumed_operator_mocks.add((target, id(mock)))
                    return self._build_operator_return(mock["return"], args, kwargs)
        try:
            return func(*args, **kwargs)
        except (DataDependentOutputException, DynamicOutputShapeException) as error:
            source = self._source_anchor()
            count_key = (target, source["file"], source["function"], source["line"], op_name)
            occurrence = self._operator_counts.get(count_key, 0)
            self._operator_counts[count_key] = occurrence + 1
            raise self._unconfigured_error(target, op_name, args, source, occurrence) from error

    def _build_operator_return(self, spec: Dict[str, Any], args: Any, kwargs: Dict[str, Any]) -> Any:
        """Construct one recursive operator-debug return value."""
        import torch  # pylint: disable=C0415

        kind, value = next(iter(spec.items()))
        if kind == "scalar":
            return value
        if kind == "by_global_rank":
            rank_key = str(self._runtime.rank)
            selected = value.get(self._runtime.rank, value.get(rank_key))
            if selected is None:
                raise ValueError(
                    f"operator_debug return has no by_global_rank entry for rank {self._runtime.rank}"
                )
            return self._build_operator_return(selected, args, kwargs)
        if kind in ("tuple", "list"):
            built = [self._build_operator_return(item, args, kwargs) for item in value]
            return tuple(built) if kind == "tuple" else built
        dtype = getattr(torch, value["dtype"], None)
        if not isinstance(dtype, torch.dtype):
            raise ValueError(f"operator_debug uses unknown torch dtype {value['dtype']!r}")
        device = self._first_tensor_device((args, kwargs))
        return torch.empty(tuple(value["shape"]), dtype=dtype, device=device)

    @staticmethod
    def _first_tensor_device(value: Any) -> Any:
        """Return the first tensor device in a nested operator argument tree."""
        import torch  # pylint: disable=C0415

        if isinstance(value, torch.Tensor):
            return value.device
        if isinstance(value, dict):
            values = value.values()
        elif isinstance(value, (tuple, list)):
            values = value
        else:
            return torch.device("cpu")
        for item in values:
            device = ValueDependencyManager._first_tensor_device(item)
            if device.type != "cpu" or isinstance(item, torch.Tensor):
                return device
        return torch.device("cpu")

    def _unconfigured_error(
            self, target: Optional[str], op_name: str, args: Any,
            source: Dict[str, Any], occurrence: int,