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

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.