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
orpmap
). - 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 1400+ tech skills courses.