jax

be0dd453 - [Pallas:MGPU] Erase swizzle from transform list before handling transforms

[Pallas:MGPU] Erase swizzle from transform list before handling transforms This lets us avoid the annoying errors when trying to swap a transpose and swizzle transform. The FragmentedArray methods always apply swizzle to the minormost array dimension so it's ok to pop it for as long as it's the first transform. PiperOrigin-RevId: 958998796