jax

cf308a84 - Use an PartitionSpec.UNCONSTRAINED to represent unconstrained dimensions in ParsedPartitionSpec, rather than None.

Use an PartitionSpec.UNCONSTRAINED to represent unconstrained dimensions in ParsedPartitionSpec, rather than None. This makes PartitionSpec and ParsedPartitionSpec more similar, and fixes some TODOs. PiperOrigin-RevId: 724927217