Chex: Testing made fun, in JAX!

Chex is a library of utilities for helping to write reliable JAX code. This includes utils to help: (1) instrument your code (e.g. assertions); (2) debug (e.g. transforming pmaps in vmaps within a context manager); and (3) test JAX code across many variants (e.g. jitted vs non-jitted).

Chex is a library of utilities for helping to write reliable JAX code. This includes utils to help: (1) instrument your code (e.g. assertions); (2) debug (e.g. transforming pmaps in vmaps within a context manager); and (3) test JAX code across many variants (e.g. jitted vs non-jitted).

OpenSource

26 Jul 2020