class_docstr: Tests for ops in horovod.tensorflow using multiple process sets.
class_docstr_tok: ["Tests","for","ops","in","horovod.tensorflow","using","multiple","process","sets","."]
docstr: Test that the allgather correctly gathers 1D, 2D, 3D tensors if restricted to non-global process sets.
code: def test_horovod_allgather_gpu_process_sets(self):
      """Test that the allgather correctly gathers 1D, 2D, 3D tensors if restricted to non-global process sets."""
  
      if not tf.test.is_gpu_available(cuda_only=True):
          self.skipTest("No GPUs available")
  
      if int(os.environ.get("HOROVOD_MIXED_INSTALL", 0)):
          # Skip if compiled with CUDA but without HOROVOD_GPU_OPERATIONS.
          self.skipTest("Not compiled with HOROVOD_GPU_OPERATIONS")
  
      rank = hvd.rank()
      local_rank = hvd.local_rank()
  
      if rank in self.even_ranks:
          set_size = len(self.even_ranks)
          set_ranks = self.even_ranks
          this_set = self.even_set
      elif rank in self.odd_ranks:
          set_size = len(self.odd_ranks)
          set_ranks = self.odd_ranks
          this_set = self.odd_set
  
      dtypes = [
          tf.uint8,
          tf.int8,
          tf.uint16,
          tf.int16,
          tf.int32,
          tf.int64,
          tf.float16,
          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)
          with tf.device("/gpu:%d" % local_rank):
              gathered = hvd.allgather(tensor, process_set=this_set)
  
          gathered_tensor = self.evaluate(gathered)
          self.assertEqual(
              list(gathered_tensor.shape),
              [17 * set_size] + [17] * (dim - 1),
          )
  
          for i in range(set_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 = set_ranks[i]
              else:
                  value = set_ranks[i] % 2
              self.assertTrue(
                  self.evaluate(
                      tf.reduce_all(
                          tf.equal(
                              tf.cast(rank_tensor, tf.int32), value
                          )
                      )
                  ),
                  "hvd.allgather produces incorrect gathered tensor",
              )
  
imports: from packaging import version
  import itertools
  import numpy
  import platform
  import tensorflow
  from horovod.tensorflow.util import _executing_eagerly
  import horovod.tensorflow
  from base_test_tensorflow import *
  from horovod.runner.common.util.env import get_env_rank_and_size
  from tensorflow.python.framework.test_util import run_all_in_graph_and_eager_modes
label_desc: Use image node_id 14 for calling the TensorFlowProcessSetsTests obj's underlying member method code with example usage: obj.test_horovod_allgather_gpu_process_sets() without return types