jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_horovod_allreduce_gpu_fused(self):
        """Test that the allreduce works on GPUs with Tensor Fusion.

        This test will crash badly if used with an MPI implementation that does
        not support GPU memory transfers directly, as it will call MPI_Send on
        a GPU data pointer."""
        # Only do this test if there are GPUs available.
        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")

        hvd.init()
        local_rank = hvd.local_rank()
        size = hvd.size()

        dtypes = [tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64]
        dims = [1, 2, 3]
        tests = []
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % local_rank):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensor = self.random_uniform([17] * dim, minval, maxval)
                tensor = t.cast(tensor, dtype=dtype)
                summed = hvd.allreduce(tensor, op=hvd.Sum)
            multiplied = tensor * size
            difference = summed - multiplied
            difference = t.cast(difference, tf.int32) if dtype == tf.uint8 else difference
            max_difference = tf.reduce_max(tf.abs(difference))

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            test = max_difference <= threshold
            tests.append(test)
        self.assertTrue(self.evaluate(tf.reduce_all(tests)),
                        "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_gpu_fused(self):
        """Test that the allreduce works on GPUs with Tensor Fusion.

        This test will crash badly if used with an MPI implementation that does
        not support GPU memory transfers directly, as it will call MPI_Send on
        a GPU data pointer."""
        # Only do this test if there are GPUs available.
        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")

        hvd.init()
        local_rank = hvd.local_rank()
        size = hvd.size()

        dtypes = [tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64]
        dims = [1, 2, 3]
        tests = []
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % local_rank):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensor = self.random_uniform([17] * dim, minval, maxval)
                tensor = tf.cast(tensor, dtype=dtype)
                summed = hvd.allreduce(tensor, op=hvd.Sum)
            multiplied = tensor * size
            difference = summed - multiplied
            difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
            max_difference = tf.reduce_max(tf.abs(difference))

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            test = max_difference <= threshold
            tests.append(test)
        self.assertTrue(self.evaluate(tf.reduce_all(tests)),
                        "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_cpu(self):
        """Test on CPU that the grouped allreduce correctly sums 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [tf.cast(self.random_uniform(
                    [17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                summed = hvd.grouped_allreduce(tensors, op=hvd.Sum)
            multiplied = [tensor * size for tensor in tensors]
            differences = [t1 - t2 for t1, t2 in zip(summed, multiplied)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_gpu_prescale(self):
        """Test on GPU that the allreduce correctly sums 1D, 2D, 3D tensors
           with prescaling"""

        # Only do this test if there are GPUs available.
        if not tf.test.is_gpu_available(cuda_only=True):
            return

        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_ALLREDUCE.
            return

        hvd.init()
        size = hvd.size()
        local_rank = hvd.local_rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32])
        int_types = [tf.uint8, tf.int8, tf.int32, tf.int64]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%s" % local_rank):
                np.random.seed(1234)
                factor = np.random.uniform()
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensor = self.random_uniform([17] * dim, minval, maxval)
                tensor = t.cast(tensor, dtype=dtype)
                summed = hvd.allreduce(tensor, op=hvd.Sum,
                                       prescale_factor=factor)

                # Scaling done in FP64 math for integer types.
                tensor = t.cast(tensor, tf.float64 if dtype in int_types else dtype)
                factor = t.convert_to_tensor(factor, tf.float64 if dtype in int_types else dtype)
                multiplied = t.cast(factor * tensor, dtype) * size
                difference = summed - multiplied
                difference = t.cast(difference, tf.int32) if dtype == tf.uint8 else difference
                max_difference = tf.reduce_max(tf.abs(difference))
            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in int_types:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold,
                            "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_gpu_prescale(self):
        """Test on GPU that the allreduce correctly sums 1D, 2D, 3D tensors
           with prescaling"""

        # Only do this test if there are GPUs available.
        if not tf.test.is_gpu_available(cuda_only=True):
            return

        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_ALLREDUCE.
            return

        hvd.init()
        size = hvd.size()
        local_rank = hvd.local_rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32])
        int_types = [tf.uint8, tf.int8, tf.int32, tf.int64]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%s" % local_rank):
                np.random.seed(1234)
                factor = np.random.uniform()
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensor = self.random_uniform([17] * dim, minval, maxval)
                tensor = tf.cast(tensor, dtype=dtype)
                summed = hvd.allreduce(tensor, op=hvd.Sum,
                                       prescale_factor=factor)

                # Scaling done in FP64 math for integer types.
                tensor = tf.cast(tensor, tf.float64 if dtype in int_types else dtype)
                factor = tf.convert_to_tensor(factor, tf.float64 if dtype in int_types else dtype)
                multiplied = tf.cast(factor * tensor, dtype) * size
                difference = summed - multiplied
                difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
                max_difference = tf.reduce_max(tf.abs(difference))
            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in int_types:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold,
                            "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_gpu_postscale(self):
        """Test on GPU that the allreduce correctly sums 1D, 2D, 3D tensors
           with postscaling"""

        # Only do this test if there are GPUs available.
        if not tf.test.is_gpu_available(cuda_only=True):
            return

        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_ALLREDUCE.
            return

        hvd.init()
        size = hvd.size()
        local_rank = hvd.local_rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32])
        int_types = [tf.uint8, tf.int8, tf.int32, tf.int64]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%s" % local_rank):
                np.random.seed(1234)
                factor = np.random.uniform()
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensor = self.random_uniform([17] * dim, minval, maxval)
                tensor = t.cast(tensor, dtype=dtype)
                summed = hvd.allreduce(tensor, op=hvd.Sum,
                                       postscale_factor=factor)

                multiplied = tensor * size
                # Scaling done in FP64 math for integer types.
                multiplied = t.cast(multiplied, tf.float64 if dtype in int_types else dtype)
                factor = t.convert_to_tensor(factor, tf.float64 if dtype in int_types else dtype)
                multiplied = t.cast(factor * multiplied, dtype)
                difference = summed - multiplied
                difference = t.cast(difference, tf.int32) if dtype == tf.uint8 else difference
                max_difference = tf.reduce_max(tf.abs(difference))

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in int_types:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold,
                            "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_gpu_postscale(self):
        """Test on GPU that the allreduce correctly sums 1D, 2D, 3D tensors
           with postscaling"""

        # Only do this test if there are GPUs available.
        if not tf.test.is_gpu_available(cuda_only=True):
            return

        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_ALLREDUCE.
            return

        hvd.init()
        size = hvd.size()
        local_rank = hvd.local_rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32])
        int_types = [tf.uint8, tf.int8, tf.int32, tf.int64]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%s" % local_rank):
                np.random.seed(1234)
                factor = np.random.uniform()
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensor = self.random_uniform([17] * dim, minval, maxval)
                tensor = tf.cast(tensor, dtype=dtype)
                summed = hvd.allreduce(tensor, op=hvd.Sum,
                                       postscale_factor=factor)

                multiplied = tensor * size
                # Scaling done in FP64 math for integer types.
                multiplied = tf.cast(multiplied, tf.float64 if dtype in int_types else dtype)
                factor = tf.convert_to_tensor(factor, tf.float64 if dtype in int_types else dtype)
                multiplied = tf.cast(factor * multiplied, dtype)
                difference = summed - multiplied
                difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
                max_difference = tf.reduce_max(tf.abs(difference))

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in int_types:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold,
                            "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_indexed_slices_cpu(self):
        """Test on CPU that the grouped allreduce correctly sums tf.IndexedSlices."""
        hvd.init()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]

        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                slice_values = [tf.cast(self.random_uniform([17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                tensors = [tf.IndexedSlices(tf.stack([x]*2),
                                            tf.convert_to_tensor([hvd.rank(), hvd.rank()+1])) for x in slice_values]
                results = hvd.grouped_allreduce(tensors, average=False)

            # Convert indexed slice to tensor (summing entries for duplicate indices)
            processed_results = []
            for x in results:
                keys = x.indices
                processed_results.append(tf.math.unsorted_segment_sum(x.values, keys, tf.size(tf.unique(keys)[0])))

            references = []
            for x in slice_values:
                slices = [x] * (hvd.size() + 1)
                for i in range(1, hvd.size()):
                    slices[i] *= 2
                references.append(tf.stack(slices))

            differences = [t1 - t2 for t1, t2 in zip(processed_results, references)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_average_cpu(self):
        """Test on CPU that the grouped allreduce correctly averages 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [tf.cast(self.random_uniform(
                    [17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                result = hvd.grouped_allreduce(tensors, op=hvd.Average)
            differences = [t1 - t2 for t1, t2 in zip(result, tensors)]
            differences = [t.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_average_cpu(self):
        """Test on CPU that the grouped allreduce correctly averages 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [tf.cast(self.random_uniform(
                    [17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                result = hvd.grouped_allreduce(tensors, op=hvd.Average)
            differences = [t1 - t2 for t1, t2 in zip(result, tensors)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_min_cpu(self):
        """Test on CPU that the grouped allreduce correctly finds minimum of 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        rank = hvd.rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                full_tensors = [tf.cast(self.random_uniform([size] + [17] * dim, -100, 100),
                                        dtype=dtype) for _ in range(5)]
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [t[rank, ...]  for t in full_tensors]
                result = hvd.grouped_allreduce(tensors, op=hvd.Min)
            reference = [tf.math.reduce_min(t, axis=0) for t in full_tensors]
            differences = [t1 - t2 for t1, t2 in zip(result, reference)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_mixed_indexed_slices_cpu(self):
        """Test on CPU that the grouped allreduce correctly sums a mix of tensors and tf.IndexedSlices."""
        hvd.init()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]

        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                # Create indexed slices
                slice_values = [tf.cast(self.random_uniform([17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                tensors = [tf.IndexedSlices(tf.stack([x]*2),
                                            tf.convert_to_tensor([hvd.rank(), hvd.rank()+1])) for x in slice_values]

                # Append additional set of regular tensors
                for x in slice_values:
                    tensors.append(x)
                results = hvd.grouped_allreduce(tensors, average=False)

            # Convert indexed slice to tensor (summing entries for duplicate indices) or append standard tensor results
            processed_results = []
            for x in results:
                if isinstance(x, tf.IndexedSlices):
                    keys = x.indices
                    processed_results.append(tf.math.unsorted_segment_sum(x.values, keys, tf.size(tf.unique(keys)[0])))
                else:
                    processed_results.append(x)

            references = []
            # First append references for indexed slices
            for x in slice_values:
                slices = [x] * (hvd.size() + 1)
                for i in range(1, hvd.size()):
                    slices[i] *= 2
                references.append(tf.stack(slices))

            # Next, append references for standard tensors
            for x in slice_values:
                references.append(x * size)

            differences = [t1 - t2 for t1, t2 in zip(processed_results, references)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_max_cpu(self):
        """Test on CPU that the grouped allreduce correctly finds maximum of 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        rank = hvd.rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                full_tensors = [t.cast(self.random_uniform([size] + [17] * dim, -100, 100),
                                        dtype=dtype) for _ in range(5)]
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [t[rank, ...]  for t in full_tensors]
                result = hvd.grouped_allreduce(tensors, op=hvd.Max)
            reference = [tf.math.reduce_max(t, axis=0) for t in full_tensors]
            differences = [t1 - t2 for t1, t2 in zip(result, reference)]
            differences = [t.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_max_cpu(self):
        """Test on CPU that the grouped allreduce correctly finds maximum of 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        rank = hvd.rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                full_tensors = [tf.cast(self.random_uniform([size] + [17] * dim, -100, 100),
                                        dtype=dtype) for _ in range(5)]
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [t[rank, ...]  for t in full_tensors]
                result = hvd.grouped_allreduce(tensors, op=hvd.Max)
            reference = [tf.math.reduce_max(t, axis=0) for t in full_tensors]
            differences = [t1 - t2 for t1, t2 in zip(result, reference)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_product_cpu(self):
        """Test on CPU that the grouped allreduce correctly finds product of 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        rank = hvd.rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                full_tensors = [tf.cast(self.random_uniform([size] + [17] * dim, -100, 100),
                                        dtype=dtype) for _ in range(5)]
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [t[rank, ...]  for t in full_tensors]
                result = hvd.grouped_allreduce(tensors, op=hvd.Product)
            reference = [tf.math.reduce_prod(t, axis=0) for t in full_tensors]
            differences = [t1 - t2 for t1, t2 in zip(result, reference)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_indexed_slices_average_cpu(self):
        """Test on CPU that the grouped allreduce correctly averages tf.IndexedSlices."""
        hvd.init()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]

        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                slice_values = [tf.cast(self.random_uniform([17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                tensors = [tf.IndexedSlices(tf.stack([x]*2),
                                            tf.convert_to_tensor([hvd.rank(), hvd.rank()+1])) for x in slice_values]
                results = hvd.grouped_allreduce(tensors, op=hvd.Average)

            # Convert indexed slice to tensor (summing entries for duplicate indices)
            processed_results = []
            for x in results:
                keys = x.indices
                processed_results.append(tf.math.unsorted_segment_sum(x.values, keys, tf.size(tf.unique(keys)[0])))

            references = []
            for x in slice_values:
                slices = [x] * (hvd.size() + 1)
                for i in range(0, hvd.size()+1):
                    if i == 0 or i == hvd.size():
                        slices[i] /= hvd.size()
                    else:
                        slices[i] *= 2
                        slices[i] /= hvd.size()
                references.append(tf.stack(slices))

            differences = [t1 - t2 for t1, t2 in zip(processed_results, references)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_gpu(self):
        """Test on GPU that the grouped allreduce correctly sums 1D, 2D, 3D tensors."""
        # Only do this test if there are GPUs available.
        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")

        hvd.init()
        local_rank = hvd.local_rank()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % local_rank):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [tf.cast(self.random_uniform(
                    [17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                summed = hvd.grouped_allreduce(tensors, op=hvd.Sum)
            multiplied = [tensor * size for tensor in tensors]
            differences = [t1 - t2 for t1, t2 in zip(summed, multiplied)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce on GPU produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_average_gpu(self):
        """Test on GPU that the grouped allreduce correctly averages 1D, 2D, 3D tensors."""
        # Only do this test if there are GPUs available.
        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")

        hvd.init()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % hvd.local_rank()):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [tf.cast(self.random_uniform(
                    [17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                result = hvd.grouped_allreduce(tensors, op=hvd.Average)
            differences = [t1 - t2 for t1, t2 in zip(result, tensors)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_min_gpu(self):
        """Test on GPU that the grouped allreduce correctly finds minimum of 1D, 2D, 3D tensors."""
        # Only do this test if there are GPUs available.
        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")

        hvd.init()
        size = hvd.size()
        rank = hvd.rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % hvd.local_rank()):
                full_tensors = [tf.cast(self.random_uniform([size] + [17] * dim, -100, 100),
                                        dtype=dtype) for _ in range(5)]
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [t[rank, ...]  for t in full_tensors]
                result = hvd.grouped_allreduce(tensors, op=hvd.Min)
            reference = [tf.math.reduce_min(t, axis=0) for t in full_tensors]
            differences = [t1 - t2 for t1, t2 in zip(result, reference)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_indexed_slices_gpu(self):
        """Test on GPU that the grouped allreduce correctly sums tf.IndexedSlices."""
        # Only do this test if there are GPUs available.
        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")

        hvd.init()
        size = hvd.size()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]

        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % hvd.local_rank()):
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                slice_values = [tf.cast(self.random_uniform([17] * dim, minval, maxval), dtype=dtype) for _ in range(5)]
                tensors = [tf.IndexedSlices(tf.stack([x]*2),
                                            tf.convert_to_tensor([hvd.rank(), hvd.rank()+1])) for x in slice_values]
                results = hvd.grouped_allreduce(tensors, average=False)

            # Convert indexed slice to tensor (summing entries for duplicate indices)
            processed_results = []
            for x in results:
                keys = x.indices
                processed_results.append(tf.math.unsorted_segment_sum(x.values, keys, tf.size(tf.unique(keys)[0])))

            references = []
            for x in slice_values:
                slices = [x] * (hvd.size() + 1)
                for i in range(1, hvd.size()):
                    slices[i] *= 2
                references.append(tf.stack(slices))

            differences = [t1 - t2 for t1, t2 in zip(processed_results, references)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_max_gpu(self):
        """Test on GPU that the grouped allreduce correctly finds maximum of 1D, 2D, 3D tensors."""
        # Only do this test if there are GPUs available.
        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")

        hvd.init()
        size = hvd.size()
        rank = hvd.rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % hvd.local_rank()):
                full_tensors = [tf.cast(self.random_uniform([size] + [17] * dim, -100, 100),
                                        dtype=dtype) for _ in range(5)]
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [t[rank, ...]  for t in full_tensors]
                result = hvd.grouped_allreduce(tensors, op=hvd.Max)
            reference = [tf.math.reduce_max(t, axis=0) for t in full_tensors]
            differences = [t1 - t2 for t1, t2 in zip(result, reference)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_product_gpu(self):
        """Test on GPU that the grouped allreduce correctly finds product of 1D, 2D, 3D tensors."""
        # Only do this test if there are GPUs available.
        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")

        hvd.init()
        size = hvd.size()
        rank = hvd.rank()
        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % hvd.local_rank()):
                full_tensors = [tf.cast(self.random_uniform([size] + [17] * dim, -100, 100),
                                        dtype=dtype) for _ in range(5)]
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                tensors = [t[rank, ...]  for t in full_tensors]
                result = hvd.grouped_allreduce(tensors, op=hvd.Product)
            reference = [tf.math.reduce_prod(t, axis=0) for t in full_tensors]
            differences = [t1 - t2 for t1, t2 in zip(result, reference)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_allgather_cpu(self):
        """Test that the allgather correctly gathers 1D, 2D, 3D tensors."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        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("/cpu:0"):
                gathered = hvd.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
                else:
                    value = i % 2
                self.assertTrue(
                    self.evaluate(tf.reduce_all(
                        tf.equal(tf.cast(rank_tensor, tf.int32), value))),
                    "hvd.allgather produces incorrect gathered tensor")