sharded_jit cleanup (#3075)
* Add sharding utilities to xla_bridge
* Change partitions argument to in_parts and out_parts
* Add unit tests
* Reuse more pxla functionality
* Remove stale translation rule
* Fail on non-TPU platforms
* Add docstring
* And more!
