jnp.allclose for Tolerance Equality

Why this matters

Floating-point arithmetic makes exact equality comparisons unreliable. jnp.allclose(a, b, atol=k) returns True when every element satisfies |a - b| <= atol + rtol * |b|. This is the standard JAX idiom for asserting numerical correctness in tests, comparing two implementations of the same function, or verifying solver convergence.

Worked mini-example

import jax.numpy as jnp

a = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([1.0001, 2.0001, 3.0001])

jnp.allclose(a, b, atol=1e-3)   # True  — within 0.001
jnp.allclose(a, b, atol=1e-5)   # False — not within 0.00001

Combine with jnp.where for a jit-safe conversion from boolean to float:

jnp.where(jnp.allclose(a, b, atol=1e-3), 1.0, 0.0)   # → 1.0

Common pitfalls

• Don’t use == for float comparisons — a == b on float32 arrays fails for any rounding mismatch, even within acceptable tolerances.
• rtol also matters — the default rtol=1e-5 adds a relative term rtol * |b|; for very small magnitudes set both atol and rtol explicitly.
• jnp.allclose returns a scalar boolean — inside jax.jit you cannot use it in a Python if statement. Use jnp.where instead.
• Shape mismatch raises an error — a and b must be broadcastable.

Problem

Implement values_equal_within_atol(a, b, atol) that:

1. Calls jnp.allclose(a, b, atol=atol).
2. Returns 1.0 if close, 0.0 otherwise — using jnp.where.

• a, b: 1-D JAX arrays of the same shape.
• atol: scalar absolute tolerance.

Returns: scalar — 1.0 or 0.0.