Multi-GPUs training with JAX and equinox Published: 2024-02-29 Multi-GPUs training with JAX and equinox