Is JAX in competition with TF? Is it something you use instead of TF, or in combination with it?
Instead of.
Using JAX to Accelerate Research
11–15 of 15 posts
Re: Using JAX to Accelerate Research
#12The description says it's autograd + XLA. So I assumed it always compiles to GPUs via XLA.
But, had a look in the code and jax has cublas and RoCm blas, and it looks like there is a flow where it uses the gpu directly, unless I'm missing something.
Definitely worth having a closer look. Autograd via function reflection should be faster than backprop. And if it's running on AMD GPUs then it's quite intriguing.
Re: Using JAX to Accelerate Research
#13I find JAX really exciting. The idea of numpy with autograd is exactly what Pythonistas want. The elephant in the room, though, is “why not Pytorch?” Everyone knows JAX is what Google realised Tensorflow should have been when they realised how much of a joy Pytorch was to use. I actually think JAX does offer some advantages, not least true numpy interoperability. However, not mentioning *torch a single time in the bl…
> why not Pytorch?
JAX enables using (parts of) existing numpy codebases in disciplines other than deep learning. Autodiff and compilation to GPUs are very useful for all kinds of algorithms and processing pipelines.
Re: Using JAX to Accelerate Research
#14I find JAX really exciting. The idea of numpy with autograd is exactly what Pythonistas want. The elephant in the room, though, is “why not Pytorch?” Everyone knows JAX is what Google realised Tensorflow should have been when they realised how much of a joy Pytorch was to use. I actually think JAX does offer some advantages, not least true numpy interoperability. However, not mentioning *torch a single time in the bl…
JAX and Pytorch have somewhat different scopes though... JAX is concerned about much more than the kinds of neural nets we write today. It's a general system for expressing and transforming numerical programs, and the devs are as genuinely excited about e.g. scientific programming, probabilistic modeling, etc. as they are about NNs. A technical reason for "why not pytorch" is that JAX was also built in part to expose…
Pytorch might be suboptimal but it's what's already there e.g. for Pyro. (PyMC3 uses Theano; can Theano drive JAX?)
OTOH I'm not sure most people know what tasks are GPU-worthy or not. I haven't the slightest idea of why MCMC/Variational Bayes is amenable to GPU speedups and Persistent Homology isn't.
Re: Using JAX to Accelerate Research
#15Earlier quoted context omitted.
JAX and Pytorch have somewhat different scopes though... JAX is concerned about much more than the kinds of neural nets we write today. It's a general system for expressing and transforming numerical programs, and the devs are as genuinely excited about e.g. scientific programming, probabilistic modeling, etc. as they are about NNs. A technical reason for "why not pytorch" is that JAX was also built in part to expose…
Pytorch might be suboptimal but it's what's already there e.g. for Pyro. (PyMC3 uses Theano; can Theano drive JAX?) OTOH I'm not sure most people know what tasks are GPU-worthy or not. I haven't the slightest idea of why MCMC/Variational Bayes is amenable to GPU speedups and Persistent Homology isn't.
> can Theano drive JAX?
Yes, that is in active development: https://pymc-devs.medium.com/the-future-of-pymc3-or-theano-i...