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)
  3. Test JAX code across many variants (e.g. jitted vs non-jitted)
No items found.