xla

afb00b06 - [SPMD][PoC] Tensor sharding util (#3962)

[SPMD][PoC] Tensor sharding util (#3962) * [SPMD] Refactor XLAShardedTensor & unit tests. * [SPMD] Add at::Tensor sharding util & sharding tests.