jax

bd0397dd - [Mosaic:TPU][Relayouts] Replicated minor large-to-small retiling (one-to-many)

[Mosaic:TPU][Relayouts] Replicated minor large-to-small retiling (one-to-many) Previously, this retiling did not preserve replication, leading to errors when both the source and target layout are replicated and retiling is needed. It is also more efficient even when the target layout does not require replication. PiperOrigin-RevId: 825818371