JAX is a huge step up from Tensorflow, even though I aboslutely don't understand why it takes Google so many iterations to eventually land on a PyTorch API. Jax might be close enough that they'll give up trying, but I feel like it still falls short in being as good as Pytorch. But Google will definitely continue using it, and they're one of the main drivers of AI research. So at least in that area it'll see non-trivial adoption. I still think Jax is inferior to PyTorch, so there's no reason to switch (the better support for TPUs might be a selling point, so if Google gets serious about pushing those onto people, there might be an uptick in Pytorch->Jax conversion).

The productization story is way worse than in PyTorch, and even within google, production still uses TF. Until that changes, I don't think Jax will make big inroads in industry.