xla

a6a77ef8 - Fix bug in ShardTensor to generate correct shards for each device (#4549)