Diff Coverage

        configs = {name: dict(config) for name, config in
                   (("muon", muon), ("sinkhorn", sinkhorn), ("adamw", adamw)) if config is not None}
        specs = {name: config.pop("param_groups", []) for name, config in configs.items()}
        if not any(specs.values()):
            self.optimizer = get_hyper_optimizer(model, **configs)
            return
        groups = _build_optimizer_groups(model, configs)
        arguments = {}
        for name, config in configs.items():
            arguments[f"{name}_params"] = _resolve_param_groups(model, groups[name], specs[name])

            family = _default_family(param, names, embedding_ids, norm_ids, configs)
        if family not in selected:
            raise ValueError(f"No enabled optimizer owns trainable parameter {names[0]}")
        if family in ("muon", "sinkhorn") and param.ndim < 2:
            raise ValueError(f"{family} requires matrix parameters, got {names[0]}")
        selected[family].append(param)
    return {family: _make_groups(family, params, aliases, norm_ids) for family, params in selected.items()}

    selectors = {}
    for family, config in configs.items():
        patterns = config.get("param_patterns", [])
        if isinstance(patterns, str):
            patterns = [patterns]
        selectors[family] = [re.compile(pattern) for pattern in patterns]
    return selectors

    """
    config = getattr(model, "config", None)
    head_dim = getattr(config, "head_dim", None)
    if head_dim is not None:
        return head_dim
    hidden = getattr(config, "hidden_size", None)
    heads = getattr(config, "num_attention_heads", None)
    if isinstance(hidden, int) and isinstance(heads, int) and heads > 0 and hidden % heads == 0:
        return hidden // heads

        name: _build_configured_optimizer(name, classes[name], family_groups, configs[name])
        for name, family_groups in groups.items() if family_groups
    }
    if not optimizers:
        raise ValueError("At least one non-empty optimizer parameter group is required")
    return ChainedOptimizer(model, optimizers=optimizers)


__all__ = [

        self._head_wise_dims: Dict[torch.nn.Parameter, int] = {}
        if head_wise and (advanced.reshape_fn is not None or advanced.ns_transform_fn is not None):
            raise ValueError("head_wise cannot be combined with reshape_fn or ns_transform_fn")
        if head_wise_patterns is not None and not head_wise:
            raise ValueError("head_wise_patterns requires head_wise=True")
        if ns_variant not in ("legacy", "asym5", "custom"):
            raise ValueError(
                f"ns_variant must be 'legacy', 'asym5', or 'custom', got {ns_variant!r}"
            )

        self._classify_parameters_for_step()

    def _new_checkpoint_state(self, param: torch.Tensor, _: str) -> torch.Tensor:
        """Muon momentum follows the parameter/gradient precision."""
        return torch.zeros_like(param)

    def _configure_head_wise(self) -> None:
        """Resolve head dimensions once, including after main-parameter replacement."""
        self._head_wise_dims = {}

    def _resolve_head_dim(param, compiled):
        """Validate the selected matrix's name and output-head geometry."""
        name = getattr(param, "model_name", None)
        if name is None:
            raise ValueError("head_wise requires model_name; use get_hyper_optimizer to bind names")
        dims = [dim for pattern, dim in compiled if pattern.search(name)]
        if not dims:
            return None
        if len(dims) != 1:
            raise ValueError(f"Overlapping head_wise_patterns for {name}")
        dim = dims[0]
        if not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0:
            raise ValueError(f"A positive integer head_dim is required for {name}")
        if param.ndim != 2 or param.shape[0] % dim:
            raise ValueError(f"{name} must be a 2D matrix with rows divisible by head_dim={dim}")
        return dim

            info = self._hsdp_assignment_batches.get(group_idx)
            linear_batches = []
            if info:
                for bg in info.get("batch_groups", []):
                    linear_batches.extend(bg.get("sub_batches", []))
            group_linear_batches[group_idx] = linear_batches
            max_num_batches = max(max_num_batches, len(linear_batches))

        return group_linear_batches, max_num_batches

        for group_idx in range(num_groups):
            info = self._hsdp_assignment_batches.get(group_idx)
            if not info:
                # No HSDP info — all params are no_comm.
                unshard_params = self.unshard_params_by_group.get(group_idx, [])
                no_comm_ns[group_idx] = self._update_muon_momentum(
                    self.param_groups[group_idx], unshard_params
                )
            else:
                no_comm_params = info.get("no_comm", [])

                    no_comm_ns[group_idx] = self._update_muon_momentum(
                        self.param_groups[group_idx], no_comm_params
                    )
                else:
                    no_comm_ns[group_idx] = {}

        return no_comm_ns

    def _run_post_update_fn(

            owner = owner_by_id.get(id(optimizer_param))
            if owner is not None:
                self.param_names_by_optimizer[owner].append(param_name)
        if owner_by_id.keys() - registered_ids:
            raise ValueError("Optimizer parameters must belong to the registered model")

    def reset_optimizer_parameters(
            self,
            model_param_by_optimizer_param: Dict[

        self._broadcast_op_fused(target="state", state_keys=list(self.state_tensor_keys))

    def _new_checkpoint_state(self, param: torch.Tensor, _: str) -> torch.Tensor:
        """Allocate a missing parameter-shaped state; subclasses may override its schema."""
        return torch.zeros_like(param, dtype=torch.float32)

    def _collect_broadcast_tensors(
            self,
            target: str,

                    if key in param_state:
                        state_tensor = param_state[key]
                        local_tensor = to_local_if_dtensor(state_tensor)
                    else:
                        local_tensor = self._new_checkpoint_state(p, key)
                        param_state[key] = local_tensor
                        local_tensor = to_local_if_dtensor(local_tensor)

                    rank_dtype_tensors[(src_coord, local_tensor.dtype, replicate_pgs)].append(local_tensor)

                batch_buffer = buffer[:batch_total_size]

            # Pack: owner rank
            if local_coord == src_coord:
                self._pack_broadcast_batch(batch_buffer, batch_tensor_offsets)

            if async_op:
                handles = BaseDistributedOptimizer._hierarchical_broadcast_buffer_async(
                    batch_buffer, src_coord, replicate_pgs, local_coord,

    @staticmethod
    def _pack_broadcast_batch(batch_buffer, batch_tensor_offsets):
        """Pack owner values and zero alignment padding before broadcast."""
        for t, offset, actual_numel, padded_numel in batch_tensor_offsets:
            batch_buffer[offset:offset + actual_numel].copy_(t.view(-1))
            if padded_numel > actual_numel:
                batch_buffer[offset + actual_numel:offset + padded_numel].zero_()

    def _wait_and_release_oldest(self) -> None:
        """Wait, unpack, and release the oldest inflight async batch."""
        batch_buffer, batch_tensor_offsets, handles = self._inflight.pop(0)

        self._group_dtensor_by_mesh()
        self._norm_groups = {}
        for _, hsdp_groups in self._hsdp_grouping.values():
            for assignment in hsdp_groups:
                row_groups, column_groups = [], []
                for (_, dimension), process_group in zip(assignment.layout_spec.shard_axes, assignment.shard_pgs):
                    (row_groups if dimension == 0 else column_groups).append(process_group)
                for param in assignment.params:
                    self._norm_groups[param] = (tuple(row_groups), tuple(column_groups))
        self._split_replicate_groups()
        self._build_hsdp_batch()
        self._build_param_broadcast_info()

    def _new_checkpoint_state(self, param: torch.Tensor, _: str) -> torch.Tensor:
        """Keep state layout metadata even inside SkipDTensorDispatch."""
        local = torch.zeros_like(to_local_if_dtensor(param), dtype=torch.float32)
        if isinstance(param, DTensor):
            return type(param).from_local(local, param.device_mesh, param.placements,
                                          shape=tuple(param.shape), stride=(param.shape[1], 1))
        return local

    def _update_parameter(self, param: torch.Tensor, group: Dict[str, Any]) -> None:

        if param.grad is None:
            return
        gradient = param.grad
        if hasattr(gradient, "placements") and any(place.is_partial() for place in gradient.placements):
            raise ValueError("Sinkhorn gradients must be reduced before optimizer.step()")
        gradient = to_local_if_dtensor(gradient)
        if gradient.is_sparse:
            raise ValueError("Sinkhorn does not support sparse gradients")
        state = self.state[param]
        if "momentum_buffer" not in state:
            state["momentum_buffer"] = self._new_checkpoint_state(param, "momentum_buffer")
        buffer = to_local_if_dtensor(state["momentum_buffer"])

            Closure loss when supplied, otherwise None.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        for group_index, group in enumerate(self.param_groups):
            self._validate_group(group)
            group["step"] = (group.get("step") or 0) + 1
            assignment = self._hsdp_assignment_batches[group_index]

            assignment = self._hsdp_assignment_batches[group_index]
            for param in assignment["no_comm"]:
                self._update_parameter(param, group)
            for batch_group in assignment["batch_groups"]:
                for batch in batch_group["sub_batches"]:
                    for param in batch.owned_params:
                        self._update_parameter(param, group)
        self._broadcast_replicate_params_after_step()
        return loss