Are there any resources going into detail about why the big players prefer JAX? I've heard this before but have never seen explanations of why/how this happened.
It's all about cost and performance. If you can train a foundation model 2x faster with JAX on the same hardware, you are effectively slashing your training costs by 2x, which is significant for a multi-million dollar training run.
“Faster” might have been true pre-Hopper, but there’s no such guarantee now.