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 :)
Probably unfair as a reaction to the unfair statements in the blog…
The future of Deep Learning frameworks
71–80 of 113 posts
Re: The future of Deep Learning frameworks
#72Earlier quoted context omitted.
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...
Python isn't hated AFAICT, though people will profess to hating building large projects in it (myself included), but many of those people also love it for shorter programs and scripts.
Re: The future of Deep Learning frameworks
#73Pushback notwithstanding, this article is 100% correct in all PyTorch criticisms. PyTorch was a platform for fast experimentation with eager evaluation, now they shoehorn "compilers" into it. "compilers", because a lot of the work is done by g++ and Triton.
It is a messy and quickly expanding codebase with many surprises like segfaults and leaks.
Is scientific experimentation really sped up by these frameworks? Everyone uses the Transformer model and uses the same algorithms over and over again.
If researchers wrote directly in C or Fortran, perhaps they'd get new ideas. The core inference (see Karparthy's llama.c) is ridiculously small. Core training does not seem much larger either.
Re: The future of Deep Learning frameworks
#74So if multi-backend is doomed, is tinygrad then doomed, too?
I haven't used TinyGrad but I'm not really sure what its goal is. To be the best autograd framework? to be a minimal one?
I'm glad they've removed the (rather arbitrary, and admittedly stupid) loc cap. And from the little I know, geohot is focusing on having its own internal compiler stack.
As much as I admire geohot, I don't think rolling your own compiler is the best way. Its not that the TinyGrad team isn't smart enough, but a compiler is a huge undertaking and that you have to support and maintain for a long time. I'm sure he's well aware of this, but no big labs would touch TG seriously because of this limitation.
XLA on the other hand is under governance seperate from Google, and is far more mature - so people trust that.
That said, I don't know much about Tinygrad so I would appreciate if someone more knowledgable can jump in here and outline the differences and key features ¯\_(ツ)_/¯
Re: The future of Deep Learning frameworks
#75the 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…
3. The project started under a Harvard affiliated Github org during the course of PhDs. These same people later joined Google where it continued to be developed and over time adopted more and more in place of TensorFlow.
Re: The future of Deep Learning frameworks
#76I 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 s…
I can personally say, I am not super concerned about it being killed. Google supported TF1 for quite a long time and all these projects have a shelf life.
What concerned me about JAX, at a small company, is that it doesn't benefit from the network effects of almost everyone developing for it. E.g. There is no Llama 3.1 implementation in JAX afaict.
So as long as there is a need to pull from the rest of the world the ecosystem will trump the framework.
Activity in the LLM space is slowing down though, so there is an opportunity to take the small set of what worked and port it to JAX and show people how good that world is.
Re: The future of Deep Learning frameworks
#77PyTorch 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)
This is exactly why I gravitated to it so quickly. The first time I looked at pytorch code it was immediately obvious what the abstractions meant and how to use them to write a model architecture.
Jax looks like something completely different to me. Maybe I’m dumb and probably not the target audience, but it occurs to me that very few people are. When I read about using Jax, I find recommendations for a handful of other libraries that make it more useable. Which of those I choose to learn is not entirely obvious because they all seem to create a very fragmented ecosystem with code that isn’t portable.
I’m still not sure why I’d spend my time learning Jax, especially when it seems like most of the complaints from the author don’t really separate out training and inference, which don’t necessarily need to occur from the same framework.
Re: The future of Deep Learning frameworks
#78Are modern NN's really just static functions, and are they going to continue to be in the future?
KV caching is directly in conflict with a purely functional approach.
Re: The future of Deep Learning frameworks
#79PyTorch 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?
Nope, changing graph shape requires recompilation: https://github.com/google/jax/discussions/17191
Re: The future of Deep Learning frameworks
#80PyTorch 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)
This is exactly why I gravitated to it so quickly. The first time I looked at pytorch code it was immediately obvious what the abstractions meant and how to use them to write a model architecture. Jax looks like something completely different to me. Maybe I’m dumb and probably not the target audience, but it occurs to me that very few people are. When I read about using Jax, I find recommendations for a handful of ot…
Honestly, when I turn to JAX, I generally do it without a framework. It’s like asking for a framework to wrap numpy to me. Just JAX plus optax is sufficient for me in the cases I turn to it.