For me the main point is that JAX is slightly lower level than Pytorch and has nicer abstraction (no need to worry about tensors that might not store a gradient

For me the main point is that JAX is slightly lower level than Pytorch and has nicer abstraction (no need to worry about tensors that might not store a gradient or wetehr you are on the GPU: it eliminates lots of newcomers bug) which makes it a great fit to build Deep learning frameworks, but also simulations, on top of it.