def _sqrt_positive_part(x):
    """
    Returns torch.sqrt(torch.max(0, x))
    but with a zero subgradient where x is 0.
    """
    ret = torch.zeros_like(x)
    positive_mask = x > 0
    ret[positive_mask] = torch.sqrt(x[positive_mask])
    return ret

    Returns torch.sqrt(torch.max(0, x))
    but with a zero subgradient where x is 0.
    
mGPT/utils/rotation_conversions.py