jax

53346b39 - [Mosaic GPU] Add support for keeping dimensions in `Reduce` constraint.

[Mosaic GPU] Add support for keeping dimensions in `Reduce` constraint. If `keep_dims=True`, the axes which are reduced are left in the result as dimensions with size one. PiperOrigin-RevId: 881945009