7 ms·

The future of Deep Learning frameworks

209 points by lairv 2y ago

casualscience 2y ago
> Multi-backend is doomed Finally! Someone says it! This is why the C programming language will never have wide adoption. /s
manlobster 2y ago
[flagged]
munchler 2y ago
This is such a hyperbolic headline it's hard to be interested in reading the actual article.
- bugglebeetle 2y ago
  [flagged]
- deleted 2y ago
  [deleted]
- deleted 2y ago
  [deleted]
  - deisteve 2y ago
    [flagged]
- crancher 2y ago
  All blanket statements are false.
- dang 2y ago
  We've unhyperbolized it via the subtitle.
  - sundarurfriend 2y ago
    I think that change was for the worse. (Just posted a comment asking for the original title: https://news.ycombinator.com/item?id=41275918 https://news.ycombinator.com/item?id=41275918) The subtitle doesn't convey the content of the article nearly as well as the title does. Perhaps you can take a sentence like "PyTorch has been a net negative for scientific computing efforts" which the article does say, or some toned down versino of the original title, but the current title makes it sound like a very different article and felt like clickbait to me.
    - dang 2y ago
      Those are good points. Probably too late to make a difference now, but I hear you.
semiinfinitely 2y ago
PyTorch 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
- bob020202 2y ago
  Nothing is even theoretically better than Javascript for its intended use cases, web frontends and backends. Mainly because it went all-in on event loop parallelism early-on, which isn't just for usability but also performance. And didn't go all-in on OOP unlike Java, and has easy imports/packages unlike Python. It has some quirks like the "0 trinity," but that doesn't really matter. No matter how good you are with something else, it still takes more dev time ($$$) than JS. Now it's been forever since I used PyTorch or TF, but I only remember TF 1.x being more like "why TF isn't this working." At some point I didn't blame myself, I blamed the tooling, which TF2 later admitted. It seemed like no matter how skilled I got with TF1, it'd always take much longer than developing with PyTorch, so I switched early.
  - josephg 2y ago
    You don't think its possible to (even theoretically) improve on javascript for its intended use case? What a terrific lack of imagination. Typescript and Elm would like a word
    - bob020202 2y ago
      No, I said there's nothing that exists right now that's theoretically better. Typescript isn't. It'd be great if TS's type inference were smart enough that it basically takes no additional dev input vs JS, but until then, it's expensive to use. It's also bolted on awkwardly, but that's changing soon. Could also imagine JS getting some nice Py features like list comp. Also, generally when people complain that JS won the web, it's not because they prefer TS, it's cause they wanted to use something else and can't. Never used Elm, but... no variables, kinda like Erlang, which I've used. That has its appeal, but you're not going to find a consensus that this is better for web.
ein0p 2y ago
Jax is dead, long live PyTorch. PyTorch has _twenty times_ as many users as Jax. Any rumors of its death are highly exaggerated
- melling 2y ago
  They used to say the same thing about Perl and Python Downvoted. Hmmm. I’m a little tired so I don’t want to go into detail. However, I was a Perl programmer when Python was rising. So, needless to say, having a big lead doesn’t matter. Please learn from history. A big lead means nothing.
  - pjmlp 2y ago
    Had the whole Parrot VM and Perl 6 distraction never happened, and I doubt Python would have taken over on UNIX scripting part. So a good lesson is not to get distracted by foolish endevours.
  - ein0p 2y ago
    It’s been years and Jax is just where it was, no growth whatsoever. And that’s with all of Google forced internally to use only Jax. Look, I like the technical side of Jax for the most part, but it’s years too late to the party and it’s harder to use than PyTorch. It just isn’t going to ever take off at this point.
- deisteve 2y ago
  lot of contrarian takes are popular but rarely implemented in reality
- tripplyons 2y ago
  It's definitely exaggerated, but I personally prefer JAX and have found it easier to use than PyTorch for almost everything. If you haven't already, I would give JAX a good try.
deleted 2y ago
[deleted]
- deleted 2y ago
  [deleted]
srush 2y ago
PyTorch 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)
- PostOnce 2y ago
  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.
  - n7g 2y ago
    > 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...
    - CuriouslyC 2y ago
      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.
      - dartos 2y ago
        Everything is hated. Python has always gotten hate for being super super slow and having an ugly syntax (subjective ofc, but I happen to agree)
    - goatlover 2y ago
      Scientific computing tends to be 1-based. Thus R, Julia, Fortran, Matlab.
- deepsquirrelnet 2y ago
  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.
hprotagonist 2y ago
> I believe that all infrastructure built on Torch is just a huge pile of technical debt, that will haunt the field for a long, long time. ... from the company that pioneered the approach with tensorflow. I've worked with worse ML frameworks, but they're by now pretty obscure; i cannot remember (and i am very happy about it) the last time i saw MXNet in the wild, for example. You'll still find Caffe on some embedded systems, but you can mostly sidestep it.
deisteve 2y ago
i like pytorch because all the academia release their code with it ive never even heard of jax nor will i have the skills to use it i literally just want to know two things: 1) how much vram 2) how to run it on pytorch
- deleted 2y ago
  [deleted]
- sva_ 2y ago
  Jax is a competing computational framework that does something similar to PyTorch, so both of your questions don't really make sense.
  - etiam 2y ago
    Maybe deisteve will answer for himself, but I don't think that's meant to mean how to run Jax on Pytorch, but rather the questions he's interested in for any published model.
    - deisteve 2y ago
      if a paper releases code in PyTorch im not going to sit there and complain Jax is more efficient its like fretting about how everything should be written in C++ instead of Python/Javascript I don't care.
  - deisteve 2y ago
    can i run pytorch code on jax? if not im not interested. i'll keep using pytorch. i hope this answer makes more sense for you.
0cf8612b2e1e 2y ago
As we all know, the technically best solution always wins.
- deleted 2y ago
  [deleted]
Hizonner 2y ago
[flagged]
- deleted 2y ago
  [deleted]
logicchains 2y ago
PyTorch 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 cuda/networking code). Needing to scale up a new architecture across many GPUs is however not a common use-case, and most teams that have the resources for large-scale multi-gpu training also have the resources for specialised engineers to do it in PyTorch.
- big-chungus4 2y ago
  Doesn't JAX support dynamic graphs as well?
  - logicchains 2y ago
    Nope, changing graph shape requires recompilation: https://github.com/google/jax/discussions/17191 https://github.com/google/jax/discussions/17191
funks_ 2y ago
I wish dex-lang [1] had gotten more traction. It’s JAX without the limitations that come from being a Python DSL. But ML researchers apparently don’t want to touch anything that doesn’t look exactly like Python. [1]: https://github.com/google-research/dex-lang https://github.com/google-research/dex-lang
- hatmatrix 2y ago
  It seems like an experimental research language. Julia also competes in this domain from a more practical standpoint and has less limitations than JAX as I understand it, but is less mature and still working on getting wider traction.
  - funks_ 2y ago
    The Julia AD ecosystem is very interesting in that the community is trying to make the entire language differentiable, which is much broader in scope than what Torch and JAX are doing. But unlike Dex, Julia is not a language built from the ground up for automatic differentiation. Shameless plug for one of my talks at JuliaCon 2024: https://www.youtube.com/live/ZKt0tiG5ajw?t=19747s https://www.youtube.com/live/ZKt0tiG5ajw?t=19747s. The comparison between Python and Julia starts at 5:31:44.
    - hatmatrix 2y ago
      Ah I had not realized I was corresponding with the author of that talk - I'd followed it back when it was happening as I'm particularly interested in adapting AD. Where do you feel Julia is at this point in time (compared to say, JAX or PyTorch) from a practitioner's standpoint?
      - funks_ 2y ago
        When it comes to general deep learning, Julia is much less mature than the JAX ecosystem. I think deep learning will be the hardest nut for Julia to crack. The field is moving incredibly fast, and network effects are strong. Julia's strength lies in scientific computing, so I think adoption will come through novel applications of AD/ML in the sciences, rather than trying to catch up with the latest LLM developments I'm positive about Julia's future because the developer experience just feels so fun and productive. I always find it impressive how much a small group of self-organized volunteers has been able to achieve. Amazing things could happen if a company like Google or Meta paid a team of full-time engineers to advance the deep learning ecosystem. Fun fact: Julia strongly influenced PyTorch's recent design decisions [1]. [1]: https://dev-discuss.pytorch.org/t/where-we-are-headed-and-why-it-looks-a-lot-like-julia-but-not-exactly-like-julia/276 https://dev-discuss.pytorch.org/t/where-we-are-headed-and-wh...
sva_ 2y ago
> I’ve personally known researchers who set the seeds in the wrong file at the wrong place and they weren’t even used by torch at all - instead, were just silently ignored, thus invalidating all their experiments. (That researcher was me) Some assert-ing won't hurt you. Seriously. It might even help keeping your sanity.
marcinzm 2y ago
My main reason to avoid Jax is Google. Google doesn't provide good support even for things you pay them for. They do things because they want to, to get their internal promotions, irrespective of their customers or the impact on them.
- deleted 2y ago
  [deleted]
smhx 2y ago
the 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 extensibility of backends (including XLA), but there's disproportionate effort into the default path. torch.compile is 2 years old, XLA is 7 years old. Compilers take a few years to mature. torch.compile will get there (and we have reasonable measures that the compiler is on track to maturity). 2. PyTorch/XLA exists, mainly to drive a TPU backend for PyTorch, as Google gives no other real way to access the TPU. It's not great to try shoe-in XLA as a backend into PyTorch -- as XLA fundamentally doesn't have the flexibility that PyTorch supports by default (especially dynamic shapes). PyTorch on TPUs is unlikely to ever have the experience of JAX on TPUs, almost by definition. 3. JAX was developed at Google, not at Deepmind.
- n7g 2y ago
  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 :)
cs702 2y ago
A more accurate title for the OP would be "I hope and wish PyTorch were dead, so Jax could become the standard." Leaving aside the fact that PyTorch's ecosystem is 10x to 100x larger, depending on how one measures it, PyTorch's biggest advantage, in my experience, is that it can be picked up quickly by developers who are new to it. Jax, despite its superiority, or maybe because of it, can not be picked up quickly. Equinox does a great job of making Jax accessible, but Jax's functional approach is in practice more difficult to learn than PyTorch's object-oriented one.
- deleted 2y ago
  [deleted]
BaculumMeumEst 2y ago
Jax is well designed? That's nice. The only thing that matters is adoption. You can make run this title when Jax's adoption surpasses PyTorch. How does someone using _python_ not understand this?
pjmlp 2y ago
The best thing about PyTorch is that we aren't stuck with Python, rather we can enjoy it alongside proper performance, by using the Java and C++ API surfaces instead.
wafngar 2y ago
PyTorch 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.
- dauertewigkeit 2y ago
  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.
- n7g 2y ago
  > 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? ;)
  - wafngar 2y ago
    Depends. PyTorch on the other hand has a large user base and well defined and tested api. So should be doable; and is already progressing and rapid speed..
  - anon389r58r58 2y ago
    So the answer is not Jax? Because JAX is not designed around a mature compiler stack. The history of Jax is more so that it matured alongside the compiler...
patrickkidger 2y ago
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 :)
- n7g 2y ago
  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/#governance-structure 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.
  - Eridrus 2y ago
    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.
- wafngar 2y ago
  Probably unfair as a reaction to the unfair statements in the blog…
auggierose 2y ago
So if multi-backend is doomed, is tinygrad then doomed, too?
- n7g 2y ago
  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 ¯\_(ツ)_/¯
0x19241217 2y ago
Pushback 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.
- dunefox 2y ago
  > If researchers wrote directly in C or Fortran... ... then they would get nothing done.
  - pklausler 2y ago
    Why not?
    - dunefox 2y ago
      1. data cleaning and other tasks for which Python excels, 2. C and Fortran don't come close to the developer speed of Python. Ideas don't come from using languages like C, but are rather enabled by implementing, prototyping, and iterating quickly and with little resistance from the language, using existing libraries, etc. Python is simply much better in these regards.
  - TheRealKing 2y ago
    Fortran cannot be placed with C in the same category of low programming productivity.
  - pjmlp 2y ago
    Apparently we could do research work before Python was invented.
FL33TW00D 2y ago
Are 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.
- visarga 2y ago
  You could sen the KV cache and even the LoRA as input args to a static model.
- spott 2y ago
  Is it? The same inputs have the same outputs, the only thing different is the performance. It is basically the same as memorization.
ianbutler 2y ago
From an eng/industry perspective, back in 2016/2017 I watched the realtime decline of Tensorflow towards Pytorch. The issue was TF had too many interfaces to accomplish the same thing and each one was rough in its own way. Along with some complexity for using serving and experiment logging via Tensorboard, but this wasn’t as bad at least for me. Keras was integrated in an attempt to help, but ultimately it wasn’t enough and people started using Torch more and more even against the perception that TF was for prod workloads and Torch was for research. TFA mentions the interface complexity as starting to be a problem with Torch, but I don’t think we’re anywhere near the critical point that would cause people to abandon it in favor of JAX. Additionally with JAX you’re just shoving the portability problems mentioned down to XLA which brings its own issues and gotchas even if it hides the immediate reality of said problems from the end user. I think the Torch maintainers should watch not to repeat the mistakes of TF, but I think theres a long way to go before JAX is a serious contender. It’s been years and JAX has stayed in relatively small usage.
- markeroon 2y ago
  Imo the biggest issue (from memory) was that Tensorflow used a static computation graph. PyTorch was so much easier to work with.
  - ianbutler 2y ago
    Yup that was also definitely one among the many issues
  - lostmsu 2y ago
    I honestly think the static graph was much better in the similar way Vulkan/DX12 are better than OpenGL/DX9. It is harder to program, but gives more explicit control of important things. E.g. who would use PyTorch if they new that for optimal performance they'd need to record CUDA graphs?
sundarurfriend 2y ago
Can we get the title changed to the actual title of the post? "The future of Deep Learning frameworks" sounds like a neutral and far wider-reaching article, and ends up being clickbait here (even if unintentionally). "PyTorch is dead. Long live JAX." conveys exactly what the article about, and is a much better title.
6gvONxR4sf7o 2y ago
One aspect of jax that’s rarely touched on is browser stuff. Completely aside from deep learning, it’s straightforward to compile jax to a graphics shader you can call in js, which in this insane world is actually my preferred way to put numerical computing or linear algebra code on a web page.
- mccoyb 2y ago
  Can you share any code or references for this? Several people I work with are interested in this.
  - 6gvONxR4sf7o 2y ago
    Sure. I’ll have to dig it up (shits boxed up as I move), but i’ll ping you or this comment when I get to it.
lostmsu 2y ago
From the article is seems that JAX is a non-starter for me as they don't have support for any kind of acceleration on Windows proper, and only experimental in WSL.
astromaniak 2y ago
The article misses multi-modal thing. Which is the future. Sure they can be considered a separate things, like today. But that's probably not the best approach. Support from framework may include partial training, easy components swap, intermediate data caching, dynamic architecture, automatic work balance and scaling.

209 points by lairv 2y ago

> Multi-backend is doomed Finally! Someone says it! This is why the C programming language will never have wide adoption. /s

manlobster 2y ago

[flagged]

munchler 2y ago

This is such a hyperbolic headline it's hard to be interested in reading the actual article.

bugglebeetle 2y ago
[flagged]
deleted 2y ago
[deleted]
deleted 2y ago
[deleted]
- deisteve 2y ago
  [flagged]
crancher 2y ago
All blanket statements are false.
dang 2y ago
We've unhyperbolized it via the subtitle.
- sundarurfriend 2y ago
  I think that change was for the worse. (Just posted a comment asking for the original title: https://news.ycombinator.com/item?id=41275918 https://news.ycombinator.com/item?id=41275918) The subtitle doesn't convey the content of the article nearly as well as the title does. Perhaps you can take a sentence like "PyTorch has been a net negative for scientific computing efforts" which the article does say, or some toned down versino of the original title, but the current title makes it sound like a very different article and felt like clickbait to me.
  - dang 2y ago
    Those are good points. Probably too late to make a difference now, but I hear you.

semiinfinitely 2y ago

PyTorch 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

bob020202 2y ago
Nothing is even theoretically better than Javascript for its intended use cases, web frontends and backends. Mainly because it went all-in on event loop parallelism early-on, which isn't just for usability but also performance. And didn't go all-in on OOP unlike Java, and has easy imports/packages unlike Python. It has some quirks like the "0 trinity," but that doesn't really matter. No matter how good you are with something else, it still takes more dev time ($$$) than JS. Now it's been forever since I used PyTorch or TF, but I only remember TF 1.x being more like "why TF isn't this working." At some point I didn't blame myself, I blamed the tooling, which TF2 later admitted. It seemed like no matter how skilled I got with TF1, it'd always take much longer than developing with PyTorch, so I switched early.
- josephg 2y ago
  You don't think its possible to (even theoretically) improve on javascript for its intended use case? What a terrific lack of imagination. Typescript and Elm would like a word
  - bob020202 2y ago
    No, I said there's nothing that exists right now that's theoretically better. Typescript isn't. It'd be great if TS's type inference were smart enough that it basically takes no additional dev input vs JS, but until then, it's expensive to use. It's also bolted on awkwardly, but that's changing soon. Could also imagine JS getting some nice Py features like list comp. Also, generally when people complain that JS won the web, it's not because they prefer TS, it's cause they wanted to use something else and can't. Never used Elm, but... no variables, kinda like Erlang, which I've used. That has its appeal, but you're not going to find a consensus that this is better for web.

ein0p 2y ago

Jax is dead, long live PyTorch. PyTorch has _twenty times_ as many users as Jax. Any rumors of its death are highly exaggerated

melling 2y ago
They used to say the same thing about Perl and Python Downvoted. Hmmm. I’m a little tired so I don’t want to go into detail. However, I was a Perl programmer when Python was rising. So, needless to say, having a big lead doesn’t matter. Please learn from history. A big lead means nothing.
- pjmlp 2y ago
  Had the whole Parrot VM and Perl 6 distraction never happened, and I doubt Python would have taken over on UNIX scripting part. So a good lesson is not to get distracted by foolish endevours.
- ein0p 2y ago
  It’s been years and Jax is just where it was, no growth whatsoever. And that’s with all of Google forced internally to use only Jax. Look, I like the technical side of Jax for the most part, but it’s years too late to the party and it’s harder to use than PyTorch. It just isn’t going to ever take off at this point.
deisteve 2y ago
lot of contrarian takes are popular but rarely implemented in reality
tripplyons 2y ago
It's definitely exaggerated, but I personally prefer JAX and have found it easier to use than PyTorch for almost everything. If you haven't already, I would give JAX a good try.

deleted 2y ago

[deleted]

deleted 2y ago
[deleted]

srush 2y ago

PyTorch 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)

PostOnce 2y ago
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.
- n7g 2y ago
  > 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...
  - CuriouslyC 2y ago
    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.
    - dartos 2y ago
      Everything is hated. Python has always gotten hate for being super super slow and having an ugly syntax (subjective ofc, but I happen to agree)
  - goatlover 2y ago
    Scientific computing tends to be 1-based. Thus R, Julia, Fortran, Matlab.
deepsquirrelnet 2y ago
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.

hprotagonist 2y ago

> I believe that all infrastructure built on Torch is just a huge pile of technical debt, that will haunt the field for a long, long time. ... from the company that pioneered the approach with tensorflow. I've worked with worse ML frameworks, but they're by now pretty obscure; i cannot remember (and i am very happy about it) the last time i saw MXNet in the wild, for example. You'll still find Caffe on some embedded systems, but you can mostly sidestep it.

deisteve 2y ago

i like pytorch because all the academia release their code with it ive never even heard of jax nor will i have the skills to use it i literally just want to know two things: 1) how much vram 2) how to run it on pytorch

deleted 2y ago
[deleted]
sva_ 2y ago
Jax is a competing computational framework that does something similar to PyTorch, so both of your questions don't really make sense.
- etiam 2y ago
  Maybe deisteve will answer for himself, but I don't think that's meant to mean how to run Jax on Pytorch, but rather the questions he's interested in for any published model.
  - deisteve 2y ago
    if a paper releases code in PyTorch im not going to sit there and complain Jax is more efficient its like fretting about how everything should be written in C++ instead of Python/Javascript I don't care.
- deisteve 2y ago
  can i run pytorch code on jax? if not im not interested. i'll keep using pytorch. i hope this answer makes more sense for you.

0cf8612b2e1e 2y ago

As we all know, the technically best solution always wins.

deleted 2y ago
[deleted]

Hizonner 2y ago

[flagged]

deleted 2y ago
[deleted]

logicchains 2y ago

PyTorch 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 cuda/networking code). Needing to scale up a new architecture across many GPUs is however not a common use-case, and most teams that have the resources for large-scale multi-gpu training also have the resources for specialised engineers to do it in PyTorch.

big-chungus4 2y ago
Doesn't JAX support dynamic graphs as well?
- logicchains 2y ago
  Nope, changing graph shape requires recompilation: https://github.com/google/jax/discussions/17191 https://github.com/google/jax/discussions/17191

funks_ 2y ago

I wish dex-lang [1] had gotten more traction. It’s JAX without the limitations that come from being a Python DSL. But ML researchers apparently don’t want to touch anything that doesn’t look exactly like Python. [1]: https://github.com/google-research/dex-lang https://github.com/google-research/dex-lang

hatmatrix 2y ago
It seems like an experimental research language. Julia also competes in this domain from a more practical standpoint and has less limitations than JAX as I understand it, but is less mature and still working on getting wider traction.
- funks_ 2y ago
  The Julia AD ecosystem is very interesting in that the community is trying to make the entire language differentiable, which is much broader in scope than what Torch and JAX are doing. But unlike Dex, Julia is not a language built from the ground up for automatic differentiation. Shameless plug for one of my talks at JuliaCon 2024: https://www.youtube.com/live/ZKt0tiG5ajw?t=19747s https://www.youtube.com/live/ZKt0tiG5ajw?t=19747s. The comparison between Python and Julia starts at 5:31:44.
  - hatmatrix 2y ago
    Ah I had not realized I was corresponding with the author of that talk - I'd followed it back when it was happening as I'm particularly interested in adapting AD. Where do you feel Julia is at this point in time (compared to say, JAX or PyTorch) from a practitioner's standpoint?
    - funks_ 2y ago
      When it comes to general deep learning, Julia is much less mature than the JAX ecosystem. I think deep learning will be the hardest nut for Julia to crack. The field is moving incredibly fast, and network effects are strong. Julia's strength lies in scientific computing, so I think adoption will come through novel applications of AD/ML in the sciences, rather than trying to catch up with the latest LLM developments I'm positive about Julia's future because the developer experience just feels so fun and productive. I always find it impressive how much a small group of self-organized volunteers has been able to achieve. Amazing things could happen if a company like Google or Meta paid a team of full-time engineers to advance the deep learning ecosystem. Fun fact: Julia strongly influenced PyTorch's recent design decisions [1]. [1]: https://dev-discuss.pytorch.org/t/where-we-are-headed-and-why-it-looks-a-lot-like-julia-but-not-exactly-like-julia/276 https://dev-discuss.pytorch.org/t/where-we-are-headed-and-wh...

sva_ 2y ago

> I’ve personally known researchers who set the seeds in the wrong file at the wrong place and they weren’t even used by torch at all - instead, were just silently ignored, thus invalidating all their experiments. (That researcher was me) Some assert-ing won't hurt you. Seriously. It might even help keeping your sanity.

marcinzm 2y ago

My main reason to avoid Jax is Google. Google doesn't provide good support even for things you pay them for. They do things because they want to, to get their internal promotions, irrespective of their customers or the impact on them.

deleted 2y ago
[deleted]

smhx 2y ago

the 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 extensibility of backends (including XLA), but there's disproportionate effort into the default path. torch.compile is 2 years old, XLA is 7 years old. Compilers take a few years to mature. torch.compile will get there (and we have reasonable measures that the compiler is on track to maturity). 2. PyTorch/XLA exists, mainly to drive a TPU backend for PyTorch, as Google gives no other real way to access the TPU. It's not great to try shoe-in XLA as a backend into PyTorch -- as XLA fundamentally doesn't have the flexibility that PyTorch supports by default (especially dynamic shapes). PyTorch on TPUs is unlikely to ever have the experience of JAX on TPUs, almost by definition. 3. JAX was developed at Google, not at Deepmind.

n7g 2y ago
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 :)

cs702 2y ago

A more accurate title for the OP would be "I hope and wish PyTorch were dead, so Jax could become the standard." Leaving aside the fact that PyTorch's ecosystem is 10x to 100x larger, depending on how one measures it, PyTorch's biggest advantage, in my experience, is that it can be picked up quickly by developers who are new to it. Jax, despite its superiority, or maybe because of it, can not be picked up quickly. Equinox does a great job of making Jax accessible, but Jax's functional approach is in practice more difficult to learn than PyTorch's object-oriented one.

deleted 2y ago
[deleted]

BaculumMeumEst 2y ago

Jax is well designed? That's nice. The only thing that matters is adoption. You can make run this title when Jax's adoption surpasses PyTorch. How does someone using _python_ not understand this?

pjmlp 2y ago

The best thing about PyTorch is that we aren't stuck with Python, rather we can enjoy it alongside proper performance, by using the Java and C++ API surfaces instead.

wafngar 2y ago

PyTorch 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.

dauertewigkeit 2y ago
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.
n7g 2y ago
> 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? ;)
- wafngar 2y ago
  Depends. PyTorch on the other hand has a large user base and well defined and tested api. So should be doable; and is already progressing and rapid speed..
- anon389r58r58 2y ago
  So the answer is not Jax? Because JAX is not designed around a mature compiler stack. The history of Jax is more so that it matured alongside the compiler...

patrickkidger 2y ago

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 :)

n7g 2y ago
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/#governance-structure 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.
- Eridrus 2y ago
  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.
wafngar 2y ago
Probably unfair as a reaction to the unfair statements in the blog…

auggierose 2y ago

So if multi-backend is doomed, is tinygrad then doomed, too?

n7g 2y ago
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 ¯\_(ツ)_/¯

0x19241217 2y ago

Pushback 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.

dunefox 2y ago
> If researchers wrote directly in C or Fortran... ... then they would get nothing done.
- pklausler 2y ago
  Why not?
  - dunefox 2y ago
    1. data cleaning and other tasks for which Python excels, 2. C and Fortran don't come close to the developer speed of Python. Ideas don't come from using languages like C, but are rather enabled by implementing, prototyping, and iterating quickly and with little resistance from the language, using existing libraries, etc. Python is simply much better in these regards.
- TheRealKing 2y ago
  Fortran cannot be placed with C in the same category of low programming productivity.
- pjmlp 2y ago
  Apparently we could do research work before Python was invented.

FL33TW00D 2y ago

Are 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.

visarga 2y ago
You could sen the KV cache and even the LoRA as input args to a static model.
spott 2y ago
Is it? The same inputs have the same outputs, the only thing different is the performance. It is basically the same as memorization.

ianbutler 2y ago

From an eng/industry perspective, back in 2016/2017 I watched the realtime decline of Tensorflow towards Pytorch. The issue was TF had too many interfaces to accomplish the same thing and each one was rough in its own way. Along with some complexity for using serving and experiment logging via Tensorboard, but this wasn’t as bad at least for me. Keras was integrated in an attempt to help, but ultimately it wasn’t enough and people started using Torch more and more even against the perception that TF was for prod workloads and Torch was for research. TFA mentions the interface complexity as starting to be a problem with Torch, but I don’t think we’re anywhere near the critical point that would cause people to abandon it in favor of JAX. Additionally with JAX you’re just shoving the portability problems mentioned down to XLA which brings its own issues and gotchas even if it hides the immediate reality of said problems from the end user. I think the Torch maintainers should watch not to repeat the mistakes of TF, but I think theres a long way to go before JAX is a serious contender. It’s been years and JAX has stayed in relatively small usage.

markeroon 2y ago
Imo the biggest issue (from memory) was that Tensorflow used a static computation graph. PyTorch was so much easier to work with.
- ianbutler 2y ago
  Yup that was also definitely one among the many issues
- lostmsu 2y ago
  I honestly think the static graph was much better in the similar way Vulkan/DX12 are better than OpenGL/DX9. It is harder to program, but gives more explicit control of important things. E.g. who would use PyTorch if they new that for optimal performance they'd need to record CUDA graphs?

sundarurfriend 2y ago

Can we get the title changed to the actual title of the post? "The future of Deep Learning frameworks" sounds like a neutral and far wider-reaching article, and ends up being clickbait here (even if unintentionally). "PyTorch is dead. Long live JAX." conveys exactly what the article about, and is a much better title.

6gvONxR4sf7o 2y ago

One aspect of jax that’s rarely touched on is browser stuff. Completely aside from deep learning, it’s straightforward to compile jax to a graphics shader you can call in js, which in this insane world is actually my preferred way to put numerical computing or linear algebra code on a web page.

mccoyb 2y ago
Can you share any code or references for this? Several people I work with are interested in this.
- 6gvONxR4sf7o 2y ago
  Sure. I’ll have to dig it up (shits boxed up as I move), but i’ll ping you or this comment when I get to it.

lostmsu 2y ago

From the article is seems that JAX is a non-starter for me as they don't have support for any kind of acceleration on Windows proper, and only experimental in WSL.

astromaniak 2y ago

The article misses multi-modal thing. Which is the future. Sure they can be considered a separate things, like today. But that's probably not the best approach. Support from framework may include partial training, easy components swap, intermediate data caching, dynamic architecture, automatic work balance and scaling.