jax

bae66000 - [pallas:mosaic_gpu] `FragmentedArray.reduce_sum` now returns a `FragmentedArray`

[pallas:mosaic_gpu] `FragmentedArray.reduce_sum` now returns a `FragmentedArray` This aligns it with the `reduce` method and also makes it clear that the reduction always produces a scalar. PiperOrigin-RevId: 703494443