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

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]:

    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

    * ``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

            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.

    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)

                    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.

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

                    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

        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,