Utilities¶

def get_supported_optimizers(filters: str | list[str] | None = None) -> list[str]:
    """List registered optimizer names in alphabetical order.

    Args:
        filters: Wildcard pattern or list of patterns, such as `'*adam*'`. `None` selects all names.

    Returns:
        list[str]: Matching names in lowercase, without duplicates.

    """
    if filters is None:
        return sorted(OPTIMIZERS.keys())

    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]

    filtered_list: set[str] = set()
    for include_filter in include_filters:
        filtered_list.update(fnmatch.filter(OPTIMIZERS.keys(), include_filter))

    return sorted(filtered_list)

def get_supported_lr_schedulers(filters: str | list[str] | None = None) -> list[str]:
    """List registered scheduler names in alphabetical order.

    Args:
        filters: Wildcard pattern or list of patterns, such as `'*cosine*'`. `None` selects all names.

    Returns:
        list[str]: Matching names in lowercase, without duplicates.
    """
    if filters is None:
        return sorted(LR_SCHEDULERS.keys())

    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]

    filtered_list: set[str] = set()
    for include_filter in include_filters:
        filtered_list.update(fnmatch.filter(LR_SCHEDULERS.keys(), include_filter))

    return sorted(filtered_list)

def get_supported_loss_functions(filters: str | list[str] | None = None) -> list[str]:
    """List registered loss function names in alphabetical order.

    Args:
        filters: Wildcard pattern or list of patterns, such as `'*focal*'`. `None` selects all names.

    Returns:
        list[str]: Matching names in lowercase, without duplicates.

    """
    if filters is None:
        return sorted(LOSS_FUNCTIONS.keys())

    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]

    filtered_list: set[str] = set()
    for include_filter in include_filters:
        filtered_list.update(fnmatch.filter(LOSS_FUNCTIONS.keys(), include_filter))

    return sorted(filtered_list)

def is_valid_parameters(parameters: ParamsT) -> bool:
    """Check for a nonempty list or tuple whose first entry is a parameter group dictionary."""
    return isinstance(parameters, (list, tuple)) and len(parameters) > 0 and isinstance(parameters[0], dict)

def has_overflow(grad_norm: torch.Tensor) -> bool:
    """Return whether a tensor contains NaN or infinite values."""
    return bool(torch.logical_or(torch.isnan(grad_norm), torch.isinf(grad_norm)).any())

def to_real(x: torch.Tensor) -> torch.Tensor:
    """Return the real part of a complex tensor, or the original real tensor."""
    return x.real if torch.is_complex(x) else x

def normalize_gradient(x: torch.Tensor, use_channels: bool = False, epsilon: float = 1e-8) -> None:
    """Divide gradients by their standard deviation in place.

    Args:
        x: Gradient tensor to normalize.
        use_channels: If True, perform channel wise normalization.
        epsilon: Small constant added for numerical stability.

    """
    size: int = x.dim()
    if size > 1 and use_channels:
        if math.prod(x.shape[1:]) <= 1:
            return

        s = x.std(dim=tuple(range(1, size)), keepdim=True).add_(epsilon)
        x.div_(s)
    elif torch.numel(x) > 2:
        s = x.std().add_(epsilon)
        x.div_(s)

def clip_grad_norm(
    parameters: ParamsT | torch.Tensor,
    max_norm: float = 0.0,
    sync: bool = False,
) -> torch.Tensor:
    """Compute the global L2 gradient norm and optionally clip gradients in place.

    Args:
        parameters: Tensor or iterable of tensors whose gradients to inspect.
        max_norm: Maximum gradient norm. A nonpositive value disables clipping.
        sync: Sum squared norms across the distributed process group for sharded gradients.

    Returns:
        torch.Tensor: Global gradient norm before clipping.

    """
    if parameters is None:
        raise ValueError('ParamsT cannot be None.')

    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]

    # make sure any generators are expanded
    parameters = cast(list, list(parameters))

    # if syncing we need to manually perform the clipping so that we aggregate properly
    if max_norm > 0 and not sync:
        return clip_grad_norm_(parameters, max_norm)

    norm_sq = get_global_gradient_norm([{'params': parameters}])
    if sync:  # pragma: no cover
        # also need to get the norms from all the other sharded works in FSDP
        all_reduce(norm_sq)

    grad_norm = norm_sq.sqrt_().squeeze_(0)
    if max_norm > 0:  # pragma: no cover
        clip_coefficient = (max_norm / (grad_norm + 1e-6)).clamp_(max=1.0)
        for p in parameters:
            if p.grad is not None:
                p.grad.detach().mul_(clip_coefficient.to(p.grad.device))

    return grad_norm

def unit_norm(x: torch.Tensor, norm: float = 2.0) -> torch.Tensor:
    """Compute parameter unit norms for adaptive gradient clipping.

    Uses the full norm for scalars and vectors, dimension 1 for 2D and 3D tensors,
    and all dimensions after the first for tensors with four or more dimensions.

    Args:
        x: Parameter or gradient tensor.
        norm: Order of the vector norm.

    Returns:
        torch.Tensor: Unit norms, with reduced dimensions retained for multidimensional inputs.

    """
    keep_dim: bool = True
    dim: int | tuple[int, ...] | None = None

    x_len: int = len(x.shape)
    if x_len <= 1:
        keep_dim = False
    elif x_len in (2, 3):
        dim = 1
    elif x_len == 4:
        dim = (1, 2, 3)
    else:
        dim = tuple(range(1, x_len))

    return x.norm(p=norm, dim=dim, keepdim=keep_dim)

def disable_running_stats(model: nn.Module):
    """Pause BatchNorm running statistic updates by setting momentum to zero."""

    def _disable(module):
        if isinstance(module, _BatchNorm):
            module.backup_momentum = module.momentum
            module.momentum = 0

    model.apply(_disable)

def enable_running_stats(model: nn.Module):
    """Restore BatchNorm momentum after pausing running statistic updates."""

    def _enable(module):
        if isinstance(module, _BatchNorm) and hasattr(module, 'backup_momentum'):
            module.momentum = module.backup_momentum

    model.apply(_enable)

@torch.no_grad()
def get_global_gradient_norm(
    param_groups: list[ParamGroup] | None,
    device: torch.device | None = None,
    *,
    weight_adaptive: bool = False,
) -> torch.Tensor:
    """Return the sum of squared L2 gradient norms across parameter groups.

    Args:
        param_groups: Optimizer parameter groups, or None for no parameters.
        device: Output device. Defaults to the first parameter's device, or CPU for empty groups.
        weight_adaptive: Weight gradients by absolute parameter values for groups with `adaptive=True`.

    Returns:
        torch.Tensor: Squared global norm as a single element float32 tensor.

    """
    param_groups = param_groups or []
    device = device or next((p.device for group in param_groups for p in group['params']), torch.device('cpu'))
    grads: list[torch.Tensor] = []
    for group in param_groups:
        adaptive = weight_adaptive and group.get('adaptive', False)
        for p in group['params']:
            if p.grad is None:
                continue

            grad = p.grad
            if adaptive:
                if grad.dtype in (torch.float16, torch.bfloat16):
                    grad = grad.float()
                grad = grad * p.abs()

            grads.append(grad)

    if not grads:
        return torch.zeros(1, dtype=torch.float32, device=device)

    norms = (
        torch._foreach_norm(grads)
        if has_foreach_support(grads) and grads[0].dtype not in (torch.float16, torch.bfloat16)
        else [(grad.float() if grad.dtype in (torch.float16, torch.bfloat16) else grad).norm() for grad in grads]
    )
    squared_norms = torch.stack([norm.to(device) for norm in norms]).float().square_()
    return squared_norms.sum().reshape(1)

def reg_noise(
    network1: nn.Module,
    network2: nn.Module,
    num_data: int,
    lr: float,
    eta: float = 8e-3,
    temperature: float = 1e-4,
) -> torch.Tensor:
    """Compute the Entropy-MCMC coupling and noise term for two networks.

    Usage example and detailed implementation can be found at:
    https://github.com/lblaoke/EMCMC/blob/master/exp/cifar10_emcmc.py

    Args:
        network1: First neural network.
        network2: Second neural network.
        num_data: Number of training data points.
        lr: Learning rate.
        eta: Eta parameter controlling auxiliary guiding variable.
        temperature: Temperature parameter for sampling.

    """
    reg_coef: float = 0.5 / (eta * num_data)
    noise_coef: float = math.sqrt(2.0 / lr / num_data * temperature)

    loss = torch.tensor(0.0, device=next(network1.parameters()).device)

    for param1, param2 in zip(network1.parameters(), network2.parameters()):
        reg = (param1 - param2).pow_(2).mul_(reg_coef).sum()

        noise = param1 * torch.randn_like(param1)
        noise.add_(param2 * torch.randn_like(param2))

        loss.add_(reg - noise.mul_(noise_coef).sum())

    return loss

@torch.no_grad()
def copy_stochastic(target: torch.Tensor, source: torch.Tensor) -> None:
    """Copy float32 values to bfloat16 with stochastic rounding.

    reference: https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905

    Args:
        target: A tensor in bfloat16 format to copy to.
        source: A tensor in float32 format to copy from.

    """
    result = torch.randint_like(
        source,
        dtype=torch.int32,
        low=0,
        high=1 << 16,
    )

    result.add_(source.view(dtype=torch.int32))

    result.bitwise_and_(-65536)

    target.copy_(result.view(dtype=torch.float32))

@torch.no_grad()
def batched_power_iteration(matrix: torch.Tensor, num_iters: int = 10, num_vectors: int = 16) -> torch.Tensor:
    """Estimate the largest eigenvalue of each positive semidefinite matrix in a batch.

    Args:
        matrix: Batch of symmetric positive semidefinite matrices with shape `(batch_size, size, size)`.
        num_iters: Number of power iterations.
        num_vectors: Number of parallel starting vectors per matrix.

    Returns:
        torch.Tensor: Largest Rayleigh quotient per matrix, with shape `(batch_size, 1, 1)` and the input data type.

    """
    vectors = torch.randn((*matrix.shape[:2], num_vectors), device=matrix.device, dtype=matrix.dtype)
    tiny = torch.finfo(vectors.dtype).tiny
    vectors.div_(torch.linalg.vector_norm(vectors, dim=1, keepdim=True).clamp_min_(tiny))

    product = torch.empty_like(vectors)
    for _ in range(num_iters):
        torch.bmm(matrix, vectors, out=product)
        torch.div(product, torch.linalg.vector_norm(product, dim=1, keepdim=True).clamp_min_(tiny), out=vectors)

    torch.bmm(matrix, vectors, out=product)

    return (vectors * product).sum(dim=1).amax(dim=1).view(-1, 1, 1)

@torch.no_grad()
def compute_power_newton_db(
    matrix: torch.Tensor, scale: torch.Tensor, num_iters: int = 10, inverse: bool = True
) -> torch.Tensor:
    """Compute batched square roots or inverse square roots with Newton-Denman-Beavers iterations.

    Args:
        matrix: Batch of symmetric positive definite matrices with shape `(batch_size, size, size)`.
        scale: Positive normalization factors with shape `(batch_size, 1, 1)` that bound the largest eigenvalues.
        num_iters: Positive number of Newton-Denman-Beavers iterations.
        inverse: Return inverse square roots when `True`, or square roots otherwise.

    Returns:
        torch.Tensor: Approximate matrix roots with the same shape and data type as `matrix`.

    """
    y = matrix / scale
    correction = y.mul(-0.5)
    correction.diagonal(dim1=-2, dim2=-1).add_(1.5)
    z = correction.clone()

    if num_iters > 1 or not inverse:
        y = y @ correction

    scratch = torch.empty_like(y)
    for _ in range(1, num_iters - 1):
        torch.bmm(z, y, out=correction)
        correction.mul_(-0.5).diagonal(dim1=-2, dim2=-1).add_(1.5)

        torch.bmm(y, correction, out=scratch)
        y, scratch = scratch, y

        torch.bmm(correction, z, out=scratch)
        z, scratch = scratch, z

    if num_iters > 1:
        torch.bmm(z, y, out=correction)
        correction.mul_(-0.5).diagonal(dim1=-2, dim2=-1).add_(1.5)

        if inverse:
            torch.bmm(correction, z, out=scratch)
        else:
            torch.bmm(y, correction, out=scratch)

        result = scratch
    else:
        result = z if inverse else y

    return result.div_(scale.sqrt()) if inverse else result.mul_(scale.sqrt())