utils

numerical

Numerical stability utilities for mixed-precision training.

General-purpose helpers for preventing NaN/Inf values in deep learning computations, particularly under FP16/BF16 autocast.

StabilityConfig

Default numerical stability bounds.

get_eps

get_eps(dtype)

Return an appropriate epsilon for the given dtype.

Source code in deltaflow/utils/numerical.py

def get_eps(dtype: torch.dtype) -> float:
    """Return an appropriate epsilon for the given dtype."""
    if dtype in (torch.float16, torch.bfloat16):
        return StabilityConfig.EPS_FP16
    return StabilityConfig.EPS_FP32

safe_normalize

safe_normalize(x, dim=-1, eps=None)

L2-normalize with numerical stability (avoids division by ~zero norms).

Source code in deltaflow/utils/numerical.py

def safe_normalize(x: torch.Tensor, dim: int = -1, eps: Optional[float] = None) -> torch.Tensor:
    """L2-normalize with numerical stability (avoids division by ~zero norms)."""
    if eps is None:
        eps = get_eps(x.dtype)
    norm = x.norm(p=2, dim=dim, keepdim=True)
    norm = torch.clamp(norm, min=eps)
    return x / norm

safe_sqrt

safe_sqrt(x, eps=None)

Compute a square root, clamping the input to be non-negative first.

Source code in deltaflow/utils/numerical.py

def safe_sqrt(x: torch.Tensor, eps: Optional[float] = None) -> torch.Tensor:
    """Compute a square root, clamping the input to be non-negative first."""
    min_val = eps if eps is not None else 0.0
    return torch.sqrt(torch.clamp(x, min=min_val))

clamp_noise

clamp_noise(noise, bound=None)

Clamp a noise tensor to prevent extreme values under FP16.

Source code in deltaflow/utils/numerical.py

def clamp_noise(noise: torch.Tensor, bound: Optional[float] = None) -> torch.Tensor:
    """Clamp a noise tensor to prevent extreme values under FP16."""
    bound = bound if bound is not None else StabilityConfig.NOISE_BOUND
    return torch.clamp(noise, -bound, bound)

clamp_prediction

clamp_prediction(pred, bound=None)

Clamp model predictions to prevent extreme values.

Source code in deltaflow/utils/numerical.py

def clamp_prediction(pred: torch.Tensor, bound: Optional[float] = None) -> torch.Tensor:
    """Clamp model predictions to prevent extreme values."""
    bound = bound if bound is not None else StabilityConfig.PREDICTION_BOUND
    return torch.clamp(pred, -bound, bound)

clamp_loss

clamp_loss(loss, max_val=None)

Clamp a scalar loss to prevent gradient explosion.

Source code in deltaflow/utils/numerical.py

def clamp_loss(loss: torch.Tensor, max_val: Optional[float] = None) -> torch.Tensor:
    """Clamp a scalar loss to prevent gradient explosion."""
    max_val = max_val if max_val is not None else StabilityConfig.LOSS_BOUND
    return torch.clamp(loss, min=0.0, max=max_val)

clamp_cosine_similarity

clamp_cosine_similarity(sim)

Clamp cosine similarity to the valid range [-1, 1].

Source code in deltaflow/utils/numerical.py

def clamp_cosine_similarity(sim: torch.Tensor) -> torch.Tensor:
    """Clamp cosine similarity to the valid range [-1, 1]."""
    return torch.clamp(sim, -StabilityConfig.COSINE_SIM_BOUND, StabilityConfig.COSINE_SIM_BOUND)