[CI] fix test_pointwise_ops.py test_mul_div_scalar_partial (pytorch#170510)

Support any world size; 2, 3 or 4.

Pull Request resolved: pytorch#170510
Approved by: https://github.com/jeffdaily

Co-authored-by: Jeff Daily <[email protected]>

Author: jeffdaily

test/distributed/tensor/test_pointwise_ops.py

@@ -467,7 +467,8 @@ def test_mul_div_scalar_partial(self):
 
         self.assertTrue(res._spec.placements[0].is_partial())
         res = res.redistribute(dt.device_mesh, placements=[Replicate()])
-        self.assertEqual(res, 12)
+        expected = sum(i * 2 for i in range(self.world_size))
+        self.assertEqual(res, expected)
 
         res = aten.div.Scalar(dt, 2)
         self.assertEqual(
@@ -478,7 +479,8 @@ def test_mul_div_scalar_partial(self):
         self.assertTrue(res._spec.placements[0].is_partial())
         res = res.redistribute(dt.device_mesh, placements=[Replicate()])
 
-        self.assertEqual(res, 3)
+        expected = expected / 4.0
+        self.assertEqual(res, expected)
 
     @with_comms
     @parametrize("op,reduce_op", [(torch.maximum, "max"), (torch.minimum, "min")])