chex (pypi): Chex: Testing made fun, in JAX!
Keywords: jax|testing|debugging|python|machine|learning