Request for developer guide: multi-node TPU distributed training with JAX

[rivershah]

Multi-node TPU Training with JAX

The multi-GPU JAX training guide is helpful, but it's unclear how to extend this to multi-node TPU setups.

Currently using tpu v4-256 with tf.distribute.cluster_resolver.TPUClusterResolver and tf.distribute.TPUStrategy for data-parallel training. We're transitioning to jax and need the equivalent approach.

Specifically:

1. How to configure TPU runtime for jax.
2. How to handle cluster resolution for TPUs (similar to TPUClusterResolver).
3. Examples for multi-node TPU data-parallel training with jax.

Detailed examples would be helpful. Thank you.