jnp.allclose for Tolerance Equality

Compare results with a tolerance

Two calculations can give slightly different floating-point results because they round at different steps. When small differences are acceptable, use a tolerance instead of requiring exact equality. Exact equality is still useful when exact agreement is the requirement.

For finite values, jnp.allclose(a, b, atol=atol) checks every pair using:

abs(a - b) <= atol + rtol * abs(b)

atol is an absolute allowance, useful near zero. rtol scales the allowance with the magnitude of b. This exercise keeps JAX's default rtol=1e-5, so the result is not an absolute-tolerance-only comparison. For that different check, you would explicitly set rtol=0.

Worked example

import jax.numpy as jnp

a = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
b = jnp.array([1.0001, 2.0001, 3.0001], dtype=jnp.float32)
loose = jnp.allclose(a, b, atol=1e-3)  # True
tight = jnp.allclose(a, b, atol=1e-5)  # False
result = jnp.where(loose, 1.0, 0.0)   # scalar 1.0

At the first element, the tight allowance is about 0.00002, smaller than the difference of about 0.0001. One failing pair is enough to make allclose false. At larger magnitudes, the relative term can allow differences greater than atol alone.

Floating-point numbers have limited range and precision

A floating-point number stores significant digits and an exponent. The exponent limits its range; the number of significant digits limits its precision. Overflow occurs when a result is too large to represent. Very small values can lose precision or round to zero through underflow. Ordinary rounding can also make nearby values indistinguishable.

Subtracting nearly equal values can leave a result dominated by earlier rounding errors. For tiny x, computing log(1 + x) can fail because 1 + x rounds to 1 before the logarithm is taken. The specialized function log1p(x) avoids that loss. Similarly, expm1(x) computes exp(x) - 1 accurately near zero.

Use formulas with safer intermediate values

A numerically stable formula computes the same mathematical quantity while reducing avoidable rounding error or overflow. For finite, nonempty inputs, log-sum-exp can be computed as

$$\log\sum_i e^{z_i}=m+\log\sum_i e^{z_i-m},\quad m=\max_i z_i.$$

After subtracting m, every exponent is at most zero. This avoids exponentiating large positive numbers. For scores [1000, 1000], the result is 1000 + log(2) without ever computing exp(1000).

Compare with a tolerance appropriate to the task

Floating-point addition is not associative: changing the order of a sum can change its rounded result. Parallel reductions and different matrix-multiplication kernels can therefore produce slightly different answers.

Approximate comparisons usually allow an absolute error near zero and a relative error that scales with the reference value. Choose tolerances suited to the problem's scale and dtype. A loose tolerance should not hide a wrong algorithm.

NaN does not compare equal to itself. Decide how invalid values should be handled rather than assuming an ordinary equality test will catch them. Replacing a NaN after an invalid operation may still leave invalid gradients. When differentiating, avoid the invalid operation itself where possible.