jax

bfe8e4f0 - Enable `all_to_all` in multi-host settings. (#3772)

Enable `all_to_all` in multi-host settings. (#3772) I tested this via `pswapaxes` and it seems to work. There may still be issues with all_to_all (e.g. https://github.com/google/jax/issues/1332) but it seems worth enabling the `pswapaxes` use case.