Support tuples in partition spec (#5488)
* Support tuples in partition spec
* Add unit test for partial replication
* yapf
* Support higher-rank tensors over lower-rank mesh
* Fix test & yapf
* Don't use partition_spec when creating group assignment
* Update documentation
* More documentation
* Translate named specs in ShardingSpec
