Using JAX in 2022(assemblyai.com)
assemblyai.com
Using JAX in 2022
https://www.assemblyai.com/blog/why-you-should-or-shouldnt-be-using-jax-in-2022/
32 コメント
This is a really great overview of JAX - the best I've seen outside of primary Google/DeepMind sources. Glad to see people besides Googlers getting familiar with its capabilities.
Thanks for your feedback! I wanted to introduce it for people who haven't worked with it before giving recommendations, but I got a little carried away because I think JAX is really cool!
I was planning on doing a more thorough introductory tutorial or deep dive into Transformations, so let me know if you (all) think that would be instructive!
I was planning on doing a more thorough introductory tutorial or deep dive into Transformations, so let me know if you (all) think that would be instructive!
> Why Should I Care About JAX?
> In short - speed.
For me personally, the magic of JAX is that it able to have this performance, while being as close as possible to having first class differentiation in Python. The latter is a far more important reason to use JAX. It can really change how you think about programming and ML. Rather than implementing a specific model, you can write up the parameterized solution to a problem then solve it.
However first class differentiation ultimately isn't really useful unless you happen to also solve the speed problem. That is what makes JAX incredible. From the programming perspective JAX is to differentiable programming what Prolog is to logic programming, however Prolog has always been limited ultimately by performance problems where JAX is not.
> In short - speed.
For me personally, the magic of JAX is that it able to have this performance, while being as close as possible to having first class differentiation in Python. The latter is a far more important reason to use JAX. It can really change how you think about programming and ML. Rather than implementing a specific model, you can write up the parameterized solution to a problem then solve it.
However first class differentiation ultimately isn't really useful unless you happen to also solve the speed problem. That is what makes JAX incredible. From the programming perspective JAX is to differentiable programming what Prolog is to logic programming, however Prolog has always been limited ultimately by performance problems where JAX is not.
Thanks for your comment - I wanted to keep it as general as possible, and I think speed touches basically every possible use of JAX, but I definitely hear what you are saying. For DL, jvps and quick/easy Hessians (along with speed) are really the reasons to use JAX imo.
I like your comment on thinking about implementing models/thinking about them as parameterized solutions
I like your comment on thinking about implementing models/thinking about them as parameterized solutions
I've been really struggling to see the benefit from autodiff... I work with DAGs, where the data processes are batches that could be processed in parallel and I would love to speed that up, or to enable some gradient descent style optimization.... but I can't figure out how autodiff plays into it.
Do you know of a good place to help learn how it could apply to my domain?
Firstly vmap() might make it easier for you from a programming standpoint. You can think of operations in terms of single instances and then autobatch them.
Also, jit() is where you'll likely see your speed increases.
As for autodiff - not every application will use it. What exactly are you trying to optimize? If you can parameterize your model somehow (e.g. parameterize edge weights) and then figure out some way to measure how "bad" your model is, you can use autodiff to tune your edge weights to minimize that metric. Not too familiar with DAGs but the first step is figuring out how to parameterize your model in a way that tuning the params can lead to your goal (and how to measure how close you are to that goal)
Also, jit() is where you'll likely see your speed increases.
As for autodiff - not every application will use it. What exactly are you trying to optimize? If you can parameterize your model somehow (e.g. parameterize edge weights) and then figure out some way to measure how "bad" your model is, you can use autodiff to tune your edge weights to minimize that metric. Not too familiar with DAGs but the first step is figuring out how to parameterize your model in a way that tuning the params can lead to your goal (and how to measure how close you are to that goal)
With the year so prominently in the title, i thought this was going to be about a technology that is obviously not trendy anymore, ie, JAX The Terrible Java XML Parser. I'm disappointed that it was just blog title spam about some non-controversial modern technology.
Wait, is the "The Terrible Java XML Parser" moniker tongue-and-cheek or sincere?
I remember JAX (the XML parser) really saving my bacon when I was parsing larger-than-memory XML files ages ago.
I remember JAX (the XML parser) really saving my bacon when I was parsing larger-than-memory XML files ages ago.
It was sincere, but I'll readily admit that the opinion was informed by a 15-year-younger me who definitely was not trying to parse larger-than-memory XML files.
I bet it was just a matter of the wrong tool for the job, but I (with some weird fondness, I might add), remember finding it extremely overengineered, even for Java standards.
I bet it was just a matter of the wrong tool for the job, but I (with some weird fondness, I might add), remember finding it extremely overengineered, even for Java standards.
If XML files fit in memory, then it totally feels like an astronomically more annoying way to parse files. Totally sympathize with the sentiment.
It’s not really spam, because it’s the opposite situation of what you describe: JAX is still in its infancy and doesn’t have high adoption yet, and as such, the user experience is changing rapidly over time. It’s valid to title your article something like “X in YEAR” when a few years ago and a few years from now the experience would be very different.
I wish JAX worked with windows natively (without using wsl). I teach a very high level intro to numpy and would _love_ to have my students try jax. These students are relatively new to programming and the idea of using a linux shell or having to compile anything themselves just wouldn't work.
I find this repo used to create Jaxlib wheels to Windows https://github.com/cloudhan/jax-windows-builder
VM with Ubuntu? If you write a guide they can get set up with VMWare pretty easily!
These benchmarks are pretty crazy, especially as I presumed NumPy to do far better. Isn’t tensorflow etc already doing GPU acceleration? Is this actually a fair comparison?
(I'm an engineer at DeepMind, and I work with JAX daily)
It's a somewhat fair comparison; in my experience, highly optimized JAX matches highly optimized Tensorflow. However, non-optimized (but JITted) JAX beats non-optimized Tensorflow, as Tensorflow requires a lot of architectural changes to make it perform well. JAX, on the other hand, tends to perform well as long as you just JIT it. So it's much easier to get to, say, 90% of optimal performance. In Tensorflow, it's much harder (in my experience- maybe I'm just bad at Tensorflow).
The JIT compilation that JAX does is really, really good, as it combines operations together in a highly performant way.
It's a somewhat fair comparison; in my experience, highly optimized JAX matches highly optimized Tensorflow. However, non-optimized (but JITted) JAX beats non-optimized Tensorflow, as Tensorflow requires a lot of architectural changes to make it perform well. JAX, on the other hand, tends to perform well as long as you just JIT it. So it's much easier to get to, say, 90% of optimal performance. In Tensorflow, it's much harder (in my experience- maybe I'm just bad at Tensorflow).
The JIT compilation that JAX does is really, really good, as it combines operations together in a highly performant way.
The first benchmark is wrong. Numpy random by default uses float64, and jax uses float32.
In [9]: x = np.random.randn(10000,10000).astype('f')
In [10]: %timeit -n5 fn(x)
623 ms ± 9.31 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)
In [17]: %timeit -n5 jax_fn(x).block_until_ready()
98.4 ms ± 1.55 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)
Thanks for this comment, I probably could've stressed JIT more. Random question - I noted in the article that you all at DeepMind announced that you're using JAX to accelerate your research.
IIRC, you guys standardized TensorFlow several years back. What does the current split look like between JAX and TF internally? Do some people use TF and some use JAX, or do you use JAX for specific tasks?
IIRC, you guys standardized TensorFlow several years back. What does the current split look like between JAX and TF internally? Do some people use TF and some use JAX, or do you use JAX for specific tasks?
[deleted]
> especially as I presumed NumPy to do far better
If you are running your benchmark on a single Numpy operation then I would expect Numpy to have equal or better speed (you are not paying for JIT compilation). However, when you are doing several operations, Numpy will do a loop on array elements for each operation while JAX can fuse everything, that can end up making a big difference.
If you are running your benchmark on a single Numpy operation then I would expect Numpy to have equal or better speed (you are not paying for JIT compilation). However, when you are doing several operations, Numpy will do a loop on array elements for each operation while JAX can fuse everything, that can end up making a big difference.
It's not quite a fair comparison, since Numpy is running with float64 while Jax is running with float32.
If you fix the benchmarks then looks like this
5 loops, best of 5: 99.2 ms per loop
10 loops, best of 5: 114 ms per loop
10 loops, best of 5: 20.2 ms per loop
5x faster is to be expected as there are 5 pointwise operations (that are bandwidth bound) that can be fused.
The leading comparison is also quite misleading, imo, since I think it's comparing Numpy on CPU vs. Jax on an accelerator.
If you fix the benchmarks then looks like this
5 loops, best of 5: 99.2 ms per loop
10 loops, best of 5: 114 ms per loop
10 loops, best of 5: 20.2 ms per loop
5x faster is to be expected as there are 5 pointwise operations (that are bandwidth bound) that can be fused.
The leading comparison is also quite misleading, imo, since I think it's comparing Numpy on CPU vs. Jax on an accelerator.
Thanks very much for posting this - I forgot JAX defaults to float32 - I'll fix that soon.
As for the other part about the leading comparison - I was trying to highlight just how much faster JAX could be in the best-case scenario. Beyond the accelerator and JIT, the function itself lends to being expedited significantly when JITted. I posted benchmarks with a comparison of JAX vs NumPy both on CPU, and then with JAX on TPU further down to control more variables. (reposted from reddit)
As for the other part about the leading comparison - I was trying to highlight just how much faster JAX could be in the best-case scenario. Beyond the accelerator and JIT, the function itself lends to being expedited significantly when JITted. I posted benchmarks with a comparison of JAX vs NumPy both on CPU, and then with JAX on TPU further down to control more variables. (reposted from reddit)
re: leading comparison, that makes sense. I recommend labeling it a bit better: JAX on TPU vs NumPy on CPU
Updated! Thanks for the feedback. Added this note in the figure description
"(n.b. JAX is using TPU and NumPy is using CPU in order to highlight that JAX's speed ceiling is much higher than NumPy's)"
"(n.b. JAX is using TPU and NumPy is using CPU in order to highlight that JAX's speed ceiling is much higher than NumPy's)"
Comparing Jax on TPU vs Numpy on CPU does not make any sense. Of course a GEMM hardware accelerator will be much faster than a general purpose CPU. What is the point of this comparison? You either run the same code on different hardware, or different code on the same hardware.
A much more interesting comparison would be CuPy or Pytorch code vs Jax code running on A100.
A much more interesting comparison would be CuPy or Pytorch code vs Jax code running on A100.
NumPy can't run on accelerators, which is part of where the difference stems from. I implemented the same operation first with both NP/JAX on CPU, then with JAX on TPU and NP again on CPU, not sure if you saw that part.
Obviously, the ability to use an accelerator makes JAX faster, but even without that on CPU it was faster. This is in part because of JIT, and in fairness the calculation in question does lend itself well to being expedited by JIT.
I actually have some preliminary benchmarks for a follow up specifically on just NumPy vs JAX, and it has become clear so far that NumPy is better in certain cases, especially for small operations where the overhead of JAX is not worth it.
In the article I mention this briefly, along with how JAX hasn't been focused on being optimized on CPU because they have bigger fish to fry, so to speak. I also link to the JAX documentation that has some comments comparing the two!
Relating to TF - I don't actually use TF at any point, but I did use PyTorch for Hessian calculation. TF and PT obviously do both work on GPU, but JAX has the benefit of being able to JIT more and implement everything in terms of XLA (although TF obviously has XLA support as well, and PT kind of does but just to get PT working on TPU).
Thinking about doing another article on a direct comparison of JAX with PT and TF - let me know if that's something you'd like to see!
Obviously, the ability to use an accelerator makes JAX faster, but even without that on CPU it was faster. This is in part because of JIT, and in fairness the calculation in question does lend itself well to being expedited by JIT.
I actually have some preliminary benchmarks for a follow up specifically on just NumPy vs JAX, and it has become clear so far that NumPy is better in certain cases, especially for small operations where the overhead of JAX is not worth it.
In the article I mention this briefly, along with how JAX hasn't been focused on being optimized on CPU because they have bigger fish to fry, so to speak. I also link to the JAX documentation that has some comments comparing the two!
Relating to TF - I don't actually use TF at any point, but I did use PyTorch for Hessian calculation. TF and PT obviously do both work on GPU, but JAX has the benefit of being able to JIT more and implement everything in terms of XLA (although TF obviously has XLA support as well, and PT kind of does but just to get PT working on TPU).
Thinking about doing another article on a direct comparison of JAX with PT and TF - let me know if that's something you'd like to see!
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.
Thanks 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!
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!
How does JAX's JIT compare to Numba?
[deleted]
I didn't get a chance to run experiments myself, but from other people I've heard JAX is superior.
Am curious esp wrt numba gpu jit...