Set up jax with multiple hosts

[Happy-to-see-you]

We read this tutorial about setting up jax in multi-host environments, which was excellent:

https://jax.readthedocs.io/en/latest/multi_process.html

But how do we get jax to recognize all of the hosts that are available. For example, suppose we have four computers on the same network, each with a GPU, and want to use them together. How can we set up jax so that jax.devices() recognizes all of them?