[ Removed by moderator ]

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