As promised, I did the experiment. Converting all floats in numpy to float32 achieves a speedup of about 20%. Writing the same code in jax (with the few small c

As promised, I did the experiment. Converting all floats in numpy to float32 achieves a speedup of about 20%. Writing the same code in jax (with the few small changes needed for example for random number generation), the speedup is of a factor of 2. But once we switch back to float64, jax becomes a tiny bit slower than numpy, maybe about 2% slower. Adding compilation to jax does not help if the code is already fully vectorized.