code: def replicate_embeddings(embeddings, replicate_times):
      """Replicates the given `embeddings` by `replicate_times`.
  
      This function is useful when comparing the same instance with multiple other
      instances. For example, given a seed and its neighbors, this function can be
      used to replicate the embeddings of the seed by the number of its neighbors,
      such that the distances between the seed and its neighbors can be computed
      efficiently.
  
      The `replicate_times` argument is either a scalar, or a 1-D tensor.
      For example, if
  
      ```
      embeddings = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
      ```
  
      then we would have the following results for different `replicate_times`
      arguments:
  
      ```
      replicate_times = 2
      result = [[0, 1, 2], [0, 1, 2], [3, 4, 5], [3, 4, 5], [6, 7, 8], [6, 7, 8]]
      ```
  
      and
  
      ```
      replicate_times = [3, 0, 1]
      result = [[0, 1, 2], [0, 1, 2], [0, 1, 2], [6, 7, 8]]
      ```
  
      Args:
        embeddings: A Tensor of shape `[batch_size, d1, ..., dN]`.
        replicate_times: An integer scalar or an integer 1-D Tensor of shape `[batch
          size]`. Each element indicates the number of times the corresponding row
          in `embeddings` should be replicated.
  
      Returns:
        A Tensor of shape `[N, d1, ..., dN]`, where `N` is the sum of all elements
          in `replicate_times`.
  
      Raises:
        InvalidArgumentError: If any value in `replicate_times` is negative.
        TypeError: If `replicate_times` contains any value that cannot be cast to
          the `int32` type.
      """
      with tf.control_dependencies(
          [tf.debugging.assert_greater_equal(replicate_times, 0)]
      ):
          replicate_times = tf.cast(replicate_times, tf.dtypes.int32)
          batch_size = tf.shape(input=embeddings)[0]
          idx_array = tf.range(batch_size, dtype="int32")
          if replicate_times.get_shape().ndims == 0:
              lookup_idx = tf.tile(
                  tf.expand_dims(idx_array, -1), [1, replicate_times]
              )
              lookup_idx = tf.reshape(
                  lookup_idx, [batch_size * replicate_times]
              )
          else:
              lookup_idx = _replicate_index(idx_array, replicate_times)
          output_embeddings = tf.gather(embeddings, lookup_idx)
          return output_embeddings
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import neural_structured_learning.configs
  import six
  import tensorflow
label_desc: Use image node_id 9 for calling a global function with example usage: replicate_embeddings(embeddings, replicate_times) and returns: output_embeddings