PyTorch is far more friendly for deep learning stuff, but sometimes all you want is pure numerical computations that can be vmapped across tensors, and this is where jax shines imho.

Personal Example: I needed to sample a bunch of datapoints, make distributions out of them, sample, and then compute the density of each sample across distributions. Doing this with pytorch was rather slow, I was probably doing something wrong with vectorization and broadcasting, but I didn't have the time to figure it out.

With jax, I wrote a function that produces the samples, then I vmapped the evaluation of a sample across all distributions, then vmapped over all samples. Took a couple of minutes to implement and seconds to execute.

PyTorch also has the advantage of a far more mature ecosystem, libraries like Lightning, Accelerate, Transformers, Evaluate, and so on make building models a breeze.