Chex

This lesson will introduce Chex.

We'll cover the following

Testing a numerical computation program can be tricky. This is especially the case when using JAX due to the parallel processing over GPU/TPU. The JAX ecosystem provides a library for it. Chex is a useful library with utilities like:

  • Assertions.
  • Debugging transformations (like vmap or pmap).
  • Testing code across JIT and non-JIT versions.

Assertions

Traditional PyType annotations do not support DeviceArray size or shape, so Chex provides assertions of its own.

Primitives

By using assert_shape() and assert_rank(), we can validate both the shape and dimension of a given JAX array.

Get hands-on with 1200+ tech skills courses.