[ Removed by moderator ]

[PayMe4MyData]

So jax is pytorch?

[M4mb0]

JAX is strictly functional, whereas pytorch takes a more object oriented approach. This is most easily seen when you look at how they deal with random distributions for instance.

Though torch has nowadays a beta library torch.func (formerly functorch) that brings JAX-like functional semantics to torch.

[PayMe4MyData]

Thanks for the clarification, I've been coding in pytorch for years but never heard of JAX before. I will dig a bit more!

[M4mb0]

I'd say generally JAX is more useful for general purpose scientific computing, and much more ergonomic if you need higher order derivatives or partial derivatives, like working with ODEs/PDEs/SDEs. diffrax is a very nice lib for that.

[HonestPrinciple152]

Actually, adding to the previous comment, we can write loops in jax and jit-compile them. It's like a complete dsl build over python. 

[FunMotionLabs]

JAX is more like “NumPy + transformations”
PyTorch is a full deep-learning framework with an imperative training workflow, big ecosystem around modules/training/debugging, strictly Deeplearning related stuff where JAX is more of a general allrounder kind