jax

4fe53bc3 - [Mosaic:GPU] Rollback: Use NCCL API on the host to collective peer addresses.

[Mosaic:GPU] Rollback: Use NCCL API on the host to collective peer addresses. Reason: Breaks oss tests. Reverts 77c0fd0c7be90c2eb0bc94b57eab9c8bb8b8fc64 PiperOrigin-RevId: 913784983