jax

f255b83f - [Mosaic:GPU] Remove supports_cross_device_collectives since collective metadata currently support multi-process usage as well.