Chex: Testing made fun, in JAX!
conda install anaconda::chex
Chex is a library of utilities for helping to write reliable JAX code.
This includes utils to help: