jax

1b3c5a72 - Unreduced/reduced can only talk about Explicit and Manual mesh axes in ShapedArray (remove auto axes from the type). This worked in v0.8.0 but fails in v0.8.1

Unreduced/reduced can only talk about Explicit and Manual mesh axes in ShapedArray (remove auto axes from the type). This worked in v0.8.0 but fails in v0.8.1 Fixes reduced bug with .get(out_sharding=...) where the input was reduced over some mesh axes. PiperOrigin-RevId: 836389460