safe_norm

Contents

safe_norm#

class braintools.metric.safe_norm(x, min_norm, ord=None, axis=None, keepdims=False)#

Compute vector or matrix norm with gradient-safe lower bound.

Calculates the norm of input arrays while ensuring the result is at least min_norm, with proper gradient handling. This prevents gradient issues that arise when using jax.numpy.maximum(jax.numpy.linalg.norm(x), min_norm) directly, which can produce NaN gradients at zero due to JAX evaluating both branches of the maximum operation.

Parameters:

  • x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Input array for which to compute the norm. Can be of any shape.

  • min_norm (float) – Minimum value for the returned norm. The result will be at least this value.

  • ord (int | float | str | None) –

    Order of the norm:

    • For vectors: any real number, inf, -inf (default: 2)

    • For matrices: ‘fro’ (Frobenius), ‘nuc’ (nuclear), inf, -inf, 1, -1, 2, -2

  • axis (None | tuple[int, ...] | int) –

    Axis or axes along which to compute the norm:

    • None: flatten array and compute single norm

    • int: compute norms along specified axis

    • tuple: for matrix norms, specify the two axes defining matrices

  • keepdims (bool) – If True, the reduced axes are left as dimensions with size one, allowing the result to broadcast against the original array.

Returns:

Array norms with values bounded below by min_norm. Shape depends on axis and keepdims parameters, following jax.numpy.linalg.norm conventions.

Return type:

Array | ndarray | bool | number | bool | int | float | complex | Quantity

Notes

This function addresses a specific issue with automatic differentiation where the gradient of max(norm(x), min_norm) is undefined (NaN) when norm(x) = 0. The implementation ensures that:

  • Forward pass: Returns max(norm(x), min_norm)

  • Backward pass: Provides well-defined gradients even at zero

The gradient handling works by masking the input when the norm would be below the threshold, ensuring the gradient computation path remains valid.

Examples

Basic usage with vector norms:

>>> import jax.numpy as jnp
>>> from braintools.metric import safe_norm
>>> x = jnp.array([0.0, 0.0, 0.0])  # Zero vector
>>> safe_norm(x, min_norm=1e-8)  # Returns 1e-8 instead of 0.0
Array(1.e-08, dtype=float32)

Compare with regular norm:

>>> jnp.linalg.norm(x)  # Returns 0.0
Array(0., dtype=float32)
>>> safe_norm(x, min_norm=0.1)  # Returns 0.1
Array(0.1, dtype=float32)

Matrix norms with axis specification:

>>> X = jnp.array([[1.0, 2.0], [3.0, 4.0]])
>>> # L2 norm along rows
>>> safe_norm(X, min_norm=0.1, axis=1)
Array([2.236068, 5.      ], dtype=float32)

See also

jax.numpy.linalg.norm

Standard norm computation without lower bound

jax.numpy.maximum

Element-wise maximum operation

Contents