Hacker News Reader

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.

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.

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

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.

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?

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.

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)

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.

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.