Comment by imtringued

I was about to say that it is literally just Jax.

It genuinely deserves to exist alongside pytorch. It's not just Google's latest framework that you're forced to use to target TPUs.