class_docstr: Tests for bluefog/torch/mpi_ops.py
class_docstr_tok: ["Tests","for","bluefog\/torch\/mpi_ops.py"]
docstr: Test neighbor allreduce works with destination weights under tensor fusion.
code: def test_neighbor_allreduce_dst_weight_fusion(self):
      """Test neighbor allreduce works with destination weights under tensor fusion."""
      size = bf.size()
      rank = bf.rank()
      K = 50  # number of tensors send in short time
      if size <= 1:
          fname = inspect.currentframe().f_code.co_name
          warnings.warn("Skip {} due to size 1".format(fname))
          return
  
      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) * K
  
      dst_weight = 0.5
      self_weight = dst_weight
      src_weights = {i: 1.0 for i in neighbor_ranks}
      dst_weights = {
          (rank + 2**i) % size: dst_weight
          for i in range(num_indegree)
      }
  
      dims = [1, 2, 3]
      for dtype, dim in itertools.product(dtypes, dims):
          tensor_list, handles, names = [], [], []
          for i in range(K):
              tensor = torch.FloatTensor(*([23] * dim)).fill_(
                  i + rank * K
              )
              tensor = self.cast_and_place(tensor, dtype)
              tensor_list.append(tensor)
              names.append("index{}_{}_{}".format(i, dtype, dim))
  
          for i in range(K):
              handle = bf.neighbor_allreduce_nonblocking(
                  tensor_list[i],
                  name=names[i],
                  self_weight=self_weight,
                  src_weights=src_weights,
                  dst_weights=dst_weights,
              )
              handles.append(handle)
  
          outputs = []
          for i in range(K):
              output = bf.synchronize(handles[i])
              outputs.append(output)
  
          for i in range(K):
              assert (
                  list(outputs[i].shape) == [23] * dim
              ), f"{names[i]} (fusion) produces incorrect reduced shape"
              output_normalized = outputs[i] / dst_weight - i * (
                  num_indegree + 1
              )
              assert (
                  output_normalized - sum_value
              ).abs().max() < EPSILON, f"{names[i]} (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 32 for calling the OpsTests obj's underlying member method code with example usage: obj.test_neighbor_allreduce_dst_weight_fusion() without return types