jaxmat.tensors.utils module

safe_fun(fun, x, norm=None, eps=1e-16)

Safely applies a function to an input, avoiding numerical issues near zero.

This function applies fun(x) only when the norm of x exceeds a small tolerance eps. Otherwise, it returns zero. This is useful for ensuring numerical stability in cases where evaluating fun at or near zero could result in undefined or unstable behavior (e.g., division by zero).

Parameters:

• fun (Callable) – The function to apply safely.

• x (array-like) – Input array or tensor.

• norm (Callable, optional) – A norm or magnitude function used to test whether x is sufficiently large. Defaults to the identity function.

• eps (float, optional) – Small threshold to determine whether x is treated as nonzero. Defaults to 1e-16.

Returns:

fun(x) if norm(x) > eps, otherwise 0 (of the same shape as x).

Return type:

array-like

safe_sqrt(x, eps=1e-16)

Computes a numerically safe square root.

Ensures the argument to the square root is greater than eps to avoid taking the square root of zero or negative values, which could cause instability or NaNs.

Parameters:

• x (array-like) – Input array or tensor.

• eps (float, optional) – Minimum threshold for x before taking the square root. Defaults to 1e-16.

Returns:

The square root of x for x > eps, otherwise eps.

Return type:

array-like

safe_norm(x, eps=1e-16, **kwargs)

Wrapper around optax.safe_norm that computes a numerically stable norm.

This function prevents numerical instability when computing vector norms for small magnitudes by internally applying a stability threshold.

Parameters:

• x (array-like) – Input vector or tensor.

• eps (float, optional) – Small constant added for numerical stability. Defaults to 1e-16.

• **kwargs – Additional arguments passed to optax.safe_norm.

Returns:

The numerically stable norm of x.

Return type:

array-like

FischerBurmeister(x, y)

Computes the scalar Fischer-Burmeister function.

The Fischer-Burmeister function is defined as:

\[\Phi(x, y) = x + y - \sqrt{x^2 + y^2}\]

and is commonly used in complementarity problem formulations to provide a smooth reformulation of the complementarity conditions

\[x \geq 0, y \geq 0, xy = 0\]

.