jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def finalize(beam_width, parent_ids, sequence_lengths, outputs, end_id, max_seq_len=None):
    maximum_lengths = tf.reduce_max(tf.reshape(
        sequence_lengths, [-1, beam_width]), axis=-1)
    
    if max_seq_len != None:
        array_shape = [max_seq_len, -1, beam_width]
    else:
        array_shape = [tf.reduce_max(maximum_lengths), -1, beam_width]
        
    step_ids = tf.reshape(outputs, array_shape)
    parent_ids = tf.reshape(parent_ids, array_shape)

    ids = tf.contrib.seq2seq.gather_tree(
        step_ids, parent_ids, maximum_lengths, end_id)

    ids = tf.transpose(ids, perm=[1, 2, 0])
    lengths = tf.not_equal(ids, end_id)
    lengths = tf.cast(lengths, tf.int32)
    lengths = tf.reduce_sum(lengths, axis=-1)
    return ids, lengths
//  [python]
def initialize_decoding_variables(decoding_args, batchxbeam):

    start_ids = tf.fill([batchxbeam], decoding_args.start_id)  # [batch_size * beam_width]

    step = tf.constant(0, dtype=tf.int32)
    # save the output ids for each step
    outputs = tf.TensorArray(tf.int32, size=0, dynamic_size=True)

    cache = init_tf_cache(batchxbeam,
                          decoding_args.decoder_args.head_num, decoding_args.decoder_args.size_per_head,
                          decoding_args.decoder_args.num_layer, dtype=decoding_args.decoder_args.dtype, num_sources=1)

    finished = tf.zeros([batchxbeam], dtype=tf.bool)  # [batch_size * beam_width], record that a sentence is finished or not
    initial_log_probs = tf.to_float(tf.tile([0.] + [-float("inf")] * (decoding_args.decoder_args.beam_width - 1),
                                        [batchxbeam / decoding_args.decoder_args.beam_width]))  # [batch_size * beam_width]
    # [batch_size * beam_width], record the lengths of all sentences
    sequence_lengths = tf.zeros([batchxbeam], dtype=tf.int32)
    # record the beam search indices, used for rebuild the whole sentence in the final
    parent_ids = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
    extra_vars = tuple([parent_ids, sequence_lengths])

    return start_ids, step, outputs, cache, finished, initial_log_probs, sequence_lengths, extra_vars
//  [python]
def initialize_decoding_variables(decoding_args, batchxbeam):

    start_ids = tf.fill([batchxbeam], decoding_args.start_id)  # [batch_size * beam_width]

    step = tf.constant(0, dtype=tf.int32)
    # save the output ids for each step
    outputs = tf.TensorArray(tf.int32, size=0, dynamic_size=True)

    cache = init_tf_cache(batchxbeam,
                          decoding_args.decoder_args.head_num, decoding_args.decoder_args.size_per_head,
                          decoding_args.decoder_args.num_layer, dtype=decoding_args.decoder_args.dtype, num_sources=1)

    finished = tf.zeros([batchxbeam], dtype=tf.bool)  # [batch_size * beam_width], record that a sentence is finished or not
    initial_log_probs = tf.cast(tf.tile([0.] + [-float("inf")] * (decoding_args.decoder_args.beam_width - 1),
                                        [batchxbeam / decoding_args.decoder_args.beam_width]), dtype=tf.float32)  # [batch_size * beam_width]
    # [batch_size * beam_width], record the lengths of all sentences
    sequence_lengths = tf.zeros([batchxbeam], dtype=tf.int32)
    # record the beam search indices, used for rebuild the whole sentence in the final
    parent_ids = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
    extra_vars = tuple([parent_ids, sequence_lengths])

    return start_ids, step, outputs, cache, finished, initial_log_probs, sequence_lengths, extra_vars
//  [python]
def create_input_pipeline(input_queue, image_size, nrof_preprocess_threads, batch_size_placeholder):
    images_and_labels_list = []
    for _ in range(nrof_preprocess_threads):
        filenames, label, control = input_queue.dequeue()
        images = []
        for filename in tf.unstack(filenames):
            file_contents = tf.read_file(filename)
            image = tf.image.decode_image(file_contents, 3)
            image = tf.cond(get_control_flag(control[0], RANDOM_ROTATE),
                            lambda:tf.py_func(random_rotate_image, [image], tf.uint8), 
                            lambda:tf.identity(image))
            image = tf.cond(get_control_flag(control[0], RANDOM_CROP), 
                            lambda:tf.random_crop(image, image_size + (3,)), 
                            lambda:tf.image.resize_image_with_crop_or_pad(image, image_size[0], image_size[1]))
            image = tf.cond(get_control_flag(control[0], RANDOM_FLIP),
                            lambda:tf.image.random_flip_left_right(image),
                            lambda:tf.identity(image))
            image = tf.cond(get_control_flag(control[0], FIXED_STANDARDIZATION),
                            lambda:(tf.to_float(image) - 127.5)/128.0,
                            lambda:tf.image.per_image_standardization(image))
            image = tf.cond(get_control_flag(control[0], FLIP),
                            lambda:tf.image.flip_left_right(image),
                            lambda:tf.identity(image))
            #pylint: disable=no-member
            image.set_shape(image_size + (3,))
            images.append(image)
        images_and_labels_list.append([images, label])

    image_batch, label_batch = tf.train.batch_join(
        images_and_labels_list, batch_size=batch_size_placeholder, 
        shapes=[image_size + (3,), ()], enqueue_many=True,
        capacity=4 * nrof_preprocess_threads * 100,
        allow_smaller_final_batch=True)
    
    return image_batch, label_batch
//  [python]
def create_input_pipeline(input_queue, image_size, nrof_preprocess_threads, batch_size_placeholder):
    images_and_labels_list = []
    for _ in range(nrof_preprocess_threads):
        filenames, label, control = input_queue.dequeue()
        images = []
        for filename in tf.unstack(filenames):
            file_contents = tf.read_file(filename)
            image = tf.image.decode_image(file_contents, 3)
            image = tf.cond(get_control_flag(control[0], RANDOM_ROTATE),
                            lambda:tf.py_func(random_rotate_image, [image], tf.uint8), 
                            lambda:tf.identity(image))
            image = tf.cond(get_control_flag(control[0], RANDOM_CROP), 
                            lambda:tf.random_crop(image, image_size + (3,)), 
                            lambda:tf.image.resize_image_with_crop_or_pad(image, image_size[0], image_size[1]))
            image = tf.cond(get_control_flag(control[0], RANDOM_FLIP),
                            lambda:tf.image.random_flip_left_right(image),
                            lambda:tf.identity(image))
            image = tf.cond(get_control_flag(control[0], FIXED_STANDARDIZATION),
                            lambda:(tf.cast(image, tf.float32) - 127.5)/128.0,
                            lambda:tf.image.per_image_standardization(image))
            image = tf.cond(get_control_flag(control[0], FLIP),
                            lambda:tf.image.flip_left_right(image),
                            lambda:tf.identity(image))
            #pylint: disable=no-member
            image.set_shape(image_size + (3,))
            images.append(image)
        images_and_labels_list.append([images, label])

    image_batch, label_batch = tf.train.batch_join(
        images_and_labels_list, batch_size=batch_size_placeholder, 
        shapes=[image_size + (3,), ()], enqueue_many=True,
        capacity=4 * nrof_preprocess_threads * 100,
        allow_smaller_final_batch=True)
    
    return image_batch, label_batch
//  [python]
def main(*args):
    cache_data_dir = '/cache/data_url'
    cache_train_url = '/cache/train_url'
    mox.file.copy_parallel(FLAGS.data_url, cache_data_dir)

    # Train model
    print('Training model...')
    mnist = input_data.read_data_sets(cache_data_dir, one_hot=True)
    serialized_tf_example = tf.placeholder(tf.string, name='tf_example')
    feature_configs = {'x': tf.FixedLenFeature(shape=[784], dtype=tf.float32),}
    tf_example = tf.parse_example(serialized_tf_example, feature_configs)
    x = tf.identity(tf_example['x'], name='x')
    y_ = tf.placeholder('float', shape=[None, 10])

    y = tf.keras.layers.Dense(128, activation='relu')(x)
    y = tf.keras.layers.Dense(10)(y)
    cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y, labels=y_))
    tf.summary.scalar('cross_entropy', cross_entropy)

    train_step = tf.train.AdamOptimizer(0.01).minimize(cross_entropy)

    correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
    accuracy = tf.reduce_mean(tf.to_float(correct_prediction))
    tf.summary.scalar('accuracy', accuracy)
    merged = tf.summary.merge_all()
    test_writer = tf.summary.FileWriter(cache_train_url, flush_secs=1)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for step in range(FLAGS.max_steps):
            batch = mnist.train.next_batch(FLAGS.batch_size)
            train_step.run(feed_dict={x: batch[0], y_: batch[1]})
            if step % 10 == 0:
                summary, acc = sess.run([merged, accuracy], 
                                        feed_dict={x: mnist.test.images, y_: mnist.test.labels})
                test_writer.add_summary(summary, step)
                print('training accuracy is:', acc)
        print('Done training!')

        builder = tf.saved_model.builder.SavedModelBuilder(os.path.join(cache_train_url, 'model'))
        tensor_info_x = tf.saved_model.utils.build_tensor_info(x)
        tensor_info_y = tf.saved_model.utils.build_tensor_info(y)

        prediction_signature = (
            tf.saved_model.signature_def_utils.build_signature_def(
                inputs={'images': tensor_info_x},
                outputs={'scores': tensor_info_y},
                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))
        builder.add_meta_graph_and_variables(
            sess, [tf.saved_model.tag_constants.SERVING],
            signature_def_map={
                'predict_images':
                    prediction_signature,
            },
            main_op=tf.tables_initializer(),
            strip_default_attrs=True)
        builder.save()
        print('Done exporting!')

    mox.file.copy_parallel(cache_train_url, FLAGS.train_url)
//  [python]
def main(*args):
    cache_data_dir = '/cache/data_url'
    cache_train_url = '/cache/train_url'
    mox.file.copy_parallel(FLAGS.data_url, cache_data_dir)

    # Train model
    print('Training model...')
    mnist = input_data.read_data_sets(cache_data_dir, one_hot=True)
    serialized_tf_example = tf.placeholder(tf.string, name='tf_example')
    feature_configs = {'x': tf.FixedLenFeature(shape=[784], dtype=tf.float32),}
    tf_example = tf.parse_example(serialized_tf_example, feature_configs)
    x = tf.identity(tf_example['x'], name='x')
    y_ = tf.placeholder('float', shape=[None, 10])

    y = tf.keras.layers.Dense(128, activation='relu')(x)
    y = tf.keras.layers.Dense(10)(y)
    cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y, labels=y_))
    tf.summary.scalar('cross_entropy', cross_entropy)

    train_step = tf.train.AdamOptimizer(0.01).minimize(cross_entropy)

    correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, 'float'))
    tf.summary.scalar('accuracy', accuracy)
    merged = tf.summary.merge_all()
    test_writer = tf.summary.FileWriter(cache_train_url, flush_secs=1)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for step in range(FLAGS.max_steps):
            batch = mnist.train.next_batch(FLAGS.batch_size)
            train_step.run(feed_dict={x: batch[0], y_: batch[1]})
            if step % 10 == 0:
                summary, acc = sess.run([merged, accuracy], 
                                        feed_dict={x: mnist.test.images, y_: mnist.test.labels})
                test_writer.add_summary(summary, step)
                print('training accuracy is:', acc)
        print('Done training!')

        builder = tf.saved_model.builder.SavedModelBuilder(os.path.join(cache_train_url, 'model'))
        tensor_info_x = tf.saved_model.utils.build_tensor_info(x)
        tensor_info_y = tf.saved_model.utils.build_tensor_info(y)

        prediction_signature = (
            tf.saved_model.signature_def_utils.build_signature_def(
                inputs={'images': tensor_info_x},
                outputs={'scores': tensor_info_y},
                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))
        builder.add_meta_graph_and_variables(
            sess, [tf.saved_model.tag_constants.SERVING],
            signature_def_map={
                'predict_images':
                    prediction_signature,
            },
            main_op=tf.tables_initializer(),
            strip_default_attrs=True)
        builder.save()
        print('Done exporting!')

    mox.file.copy_parallel(cache_train_url, FLAGS.train_url)
//  [python]
def test_horovod_allreduce_cpu(self):
        """Test on CPU that the 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"):
                # prevent underflows/overflows in uint8, int8
                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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_average_cpu(self):
        """Test on CPU that the 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
                tensor = self.random_uniform([17] * dim, minval, maxval)
                tensor = tf.cast(tensor, dtype=dtype)
                averaged = hvd.allreduce(tensor, op=hvd.Average)
            # handle int8, uint8 overflows when allreduce sums up and averages the values
            tensor = tf.cast((tensor*size)/size, dtype=dtype)
            difference = tf.cast(averaged, dtype=dtype) - tensor
            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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_min_cpu(self):
        """Test on CPU that the allreduce correctly minimizes 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"):
                tensors = self.random_uniform([size] + [17] * dim, -100, 100)
                tensors = tf.cast(tensors, dtype=dtype)
                tensor = tensors[rank,...]
                result = hvd.allreduce(tensor, op=hvd.Min)
            reference = tf.math.reduce_min(tensors, axis=0)
            difference = tf.cast(result, dtype=dtype) - reference
            difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
            max_difference = tf.reduce_max(difference)

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results for min")
//  [python]
def test_horovod_allreduce_max_cpu(self):
        """Test on CPU that the allreduce correctly maximizes 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"):
                tensors = self.random_uniform([size] + [17] * dim, -100, 100)
                tensors = tf.cast(tensors, dtype=dtype)
                tensor = tensors[rank,...]
                result = hvd.allreduce(tensor, op=hvd.Max)
            reference = tf.math.reduce_max(tensors, axis=0)
            difference = tf.cast(result, dtype=dtype) - reference
            difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
            max_difference = tf.reduce_max(difference)

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results for max")
//  [python]
def test_horovod_allreduce_product_cpu(self):
        """Test on CPU that the allreduce correctly multiplies 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"):
                tensors = self.random_uniform([size] + [17] * dim, -100, 100)
                tensors = tf.cast(tensors, dtype=dtype)
                tensor = tensors[rank,...]
                result = hvd.allreduce(tensor, op=hvd.Product)
            reference = tf.math.reduce_prod(tensors, axis=0)
            difference = tf.cast(result, dtype=dtype) - reference
            difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
            max_difference = tf.reduce_max(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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results for product")
//  [python]
def test_horovod_allreduce_cpu_fused(self):
        """Test on CPU that the allreduce correctly sums 1D, 2D, 3D tensors
        with Tensor Fusion."""
        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]
        tests = []
        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
                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_allreduce_cpu_postscale(self):
        """Test on CPU that the allreduce correctly sums 1D, 2D, 3D tensors
           with postscaling"""
        hvd.init()
        size = hvd.size()
        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("/cpu:0"):
                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, FP32 math for FP16 on CPU
                multiplied = tf.cast(multiplied, tf.float32 if dtype == tf.float16 else
                                     tf.float64 if dtype in int_types else dtype)
                factor = tf.convert_to_tensor(factor, tf.float32 if dtype == tf.float16 else
                                              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_allreduce_cpu_prescale(self):
        """Test on CPU that the allreduce correctly sums 1D, 2D, 3D tensors
           with prescaling"""
        hvd.init()
        size = hvd.size()
        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("/cpu:0"):
                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, FP32 math for FP16 on CPU
                tensor = tf.cast(tensor, tf.float32 if dtype == tf.float16 else
                                 tf.float64 if dtype in int_types else dtype)
                factor = tf.cast(factor, tf.float32 if dtype == tf.float16 else
                                              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_cpu_prescale(self):
        """Test on CPU that the allreduce correctly sums 1D, 2D, 3D tensors
           with prescaling"""
        hvd.init()
        size = hvd.size()
        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("/cpu:0"):
                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, FP32 math for FP16 on CPU
                tensor = tf.cast(tensor, tf.float32 if dtype == tf.float16 else
                                 tf.float64 if dtype in int_types else dtype)
                factor = tf.convert_to_tensor(factor, tf.float32 if dtype == tf.float16 else
                                              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_indexed_slices_cpu(self):
        """Test on CPU that the 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"):
                # prevent underflows/overflows in uint8, int8
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                slice_values = self.random_uniform([17] * dim, minval, maxval)
                slice_values = tf.cast(slice_values, dtype=dtype)
                tensor = tf.IndexedSlices(tf.stack([slice_values] * 2),
                                          tf.convert_to_tensor([hvd.rank(), hvd.rank()+1]))
                result = hvd.allreduce(tensor, average=False)

            # Convert indexed slice to tensor (summing entries for duplicate indices)
            keys = result.indices
            result = tf.math.unsorted_segment_sum(result.values, keys, tf.size(tf.unique(keys)[0]))

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

            difference = result - reference
            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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_indexed_slices_average_cpu(self):
        """Test on CPU that the 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"):
                # prevent underflows/overflows in uint8, int8
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                slice_values = self.random_uniform([17] * dim, minval, maxval)
                slice_values = tf.cast(slice_values, dtype=dtype)
                tensor = tf.IndexedSlices(tf.stack([slice_values] * 2),
                                          tf.convert_to_tensor([hvd.rank(), hvd.rank()+1]))
                result = hvd.allreduce(tensor, op=hvd.Average)

            # Convert indexed slice to tensor (summing entries for duplicate indices)
            keys = result.indices
            result = tf.math.unsorted_segment_sum(result.values, keys, tf.size(tf.unique(keys)[0]))

            slices = [slice_values] * (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()
            reference = tf.stack(slices)

            difference = result - reference
            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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_indexed_slices_gpu(self):
        """Test on GPU that the 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()
        local_rank = hvd.local_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" % local_rank):
                # prevent underflows/overflows in uint8, int8
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                slice_values = self.random_uniform([17] * dim, minval, maxval)
                slice_values = tf.cast(slice_values, dtype=dtype)
                tensor = tf.IndexedSlices(tf.stack([slice_values] * 2),
                                          tf.convert_to_tensor([hvd.rank(), hvd.rank()+1]))
                result = hvd.allreduce(tensor, average=False)

            # Convert indexed slice to tensor (summing entries for duplicate indices)
            keys = result.indices
            result = tf.math.unsorted_segment_sum(result.values, keys, tf.size(tf.unique(keys)[0]))

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

            difference = result - reference
            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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_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):
            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]
        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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce on GPU produces incorrect results")
//  [python]
def test_horovod_allreduce_indexed_slices_average_gpu(self):
        """Test on GPU that the allreduce correctly averages 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()
        local_rank = hvd.local_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" % local_rank):
                # prevent underflows/overflows in uint8, int8
                maxval = 100 if dtype not in [tf.uint8, tf.int8] else 1
                minval = -maxval if dtype not in [tf.uint8] else 0
                slice_values = self.random_uniform([17] * dim, minval, maxval)
                slice_values = tf.cast(slice_values, dtype=dtype)
                tensor = tf.IndexedSlices(tf.stack([slice_values] * 2),
                                          tf.convert_to_tensor([hvd.rank(), hvd.rank()+1]))
                result = hvd.allreduce(tensor, op=hvd.Average)

            # Convert indexed slice to tensor (summing entries for duplicate indices)
            keys = result.indices
            result = tf.math.unsorted_segment_sum(result.values, keys, tf.size(tf.unique(keys)[0]))

            slices = [slice_values] * (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()
            reference = tf.stack(slices)

            difference = result - reference
            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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_min_gpu(self):
        """Test on GPU that the allreduce correctly minimizes 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()
        rank = hvd.rank()
        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):
                tensors = self.random_uniform([size] + [17] * dim, -100, 100)
                tensors = tf.cast(tensors, dtype=dtype)
                tensor = tensors[rank,...]
                result = hvd.allreduce(tensor, op=hvd.Min)
            reference = tf.math.reduce_min(tensors, axis=0)
            difference = tf.cast(result, dtype=dtype) - reference
            difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
            max_difference = tf.reduce_max(difference)

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results for min")
//  [python]
def test_horovod_allreduce_max_gpu(self):
        """Test on GPU that the allreduce correctly maximizes 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()
        rank = hvd.rank()
        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):
                tensors = self.random_uniform([size] + [17] * dim, -100, 100)
                tensors = tf.cast(tensors, dtype=dtype)
                tensor = tensors[rank,...]
                result = hvd.allreduce(tensor, op=hvd.Max)
            reference = tf.math.reduce_max(tensors, axis=0)
            difference = tf.cast(result, dtype=dtype) - reference
            difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
            max_difference = tf.reduce_max(difference)

            threshold = 0
            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results for max")
//  [python]
def test_horovod_allreduce_average_gpu(self):
        """Test that the allreduce with average works on GPUs."""
        # 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]
        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)
                averaged = hvd.allreduce(tensor, op=hvd.Average)
            # handle int8, uint8 overflows when allreduce sums up and averages the values
            tensor = tf.cast((tensor*size)/size, dtype=dtype)
            difference = tf.cast(averaged, dtype=dtype) - tensor
            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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce on GPU produces incorrect results")
//  [python]
def test_horovod_allreduce_product_gpu(self):
        """Test on GPU that the allreduce correctly multiplies 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()
        rank = hvd.rank()
        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):
                tensors = self.random_uniform([size] + [17] * dim, -100, 100)
                tensors = tf.cast(tensors, dtype=dtype)
                tensor = tensors[rank,...]
                result = hvd.allreduce(tensor, op=hvd.Product)
            reference = tf.math.reduce_prod(tensors, axis=0)
            difference = tf.cast(result, dtype=dtype) - reference
            difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
            max_difference = tf.reduce_max(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")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results for product")