9 ms·
JAX – NumPy on the CPU, GPU, and TPU
- p1esk 3y agoWhat happened to Jax? Is it still alive?
- maegul 3y agoDevelopment seems not to have dropped at all from the contributions page: https://github.com/google/jax/graphs/contributors https://github.com/google/jax/graphs/contributors Don’t know about usage and uptake though.
- antognini 3y agoVery much alive. From what I can tell it has more or less replaced Tensorflow for research purposes. (A lot of researchers use PyTorch though.)
- technocratius 3y agoWhy would it be dead?
- esafak 3y agoAnybody using it in production? Is it, or its derivatives like Flax, worth using over pyTorch for anything? edit: Made comparison more fair.
- _delirium 3y agoI’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.
- mvuksano 3y agoAgreed. I used Jax about a year ago to estimate some diode parameters for a side project of mine.
- mvuksano 3y agoNot 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.
- fnbr 3y agoA number of large AI companies use it to train their large models; Midjourney, Stability, Anthropic, DeepMind, among others.
- technocratius 3y agoWe recently switched to Jax to boost performance as we scale up our algorithmic core. The nice thing is that it presents only a minor jump in capabilities to get developers working with it, if they have prior exposure to numpy ofcourse. Quite nice :)
- patrickkidger 3y agoIf you like PyTorch then you might like Equinox, by the way. (https://github.com/patrick-kidger/equinox https://github.com/patrick-kidger/equinox ; 1.4k GitHub stars now!) Basically designed to offer PyTorch-like syntax for working with JAX. The latter is excellent for the reasons the sibling replies have stated, but PyTorch absolutely got the usability story correct.
- subhrm 3y agoJAX GPU support is limited to Linux only. Even the WSL2 support is experimental. https://jax.readthedocs.io/en/latest/installation.html#supported-platforms https://jax.readthedocs.io/en/latest/installation.html#suppo...
- pjtr 3y agoIs there a specific reason why Windows is not supported?
- sdenton4 3y agoBecause no one has done the work to add it... Could be you!
- KeplerBoy 3y agoPresumably because the Google cloud doesn't run on Windows. Well, nothin HPC related runs Windows.
- pjmlp 3y agoLife science industry uses plenty of Windows, including HPC workloads.
- emu 3y agoWe ship Windows CPU only at the moment. We don't support Windows GPU because we haven't had the engineer bandwidth to support it well. We recommend WSL2 for GPU on Windows at the moment because that is a compromise: it allows CUDA support, without us having to support another release variant. But we welcome community contributions!
- marmaduke 3y agoThey don't build on Windows at all, as well.
- emu 3y agoNot true! We release Windows CPU wheels (https://pypi.org/project/jaxlib/#files https://pypi.org/project/jaxlib/#files). So JAX on CPU works great on Windows. We don't release Windows GPU wheels at the moment, but that's because we're a small team and none of us use Windows personally. We welcome contributions! (I verified that the Windows CUDA GPU support built as recently as two weeks ago, but I don't have the ability to test that it works.) We recommend WSL2 because that's just using our existing Linux CUDA release.
- antognini 3y agoIt 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.
- adgjlsfhk1 3y agothe issue here is that if your ideal algorithm isn't simply expressible in numpy (which many aren't), you're pretty much out of luck. As a result, imo the better approach is to use a fast language that also compiles to GPU (e.g. Julia)
- krasin 3y agoAn alternative is to write most of the program in Python + JAX + implement a few custom XLA ops in CUDA / Triton. That way, the program is very readable and can interoperate with the larger ecosystem, while still being fast to run.
- nestorD 3y agoHaving used JAX quite a bit for numerical computing (and having lectured on this use-case) I would say that a surprisingly large number of algorithms can be expressed as array[0] operations (even if it sometimes takes a bit of thinking). And, more importantly, things that cannot be expressed that way tend to not be a good fit for GPU computing anyway (independently of the language / framework you are using). [0]: `array` is a shortcut here, JAX is not limited to operations on arrays.
- nharada 3y agoFor large scale ML Jax is pretty nice. It makes the multi-host computation feel like a first class citizen and your programs tend to be designed with that in mind.
- amir734jj 3y agoVery unrelated but I did job interview with Nvidia JAX team for a compiler engineer role some time ago, not very friendly and very opinionated.
- aportnoy 3y agoWhat happened?
- amir734jj 3y agoNothing happened. It was an informal technical interview with the program manager at JAX. 1 hour call and the interview was remote but describing him as opinionated and entitled is an understatement. Best of luck to them.
- Hendrikto 3y agoWould you mind sharing some details? It sounds like an interesting peek behind the curtain.
- aportnoy 3y agoAsking because I am on that team :) I went through the same hiring process and had a positive experience at every stage. I had a strong competing offer but went with the JAX team at NVIDIA. I'll pass it along as feedback.
- tasubotadas 3y agoIs there any benefit using it instead of pytorch?
- fnbr 3y agoThe main benefit in my experience is that it’s much easier to do distributed computations in JAX. It has a much nicer API. For single device computing there’s no advantage either way.
- marmaduke 3y agoIf you like functional languages, then Jax will fit better for you. It provides a bunch of function transformations to implement eg grad, JIT etc.
- missingET 3y agoJax has a much nicer handling of higher order differentiation. PyTorch has functions to compute Hessians and there are libraries to keep differentiability through optimizers, but going out of their standard use-cases becomes tricky very fast. In contrast, JAX can compute nth-derivatives of things very easily.
- ayakang31415 3y agoI have been working on my DNN model using TensorFlow even though ML is not my main research. But it is a substantial part of my research, so I have to figure things out on my own, and I have done so over the past 3 years. However, I spend so much time on figuring out how any of TF methods works and debugging them. I never used JAX but I am not sure if this sort of grinding is normal when you use JAX as well (I always hear great things about JAX). I have built so many things using TF I don't think it is wise for me to learn JAX and migrate my work into JAX code base.
- polygamous_bat 3y agoIs there a reason you are limited to JAX/tensorflow and can't use pytorch? For a lot of people I know whose main job was not to write code, switching from tensorflow to pytorch was something that saved them ten to hundreds of hours in the long run, even accounting for the initial learning time.
- ayakang31415 3y agoI always feel like I am a prisoner of sunk cost fallacy
- rg111 3y agoYou should definitely switch to PyTorch. Or even JAX. TF is not worth it anymore.
- jphoward 3y agoI think JAX is cool, but I do find it slightly disingenuous when it claims to be "numpy by on the GPU" (as opposed to PyTorch), when actually there's a fundamental difference; it's functional. So if I have an array `x` and want to set index 0 to 10, I can't do: x[0] = 10 Instead I have to do: y = x.at[0].set(10) Of course this has advantages, but you can't then go and claim that JAX is a drop in replacement for numpy, because this such a fundamental change to how numpy developers think (and in this regard, PyTorch is closer to numpy than JAX).
- j7ake 3y agoAlso conditionals can be tricky (greater, if else) and often need rewriting.
- thrdbndndn 3y agoSorry for my potentially VERY ignorant question, I only know functional programming at average joe level. Why can't you do the first in functional programming (not in this specific case because it's just how it is, but in general)? And even if you can't do so for any reasonable reason in functional (again, in general), what stops us to just add syntactic sugar to equal it to the second to make programmer's life easier?
- Polygator 3y agoThe fundamental reason why many functional languages won't allow you to do the first is that they use immutable data structures. We could indeed introduce syntactic sugar (`y= (x[0]:=10)` maybe), but you'll still need to introduce a new variable to hold the modified list.
- thrdbndndn 3y agoIt's my understanding that, at least in Python, you can't change immutable data type but you can just assign a new data to the same variable and therefore overwrite it, right? So even if JAX makes list type immutable, you can still just re-use `x` to save the new modified list.
- albertzeyer 3y agoDoes it support arrays of variable lengths now? Last time I looked, I think this was not supported. So it means, for every variable dimension, you need to use an upper bound, and then use masking properly, and hope that it would not waste computation too much on the unused part (e.g. when running a loop over it). I'm working with sequences, e.g. speech recognition, machine translation, language modeling. This is a quite fundamental property for this type of models, that we have variable lengths sequences. In those cases, for some example code, I have seen that training also used only fixed size dimensions. And at inference time, they had some non-JAX code for the loop over the sequence around the JAX code with fixed-size dimensions. This seems like a quite fundamental issue to me? I wonder a bit that this is not an issue for others.
- karterk 3y agoFor JIT-ing you need to know the sizes upfront. There was an experimental branch for introducing jagged tensors, but as far as I know, it has been abandoned.
- ipunchghosts 3y agoJax is ceres but for python!
- nologic01 3y agoImho the project should emphasize more the potential for simple and uniquitous multi-core acceleration of vector compute that is available by definition to anybody having any modern cpu. nvidia dropped cuda support for perfectly good gpu's, showing the perils and waste of being locked-in in a profit-maximazing monopoly.
- jarym 3y agoNow if only it supported ONNX export or was cross-platform so I could run it from Java and .NET land
- BenoitP 3y agohttps://github.com/google/jaxonnxruntime https://github.com/google/jaxonnxruntime
- jarym 3y agoI need the opposite - JAX to ONNX
- agibsonccc 3y agoI've been looking in to this for the java world. What's your use case? Deployment in to existing applications?
- jarym 3y agoYea exactly - Python for training, Java/.NET for inference at production. I looked at approaches like GRPC and things but my case is a bit more time-sensitive and the latency added by going over a network layer was too much. For now I'm happy with Pytorch->ONNX and then running the ONNX model directly. But as I said, that means I can't easily train using JAX :-(
- patrickkidger 3y agoYou can do JAX->TF->ONNX I believe: https://github.com/patrick-kidger/equinox/blob/main/equinox/internal/_onnx.py https://github.com/patrick-kidger/equinox/blob/main/equinox/...
- jarym 3y agoOhh, I'll check that out!
- hashtag-til 3y agoIs there a path to compile and deploy JAX models on Android apps locally?
- woadwarrior01 3y agoI haven't tried it myself, but perhaps it's worth looking into the jax2tf -> tflite route.
- sampo 3y ago> With its updated version of Autograd, JAX can automatically differentiate native Python and NumPy code. Nice. When did they make this change? Here is the old way in the docs, where you needed to define functions for the if-true branch and the if-false branch, and feed them to a conditional function, to get the normal if-then-else conditional. https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#cond https://jax.readthedocs.io/en/latest/notebooks/Common_Gotcha...
- cl3misch 3y agoThis is still the case afaik. For vanilla "if", the condition must be known at compile time. For runtime, you have to use "cond", "where", or "select" (which may be analogous).
- mattjjatgoogle 3y agoActually, that's never been a constraint for JAX autodiff. JAX grew out of the original Autograd (https://github.com/hips/autograd https://github.com/hips/autograd), so differentiating through Python control flow always worked. It's jax.jit and jax.vmap which place constraints on control flow, requiring structured control flow combinators like those.
- mattjjatgoogle 3y agoActually that never changed. The README has always had an example of differentiating through native Python control flow: https://github.com/google/jax/commit/948a8db0adf233f333f3e5f64d324f308c277773#diff-b335630551682c19a781afebcf4d07bf978fb1f8ac04c6bf87428ed5106870f5R309-R322 https://github.com/google/jax/commit/948a8db0adf233f333f3e5f... The constraints on control flow expressions come from jax.jit (because Python control flow can't be staged out) and jax.vmap (because we can't take multiple branches of Python control flow, which we might need to do for different batch elements). But autodiff of Python-native control flow works fine!
- tinytera 3y agoI am running JAX version of "ctoraman/hate-speech-bert" over here https://news.ycombinator.com/item?id=37696033 https://news.ycombinator.com/item?id=37696033 and its pretty efficient
- wjessup 3y agoGreat basics tutorial: https://colinraffel.com/blog/you-don-t-know-jax.html https://colinraffel.com/blog/you-don-t-know-jax.html
- __rito__ 3y agoHere's mine [0]. It's a very gentle introduction. [0]: https://www.kaggle.com/code/truthr/jax-0 https://www.kaggle.com/code/truthr/jax-0
- sheepscreek 3y agoI’ve used Numpy with Numba primarily (on the CPU) and it had been a game changer for my data science workloads. Naturally, I pay close attention to Jax and give it a look every now and then. So, I’ll focus my observations below on Jax’s Numpy API support. At a glance, Jax code looks like regular Python, but it’s a very different style of programming. Two big differences I’ve found are: - All Jax functions must be pure. You can’t pass references. - ndarrays cannot be created with dynamic shapes. You have to hardcode the shape tuples. One possible workaround can be to create a buffer much bigger than you need and return that along with actual shape. Then there are many small things that are very well documented[1] by the Jax team. Overall, if you are training ML models, the trouble might be worth it (Autograd). But for accelerating Numpy alone, it is no Numba replacement - which will happily work in the above mentioned use-cases. [1] https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html https://jax.readthedocs.io/en/latest/notebooks/Common_Gotcha...
- cl3misch 3y ago> You have to hardcode the shape tuples. This is not true. Rather, all shapes have to be known at compile time. That means that output shapes must not depend on input values, but may depend on input shapes -- also explicitely. Furthermore, there are two useful additions: 1. You can use vanilla numpy for compile-time computations. An example would be computing an array of indices for some moving-window filter, depending on the input shape and a stride parameter. 2. You can mark function arguments as "static". Then their values may change shapes of the output, but accordingly the function is compiled for each value of those arguments.
- sheepscreek 3y agoYou’re right, I apologize for wording it incorrectly. It might have been a restriction with the # of dimensions. Either way, it wasn’t cut out for my use case.
- brap 3y agoJust going over the docs and there’s something I don’t understand, hoping someone here can explain: What’s the point of having to explicitly call grad(jit(f))? Why doesn’t grad just call jit internally? Is there a usecase where you want the grad without jit?
- patrickkidger 3y agoIndeed, it is much better to use jit(grad(f)) in general. Supporting the opposite composition is still useful in some edge cases -- for example when debugging, you want to step through a computation without jit, and simply not crash when differentiating any inner functions also decorated with jit.
- brap 3y agoWouldn't you use jax.disable_jit() for that?
- patrickkidger 3y agoYou can do that too! That will disable every JIT though. In practice you might only want to disable just one. To add some colour to my answer. When writing a library, it's typical to a put a JIT statement on everything in the public API. This means you get the benefits of JIT compilation even when you're just hacking around in the REPL, and mitigates the new-user-footgun in which they forget to use JIT themselves. Meanwhile, good practice is always to JIT your whole computation. Combined, this mean that it's fairly common to go jit (at the top level) -> grad (of your operation) -> jit (of some library call). When debugging your code, the JIT'd library call is _probably_ not the culprit. So you only want to disable the top-level JIT when stepping through, and still take advantage of JIT compilation where you can. Overall one obtains a composition of the form grad(jit(...)). TL;DR: even if use case doesn't come up super frequently, it's more user-friendly to support grad(jit(...)) than it is to just crash.
- mi_lk 3y agoLooking for some numerical problems to practice JAX, any suggestion or resource?
- cl3misch 3y agoAlthough not terribly complicated, you can compute fractals like the Mandelbrot set or Julia sets. They look really cool and you can play with visualization.
- jjtheblunt 3y agothere's a comment above : search for kaggle and you'll find such. also: https://jax.readthedocs.io/en/latest/advanced_guide.html https://jax.readthedocs.io/en/latest/advanced_guide.html
- munro 3y agoI love JAX. It can be a great replacement for Numba or whatever, the @jit works really well. And vmap is amazing... I often get lost in the matrix sauce when batching, using JAX I develop as if it's just a single instance, then vmap that shit. The ecosystem I was using was: jax, optax, haiku. The big issue I had was: I was developing on the CPU, then moved to running it on a GPU, and it wasn't as fast as I expected-- I started debugging, and saw there was still lots of communication between the CPU and GPU even tho it was all jit'd. I think PyTorch is a more user friendly for writing high performance models if you're not straying too far from the beaten path. But I really love JAX would like to play around with it more to understand these pits I'm falling into. And another complaint is I can't run it on my Macbook M1 GPU... but I'm seeing this page now, so maybe that's not true anymore: https://developer.apple.com/metal/jax/ https://developer.apple.com/metal/jax/
- joey_bob 3y agoI honestly had no idea JAX was useful outside autograd, due to my own tunnel vision. I even used it with other libraries to do this kind of work. Is there a term for this type of mistake?
- __rito__ 3y agoI wrote a Notebook [0] that gets you introduced to JAX in a very gentle manner. It also covers things like functional purity in Deep Learning, and handling of random numbers in JAX. [0]: Learn JAX: From Linear Regression to Neural Networks - https://www.kaggle.com/code/truthr/jax-0 https://www.kaggle.com/code/truthr/jax-0
- gorenb 3y agoThis is pretty well known, so I don't know why it would get so popular.
- robertnishihara 3y agoI'm a huge fan of Jax. The Jax team is incredibly strong! Just want to share that Ray (an open source project we're developing at Anyscale), can be used to scale Jax (e.g., across TPUs). Some docs from Google on how to do this https://cloud.google.com/tpu/docs/ray-guide https://cloud.google.com/tpu/docs/ray-guide Alpa is an open source project scaling Jax on 1000+ GPUs https://www.anyscale.com/blog/training-175b-parameter-language-models-at-1000-gpu-scale-with-alpa-and-ray https://www.anyscale.com/blog/training-175b-parameter-langua... Cohere uses Ray + Jax + TPUs to build their LLMs https://www.youtube.com/watch?v=For8yLkZP5w https://www.youtube.com/watch?v=For8yLkZP5w A demo from Matt Johnson on the Jax team https://www.youtube.com/watch?v=hyQ-tgD5sgc https://www.youtube.com/watch?v=hyQ-tgD5sgc
- dekhn 3y agoI'll remember this always: when DeepMind solved a subset of the protein structure prediction problem, they used Jax as the framework. PSPP was a long-standing issue and to see a fairly new computational tool used to significantly aid in the process of "solving" it speaks greatly towards its general utility in the sciences.