AlphaZero did not run game logic on TPUs (neither chess nor other games), implementing it in C++ is more than fast enough and much simpler.

TPUs were used for neural network inference and training, but game logic as well as MCTS was on the CPU using C++.

JAX is awesome though, I use it for all my neural network stuff!