jax

f2edc83a - [Pallas:MGPU] Properly commute indexing with other transforms

[Pallas:MGPU] Properly commute indexing with other transforms Doing so requires us to modify the other transforms when we attempt to move indexing before them. PiperOrigin-RevId: 687240515