Fix spmd reduce scatter python test

[bhavya01]

These tests were failing on v5e-8 with the following error:

$ python test/spmd/test_xla_sharding.py
.....................................................F
Actual: tensor([[8.],
        [8.],
        [8.],
        [8.],
        [8.],
        [8.],
        [8.],
        [8.]]) 
        
Expected: tensor([[4.],
        [4.],
        [4.],
        [4.],
        [4.],
        [4.],
        [4.],
        [4.]])
F........
======================================================================
FAIL: test_spmd_reduce_scatter (__main__.BasicXlaShardingTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/mnt/disks/bbahl/pytorch/xla/test/spmd/test_xla_sharding.py", line 1235, in test_spmd_reduce_scatter
    torch.testing.assert_close(x.cpu(), expected_x)
  File "/mnt/disks/bbahl/miniconda3/envs/torchnightly/lib/python3.10/site-packages/torch/testing/_comparison.py", line 1524, in assert_close
    raise error_metas[0].to_error(msg)
AssertionError: Tensor-likes are not close!

Mismatched elements: 8 / 8 (100.0%)
Greatest absolute difference: 4.0 at index (0, 0) (up to 1e-05 allowed)
Greatest relative difference: 1.0 at index (0, 0) (up to 1.3e-06 allowed)

======================================================================
FAIL: test_spmd_reduce_scatter_canonical_index (__main__.BasicXlaShardingTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/mnt/disks/bbahl/pytorch/xla/test/spmd/test_xla_sharding.py", line 1257, in test_spmd_reduce_scatter_canonical_index
    torch.testing.assert_close(x.cpu(), expected_x)
  File "/mnt/disks/bbahl/miniconda3/envs/torchnightly/lib/python3.10/site-packages/torch/testing/_comparison.py", line 1524, in assert_close
    raise error_metas[0].to_error(msg)
AssertionError: Tensor-likes are not close!

Mismatched elements: 8 / 8 (100.0%)
Greatest absolute difference: 4.0 at index (0, 0) (up to 1e-05 allowed)
Greatest relative difference: 1.0 at index (0, 0) (up to 1.3e-06 allowed)

----------------------------------------------------------------------
Ran 63 tests in 2.984s

FAILED (failures=2)

It seems like the content of the tensor should depend on the number of devices that we shard across.

[alanwaketan]

LGTM. Sorry for late review.