Details of using TPU slices from Jax can be found on this page:

• https://cloud.google.com/tpu/docs/jax-pods

There is some introspection on available devices in the Jax API:

• https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap

If you are going multi-slice, you could consider using Google's Pax (aka PaxML) library as it is built for this:

• https://github.com/google/paxml/blob/main/paxml/docs/about-pax.md