Fault-tolerant training with jax.distributed

[superMDguy]

The docs say that jax will "ensure that all processes shut down if any process dies". I'm working on getting a training run setup with preemptible TPUs, and it's likely that nodes will go offline during training. It would be great if Jax could pause operations when a node dies so I can automatically create a new node from the latest checkpoint, instead of having to reload all of the running instances.