xla

0025ca76 - [SPMD] Fix reduce_scatter for TPUv2 (#7248)

[SPMD] Fix reduce_scatter for TPUv2 (#7248) Summary: Fix reduce_scatter test for TPUv2 Test Plan: python test/spmd/test_xla_sharding.py -v -k test_spmd_reduce_scatter