5 ms·
NumPy does not automatically use GPU, so my guess is that the speedup is due to the CPU/GPU differential. This is specifically noted later on in the "Should I u
by brilee 5y ago
NumPy does not automatically use GPU, so my guess is that the speedup is due to the CPU/GPU differential. This is specifically noted later on in the "Should I use JAX" section - if you just want numpy for accelerators, JAX is a no-brainer.
- SleekEagle 5y agoThanks for your comment - in the JIT section under "Jax Transformations" I implement the same calculation on NP and JAX on CPU (and TPU) to get a direct comparison for CPU. I mentioned in another comment that the calculation lends itself well to being JITted, but I have some preliminary data for a follow-up JAX vs NumPy performance comparison. It's clear so far that NumPy can outperform JAX on CPU for small computations (unsurprisingly). Please let me know if you'd be interested in seeing a more thorough analysis!