jax

6400a754 - [Mosaic GPU] Respect user provided swizzle.

[Mosaic GPU] Respect user provided swizzle. Now we specify swizzle upon construction and exercise previously added paths in full. We lift the constraint that the user provided swizzle should be equal to inferred swizzle. PiperOrigin-RevId: 935051449