I think a lot of the commenters here are being rather unfair.
PyTorch has better adoption / network effects. JAX has stronger underlying abstractions.
I use both. I like both :)
The future of Deep Learning frameworks
61–70 of 113 posts
Re: The future of Deep Learning frameworks
#62PyTorch is developed by multiple companies / stake holders while jax is google only with internal tooling they don’t share with the world. This alone is a major reason not to use jax. Also I think it is more the other way around: with torch.compile the main advantage of jax is disappearing.
It's the old age question in programming: Do you use a highly constrained paradigm that allows easy automatic optimization or do you use a very flexible and more user intuitive paradigm that makes automatic optimization harder?
If the future is going to be better more intelligent compilers, then that settles the question in my opinion.
Re: The future of Deep Learning frameworks
#63PyTorch beat Tensorflow because it was much easier to use for research. Jax is much harder to use for exploratory research than PyTorch, due to requiring a fixed shape computation graph, which makes implementing many custom model architectures very difficult. Jax's advantages shine when it comes to parallelizing a new architecture across multiple GPU/TPUs, which it makes much easier than PyTorch (no need for custom c…
Doesn't JAX support dynamic graphs as well?
Re: The future of Deep Learning frameworks
#64So if multi-backend is doomed, is tinygrad then doomed, too?
Re: The future of Deep Learning frameworks
#65PyTorch is the javascript of ML. sadly "worse is better" software has better survival characteristics even when there is consensus that technology X is theoretically better
How is Jax theoretically better than PyTorch? The author is ignorant of torch.compile and biased as other commenters have pointed out.
Curious to know how the OP is biased when they have no conflict of interest, and they explicitly mention the `torch.compile` stack like... a few dozen times in the blog?
Re: The future of Deep Learning frameworks
#66PyTorch is a generationally important project. I've never seen a tool that is so inline with how researchers learn and internalize a subject. Teaching Machine Learning before and after its adoption has been a completely different experience. Never can be said enough how cool it is that Meta fosters and supports it. Viva PyTorch! (Jax rocks too)
Torch was originally a Lua project, hence why pytorch is called pytorch and not just torch.
In another timeline AI would have made Lua popular.
The best part is it trampled TensorFlow which I personally find obtuse.
Re: The future of Deep Learning frameworks
#67the author got a couple of things wrong, that are worth pointing out: 1. PyTorch is going all-in on torch.compile -- Dynamo is the frontend, Inductor is the backend -- with a strong default Inductor codegen powered by OpenAI Triton (which now has CPU, NVIDIA GPU and AMD GPU backends). The author's view that PyTorch is building towards a multi-backend future isn't really where things are going. PyTorch supports extens…
Hey, thanks for actually engaging with the blog's points instead of "Google kills everything it touches" :)
1. I'm well aware of the PyTorch stack, but this point:
> PyTorch is building towards a multi-backend future isn't really where things are going
>PyTorch supports extensibility of backends (including XLA)
Is my problem. Those backends just never integrate well as I mentioned in the blogpost. I'm not sure if you've ever gone into the weeds, but there are so many (often undocumented) sharp edges when using different backends that they never really work well. For example, how bad Torch:XLA is and the nightmare inducing bugs & errors with it.
> torch.compile is 2 years old, XLA is 7 years old. Compilers take a few years to mature
That was one of my major points - I don't think leaning on torch.compile is the best idea. A compiler would inherently place restrictions that you have to work-around.
This is not dynamic, nor flexible - and it flies in the face of torch's core philosophies just so they can offer more performance to the big labs using PyTorch. For various reasons, I dislike pandering to the rich guy instead of being an independent, open-source entity.
2. Torch/XLA is indeed primarily meant for TPUs - like the quoted announcement, where they declare to be ditching TF:XLA in favour of OpenXLA. But there's still a very real effort to get it working on GPUs - infact, a lab on twitter declared that they're using Torch/XLA on GPUs and will soon™ release details.
XLA's GPU support is great, its compatible across different hardware, its optimized and mature. In short, its a great alternative to the often buggy torch.compile stack - if you fix the torch integration.
So I won't be surprised if in the long-term they lean on XLA. Whether that's a good direction or not is upto the devs to decide unfortunately - not the community.
3. Thank you for pointing that out. I'm not sure about the history of JAX (maybe might make for a good blogpost for JAX devs to write someday), but it seems that it was indeed developed at Google research, though also heavily supported + maintained by DeepMind.
Appreciate you giving the time to comment here though :)
Re: The future of Deep Learning frameworks
#68PyTorch is a generationally important project. I've never seen a tool that is so inline with how researchers learn and internalize a subject. Teaching Machine Learning before and after its adoption has been a completely different experience. Never can be said enough how cool it is that Meta fosters and supports it. Viva PyTorch! (Jax rocks too)
Torch was originally a Lua project, hence why pytorch is called pytorch and not just torch. In another timeline AI would have made Lua popular. The best part is it trampled TensorFlow which I personally find obtuse.
> In another timeline AI would have made Lua popular.
I wonder if it'd have been hated more than Python is - especially with the 1-based indexing...
Re: The future of Deep Learning frameworks
#69PyTorch is developed by multiple companies / stake holders while jax is google only with internal tooling they don’t share with the world. This alone is a major reason not to use jax. Also I think it is more the other way around: with torch.compile the main advantage of jax is disappearing.
> with torch.compile the main advantage of jax is disappearing.
Interesting take - I agree here somewhat.
But also, wouldn't you think a framework that has been from the ground-up designed around a specific, mature compiler stack be better able to integrate compilers in a more stable fashion than just shoe-horning static compilers into a very dynamic framework? ;)
Re: The future of Deep Learning frameworks
#70I think a lot of the commenters here are being rather unfair. PyTorch has better adoption / network effects. JAX has stronger underlying abstractions. I use both. I like both :)
Hey patrick, love your work!
I think the biggest, well "con" I've seen is non-technical - the fear of JAX being killed by Google.
I mention in the blog as well [here](https://neel04.github.io/my-website/blog/pytorch_rant/#gover...) how important having an independent governance structure is. I'm sure for many big companies and labs, the lack of a promise of long-term, stable support is a huge dealbreaker.
I'm not sure how much Google bureaucracy would limit this, but have you raised the subject of forming an independent entity to govern JAX, very much like PyTorch? I believe XLA is protected, as its with TF governance. But perhaps, there could be one for JAX's ecosystem as well, encompassing optax, equinox, flax etc.