2 ms·

Like what?

• sudosysgen 3y ago
  Anytime you want to make something GPU accelerated that doesn't fit as standard operations on tensors. For example, I often write RL environments in Jax, which is something you can't do in PyTorch. There's also things you can do in PyTorch but that would be far more difficult, for example an efficient implementation of MCTS. I also used Jax a lot for differential equations, not even sure how I would do that with PyTorch. Basically, Torch is a lot more like a specialization of Numpy for neural networks, while Jax feels a lot more like if you could just write CUDA as Python, and also get the Jacobians (jacs! jax!) and jvp for free (of everything, you can even differentiate hyperparameters through your optimizer which is crazy). At the end, when you're doing fundamental research and coming up with something new, I think Jax is just better. If all I had to do was implementation, then I would be a happy PyTorch user.