jax

3e8cb42e - [Mosaic GPU] Add support for reductions of arrays that have untiled dimensions

[Mosaic GPU] Add support for reductions of arrays that have untiled dimensions Previously we would reject those cases and would only support the ones where the rank of the array equals its tiling rank. PiperOrigin-RevId: 872333005