Numerics¶

The numerical helpers define finite behavior for common operations whose plain JAX derivatives are undefined at singular inputs. They are useful when a valid model configuration can contain a zero vector, zero distance, or removable quotient.

The default conventions are:

use the analytic limiting value where it exists, and zero otherwise;
define the derivative to be zero exactly at the singular point; and
match the corresponding JAX operation in value and derivative away from the singular point.

For example:

import jax
import jax.numpy as jnp

from soromox.utils import safe_norm, safe_normalize

zero = jnp.zeros(3)
gradient = jax.grad(safe_norm)(zero)  # [0., 0., 0.]
direction = safe_normalize(zero)     # [0., 0., 0.]

If only the squared norm is needed, prefer jnp.sum(x**2, axis=...); it is already finite everywhere and avoids an unnecessary square root. Use safe_norm when the norm itself is required.

Diagnosing Singularities¶

The finite defaults can hide where a model reaches a singular configuration. Use soromox.autodiff.strict_singularities_mode() while diagnosing a model to restore the underlying JAX operations and expose non-finite results.

API Reference¶

soromox.utils.numerics ¶

Numerical helpers shared across JAX utilities.

Singular-point conventions used by the safe_* helpers:

The value is the analytic limit where one exists, otherwise zero.
The derivative is defined to be zero at the singular point.
Away from the singular point the helpers match their plain counterparts (jnp.sqrt, jnp.linalg.norm, /) in both value and derivative.

Choosing between them:

jnp.linalg.norm(x, axis=k) ** 2 is exactly jnp.sum(x**2, axis=k). Prefer the latter whenever a norm is only ever squared: it needs no threshold and is finite everywhere.
Use :func:safe_norm only when the norm's own value is needed.

Guarding a singular expression only works on its input, never on its output. jnp.where and lax.select evaluate both sides and their transpose multiplies the unused side by a materialized zero, so a NaN produced there survives as 0 * NaN::

jnp.where(x > 0, jnp.sqrt(x), 0.0)       # broken at x == 0
x_safe = jnp.where(x > 0, x, 1.0)        # safe: no branch is ever singular
jnp.where(x > 0, jnp.sqrt(x_safe), 0.0)

For a scalar predicate, lax.cond evaluates and differentiates only the taken branch. Under vmap, however, batched predicates may lower to a select operation and both branch computations can execute. Branch inputs must therefore still be sanitized when a condition is vectorized.

The fallbacks hide where a model reaches a singular configuration. To find that out, wrap the call in soromox.autodiff.strict_singularities_mode(): the helpers then behave like the plain operations they replace, so the singular input surfaces as NaN or inf.

eps_for_dtype ¶

eps_for_dtype(eps: float | Array | None, dtype: dtype) -> Array

Return an epsilon scalar as a JAX array with the requested dtype.

Parameters:

Name	Type	Description	Default
`eps`	`float \| Array \| None`	Optional epsilon value. If `None`, machine epsilon for `dtype` is used.	required
`dtype`	`dtype`	JAX floating-point dtype for the returned scalar.	required

Returns:

Type	Description
`Array`	Scalar JAX array with dtype `dtype`.

safe_sqrt ¶

safe_sqrt(x: Array) -> Array

Return sqrt(max(x, 0)) with a zero derivative at the origin.

jnp.sqrt has an infinite derivative at zero, which becomes NaN as soon as it is multiplied by a cotangent.

Under strict_singularities_mode this is plain jnp.sqrt, so a negative or zero argument surfaces instead of being clamped.

Parameters:

Name	Type	Description	Default
`x`	`Array`	Real-valued array. Negative entries are treated as zero.	required

Returns:

Type	Description
`Array`	Array with the same shape as `x`.

Example

import jax
import jax.numpy as jnp
from soromox.utils import safe_sqrt

jax.grad(safe_sqrt)(jnp.array(0.0))  # 0.0 rather than inf

safe_norm ¶

safe_norm(x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array

Return the Euclidean norm of x with a zero derivative at the origin.

jnp.linalg.norm differentiates to x / norm(x), which is 0 / 0 at the zero vector. Singular configurations such as a straight backbone or a collapsed routing path reach that point exactly.

Under strict_singularities_mode the zero vector differentiates to NaN, as the plain norm does.

Parameters:

Name	Type	Description	Default
`x`	`Array`	Real-valued array.	required
`axis`	`int \| tuple[int, ...] \| None`	Static axis or axes to reduce over. If `None`, the norm is taken over the flattened array.	`None`
`keepdims`	`bool`	Static flag. If `True`, reduced axes are retained with length one, which is convenient when dividing `x` by its norm.	`False`

Returns:

Type	Description
`Array`	Array containing the Euclidean norm, reduced over `axis`.

Example

import jax
import jax.numpy as jnp
from soromox.utils import safe_norm

jax.grad(safe_norm)(jnp.zeros(3))  # [0. 0. 0.] rather than [nan nan nan]

safe_divide ¶

safe_divide(numerator: Array, denominator: Array, eps: float | Array, fallback: float | Array = 0.0) -> Array

Divide by a denominator that may vanish, keeping derivatives finite.

The denominator is sanitized before the division, so neither branch ever evaluates a singular quotient.

Parameters:

Name	Type	Description	Default
`numerator`	`Array`	Dividend array.	required
`denominator`	`Array`	Divisor array. Entries with magnitude at or below `eps` select `fallback` instead of the quotient.	required
`eps`	`float \| Array`	Small nonnegative threshold below which the denominator is treated as zero.	required
`fallback`	`float \| Array`	Value returned where the denominator vanishes. Broadcast against the quotient shape.	`0.0`

Under strict_singularities_mode this is a plain division, so a vanishing denominator surfaces as inf or NaN.

Returns:

Type	Description
`Array`	Array containing `numerator / denominator` away from the singularity
`Array`	and `fallback` at it.

Notes

The two jnp.where calls are safe here because the denominator is sanitized before division. Thus both evaluated branches remain finite.

safe_normalize ¶

safe_normalize(x: Array, axis: int | tuple[int, ...] | None = None, eps: float | Array | None = None, fallback: Array | None = None) -> Array

Scale x to unit length, keeping the zero vector finite.

Parameters:

Name	Type	Description	Default
`x`	`Array`	Real-valued array to normalize.	required
`axis`	`int \| tuple[int, ...] \| None`	Static axis or axes holding the vector components. If `None`, the whole array is treated as a single vector.	`None`
`eps`	`float \| Array \| None`	Small nonnegative norm threshold. Defaults to machine epsilon for the dtype of `x`.	`None`
`fallback`	`Array \| None`	Value returned where the norm vanishes. Defaults to zeros, which keeps the result finite without inventing a direction.	`None`

Returns:

Type	Description
`Array`	Array with the same shape as `x`.

Note

Under strict_singularities_mode the zero vector surfaces as NaN rather than fallback.