jax

ee832c1e - [Mosaic GPU] Fix bug in `Tiling.remove_dimension`.

[Mosaic GPU] Fix bug in `Tiling.remove_dimension`. We forgot to update `last_tile_rank` when looping over tiles, which could lead to incorrect dimension offsets. PiperOrigin-RevId: 777473496