jax

8f232c94 - [Mosaic GPU] Add a `rank` field to the `Reduce` constraint.

[Mosaic GPU] Add a `rank` field to the `Reduce` constraint. This is necessary to figure out which dimensions must be reduced in a layout when the corresponding `vector` has untiled dimensions. PiperOrigin-RevId: 879640792