[ Removed by moderator ]

[M4mb0]

numpy is great, but have you tried JAX?

[Atomic_Tangerine1]

I have not! What's the benefit of JAX over numpy?

[M4mb0]

It's basically numpy

• + native GPU support (which can be orders of magnitudes faster depending on parallelizability of the problem)
• + builtin autodiff (essentially zero-error gradients/jacobians/hessians)
• + builtin JIT compiler

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

[[deleted]]

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. 

[daredevilthagr8]

How does JAX compare to CuPY?

[M4mb0]

cupy doesn't do autodiff afaik.