[rawtxapp]

TPUs are amazing, but in my experience, debugging issues with them can be a bit tricky. Since nvidia's gpus are more common place (especially outside gcp), you can find a lot more information when you get stuck, it's also more battle tested, etc.

[6gvONxR4sf7o]

For what it's worth, jax is helpful to me here. You can drop out of the jit to debug it as if it were numpy.

Of course that assumes your issues aren't with the jit itself or inside pmap, etc. That shit's hard.