jax

f59d8213 - [mosaic] Verify that `tpu.memref_slice` always produces tile-aligned slices

[mosaic] Verify that `tpu.memref_slice` always produces tile-aligned slices Unlike the original version, this one checks both the indices and sizes. Reverts 5abf880a821108f5be1c2f6ecd556ce6179594f3 PiperOrigin-RevId: 874778201