3 ms·
JAX can be used as a drop-in replacement for numpy but it enables JIT compilation (similarly to numba) and, more important, very efficient GPU compilation. Num
by nestorD 6y ago
JAX can be used as a drop-in replacement for numpy but it enables JIT compilation (similarly to numba) and, more important, very efficient GPU compilation.
Numba is great for CPU-only code, JAX would be my goto solution to add GPU on top of a numpy code.
Theano had a numpy backend so moving to JAX was very straightforward and enabled GPU computations which is now a must for fast deep learning.
- sandGorgon 6y agounderstood. and what about the cudf/cupy libraries ? is jax better ?
- nestorD 6y agoJAX does JIT compilation which is not done by Cupy so you get better performance by fusing operations and compiling them cleverly ([0] is a good benchmark to get a feeling for the performance of the various options). However, it is a trade-off, the jitting time will increase your run time (you can expect at least 10 seconds of compiling time) and JAX does not let you do in-place operations (however, it has workarounds). If you can afford the time needed to jit your code then I highly recommend JAX. [0]: https://github.com/dionhaefner/pyhpc-benchmarks https://github.com/dionhaefner/pyhpc-benchmarks