Yes, I think Jax is indeed the nail for tensorflow. It's not there yet, but the part of the research community that did not go to pytorch is going to jax now.

Yes, I think Jax is indeed the nail for tensorflow. It's not there yet, but the part of the research community that did not go to pytorch is going to jax now. And jax is made by the autograd people.