ML Katas

by: newest upvotes saves

Debugging JAX code with `jax.debug.print`

yesterday easy (<10 mins) | debugging jax jit

In JAX, standard Python `print` statements don't always work as expected within `jit`-compiled functions because they execute at trace time. [11] The solution is to use `jax.debug.print`. [11, 23]...

1