jax

f96c5661 - [Mosaic TPU][NFC] Refactor tpu matmul rule.

[Mosaic TPU][NFC] Refactor tpu matmul rule. * Separate MXU size to MXU contracting size and MXU non-contracting size. * Rename tile to group for MXU shaped tiling since tile is overused in Mosaic. PiperOrigin-RevId: 684116306