jax

7ba811eb - Support auto in shard_map.

Support auto in shard_map. - Pull mesh from NamedSharding when rewriting manual axes. - Properly set manual axes in SPMDAxisContext in shard_map. - Properly set dims as unspecified inside shard_map. PiperOrigin-RevId: 627156892