Function _norm_no_nan

s3f/gvp_layer.py:55–62

L2 norm of tensor clamped above a minimum value `eps`. :param sqrt: if `False`, returns the square of the L2 norm

(x, axis=-1, keepdims=False, eps=1e-8, sqrt=True)

def _norm_no_nan(x, axis=-1, keepdims=False, eps=1e-8, sqrt=True):
    '''
    L2 norm of tensor clamped above a minimum value `eps`.
    
    :param sqrt: if `False`, returns the square of the L2 norm
    '''
    out = torch.clamp(torch.sum(torch.square(x), axis, keepdims), min=eps)
    return torch.sqrt(out) if sqrt else out


def _split(x, nv):