xla

3b121151 - Merge xm all_gather patch (#3416)

Merge xm all_gather patch (#3416) * Set proper shard_count for all_gather, when replica groups are non-empty. * Update test_mp_all_gather.py