Hacker News Reader

Comment by whimsicalism | 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.

View on HN

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

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.

Reply on news.ycombinator.com