3 ms·

I don't get it. If code written in numpy is slow because the algorithm is bad, and you need to rewrite it, we can't say "numpy is too slow". A good chunk of the

I don't get it. If code written in numpy is slow because the algorithm is bad, and you need to rewrite it, we can't say "numpy is too slow". A good chunk of the blogpost talks about this particular case, which is a non-case. The rest is some vague suggestions to use numba, or jax. But if your code is vectorizable in my experience you don't get any benefit from numba and jax. With the exception that jax is by default single-precision so you get some speed up from that. If you make it double precision, you get back to numpy speed. I suppose the author's conclusion is not all that different, otherwise they'd put a code snippet and show some numbers.

tnecniv 3y ago
I have definitely gotten significant speed ups using Jax. I try to vectorize my code as much as possible but that’s not always feasible for every aspect of an algorithm. Even for the vectorized operations, the JIT compilation gets me a speed up when used correctly (I.e. structuring code to minimize recompiling). Jax also has the added bonus of trivially running your algorithm on the GPU.
akasakahakada 3y ago
> 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 existence of jax.
- 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.

tnecniv 3y ago

I have definitely gotten significant speed ups using Jax. I try to vectorize my code as much as possible but that’s not always feasible for every aspect of an algorithm. Even for the vectorized operations, the JIT compilation gets me a speed up when used correctly (I.e. structuring code to minimize recompiling). Jax also has the added bonus of trivially running your algorithm on the GPU.

akasakahakada 3y ago

> 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 existence of jax.

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.