[SPMD] Support manual all-reduce (#7576)
Summary:
This is to add manual all-reduce support to SPMD and it currently only supports one input tensor. For array support, we can do that in python layer instead.
Test Plan:
python ./test/spmd/test_xla_sharding.py -v -k test_spmd_all_reduce
