[goingtosleep]

You said from the beginning:

>> "Not sure why one would bother with this".

Then one provided the reasons for that, and you brought out your own opinions to defense that you don't wanna use Jax (without even trying)!? With that in mind, a thousand more reasons would not satisfy you.

From the perspective of a person who does DL research stuff with math background, I find Jax way more intuitive than any of other frameworks, Pytorch included. In math, you just don't "zero-ing the gradient" every iteration. In the same way, you just don't have to "forward" first to get the gradient. The more one does experiments with Jax, the more she / he would find it's faster / less steps / more intuitive to test new ideas. And I don't like the object-oriented design of Pytorch, functional seems more intuitive for me. That's why I would bother with Jax.