Move some sharded_jit functionality into pxla.py (#3142)
Specifically:
* Move `_inner_partitions`
* Move `get_num_partitions`
* Move and slightly modify the logic for finding and validating inner
partitions to a new function, `reconcile_num_partitions`
* Move `_partitioned_sharding_spec` and rename to `partitioned_sharding_spec`
This is in preparation for enabling pmap-of-sharded_jit, since pmap
will need access to this functionality as well.
