jax

45e42c27 - [pallas/mgpu] Add a copy_smem_to_gmem scatter operation.