🧪 testax: Jit-able runtime assertions for JAX in NumPy style.

[tillahoffmann]

testax provides jit-able runtime assertions for JAX through the testing interface familiar to NumPy users.

Overview

>>> import jax
>>> from jax import numpy as jnp
>>> import testax
>>>
>>> def safe_log(x):
...     testax.assert_array_less(0, x)
...     return jnp.log(x)
>>>
>>> safe_log(jnp.arange(2))
Traceback (most recent call last):
    ...
jax._src.checkify.JaxRuntimeError:
Arrays are not less-ordered

Mismatched elements: 1 / 2 (50%)
Max absolute difference: 1
Max relative difference: 1
 x: Array(0, dtype=int32, weak_type=True)
 y: Array([0, 1], dtype=int32)

testax assertions are jit-able, although errors need to be functionalized to conform to JAX’s requirement that functions are pure and do not have side effects (see the checkify guide for details). In short, a checkify-d function returns a tuple (error, value). The first element is an error that may have occurred, and the second is the return value of the original function.

>>> jitted = jax.jit(safe_log)
>>> checkified = testax.checkify(jitted)
>>> error, y = checkified(jnp.arange(2))
>>> error.throw()
Traceback (most recent call last):
    ...
jax._src.checkify.JaxRuntimeError:
Arrays are not less-ordered

Mismatched elements: 1 / 2 (50%)
Max absolute difference: 1
Max relative difference: 1
 x: Array(0, dtype=int32, weak_type=True)
 y: Array([0, 1], dtype=int32)
>>> y
Array([-inf,   0.], dtype=float32)

[patrick-kidger]

equinox.error_if may also be of interest here.

[tillahoffmann]

Neat! Is there an option to apply the string formatting only once an error is raised? E.g., using the current implementation, all the formatting happens in __str__ of the exception class which is only called if an error occurred.

[patrick-kidger]

Not right now, but it's come up before, and I'd be happy to take a PR adding this: patrick-kidger/equinox#644