class_docstr: Tests for bluefog/tensorflow/mpi_ops.py
class_docstr_tok: ["Tests","for","bluefog\/tensorflow\/mpi_ops.py"]
docstr: Test the correctness of the broadcast gradient on CPU.
docstr_tok: ["Test","the","correctness","of","the","broadcast","gradient","on","CPU","."]
code: def test_bluefog_broadcast_grad_cpu(self):
      """Test the correctness of the broadcast gradient on CPU."""
      rank = bf.rank()
      size = bf.size()
  
      # This test does not apply if there is only one worker.
      if size == 1:
          return
  
      # As of TensorFlow v1.9, gradients are not supported on
      # integer tensors
      dtypes = [tf.float32, tf.float64]
      dims = [1, 2, 3]
      root_ranks = list(range(size))
      for dtype, dim, root_rank in itertools.product(
          dtypes, dims, root_ranks
      ):
          if _executing_eagerly():
              tensor = self.tfe.Variable(tf.ones([5] * dim) * rank)
          else:
              tensor = tf.ones([5] * dim) * rank
          if dtype == tf.bool:
              tensor = tensor % 2
          if _executing_eagerly():
              with tf.GradientTape() as tape:
                  tensor = tf.cast(tensor, dtype=dtype)
                  broadcasted_tensor = bf.broadcast(tensor, root_rank)
              with tf.device("/cpu:0"):
                  grad_out = tape.gradient(broadcasted_tensor, tensor)
          else:
              tensor = tf.cast(tensor, dtype=dtype)
              broadcasted_tensor = bf.broadcast(tensor, root_rank)
  
              grad_ys = tf.ones([5] * dim)
              with tf.device("/cpu:0"):
                  grad = tf.gradients(
                      broadcasted_tensor, tensor, grad_ys
                  )[0]
              grad_out = self.evaluate(grad)
  
          c = size if rank == root_rank else 0
          expected = np.ones([5] * dim) * c
          err = np.linalg.norm(expected - grad_out)
          self.assertLess(
              err,
              0.00000001,
              "gradient %s differs from expected %s, "
              "error: %s" % (grad_out, expected, str(err)),
          )
  
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 9 for calling the OpsTests obj's underlying member method code with example usage: obj.test_bluefog_broadcast_grad_cpu() without return types