xla

c74deb63 - Fix xm.all_gather for non-empty groups (#3413)

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