jax

e4a381c1 - [pallas:mgpu] Check that swizzle dim is not transposed in copy_smem_to_gmem()