jax

a231d35f - [Shardy MPMD] Add support for unreduced axes to NamedSpmdShardingSpec.

[Shardy MPMD] Add support for unreduced axes to NamedSpmdShardingSpec. This change is needed in order to pass information about unreduced axes after MPMD partitioning to the backend compiler. The MPMD partitioner returns NamedSpmdShardingSpec for inputs/outputs of the module's main function as part of PartitioningResult, and these specs are used to create the input/output NamedShardings passed to the backend compiler. This change also includes updates to the Shardy MLIR utilities to extract unreduced axes from MeshTensorTypes as well as updates to the Python bindings in Shardy and jaxlib PiperOrigin-RevId: 886852988