Hacker News Reader

Comment by chazeon | Hacker News Reader

How is PyTorch compares to JAX and its stack?

I find that Jax tends to result in messy code unless you build good abstractions. I personally don't like Flax and Haiku, I prefer stax and Equinox as they are more transparent on what is happening, feel a lot less like magic, and more pythonic (explicit is better than implicit etc).

PyTorch is far more friendly for deep learning stuff, but sometimes all you want is pure numerical computations that can be vmapped across tensors, and this is where jax shines imho.

Personal Example: I needed to sample a bunch of datapoints, make distributions out of them, sample, and then compute the density of each sample across distributions. Doing this with pytorch was rather slow, I was probably doing something wrong with vectorization and broadcasting, but I didn't have the time to figure it out.

With jax, I wrote a function that produces the samples, then I vmapped the evaluation of a sample across all distributions, then vmapped over all samples. Took a couple of minutes to implement and seconds to execute.

PyTorch also has the advantage of a far more mature ecosystem, libraries like Lightning, Accelerate, Transformers, Evaluate, and so on make building models a breeze.

> Personal Example: I needed to sample a bunch of datapoints, make distributions out of them, sample, and then compute the density of each sample across distributions. Doing this with pytorch was rather slow, I was probably doing something wrong with vectorization and broadcasting, but I didn't have the time to figure it out.

You probably were not doing anything wrong. I spent a lot of time trying to be clever in order to parallelize things like this and it just wasn't possible without doing CUDA extensions. But it is now! PyTorch now has vmap through functorch and it works.

I found that there still some limitations with functorch's vmap, but I can't recall what it was.

Probably not your issue, but one kind of annoying bit is that the inputs need to be tensors. I ended up calling partial on the function I was messing around with and then vmapping the partial, which seemed to work.

Pretty similar actually, I needed to pass in some parameters as tuples, lists, etc. Partial would have probably worked, but tbh I wasn't in the mood to try harder at 2am.

PyTorch and JAX are both open-source libraries for developing machine learning models, but they have some important differences. PyTorch is a more general-purpose library that provides a wide range of functionalities for developing and training machine learning models. It also has strong support for deep learning and is used by many researchers and companies in production environments.

JAX, on the other hand, is designed specifically for high-performance machine learning research. It is built on top of the popular NumPy library and provides a set of tools for creating, optimizing, and executing machine learning algorithms with high performance. JAX also integrates with the popular Autograd library, which allows users to automatically differentiate functions for training machine learning models.

Overall, the choice between PyTorch and JAX will depend on the specific requirements and goals of the project. PyTorch is a good choice for general-purpose machine learning development and is widely used in industry, while JAX is a better choice for high-performance research and experimentation.

https://chat.openai.com/chat

I was reading this and thinking it was a pretty terrible answer - glad it is just generated by an AI and not you personally so I'm not insulting you.

JAX is basically numpy on steroids and lets you do a lot of non-standard things (like a differentiable physics simulation or something) that would be harder with Pytorch.

They are both "high-performance."

Pytorch is more geared towards traditional deep learning and has the utilities and idioms to support it.

jax is not numpy on steroids. jax is "use python idiomatically to generate optimized XLA code for evaluating functions both forward and backward."

Probably the primary use of jax is `jax.numpy` which is XLA accelerated and differentiable numpy.

I'll admit that saying "basically numpy on steroids" might have been an overreduction. It is a system for function transformations that is built on XLA and oriented towards science & ML applications.

It's not just me saying stuff like this.

François Chollet (creator of Keras): "[jax is] basically Numpy with gradients. And it can compile to XLA, for strong GPU/TPU acceleration. It's an ideal fit for researchers who want maximum flexibility when implementing new ideas from scratch."

Yes- and that gradient part is a key detail that makes it more than "numpy on steroids". numpy on steroids would be a hardware accelerator that took numpy calls and made them return more quickly, but without the command-and-control and compile-python-to-xla aspects.

Well clearly I meant steroids of the gradient-developing variety.

I think you are being far too pedantic about what a biological compound would analogously do to a software library, especially given that I mention the differentiability property in the same sentence you are taking issue with.

OK, actually as long as it's gradient-developing steroids, I'll allow it.

Can someone comment more on what makes JAX that much better for differentiable simulations than PyTorch?

I'm working on a new module for work and none of my colleagues have much experience developing ML per se. I'm trying to decide whether to force their hand by implementing v1 in PyTorch or JAX and differentiable physics simulations is a likely future use case. Why is PyTorch harder?

Because the `jax.numpy` operations & primitives are almost 1:1 with numpy, many working scientists who already have experience working with numpy will be able to figure out jax faster.

It is also easier to rewrite existing code/snippets (say you were working on a non-differentiable simulator before) into jax if you already have them in numpy then to do the whole rewrite in pytorch.

I will say that I think pytorch has improved its numpy compatability a lot in recent years, functions that I was convinced didn't exist with pytorch (like eigh) apparently actually do.

I have seen JAX-MD[1] but not sure about "much better". On the other hand, there is just no MD implemented with PyTorch.

[1]: https://github.com/jax-md/jax-md

There is [torch-md](https://github.com/torchmd/torchmd)

At least prior to this announcement: JAX was much faster than PyTorch for differentiable physics. (Better JIT compiler; reduced Python-level overhead.)

E.g for numerical ODE simulation, I've found that Diffrax (https://github.com/patrick-kidger/diffrax) is ~100 times faster than torchdiffeq on the forward pass. The backward pass is much closer, and for this Diffrax is about 1.5 times faster.

It remains to be seen how PyTorch 2.0 will compare, of course!

Right now my job is actually building out the scientific computing ecosystem in JAX, so feel free to ping me with any other questions.

If you care about performance of differential physics you shouldn't use python. Diffrax is almost OKish, but is missing a ton of features (e.g. good stiff solvers, arbitrary precision support, events for anything other than stopping the simulation, ability to control the linear solve which are needed for large problems). For simple cases it can come close to the C++/Julia solvers, but for anything complicated, you either won't be able to formulate the model, or you won't be able to solve it efficiently.

> If you care about performance

This definitely isn't true. On any benchmark I've tried, JAX and Julia basically match each other. Usually I find JAX to be a bit faster, but that might just be that I'm a bit more skilled at optimising that framework.

Anyway I'm not going to try and debunk things point-by-point, I'd rather avoid yet another unpleasant Julia flame-war.

I’m not sure why, but I realized it was AI from the very first sentence, not exaggerating. It’s just not something someone on HN would write.

It reminded me of the sort of lazy Wikipedia regurgitation that a lot of undergrads used to give when I was teaching. So it is a bit jarring to see a response like that in a non-compulsory setting.

Yup. Reminds me of an article you'd find in the top 10 Google search results...

It's funny, cause already after the first sentence it felt like ChatGPT, probably because I've played with it a lot these past few days, and expectedly I found a disclaimer at the end.

That said, the answer isn't really useful, as it's very generic, without anything concrete (other than the mention of Autograd) imo.

Though a follow up question might improve on that.

It seems to use the same type of template for comparisons:

React and Vue are both JavaScript libraries for building user interfaces. The main difference between the two is that React is developed and maintained by Facebook, while Vue is an independent open-source project.

React uses a virtual DOM (Document Object Model) to update the rendered components efficiently, while Vue uses a more intuitive and straightforward approach to rendering components. This makes Vue easier to learn and use, especially for developers who are new to front-end development.

React also has a larger community and ecosystem, with a wider range of available libraries and tools. This can make it a better choice for larger, more complex projects, while Vue may be a better fit for smaller projects or teams that prefer a more lightweight and flexible approach.

Overall, the choice between React and Vue will depend on your specific project requirements and personal preferences. It's worth trying out both to see which one works better for you.

Reply on news.ycombinator.com