A small addendum: the only people I know who uses Jax are people who work at Google, or people who had a big GCP grant and needed to use TPUs as a result.