jax

f96339be - [Mosaic TPU] Be much more aggressive in inferring large 2nd minor layouts for 16-bit types on v6

[Mosaic TPU] Be much more aggressive in inferring large 2nd minor layouts for 16-bit types on v6 This often lets us avoid ambiguities between selecting the (8, 128) and (16, 128) tiling, by biasing the layout inference to prefer the latter. PiperOrigin-RevId: 713270421