xla

56a6a02a - Partially replicate lower-rank tensors (#5409)

Partially replicate lower-rank tensors (#5409) * Partially replicate lower-rank tensors * Fix unit test * Remove unnecessary device count check * Fix unordered partition spec test * yapf