jax

a1de0c73 - Add heartbeat argument to distributed.initialize.

Add heartbeat argument to distributed.initialize. This PR adds a new `heartbeat_timeout_seconds` argument to `jax.distributed.initialize` that allows users to customize how long a process can fail to send a heartbeat before being considered dead.