class_docstr: Tests for bluefog/tensorflow/mpi_ops.py
class_docstr_tok: ["Tests","for","bluefog\/tensorflow\/mpi_ops.py"]
docstr: Test that the allreduce works on GPUs.
docstr_tok: ["Test","that","the","allreduce","works","on","GPUs","."]
code: def test_bluefog_allreduce_gpu(self):
      """Test that the allreduce works on GPUs."""
      # Only do this test if there are GPUs available.
      if not tf.test.is_gpu_available(cuda_only=True):
          return
  
      local_rank = bf.local_rank()
      size = bf.size()
  
      dtypes = [tf.int32, tf.int64, tf.float32, tf.float64]
      dims = [1, 2, 3]
      for dtype, dim in itertools.product(dtypes, dims):
          with tf.device("/gpu:%d" % local_rank):
              tensor = random_uniform(
                  [17] * dim, -100, 100, dtype=dtype
              )
              summed = bf.allreduce(tensor, average=False)
          multiplied = tensor * size
          max_difference = tf.reduce_max(tf.abs(summed - multiplied))
  
          # Threshold for floating point equality depends on number of
          # ranks, since we're comparing against precise multiplication.
          if size <= 3 or dtype in [tf.int32, tf.int64]:
              threshold = 0
          elif size < 10:
              threshold = 1e-4
          elif size < 15:
              threshold = 5e-4
          else:
              return
  
          diff = self.evaluate(max_difference)
          self.assertTrue(
              diff <= threshold,
              "bf.allreduce on GPU produces incorrect results",
          )
  
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 5 for calling the OpsTests obj's underlying member method code with example usage: obj.test_bluefog_allreduce_gpu() without return types