jax

0aa31d19 - [Mosaic GPU] Minor logic cleanup in FragmentedArray.reduce

[Mosaic GPU] Minor logic cleanup in FragmentedArray.reduce `bool(reduced_shape)` is always true, since `reduced_shape` is a tuple of the same rank as the input array. PiperOrigin-RevId: 880786644