jax

220ec2aa - [Mosaic TPU] (8,128),-2 -> (8,128) for non-zero and replicated 2nd minor offset

[Mosaic TPU] (8,128),-2 -> (8,128) for non-zero and replicated 2nd minor offset Also fix bug where relayouts for fully replicated source assumed it was a no-op without checking implicit dims PiperOrigin-RevId: 655746766