class_docstr: Tests for bluefog/torch/mpi_ops.py
class_docstr_tok: ["Tests","for","bluefog\/torch\/mpi_ops.py"]
docstr: Test neighbor_allgather 1D, 2D, 3D tensors with variable size in dynamic setting.
code: def test_neighbor_allgather_dynamic_variable_size(self):
      """Test neighbor_allgather 1D, 2D, 3D tensors with variable size in dynamic setting."""
      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
      dtypes = [
          torch.FloatTensor,
          torch.IntTensor,
          torch.DoubleTensor,
          torch.LongTensor,
          torch.ByteTensor,
          torch.CharTensor,
          torch.ShortTensor,
          torch.HalfTensor,
      ]
      if TEST_ON_GPU:
          dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
  
      # Connect to all other ranks
      neighbor_ranks = [i for i in range(size) if i != rank]
      dims = [1, 2, 3]
      for dtype, dim in itertools.product(dtypes, dims):
          # Support tests up to MPI Size of 35
          if size > 35:
              break
  
          tensor_sizes = [17, 32, 81, 12, 15, 23, 22] * 5
          tensor_sizes = tensor_sizes[:size]
  
          tensor = (
              torch.FloatTensor(
                  *([tensor_sizes[rank]] + [17] * (dim - 1))
              )
              .fill_(1)
              .mul_(rank)
          )
          tensor = self.cast_and_place(tensor, dtype)
          gathered = bf.neighbor_allgather(
              tensor, dst_ranks=neighbor_ranks, src_ranks=neighbor_ranks
          )
          tensor, gathered = self.convert_cpu_fp16_to_fp32(
              tensor, gathered
          )
  
          tensor_sizes[
              rank
          ] = 0  # remove self since neighbor_allgather does not include self
          expected_size = sum(tensor_sizes)
          assert list(gathered.shape) == [expected_size] + [17] * (
              dim - 1
          )
  
          for i in range(size):
              if i == rank:
                  continue
              rank_size = [tensor_sizes[i]] + [17] * (dim - 1)
              rank_tensor = gathered[
                  sum(tensor_sizes[:i]) : sum(tensor_sizes[: i + 1])
              ]
              assert (
                  list(rank_tensor.shape) == rank_size
              ), "bf.neighbor_allgather(dynamic, var) produces incorrect gathered shape"
              assert (
                  rank_tensor.data.min() == i
              ), "bf.neighbor_allgather(dynamic, var) produces incorrect gathered tensor"
              assert (
                  rank_tensor.data.max() == i
              ), "bf.neighbor_allgather(dynamic, var) produces incorrect gathered 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 35 for calling the OpsTests obj's underlying member method code with example usage: obj.test_neighbor_allgather_dynamic_variable_size() without return types