Diff Coverage

        cls = self.__class__
        result = cls.__new__(cls)
        memo[id(self)] = result
        for k, v in self.__dict__.items():
            setattr(result, k, v if k == "_mesh" else copy.deepcopy(v, memo))
        return result

    @staticmethod
    def _process_placement_layout(obj, placements):

        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."""

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

    def __init__(self, fake_mode: Any, device: Any) -> None:
        """Store the FakeTensor mode and logical simulation device."""
        self.fake_mode = fake_mode
        self.device = device

    def mock(self, value: Any) -> Any:
        """Recursively replace tensor leaves with value-free FakeTensors."""
        import torch  # pylint: disable=C0415

        if isinstance(value, torch.Tensor):
            with self.fake_mode:
                return torch.empty_like(
                    value,
                    device=self.device,
                    memory_format=torch.preserve_format,
                )
        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()]
        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

                ) from exc

    def _mock_sequence(self, value: Any) -> Any:
        """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],

        vocab_size: int,
        tp_size: int,
) -> tuple[int, tuple[int, ...]]:
    """Derive valid-label ownership using the real loss-parallel vocab split."""
    import torch  # pylint: disable=C0415

    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(
        int(((valid_labels >= rank * chunk_size) &
             (valid_labels < min((rank + 1) * chunk_size, vocab_size))).sum().item())
        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)
class _DryRunMoERoutingPlan:

    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
        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."""
    experts = getattr(module, "experts", None)
    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 any(all(hasattr(experts, name) for name in layout) for layout in layouts)


class MoERoutingValueDependencyHandler(ValueDependencyHandler):
    """Validate and describe configured routed-expert replay."""

    paths = ("balanced", "hotspot", "explicit")

    def discover(self, model: Any, context: Any) -> list[str]:
        """Return supported routed-expert module FQNs."""
        del context
        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)
        allowed = {
            "balanced": set(),
            "hotspot": {"hotspot_expert"},
            "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
        ):
            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,

    """Parse, expand, and report unified value-dependency rules."""

    def __init__(self, config: Any) -> None:
        """Parse statically valid rules from ``dry_run.value_dependencies``."""
        self._config = config
        self._rules = self._parse_rules(config.value_dependencies)
        self._handlers = {
            name: factory()
            for name, factory in _VALUE_DEPENDENCY_HANDLER_FACTORIES.items()
        }
        self._resolved: Dict[str, Dict[str, _ValueDependencyRule]] = {}
        self._target_modules: Dict[str, Any] = {}
        self._resolved_moe_plans: Dict[str, Dict[str, Any]] = {}
        self._runtime_metadata: Dict[str, Any] = {}
        self._tp_cross_entropy_counts: Optional[tuple[int, tuple[int, ...]]] = None

    @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", {})
            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:
        """Expand all configured rules against one model and logical scopes."""
        self._target_modules = dict(model.named_modules())
        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)
            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 target in matches:
                candidates.setdefault((rule.handler, target), []).append((exact, rule))

        resolved: Dict[str, Dict[str, _ValueDependencyRule]] = {}
        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]
            self._handlers[handler_name].validate(target, rule.path, rule.inputs, context)
            resolved.setdefault(handler_name, {})[target] = rule
        self._resolved = resolved

    def project_model(self, model: Any, context: Any = None) -> None:
        """Keep validated value-dependency rules owned by one pipeline stage.

        ``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, {}))

    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."""
        return {
            target: {"routing_policy": rule.path, **rule.inputs}
            for target, rule in self._resolved.get("moe_routing", {}).items()
        }

    @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 {
            "rules": [
                {
                    "index": rule.index,
                    "match": rule.match,

        }

    def set_runtime_metadata(self, metadata: Dict[str, Any]) -> None:
        """Store JSON-safe runtime consumption details for the final report."""
        self._runtime_metadata = metadata

    def _moe_routing_matrix(
            self,
            layer: Dict[str, Any],

            expected: int,
            module_name: str,
    ) -> tuple[tuple[int, ...], ...]:
        """Build the configured or synthetic source-to-expert load matrix."""
        policy = layer["routing_policy"]
        if policy == "explicit":
            return self._explicit_moe_matrix(
                layer.get("source_expert_loads"),
                ep_size,
                num_experts,
                expected,

                expected,
                module_name,
            )

        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))

    @staticmethod
    def _local_expert_start(
            module_name: str,

            num_experts: int,
            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(
            ep_size: int,

            num_experts: int,
            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)
        )

            self, module_name: str, module: Any, ep_size: int, ep_rank: int,
            local_tokens: int, local_expert_count: int,
    ) -> _DryRunMoERoutingPlan:
        """Build aggregate EP source-to-expert load for one MoE layer."""
        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,
            ep_rank,
            num_experts,

            ep_rank,
            num_experts,
            local_expert_count,
        )
        local_loads = tuple(
            sum(row[index] for row in matrix)
            for index in range(start, start + local_expert_count)
        )
        received_loads = tuple(
            matrix[source_rank][expert_index]
            for source_rank in range(ep_size)
            for expert_index in range(start, start + local_expert_count)
        )
        expert_ranges = self._expert_ranges(ep_size, num_experts, local_expert_count)
        input_splits = tuple(
            sum(matrix[ep_rank][start_index:end_index])
            for start_index, end_index in expert_ranges
        )
        output_splits = tuple(
            sum(matrix[source_rank][start:start + local_expert_count])
            for source_rank in range(ep_size)
        )
        plan = _DryRunMoERoutingPlan(
            matrix,
            local_loads,
            received_loads,
            input_splits,

            input_splits,
            output_splits,
            sum(matrix[ep_rank]),
        )
        self._resolved_moe_plans[module_name] = {
            "source_expert_loads": [list(row) for row in plan.source_expert_loads],
            "local_expert_loads": list(plan.local_expert_loads),
            "received_expert_loads": list(plan.received_expert_loads),
            "input_split_sizes": list(plan.input_split_sizes),

            "input_split_sizes": list(plan.input_split_sizes),
            "output_split_sizes": list(plan.output_split_sizes),
            "outgoing_tokens": plan.outgoing_tokens,
        }
        return plan

    @staticmethod
    def _explicit_moe_matrix(
            source_loads: Any, ep_size: int, num_experts: int,

            source_loads: Any, ep_size: int, num_experts: int,
            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(
                    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."""

    _SEMANTIC_HANDLERS = {"moe_routing"}

    def __init__(self, profile: _DryRunValueProfile, model: Any, runtime: DryRunRuntime) -> None:
        """Bind runtime tracking to one expanded value-dependency profile."""
        self._profile = profile
        self._model = model
        self._runtime = runtime
        self._base = None
        self._stack = ExitStack()
        self._hook_handles = []
        self._active_token = None
        self._consumed_decisions = set()
        self._operator_counts: Dict[tuple[Any, ...], int] = {}
        self._consumed_operator_mocks = set()
        self._module_by_name = dict(model.named_modules())

    @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
        return self

    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),
        )
        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

    def _install_handler_contexts(
            self, stack: ExitStack, stage: ValueDependencyStage, context: Any,
    ) -> None:

    def _install_handler_contexts(
            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
        try:
            if exc_type is None:
                self._validate_consumption()
        finally:
            self._profile.set_runtime_metadata({
                "branch": {
                    "consumed_decisions": sorted(
                        f"{target}:{name}" for target, name in self._consumed_decisions
                    ),