jax

32add9e6 - [Pallas:MGPU] Add a small benchmark script to the reduce scatter kernel