jax

4e35adb4 - [pallas:mgpu] Slightly reworded the docstring of `copy_gmem_to_smem`