jax

d0d53dcb - [pallas:mgpu] Allow multi-dimensional barriers.