4 ms·
I think float16 is called half precision. Jax uses as default float32 and numpy uses float64. I didn't check how fast numpy is if you force the calculations to
by credit_guy 3y ago
I think float16 is called half precision. Jax uses as default float32 and numpy uses float64. I didn't check how fast numpy is if you force the calculations to be a different precision. I'll give it a try and come back.
- credit_guy 3y agoAs 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.