Add an OpShardingSharding codepath for vmap for pjit. In xmap cases, still round trip through MeshPspecSharding.
Add an OpShardingSharding codepath for vmap for pjit. In xmap cases, still round trip through MeshPspecSharding.

PiperOrigin-RevId: 465629138

google/jax,google/jax,google/jax,google/jax