jax

6f67db57 - [Mosaic TPU] Support reshapes of form `(*src_major, src_minor) <-> (*dst_major, dst_minor)`, where `max(src_minor, dst_minor) <= 128` and one divides the other (row merging and splitting).

[Mosaic TPU] Support reshapes of form `(*src_major, src_minor) <-> (*dst_major, dst_minor)`, where `max(src_minor, dst_minor) <= 128` and one divides the other (row merging and splitting). This change is targeting at small lane dims, since when src_minor/dst_minor are multiples of 128, it is already supported. PiperOrigin-RevId: 966129068