Diff Coverage

        _prepare_model_source_compat(context["model_spec"])
        with _codegen_model_init_context():
            if factory is not None:
                if model_args:
                    raise ValueError("model_factory uses config and context, not additional positional model arguments")
                model = auto_model.construct_family_model(factory, hf_config, context=context)
            else:
                _, model = auto_model._init_model(  # pylint: disable=protected-access
                    cls, None, hf_config, attn_implementation, torch_dtype, is_hf_model,

    """Check complete, unique whitelist entries before registration."""
    defaults = _validate_selector_specs(family)
    _validate_variant_paths(family)
    if family.parallel_axis not in (None, "cp", "ep"):
        raise ValueError("parallel_axis must be cp, ep or None")
    if family.default_variant is not None:
        if family.default_variant not in family.variants:
            raise ValueError("default_variant must name a declared concrete variant")
        if "parallel" in family.selector_specs:
            raise ValueError("parallel is derived from context, not a public selector")
        defaults = {**defaults, "parallel": False}
    conditions = [defaults]  # The default combination is an explicit baseline allowance.
    for rule in family.variant_rules:
        if not isinstance(rule, VariantRule) or not isinstance(rule.when, Mapping):

        for name, value in rule.when.items():
            if name == "parallel" and family.default_variant is not None:
                actual_type = type(value)
                if actual_type is not bool:
                    raise TypeError("derived parallel conditions must be bool")
            else:
                _validate_selector_value(family, name, value)
        if rule.variant not in family.variants:
            raise ValueError(f"ModuleFamily {family.name}: undeclared variant {rule.variant!r}")

    def __post_init__(self) -> None:
        family = get_module_family(self.target)
        if family is None or family.default_variant is None:
            raise ValueError("FamilyRequest requires a family with a concrete default_variant")
        _validate_deferred_arguments(family, self.kwargs)
        object.__setattr__(self, "kwargs", MappingProxyType(deepcopy(dict(self.kwargs))))

                raise ValueError(f"{axis}_size is required when {axis} is active")
            size = 1
        actual_type = type(size)
        if actual_type is not int or size < 1:
            raise ValueError(f"{axis}_size must be a positive integer")
        result[axis] = size
    return MappingProxyType(result)

        if getattr(mesh, "device_mesh", None) is None:
            return result
    device_mesh = getattr(mesh, "device_mesh", None)
    if device_mesh is None:
        raise ValueError("mesh requires a prebuilt device_mesh")
    previous = result.get("mesh")
    if previous is not None and not _same_mesh(previous, mesh):
        raise ValueError("context and mesh belong to different jobs")
    for axis in ("tp", "cp", "ep", "pp"):

    args = deepcopy(dict(request.kwargs))
    parallel = family.parallel_axis is not None and axes[family.parallel_axis] > 1
    constructor, variant, selectors = _select_variant(family, args, parallel=parallel)
    if constructor is family.baseline:
        raise ValueError("a concrete family member cannot be its public factory")
    return ResolvedFamily(
        family, variant, constructor, MappingProxyType(args), MappingProxyType(selectors), axes,
    )

    """Bind source or config construction arguments accepted by the selected member."""
    signature = inspect.signature(member.constructor)
    required = "config" if module is None else "module"
    if required not in signature.parameters:
        raise TypeError(f"member {member.variant!r} must explicitly accept {required}")
    kwargs = dict(member.arguments)
    unknown = set(kwargs) - set(signature.parameters)
    if unknown:
        raise TypeError(f"member {member.variant!r} does not accept {sorted(unknown)}")
    runtime = {"config": config, "module": module, "module_fqn": fqn, "context": context}
    kwargs.update({name: value for name, value in runtime.items() if name in signature.parameters})
    signature.bind(**kwargs)
    return kwargs

def _normalize_family_request(target, kwargs):
    """Resolve a public name without selecting a topology-dependent member."""
    if isinstance(target, str):
        matches = [family.baseline for family in _FAMILIES.values() if family.name == target]
        if len(matches) != 1:
            raise ValueError(f"family name {target!r} must identify exactly one registration")
        target = matches[0]
    if isinstance(target, FamilyRequest):
        if kwargs:
            raise ValueError("request arguments cannot be combined with additional keywords")
        return target
    return FamilyRequest(target, kwargs)

        member.family.validate(config=config, module=module, context=context, member=member)
    config = deepcopy(config)
    result = member.constructor(**_constructor_arguments(member, config, module, module_fqn, context))
    if not isinstance(result, nn.Module):
        raise TypeError("a family constructor must return nn.Module")
    info = FamilyBuildInfo(request, member, deepcopy(config), "direct" if module is None else "replacement")
    result._hp_family_build_info = info  # pylint: disable=protected-access
    if config is not None and callable(getattr(result, "reset_parameters", None)):
        result._hp_reset_after_materialization = True  # pylint: disable=protected-access

    """
    request = _normalize_family_request(target, kwargs)
    arguments = dict(request.kwargs)
    if "config" in arguments:
        if config is not None:
            raise ValueError("config was supplied twice")
        config = arguments.pop("config")
        request = FamilyRequest(request.target, arguments)
    if (config is None) == (module is None):
        raise ValueError("provide exactly one of config or module")
    context = _context_from_mesh(mesh, context)
    member = resolve_family(request, context=context)

        if layout is not None:
            for name, shape in layout.parameter_shapes.items():
                identity = id(boundary.module.get_parameter(name))
                if by_identity.setdefault(identity, shape) != shape:
                    raise ValueError(f"shared family parameter {name!r} has conflicting logical shapes")
    return {name: by_identity.get(id(value), tuple(value.shape))
            for name, value in model.named_parameters(remove_duplicate=False)}

    """
    axes = _axis_sizes(_context_from_mesh(None, context))
    for boundary in collect_family_modules(model):
        if boundary.info.layout is not None and boundary.info.layout.applied:
            _context_from_mesh(boundary.info.layout.mesh, context)
        member = boundary.info.resolved
        for axis in ("cp", "ep"):
            if member.axis_sizes[axis] != axes[axis]:
                raise ValueError(

    if any(name not in original.params or named != original.params[name] for name, named in spec.params.items()):
        raise ValueError(f"layout at {fqn!r} conflicts with its constructed family layout")
    for field_name in ("in_src", "in_dst", "out_src", "out_dst", "out_names", "is_boundary"):
        if getattr(spec, field_name) != getattr(original, field_name):
            raise ValueError(f"{field_name} at {fqn!r} conflicts with its constructed family layout")


def _record_family_layouts(model, plan, mesh, shapes, source_shard_info, *, applied=True):
    """Retain relative declarations once for constructor and model-level application."""

    """Retain relative declarations once for constructor and model-level application."""
    for boundary in collect_family_modules(model):
        fqn = boundary.fqns[0]
        if fqn not in plan.modules or boundary.info.resolved.sharding is None:
            continue
        existing = boundary.info.layout
        if existing is not None and existing.applied:
            continue
        prefix = fqn + "." if fqn else ""

               if boundary.info.layout is not None and boundary.info.layout.applied]
    if not hasattr(mesh, "device_mesh"):
        mesh = next((layout.mesh for layout in layouts if layout.mesh.device_mesh is mesh), mesh)
    if any(not _same_mesh(layout.mesh, mesh) for layout in layouts):
        raise ValueError("constructed families cannot be applied with a different job mesh")
    return mesh

def _optimizer_tensor_result(tensor: DTensor, stores_dtensor: bool) -> Union[torch.Tensor, DTensor]:
    """Keep global layout only when the target optimizer stores DTensor values."""
    return tensor if stores_dtensor else tensor.to_local()


def _convert_input_tensor_to_target(
    tensor: torch.Tensor,

    target_device = param.to_local().device

    if isinstance(tensor, DTensor):
        if tensor.device_mesh == mesh and tensor.placements == placements:
            return _optimizer_tensor_result(tensor, stores_dtensor)
        redistributed = tensor.redistribute(mesh, placements)
        return _optimizer_tensor_result(redistributed, stores_dtensor)

    if full_state_dict:
        if cpu_offload:
            tensor = tensor.to(target_device)

    if full_state_dict:
        if cpu_offload:
            tensor = tensor.to(target_device)
        dt = distribute_tensor(tensor, mesh, placements)
        return _optimizer_tensor_result(dt, stores_dtensor)

    if cpu_offload and tensor.device != target_device:
        tensor = tensor.to(target_device)

    cpu_offload: bool,
) -> Optional[Tuple[str, str, torch.Tensor]]:
    """Parse a state.* metadata entry and return (fqn, state_key, tensor) or None."""
    # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.
    from hyper_parallel.core.distributed_checkpoint.metadata import (  # pylint: disable=import-outside-toplevel
        TensorStorageMetadata,
    )

    remainder = meta_key[len("state."):]

    meta_val: Any,
) -> Optional[Tuple[int, str, torch.Tensor]]:
    """Parse a param_group.* metadata entry and return (group_idx, field_name, tensor) or None."""
    # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.
    from hyper_parallel.core.distributed_checkpoint.metadata import (  # pylint: disable=import-outside-toplevel
        TensorStorageMetadata,
    )

    remainder = meta_key[len("param_group."):]

        state.<FQN>.<state_key>  -> TensorStorageMetadata
        param_group.<FQN>.<field> -> TensorStorageMetadata or BytesStorageMetadata
    """
    # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.
    from hyper_parallel.core.distributed_checkpoint.metadata import (  # pylint: disable=import-outside-toplevel
        TensorStorageMetadata,
    )

    flat: Dict[str, Any] = {}

    if cp_mesh is None or cp_mesh.size() <= 1:
        return
    for fqn, spec in plan.modules.items():
        if fqn in plan.parallel_modules:
            continue
        if not spec.is_boundary or not getattr(spec, "_needs_cp_attn", False):
            continue
        if spec.inner_wrapper is not None or spec.local_compute_fn is not None:
            continue

            continue
        module = _resolve_module(model, module_fqn)
        family = get_family_build_info(module)
        if family is not None and family.layout is not None:
            continue
        boundary = PrecompiledBoundary(
            spec,
            mesh,
            mesh_dim_names,

def _module_parallel(module, member):
    provider = getattr(module, "parallel_module_spec", None)
    declaration = provider() if provider is not None else None
    if declaration is not None and not isinstance(declaration, ParallelModuleSpec):
        raise TypeError("parallel_module_spec() must return ParallelModuleSpec")
    if member is not None and member.selectors["parallel"] and declaration is None:
        raise ValueError("a parallel family member must declare its explicit execution protocol")
    return declaration


def _add_provider(entries, model, fqn, sharding, parallel):

def _add_provider(entries, model, fqn, sharding, parallel):
    module = model.get_submodule(fqn)
    identity = id(module)
    if identity in entries:
        raise ValueError(f"multiple layout/communication owners for module {fqn!r}")
    if not isinstance(sharding, ModuleShardingSpec):
        raise TypeError("layout providers must return ModuleShardingSpec")
    sharding = copy.deepcopy(sharding)
    _normalize_out_fields(sharding)
    entries[identity] = (fqn, module, sharding, parallel)

    all_modules = tuple(model.named_modules(remove_duplicate=False))
    family_subtrees = {id(child) for boundary in boundaries for child in boundary.module.modules()}
    for request in requests:
        if not isinstance(request, NativeParallelBindingRequest):
            raise TypeError("native_parallel must return NativeParallelBindingRequest objects")
        matched = set()
        for pattern in request.match:
            hits = [(fqn, module) for fqn, module in all_modules if fnmatch.fnmatchcase(fqn, pattern)]
            if not hits:
                raise ValueError(f"native binding pattern {pattern!r} matched no module")
            for fqn, module in hits:
                if id(module) in matched:
                    continue
                if _overlaps_family(module, family_subtrees):
                    raise ValueError(f"native binding at {fqn!r} overlaps a family-owned subtree")
                matched.add(id(module))
                layout = request.sharding(module=module, module_fqn=fqn, member=None, context=context)

def _validate_ports(module, spec):
    """Validate named input and output layout declarations against the module interface."""
    injection_fields = ("local_compute_fn", "inner_wrapper", "inner_target", "region_dispatch")
    if any(getattr(spec, name) is not None for name in injection_fields):
        raise ValueError("module-owned layouts cannot install legacy compute injection")
    if spec.params is None or spec.in_src is None or spec.in_dst is None:
        raise ValueError("module-owned providers must declare params, in_src and in_dst explicitly")
    if set(spec.in_src) != set(spec.in_dst):
        raise ValueError("module-owned in_src and in_dst must declare the same input names")
    signature = inspect.signature(module.forward)
    accepts_kwargs = any(item.kind == inspect.Parameter.VAR_KEYWORD for item in signature.parameters.values())
    unknown = set(spec.in_src) - set(signature.parameters)
    if unknown and not accepts_kwargs:
        raise ValueError(f"layout declares unknown forward inputs {sorted(unknown)}")
    if spec.out_names is not None and len(set(spec.out_names)) != len(spec.out_names):
        raise ValueError("output names must be unique")


def _validate_placements(spec, context):
    """Reject unknown mesh axes and incompatible parameter placement declarations."""

        for axis, placement in named.items():
            if axis not in context.axis_sizes:
                raise ValueError(f"module layouts cannot declare storage/unknown axis {axis!r}")
            if not hasattr(placement, "is_shard"):
                raise TypeError("module providers must return concrete Placement objects")
    if any(isinstance(placement, Partial) for named in spec.params.values() for placement in named.values()):
        raise ValueError("parameter source layouts cannot be Partial")


def _claim_layout_parameters(fqn, parameters, spec, children, owner_by_id, canonical_by_id):
    """Assign each Parameter identity once and verify every alias has the same layout."""

    declared = set()
    for name, placement in tuple(spec.params.items()):
        parameter = parameters.get(name)
        if parameter is None:
            raise ValueError(f"layout parameter {_full(fqn, name)!r} does not exist")
        full_name = _full(fqn, name)
        if any(full_name.startswith(child + ".") for child in children):
            raise ValueError(f"parent layout duplicates nested module ownership at {full_name!r}")
        identity = id(parameter)

        identity = id(parameter)
        declared.add(identity)
        if identity in owner_by_id:
            if owner_by_id[identity] != placement:
                raise ValueError(f"shared parameter {full_name!r} has conflicting layouts")
            del spec.params[name]
        else:
            owner_by_id[identity] = placement
            canonical_by_id[identity] = full_name

        expected = {id(parameter) for name, parameter in parameters.items()
                    if not any(_full(fqn, name).startswith(child + ".") for child in children)}
        declared = _claim_layout_parameters(fqn, parameters, spec, children, owner_by_id, canonical_by_id)
        if expected - declared:
            missing = [name for name, parameter in parameters.items() if id(parameter) in expected - declared]
            raise ValueError(f"module-owned layout {fqn!r} omits parameters {missing}")
        result.expected_ids.update(expected)
        result.specs[fqn] = spec
        result.required_by_fqn[fqn] = {id(parameters[name]) for name in spec.params}
    for name, parameter in all_parameters.items():

    boundaries = collect_family_modules(model)
    for boundary in boundaries:
        member = boundary.info.resolved
        if member.sharding is None:
            continue
        fqn = boundary.fqns[0]
        cached = boundary.info.layout
        layout = (cached.sharding if cached is not None else
                  member.sharding(module=boundary.module, module_fqn=fqn, member=member, context=context))

              if parallel is not None and parallel.binding_kind != "none"]
    for fqn, kind in active:
        for parent, parent_kind in active:
            if _is_child(fqn, parent) and (kind == parent_kind or "mesh" in (kind, parent_kind)):
                raise ValueError(f"overlapping communication owners at {parent!r} and {fqn!r}")


def validate_owned_layouts(model: nn.Module, plan: Any, owned: OwnedLayouts, context: Any) -> None:
    """Check ownership and member capabilities after the final explicit overrides.

        parameters = dict(module.named_parameters(remove_duplicate=False))
        for name, placement in spec.params.items():
            parameter = parameters.get(name)
            if parameter is None:
                continue
            identity = id(parameter)
            if identity in declarations and declarations[identity] != placement:
                raise ValueError(f"explicit overrides give shared parameter {_full(fqn, name)!r} conflicting layouts")
            declarations[identity] = placement

                _validate_applied_layout(cached, spec, fqn)
            elif parallel is not None:
                validate_parallel_module(module, parallel, spec, context)
    if owned.expected_ids - covered:
        raise ValueError("final layout overrides removed required module-owned parameter coverage")


def _prepare_family_layout(module: nn.Module, mesh: Any) -> None:
    """Apply a self-described component through the existing planner and executor."""

    from hyper_parallel.distributed.apply import apply_sharding_plan  # pylint: disable=C0415,R0401

    info = get_family_build_info(module)
    if info.resolved.sharding is None:
        if info.resolved.selectors["parallel"]:
            raise ValueError("a mesh-constructed parallel family must declare sharding")
        return
    # A component's config identifies its containing model, whose providers do not apply here.
    module._hp_model_adapter_spec = None  # pylint: disable=protected-access
    planner = ShardingPlanner()
    plan = planner.plan(module, mesh.device_mesh, tp_size=mesh.tp_size, cp_size=mesh.cp_size,

def _declared_expert_unit_fqns(model, expert_parameters):
    """Find family/native execution units by identity rather than an experts path suffix."""
    expert_module_fqns = set()
    for boundary in collect_family_modules(model):
        parameters = set(boundary.module.parameters())
        if parameters and parameters <= expert_parameters:
            expert_module_fqns.add(boundary.fqns[0])
    for fqn, module in model.named_modules():
        binding = getattr(module, "_hp_parallel_binding", None)
        if binding is not None and binding.ep_context is not None:
            expert_module_fqns.add(fqn)
    expert_module_fqns.discard("")
    return expert_module_fqns

    @staticmethod
    def _get_model_adapter_spec(model: ModuleClass) -> Any | None:
        """Resolve a model adapter from the final model configuration."""
        if hasattr(model, "_hp_model_adapter_spec"):
            return model._hp_model_adapter_spec  # pylint: disable=protected-access
        config = getattr(model, "config", None)
        if config is None:
            return None
        identities = (

        existing.update(_family_source_shard_info(model))
    if not source_shard_records and not existing:
        return None
    if dense_source_mesh is None and expert_source_mesh is None:
        return existing or None
    result = build_source_shard_info(
        plan,
        dense_source_mesh,
        expert_source_mesh=expert_source_mesh,
    )
    result.update(existing)
    return result


def _shard_module_params(module, param_specs, mesh, mesh_dim_names):
    """distribute_tensor() converts parameters into DTensors.

            contains ``Partial``, or gives tied aliases conflicting layouts.
    """
    constructed = _family_source_shard_info(model)
    if constructed:
        for name, layout in constructed.items():
            if name in (source_shard_info or {}) and source_shard_info[name] != layout:
                raise ValueError(f"source_shard_info conflicts with the constructed family layout at {name!r}")
        source_shard_info = {**constructed, **(source_shard_info or {})}
    if source_shard_info is None:
        return _build_dtensor_source_shard_info(manager, model)
    if manager.mesh_context.tp_size > 1 and not source_shard_info:
        raise ValueError("TP is enabled but source_shard_info is empty")

        validate_mode,
    )
    for part in models:
        for canonical, alias in plan.parameter_aliases:
            _set_param_by_path(part, alias, part.get_parameter(canonical))

    # ====== Phase C: wrap forward ======
    for part in models:
        _apply_phase_c(part, plan, mesh, validate_mode, expert_mesh=expert_mesh)

    def __init__(self, mesh: Any) -> None:
        """Borrow a one-dimensional CP mesh without creating process groups."""
        if mesh is None or mesh.ndim != 1:
            raise ValueError("CPContext requires an existing one-dimensional HP mesh")
        self.mesh = mesh
        self.size = mesh.size()
        self.rank = mesh.get_local_rank()

        Args:
            local_length: Sequence length on each rank of the CP group.
        """
        if local_length < 0:
            raise ValueError("local sequence length cannot be negative")
        return SequencePartition((local_length,) * self.size, self.rank)

    def launch_gather_sequence(self, tensor: Tensor, *, dim: int = 1) -> PendingSequence:
        """Launch differentiable sequence gathering with explicit readiness.

            tensor: Local sequence shard with divisible head count.
            seq_dim: Sequence dimension.
            head_dim: Head dimension.
        """
        return ulysses_seq_to_head(tensor, seq_dim=seq_dim, head_dim=head_dim, cp_mesh=self.mesh)

    def head_to_sequence(self, tensor: Tensor, *, seq_dim: int = 1, head_dim: int = 2) -> Tensor:
        """Apply the inverse head-to-sequence exchange.

            tensor: Local head shard with divisible sequence length.
            seq_dim: Sequence dimension.
            head_dim: Head dimension.
        """
        return ulysses_head_to_seq(tensor, seq_dim=seq_dim, head_dim=head_dim, cp_mesh=self.mesh)

def _validate_slots(hidden: Tensor, indices: Tensor, weights: Tensor, experts: int) -> None:
    """Validate routed token slots, expert indices and optional routing probabilities."""
    if hidden.ndim != 2 or indices.ndim != 2 or weights.shape != indices.shape:
        raise ValueError("expected hidden[T,H], expert_indices[T,K], routing_weights[T,K]")
    if hidden.shape[0] != indices.shape[0] or indices.shape[1] < 1:
        raise ValueError("token dimensions must agree and top-k must be positive")
    if indices.dtype not in (torch.int32, torch.int64) or not weights.is_floating_point():
        raise ValueError("expert IDs must be integers and routing weights floating point")
    if hidden.device != indices.device or hidden.device != weights.device:
        raise ValueError("all routing inputs must be on the same device")
    if bool(((indices < 0) | (indices >= experts)).any()):
        raise ValueError("every supplied slot must have an expert ID in [0, num_experts); drop is unsupported")

    def __init__(self, mesh: Any, *, num_experts: int) -> None:
        """Borrow the EP mesh; never create a group or reduce parameter gradients."""
        if mesh is None or mesh.ndim != 1:
            raise ValueError("EPContext requires an existing one-dimensional HP mesh")
        if num_experts < 1 or num_experts % mesh.size():
            raise ValueError("num_experts must be positive and divisible by EP size")
        self.mesh = mesh
        self.num_experts = num_experts
        self.local_experts = num_experts // mesh.size()
        self.rank = mesh.get_local_rank()

    def _exchange(self, tensor, send, receive):
        if self.mesh.size() == 1:
            return tensor
        return ep_all_to_all(tensor, send, receive, self.mesh.get_group())

    def dispatch(self, hidden_states: Tensor, expert_indices: Tensor, routing_weights: Tensor) -> DispatchedExperts:
        """Transport every already-selected slot, including zero-probability slots.

        """
        if receipt.owner is not self._identity or receipt.consumed:
            raise ValueError("dispatch receipt belongs to another context or was already consumed")
        if values.ndim != 2 or values.shape[0] != sum(receipt.recv_counts):
            raise ValueError("expert output must have one row per received slot")
        peer_major = values[torch.argsort(receipt.expert_order)]
        source_slots = self._exchange(peer_major, receipt.recv_counts, receipt.send_counts)
        source_slots = source_slots[torch.argsort(receipt.send_order)]
        output = source_slots.reshape(receipt.tokens, receipt.top_k, values.shape[-1]).sum(dim=1)

    validate_materialized: Callable | None = None

    def __post_init__(self) -> None:
        if not self.protocol or not self.version:
            raise ValueError("parallel protocol and version must be non-empty")
        if self.binding_kind not in ("none", "cp", "ep", "mesh"):
            raise ValueError("binding_kind must be none, cp, ep or mesh")
        if self.binding_kind != "none" and (not self.stages or self.validate is None):
            raise ValueError("parallel members must declare stage signatures and a layout validator")
        object.__setattr__(self, "stages", MappingProxyType({name: tuple(args) for name, args in self.stages.items()}))
        object.__setattr__(self, "options", MappingProxyType(deepcopy(dict(self.options))))

    sharding: Callable

    def __post_init__(self) -> None:
        if isinstance(self.match, str) or not self.match or not all(isinstance(name, str) for name in self.match):
            raise ValueError("native matches must be a non-empty tuple of module FQN patterns")
        if not callable(self.sharding) or not isinstance(self.parallel, ParallelModuleSpec):
            raise ValueError("native binding requires a complete layout provider and ParallelModuleSpec")
        object.__setattr__(self, "match", tuple(self.match))


@dataclass(frozen=True)

        Args:
            axis: cp or ep, in the runtime's established rank order.
        """
        if axis == "ep":
            region = self.mesh_context.fsdp_moe_mesh
            if region is not None:
                return region["ep"]
        region = self.mesh_context.device_mesh
        if region is None or axis not in tuple(region.mesh_dim_names or ()):
            raise ValueError(f"no prebuilt {axis!r} mesh is available for parallel binding")
        return region[axis]


@dataclass(frozen=True)

    """
    for name, arguments in parallel.stages.items():
        implementation = inspect.getattr_static(type(module), name, None)
        if not callable(implementation):
            raise ValueError(f"parallel stage {name!r} must be implemented in the module class")
        inspect.signature(getattr(module, name)).bind(**dict.fromkeys(arguments))
    if parallel.binding_kind != "none":
        implementation = inspect.getattr_static(type(module), "bind_parallel", None)
        if not callable(implementation):
            raise ValueError("parallel modules must define bind_parallel(binding) in their class")
        inspect.signature(module.bind_parallel).bind(None)
    if getattr(sharding, "_ep_stack", None):
        raise ValueError("explicit expert stages cannot inherit legacy whole-MoE stacking")
    if sharding.in_src != sharding.in_dst:
        raise ValueError("explicit stage modules must receive their declared computation layout directly")
    if sharding.out_dst is not None and sharding.out_src != sharding.out_dst:
        raise ValueError("explicit stage modules must restore their declared output layout themselves")
    if parallel.binding_kind == "ep":
        _validate_expert_ownership(module, parallel, sharding, context)
    if parallel.validate is not None:
        parallel.validate(module=module, sharding=sharding, context=context)

    """Check the bound EP mesh against local and global expert ownership."""
    experts = parallel.options.get("num_experts")
    expert_type = type(experts)
    if expert_type is not int or experts < 1 or experts % context.axis_sizes["ep"]:
        raise ValueError("EP protocols must declare a divisible global num_experts")
    sharded = {name: named["ep"] for name, named in sharding.params.items()
               if isinstance(named.get("ep"), Shard)}
    if not sharded:
        raise ValueError("EP transport requires explicit DTensor expert ownership")
    for name, placement in sharded.items():
        parameter = module.get_parameter(name)
        if placement.dim != 0 or parameter.shape[0] != experts:
            raise ValueError("EPContext supports balanced contiguous expert Shard(0) ownership")


def bind_parallel_modules(model: nn.Module, plan: Any, mesh_context: Any) -> None:
    """Inject shared resources after layout and before activation wrappers/FSDP2.

    mesh_binding = MeshBinding(mesh_context)
    prepared = []
    for fqn, parallel in plan.parallel_modules.items():
        if parallel.binding_kind == "none":
            continue
        module = model.get_submodule(fqn)
        existing = getattr(module, "_hp_parallel_binding", None)
        if existing is not None:
            info = get_family_build_info(module)
            cached = info.layout if info is not None else None
            if (cached is None or not cached.applied or cached.parallel != parallel
                    or not _same_mesh(existing.mesh_binding.mesh_context, mesh_context)):
                raise ValueError(f"module {fqn!r} already has a communication owner")
            continue
        cp = CPContext(mesh_binding.axis_mesh("cp")) if parallel.binding_kind == "cp" else None
        ep = (EPContext(mesh_binding.axis_mesh("ep"), num_experts=parallel.options["num_experts"])
              if parallel.binding_kind == "ep" else None)
        prepared.append((module, ParallelBinding(mesh_binding, cp, ep)))

    Args:
        model: Final loaded model whose class protocols may validate physical resources.
    """
    for fqn, parallel in getattr(model, "_hp_parallel_module_specs", {}).items():
        module = model.get_submodule(fqn)
        if parallel.validate_materialized is not None:
            parallel.validate_materialized(module=module, binding=getattr(module, "_hp_parallel_binding", None))

    def _validate_override_match(self):
        """Reject method-only and empty declarations before creating a boundary."""
        if self.replace is not None and self.replace.kind == "method" and (
                self.match.rpartition(".")[2] != "forward"):
            raise ValueError("non-forward method replacements do not create planner boundaries")

        if not self.match:
            raise ValueError(
                "plan_overrides entry is missing 'match' (an fqn or "

    for entry_index, entry in enumerate(entries):
        if entry.replace is None or entry.replace.kind != "module":
            continue
        if entry.when == "low_precision" and not low_precision_enabled:
            continue
        _validate_replacement_entry(entry, effective)
        target = entry.replace.implementation
        baseline = target.requested_target if isinstance(target, ReplacementTarget) else target.callable
        arguments = {name: value for name, value in target._kwargs.items()  # pylint: disable=protected-access

        if isinstance(target, ReplacementTarget):
            arguments.update(target.selectors)
        if entry.when == "low_precision":
            # Legacy policy-bearing closures remain on the explicit compatibility path.
            continue
        result.append(ModuleReplacementRequest(
            match=(entry.match,) if isinstance(entry.match, str) else tuple(entry.match),
            target=baseline, kwargs=arguments, origin="user",
            declaration_id=f"plan_override:{entry_index}:replace",

    if entry.replace is not None and entry.replace.kind == "module" and entry.when not in (
        None,
        "low_precision",
    ):
        raise ValueError(
            f"plan_overrides match={entry.match!r} uses replace.kind='module' and "
            "does not support 'when' values other than 'low_precision' "
            "(module replacements are not conditioned on parallel topology)"
        )