jax

7e2f487a - [Mosaic TPU] Canonicalize arith.select's condition to vector if other types are vector.

[Mosaic TPU] Canonicalize arith.select's condition to vector if other types are vector. This fixes the failure in elementwise rule of apply vector layout pass. If the condition scalar is static, it will be simplified to corresponding vector from true value and false value by MLIR. If the condition scalar is dynamic, we want to use vselect over scf.if anyway. Because latter creates a inner region. PiperOrigin-RevId: 680674560