jax

f899d515 - [Mosaic TPU] Fold sublane offset to indices when storing to untiled ref.

[Mosaic TPU] Fold sublane offset to indices when storing to untiled ref. This optimization avoids unnecessary retiling when storing to untiled ref but adds at most one extra store op for sublane offset (since sublane offset is limieted to < VregSlice[0]). PiperOrigin-RevId: 698896373