jax

e5b542ac - [Pallas/Mosaic GPU] Add cross-device barrier before running multi-device kernel with collective metadata