3 ms·

> default single-precision Did't know about that. So basically attach .astype(np.float16) to any native numpy array can have the same speed up without any exis

• 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 ago
    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.