class_docstr: Tests for bluefog/tensorflow/mpi_ops.py
class_docstr_tok: ["Tests","for","bluefog\/tensorflow\/mpi_ops.py"]
docstr: Test that the allgather correctly gathers 1D, 2D, 3D tensors.
docstr_tok: ["Test","that","the","allgather","correctly","gathers","1D",",","2D",",","3D","tensors","."]
code: def test_bluefog_allgather(self):
      """Test that the allgather correctly gathers 1D, 2D, 3D tensors."""
      rank = bf.rank()
      size = bf.size()
  
      dtypes = [tf.int32, tf.int64, tf.float32, tf.float64, tf.bool]
      dims = [1, 2, 3]
      for dtype, dim in itertools.product(dtypes, dims):
          tensor = tf.ones([17] * dim) * rank
          if dtype == tf.bool:
              tensor = tensor % 2
          tensor = tf.cast(tensor, dtype=dtype)
          gathered = bf.allgather(tensor)
  
          gathered_tensor = self.evaluate(gathered)
          self.assertEqual(
              list(gathered_tensor.shape),
              [17 * size] + [17] * (dim - 1),
          )
  
          for i in range(size):
              rank_tensor = tf.slice(
                  gathered_tensor,
                  [i * 17] + [0] * (dim - 1),
                  [17] + [-1] * (dim - 1),
              )
              self.assertEqual(list(rank_tensor.shape), [17] * dim)
              # tf.equal() does not support tf.uint16 as of TensorFlow 1.2,
              # so need to cast rank_tensor to tf.int32.
              if dtype == tf.bool:
                  value = i % 2
              else:
                  value = i
              self.assertTrue(
                  self.evaluate(
                      tf.reduce_all(
                          tf.equal(
                              tf.cast(rank_tensor, tf.int32), value
                          )
                      )
                  ),
                  "bf.allgather produces incorrect gathered tensor",
              )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import itertools
  import numpy
  import os
  import tensorflow
  from bluefog.tensorflow.util import _executing_eagerly, _has_eager
  from tensorflow.python.framework import ops
  import warnings
  import bluefog.tensorflow
label_desc: Use image node_id 10 for calling the OpsTests obj's underlying member method code with example usage: obj.test_bluefog_allgather() without return types