Diff Coverage

        _device_has_foreach_support,
        _group_tensors_by_device_and_dtype,
        _has_foreach_support,
    )
except ImportError:
    _device_has_foreach_support = None  # type: ignore[assignment]
    _group_tensors_by_device_and_dtype = None  # type: ignore[assignment]
    _has_foreach_support = None  # type: ignore[assignment]

__all__: list[str] = ["clip_grad_norm_"]

    manually divide by the group size.
    """
    lower = op_str.lower()
    if lower == "avg" and not _REDUCE_OP_AVG_SUPPORTED:
        return dist.ReduceOp.SUM, True
    op = _STR_TO_REDUCE_OP.get(lower)
    if op is None:
        raise ValueError(
            f"Unsupported Partial reduce_op: {op_str!r}. "
            f"Supported: {sorted(set(list(_STR_TO_REDUCE_OP) + ['avg']))}"
        )
    return op, False

    * single ``torch.Tensor`` -> ``[tensor]``
    * iterable of tensors    -> ``list(iterable)``
    """
    if isinstance(parameters, torch.nn.Module):
        return list(parameters.parameters())
    if isinstance(parameters, torch.Tensor):
        return [parameters]
    return list(parameters)


def _param_device(param: torch.Tensor) -> torch.device:

def _param_device(param: torch.Tensor) -> torch.device:
    """Return the local device of *param* (unwrap DTensor if needed)."""
    if _is_dtensor(param):
        return param._local_tensor.device  # pylint: disable=protected-access
    return param.device


def _get_grad_obj(param: torch.nn.Parameter) -> Optional[torch.Tensor]:

    falling back to ``param.grad``.
    """
    grad = getattr(param, "main_grad", None)
    if grad is not None:
        return grad
    return param.grad


def _get_local_grad(param: torch.nn.Parameter) -> Optional[torch.Tensor]:

    Supports ``main_grad`` for fp32 mixed-precision training.
    """
    if not param.requires_grad:
        return None
    grad = _get_grad_obj(param)
    if grad is None:
        return None
    if _is_dtensor(grad):
        return grad._local_tensor  # pylint: disable=protected-access
    return grad


def _get_param_mesh_info(
    param: torch.nn.Parameter,

    grad = _get_grad_obj(param)
    # Prefer grad's spec (most accurate); fall back to param's.
    spec_source = grad if _is_dtensor(grad) else param
    if not _is_dtensor(spec_source):
        return None, (), ()

    shard_dims = tuple(
        i for i, p in enumerate(spec_source.placements)
        if p.is_shard()

    device: torch.device,
    total: torch.Tensor,
) -> None:
    """Accumulate sum-of-p-th-powers for *dev_grads* into *total*."""
    for g in dev_grads:
        n = torch.linalg.vector_norm(g, norm_type)
        total.add_(n.to(device=device) ** norm_type)


def _foreach_p_norms(
    grads: List[torch.Tensor],

    precision as ``vector_norm(dtype=float32)``.  Non-float32 tensors
    and backends that raise ``RuntimeError`` fall back to per-tensor
    ``vector_norm``.
    """
    total = torch.tensor(0.0, device=device, dtype=torch.float32)
    grouped = _group_tensors_by_device_and_dtype(
        [[g.detach() for g in grads]],
    )
    for (dev, _), ([dev_grads], _) in grouped.items():
        if (
            dev_grads[0].dtype == torch.float32
            and _has_foreach_support(dev_grads, dev)
        ):
            try:
                per_norms = torch._foreach_norm(  # pylint: disable=W0212
                    dev_grads, norm_type,
                )
            except RuntimeError:
                per_norms = None
            if per_norms is not None:
                total.add_(
                    torch.stack([
                        n.to(device=device) ** norm_type
                        for n in per_norms
                    ]).sum(),

                        for n in per_norms
                    ]).sum(),
                )
            else:
                _sum_p_norms(dev_grads, norm_type, device, total)
        else:
            _sum_p_norms(dev_grads, norm_type, device, total)
    return total


def _per_tensor_norms(
    grads: List[torch.Tensor],

    device: torch.device,
) -> List[torch.Tensor]:
    """Return per-tensor norms as a list of scalar tensors on *device*."""
    if not grads:
        return []

    if _group_tensors_by_device_and_dtype is None or not hasattr(torch, "_foreach_norm"):
        return [
            torch.linalg.vector_norm(g.detach(), norm_type).to(device=device)
            for g in grads
        ]

            try:
                per_norms = torch._foreach_norm(  # pylint: disable=W0212
                    dev_grads, norm_type,
                )
            except RuntimeError:
                per_norms = None
            if per_norms is not None:
                norms.extend(
                    [n.to(device=device) for n in per_norms],
                )

                norms.extend(
                    [n.to(device=device) for n in per_norms],
                )
                continue
        norms.extend([
            torch.linalg.vector_norm(g, norm_type).to(device=device)
            for g in dev_grads
        ])
    return norms

    * ``-inf``  -> +inf (neutral for MIN)
    * ``0``     -> 0    (neutral for SUM)
    * finite    -> 0    (neutral for SUM)
    """
    if not grads:
        if norm_type == -math.inf:
            return torch.tensor(
                float("inf"), device=device, dtype=torch.float32,
            )
        return torch.tensor(0.0, device=device, dtype=torch.float32)

    if norm_type == math.inf:
        norms = [
            torch.linalg.vector_norm(g.detach(), math.inf)
            for g in grads
        ]
        return torch.stack(norms).max().to(device)

    if norm_type == -math.inf:
        norms = [
            torch.linalg.vector_norm(g.detach(), -math.inf)
            for g in grads
        ]
        return torch.stack(norms).min().to(device)

    if norm_type == 0:
        norms = [
            torch.linalg.vector_norm(g.detach(), 0)
            for g in grads
        ]
        return torch.stack(norms).sum().to(device)

    # Finite p-norm: return sum of p-th powers.
    if (
        len(grads) > 1
        and _group_tensors_by_device_and_dtype is not None
        and hasattr(torch, "_foreach_norm")
    ):
        return _foreach_p_norms(grads, norm_type, device)

    # Scalar fallback when foreach utilities are unavailable.
    norms = [
        torch.linalg.vector_norm(g.detach(), norm_type)
        for g in grads
    ]
    norm_powers = [n.to(device=device) ** norm_type for n in norms]
    return torch.stack(norm_powers).sum()


# ---------------------------------------------------------------------------
# Total norm aggregation with collectives

    ``norm_grads`` (parallel to ``key_per_grad``) holds the tensor whose
    norm to take per parameter; only the finite p-norm path consumes it.
    """
    if norm_type == math.inf:
        return _total_norm_inf(
            grad_groups, norm_type, mesh_cache, device,
            dist.ReduceOp.MAX,
        )

            dist.ReduceOp.MAX,
        )

    if norm_type == -math.inf:
        return _total_norm_inf(
            grad_groups, norm_type, mesh_cache, device,
            dist.ReduceOp.MIN,
        )

            dist.ReduceOp.MIN,
        )

    if norm_type == 0:
        return _total_norm_sum(
            grad_groups, norm_type, mesh_cache, device,
        )

    # Finite p-norm: FSDP2-aligned sequence.

def _total_norm_inf(  # pylint: disable=R0913,R0917
    grad_groups, norm_type, mesh_cache, device, reduce_op,
):
    """Shared logic for inf / -inf norms."""
    group_norms: List[torch.Tensor] = []
    for (mesh_id, shard_dims), grads in grad_groups.items():
        local_norm = _compute_local_norm(grads, norm_type, device)
        if mesh_id is not None:
            mesh = mesh_cache[mesh_id]
            for dim in shard_dims:
                dist.all_reduce(
                    local_norm, op=reduce_op,
                    group=mesh.get_group(dim),
                )
        group_norms.append(local_norm)
    if not group_norms:
        if norm_type == -math.inf:
            return torch.tensor(float("inf"), device=device)
        return torch.tensor(0.0, device=device)
    stacked = torch.stack(group_norms)
    return stacked.max() if reduce_op == dist.ReduceOp.MAX else stacked.min()


def _total_norm_sum(grad_groups, norm_type, mesh_cache, device):
    """Shared logic for finite norms and L0 (all use SUM all-reduce)."""
    total = torch.tensor(0.0, device=device)
    for (mesh_id, shard_dims), grads in grad_groups.items():
        local_val = _compute_local_norm(grads, norm_type, device)
        if mesh_id is not None:
            mesh = mesh_cache[mesh_id]
            for dim in shard_dims:
                dist.all_reduce(
                    local_val, op=dist.ReduceOp.SUM,
                    group=mesh.get_group(dim),
                )
        total.add_(local_val)
    return total


def _reduction_signature(grad_groups, mesh_cache):
    """Bucket grad-group keys by the *process group(s)* they reduce over.

            local_p = torch.linalg.vector_norm(
                torch.stack(norms).to(torch.float32), norm_type,
            ) ** norm_type
        else:
            local_p = torch.tensor(0.0, device=device, dtype=torch.float32)

        for group in sig_groups[sig]:
            dist.all_reduce(local_p, op=dist.ReduceOp.SUM, group=group)

            )
            chunk_sizes.append(local_grad.numel())
            has_grad.append(True)
            active_indices.append(idx)
        elif param.requires_grad:
            local_p = (
                param._local_tensor  # pylint: disable=W0212
                if _is_dtensor(param) else param.data
            )
            numel = local_p.numel()
            chunks.append(
                torch.zeros(
                    numel, device=local_p.device,
                    dtype=torch.float32,
                ),

                    numel, device=local_p.device,
                    dtype=torch.float32,
                ),
            )
            chunk_sizes.append(numel)
            has_grad.append(False)
            active_indices.append(idx)

    return chunks, chunk_sizes, has_grad, active_indices

    for idx, info in enumerate(param_infos):
        mesh, partial_info = info[2], info[3]
        if partial_info:
            if mesh is None:
                raise RuntimeError(
                    "clip_grad_norm_: parameter has Partial placements "
                    "but no DeviceMesh. This is a DTensor invariant "
                    "violation."
                )

            _build_coalesce_buffer(param_infos, indices)
        )

        if not chunks:
            continue  # all params frozen, no collective needed

        # Sanity check: same mesh → same device.  Fail fast on
        # misconfigured inputs rather than silent NCCL errors.
        buf_device = chunks[0].device

        # Sanity check: same mesh → same device.  Fail fast on
        # misconfigured inputs rather than silent NCCL errors.
        buf_device = chunks[0].device
        for chunk in chunks[1:]:
            if chunk.device != buf_device:
                raise RuntimeError(
                    f"clip_grad_norm_: parameters in the same Partial "
                    f"coalesce group are on different devices "
                    f"({buf_device} vs {chunk.device}). All parameters "
                    f"sharing the same DeviceMesh must reside on the "

        for pdim, reduce_op, needs_avg in partial_info:
            group = mesh.get_group(pdim)
            dist.all_reduce(buf, op=reduce_op, group=group)
            if needs_avg:
                buf /= dist.get_world_size(group=group)

        # Extract views for params with actual gradients.
        offset = 0
        for i, idx in enumerate(active_indices):

            (param, local_grad, mesh, partial_info, key),
        )

    if device is None:
        device = torch.device("cpu")

    # --- Phase 2: coalesced Partial reduction (O(N) → O(G)) ---
    reduced = _coalesce_partial_reduce(param_infos, mesh_cache)

        param, local_grad, key = info[0], info[1], info[4]
        if local_grad is None:
            # Ensure the key exists so the Shard norm all-reduce is
            # entered even when this rank has no grads for the group.
            if key not in grad_groups:
                grad_groups[key] = []
            continue

        grad_obj = _get_grad_obj(param)
        if _is_dtensor(grad_obj):
            has_dtensor_grad = True

        if idx in reduced:
            grad_groups[key].append(reduced[idx])
            norm_grads.append(reduced[idx])
        else:
            grad_groups[key].append(local_grad)
            norm_grads.append(local_grad)

    return _GradGroups(
        grad_groups, all_grads, norm_grads, key_per_grad,
        mesh_cache, device, has_dtensor_grad,

            use_foreach = (
                foreach is None and _has_foreach_support(device_grads, device)
            ) or (foreach and _device_has_foreach_support(device))
            if use_foreach:
                torch._foreach_mul_(  # pylint: disable=W0212
                    device_grads,
                    clip_coef_clamped.to(device=device, dtype=dtype),
                )
            elif foreach:
                raise RuntimeError(
                    f"foreach=True was passed, but can't use the "
                    f"foreach API on {device.type} tensors"
                )
            else:

                for g in device_grads:
                    g.mul_(clip_coef_clamped_cast)
    else:
        # Fallback when _foreach_utils is unavailable.
        if foreach:
            raise RuntimeError(
                "foreach=True was passed, but "
                "torch.utils._foreach_utils is not available"
            )
        for grad in all_grads:
            grad.mul_(clip_coef_clamped.to(grad.device, grad.dtype))


# ---------------------------------------------------------------------------
# Public API

    if error_if_nonfinite and torch.logical_or(
        total_norm.isnan(), total_norm.isinf()
    ):
        raise RuntimeError(
            f"The total norm of order {norm_type} for gradients from "
            "`parameters` is non-finite, so it cannot be clipped. To "
            "disable this error and scale the gradients by the "
            "non-finite norm anyway, set "

    # Promote return dtype to match gradient dtypes (FSDP1 convention).
    # When this rank has no gradients, return in the default FP32 dtype
    # (same as FSDP1's behavior to avoid extra communication).
    if not all_grads:
        warnings.warn(
            "clip_grad_norm_ called on this rank with no gradients -- "
            "returning the local norm in the default dtype "
            f"{total_norm.dtype}",
            stacklevel=2,

            "returning the local norm in the default dtype "
            f"{total_norm.dtype}",
            stacklevel=2,
        )
        return total_norm

    total_norm_dtype = functools.reduce(
        torch.promote_types,
        [g.dtype for g in all_grads],

# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Compatibility entry point for the relocated gradient clipping utility."""
from hyper_parallel.core.utils.clip_grad import clip_grad_norm_

__all__: list[str] = ["clip_grad_norm_"]

        parameters, max_norm, norm_type=2.0,
        error_if_nonfinite=False, foreach=None,
    ):
        # pylint: disable=C0415
        from hyper_parallel.core.utils.clip_grad import (
            clip_grad_norm_ as _clip_grad_norm,
        )
        return _clip_grad_norm(
            parameters, max_norm, norm_type,