I've found jax's debugging to be in different ways better and worse. The fact that the function transformations are traced is great. It means you can step debug

I've found jax's debugging to be in different ways better and worse. The fact that the function transformations are traced is great. It means you can step debug in the tracing steps just as well as the actual eval steps, and you just have jaxpr.Tracers instead of jnp.ndarrays, or whatever. Outside of the transformations, it's just as easy to debug as numpy, which is a blessing. That's one of the biggest selling points. Debugging jitted and pmapped code, on the other hand, is a pain. Since you can always step out of them to debug, it means that it's debugging performance issues that sucks. And boy does it suck. If anyone knows a good story for figuring out why my jitted thing is slow as hell on TPU, I'm all ears. The profiling section of the official docs is one of their weaker sections. (but big props to the overall documentation quality!)