Hacker News Reader

Comment by sandGorgon | Hacker News Reader

I have a Jax question here - could someone explain what the PyMC developers exactly did to move their C/C++ backend to Jax ?

Because this sounds very numba-esque. I always thought Jax was just a math library, that was slightly more usable than Numpy. This article makes it seem that code written in Jax ends up being significantly faster than Numpy...almost close to C

That's NUMBA territory

Also, Tensorflow probability is moving to JAX.

View on HN

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.

understood. and what about the cudf/cupy libraries ? is jax better ?

JAX 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/pymc-devs/Theano-PyMC/issues/10

Reply on news.ycombinator.com