class_docstr: Tests for bluefog/torch/mpi_ops.py
class_docstr_tok: ["Tests","for","bluefog\/torch\/mpi_ops.py"]
docstr: Test that the neighbor allreduce works under tensor fusion.
code: def test_neighbor_allreduce_fusion(self):
      """Test that the neighbor allreduce works under tensor fusion."""
      size = bf.size()
      rank = bf.rank()
      if size <= 1:
          fname = inspect.currentframe().f_code.co_name
          warnings.warn("Skip {} due to size 1".format(fname))
          return
      # TODO(hanbinhu) torch Half have problem in accuracy
      dtypes = [torch.FloatTensor, torch.DoubleTensor]
      if TEST_ON_GPU:
          dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
  
      # By default, we use exponential two ring topology.
      num_indegree = int(np.ceil(np.log2(size)))
      neighbor_ranks = [
          (rank - 2**i) % size for i in range(num_indegree)
      ]
      sum_value = np.sum(neighbor_ranks) + rank
  
      dims = [1, 2, 3]
      for dtype, dim in itertools.product(dtypes, dims):
          tensor_1 = torch.FloatTensor(*([23] * dim)).fill_(rank)
          tensor_2 = torch.FloatTensor(*([23] * dim)).fill_(rank + 1)
          name_1 = "neighbor_allreduce_fusion_tensor_{}_{}_1".format(
              dim, dtype
          )
          name_2 = "neighbor_allreduce_fusion_tensor_{}_{}_2".format(
              dim, dtype
          )
          tensor_1 = self.cast_and_place(tensor_1, dtype)
          tensor_2 = self.cast_and_place(tensor_2, dtype)
  
          handle_1 = bf.neighbor_allreduce_nonblocking(
              tensor_1, name=name_1
          )
          handle_2 = bf.neighbor_allreduce_nonblocking(
              tensor_2, name=name_2
          )
  
          output_1 = bf.synchronize(handle_1)
          output_2 = bf.synchronize(handle_2)
          tensor_1, output_1 = self.convert_cpu_fp16_to_fp32(
              tensor_1, output_1
          )
          tensor_2, output_2 = self.convert_cpu_fp16_to_fp32(
              tensor_2, output_2
          )
  
          eps = (
              EPSILON
              if tensor_1.dtype != torch.float16
              else LOOSE_EPSILON
          )
          assert (
              list(output_1.shape) == [23] * dim
          ), "bf.neighbor_allreduce_1 (fusion) produces incorrect reduced shape"
          assert (
              output_1.data.mul_(num_indegree + 1) - sum_value
          ).abs().max() < eps, "bf.neighbor_allreduce_1 (fusion) produces incorrect reduced tensor"
          assert (
              list(output_2.shape) == [23] * dim
          ), "bf.neighbor_allreduce_2 (fusion) produces incorrect reduced shape"
          assert (
              (output_2.data - 1).mul_(num_indegree + 1) - sum_value
          ).abs().max() < eps, "bf.neighbor_allreduce_2 (fusion) produces incorrect reduced tensor"
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import inspect
  import itertools
  import unittest
  import warnings
  import numpy
  import pytest
  import torch
  import networkx
  import bluefog.torch
  from bluefog.common import topology_util
  from common import is_openmpi_built
label_desc: Use image node_id 29 for calling the OpsTests obj's underlying member method code with example usage: obj.test_neighbor_allreduce_fusion() without return types