JAX – NumPy on the CPU, GPU, and TPU
jax.readthedocs.io
JAX – NumPy on the CPU, GPU, and TPU
1–10 of 147 posts
Re: JAX – NumPy on the CPU, GPU, and TPU
#2What happened to Jax? Is it still alive?
Re: JAX – NumPy on the CPU, GPU, and TPU
#3What happened to Jax? Is it still alive?
Development seems not to have dropped at all from the contributions page: https://github.com/google/jax/graphs/contributors
Don’t know about usage and uptake though.
Re: JAX – NumPy on the CPU, GPU, and TPU
#4Anybody using it in production? Is it, or its derivatives like Flax, worth using over pyTorch for anything?
edit: Made comparison more fair.
Re: JAX – NumPy on the CPU, GPU, and TPU
#5What happened to Jax? Is it still alive?
Very much alive. From what I can tell it has more or less replaced Tensorflow for research purposes. (A lot of researchers use PyTorch though.)
Re: JAX – NumPy on the CPU, GPU, and TPU
#6Anybody using it in production? Is it, or its derivatives like Flax, worth using over pyTorch for anything? edit: Made comparison more fair.
I’m a researcher, not using anything in production, but I find jax more usable as a general GPU-accelerated tensor math library. PyTorch is more specifically targeted at the neural network use case. It can be shoehorned into other use cases, but is clearly designed & documented for NN training & inference.
Re: JAX – NumPy on the CPU, GPU, and TPU
#7JAX GPU support is limited to Linux only. Even the WSL2 support is experimental. https://jax.readthedocs.io/en/latest/installation.html#suppo...
Re: JAX – NumPy on the CPU, GPU, and TPU
#8It took me a while to realize it, but Jax is actually a huge opportunity for a lot of scientific computing. Jax was originally developed as a more flexible platform for doing machine learning research. But Jax's real superpower is that it bundles XLA and makes it really easy to run computations on GPU or TPU. And huge swathes of scientific computation basically run large scale vectorized computations.
When I was in astronomy (about a decade ago) I did large scale simulations of gravitational interactions. But at the time all these simulations were done on CPU. Some of the really big efforts used more specialized chips, but it was a huge effort to write the code for it.
But today with Jax, if you want to write an N-body simulation of a globular cluster, you can just code it up in numpy and it'll run on a GPU for free and be about 1000x faster. From what I can tell though, very few people in the sciences have caught on yet.
Re: JAX – NumPy on the CPU, GPU, and TPU
#9Anybody using it in production? Is it, or its derivatives like Flax, worth using over pyTorch for anything? edit: Made comparison more fair.
Not a fair comparison IMO. Jax is low level library used to make ML frameworks while pytorch is a full blow ML framework.
In terms of is it worth using it - that depends on what you're doing. If you just want to start with ML training probably not. If you have something already and you want to take it to next level (e.g. influence how training and inference work) than it's a good choice. You might be interested in looking into flax or haiku instead of using vanilla Jax. These are closer to pytorch.
Re: JAX – NumPy on the CPU, GPU, and TPU
#10Anybody using it in production? Is it, or its derivatives like Flax, worth using over pyTorch for anything? edit: Made comparison more fair.
I’m a researcher, not using anything in production, but I find jax more usable as a general GPU-accelerated tensor math library. PyTorch is more specifically targeted at the neural network use case. It can be shoehorned into other use cases, but is clearly designed & documented for NN training & inference.
Agreed. I used Jax about a year ago to estimate some diode parameters for a side project of mine.