jax

307c8d3a - [Mosaic:TPU] For trunc, expand supported tilings, offsets and bitwidths

[Mosaic:TPU] For trunc, expand supported tilings, offsets and bitwidths infer-vector-layout won't use the full generality anytime soon, but we could reuse this logic for relayouts PiperOrigin-RevId: 708011538