def METHOD_NAME(self):
    device_mesh = self.build_device_mesh()
    shard_spec = [Shard(0)]
    tensor_to_sum = torch.randn(12, 8, 8)
    mat1 = distribute_tensor(tensor_to_sum, device_mesh, shard_spec)
    keep_dim_or_not = [True, False, None]
    for dim in range(tensor_to_sum.ndim):
        for keep_dim in keep_dim_or_not:
            sum_args = (dim, keep_dim) if keep_dim is not None else (dim,)
            dim_sumed_tensor = tensor_to_sum.sum(*sum_args)
            dt_dim_sumed_tensor = mat1.sum(*sum_args).redistribute(
                device_mesh, [Replicate()] * device_mesh.ndim
            )
            self.assertEqual(dt_dim_sumed_tensor.to_local(), dim_sumed_tensor)
    full_sumed_tensor = tensor_to_sum.sum()
    dt_sum = mat1.sum().redistribute(device_mesh, [Replicate()] * device_mesh.ndim)
    self.assertEqual(dt_sum.to_local(), full_sumed_tensor)