def repeat(tensor, repeats, axis=0, name="repeat", sum_repeats_hint=None):
  """Repeats a `tf.Tensor`'s elements along an axis by custom amounts.

  Equivalent to Numpy's `np.repeat`.
  `tensor and `repeats` must have the same numbers of elements along `axis`.

  Args:
    tensor: A `tf.Tensor` to repeat.
    repeats: A 1D sequence of the number of repeats per element.
    axis: An axis to repeat along. Defaults to 0.
    name: (string, optional) A name for the operation.
    sum_repeats_hint: Integer with the total sum of repeats in case it is
      known at graph definition time.

  Returns:
    The `tf.Tensor` with repeated values.
  """
  with tf.name_scope(name):
    if sum_repeats_hint is not None:
      sum_repeats = sum_repeats_hint
    else:
      sum_repeats = tf.reduce_sum(repeats)

    # This is TPU compatible.
    # Create a tensor consistent with output size indicating where the splits
    # between the different repeats are. For example:
    #   repeats = [2, 3, 6]
    # with cumsum(exclusive=True):
    #   scatter_indices = [0, 2, 5]
    # with scatter_nd:
    #   block_split_indicators = [1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0]
    # cumsum(exclusive=False) - 1
    #   gather_indices =         [0, 0, 1, 1, 1, 2, 2, 2, 2, 2, 2]

    # Note that scatter_nd accumulates for duplicated indices, so for
    # repeats = [2, 0, 6]
    # scatter_indices = [0, 2, 2]
    # block_split_indicators = [1, 0, 2, 0, 0, 0, 0, 0]
    # gather_indices =         [0, 0, 2, 2, 2, 2, 2, 2]

    # Sometimes repeats may have zeros in the last groups. E.g.
    # for repeats = [2, 3, 0]
    # scatter_indices = [0, 2, 5]
    # However, the `gather_nd` only goes up to (sum_repeats - 1) index. (4 in
    # the example). And this would throw an error due to trying to index
    # outside the shape. Instead we let the scatter_nd have one more element
    # and we trim it from the output.
    scatter_indices = tf.cumsum(repeats, exclusive=True)
    block_split_indicators = tf.scatter_nd(
        indices=tf.expand_dims(scatter_indices, axis=1),
        updates=tf.ones_like(scatter_indices),
        shape=[sum_repeats + 1])[:-1]
    gather_indices = tf.cumsum(block_split_indicators, exclusive=False) - 1

    # An alternative implementation of the same, where block split indicators
    # does not have an indicator for the first group, and requires less ops
    # but requires creating a matrix of size [len(repeats), sum_repeats] is:
    # cumsum_repeats = tf.cumsum(repeats, exclusive=False)
    # block_split_indicators = tf.reduce_sum(
    #     tf.one_hot(cumsum_repeats, sum_repeats, dtype=tf.int32), axis=0)
    # gather_indices = tf.cumsum(block_split_indicators, exclusive=False)

    # Now simply gather the tensor along the correct axis.
    repeated_tensor = tf.gather(tensor, gather_indices, axis=axis)

    shape = tensor.shape.as_list()
    shape[axis] = sum_repeats_hint
    repeated_tensor.set_shape(shape)
    return repeated_tensor
Repeats a `tf.Tensor`'s elements along an axis by custom amounts.

  Equivalent to Numpy's `np.repeat`.
  `tensor and `repeats` must have the same numbers of elements along `axis`.

  Args:
    tensor: A `tf.Tensor` to repeat.
    repeats: A 1D sequence of the number of repeats per element.
    axis: An axis to repeat along. Defaults to 0.
    name: (string, optional) A name for the operation.
    sum_repeats_hint: Integer with the total sum of repeats in case it is
      known at graph definition time.

  Returns:
    The `tf.Tensor` with repeated values.
  
graph_nets/utils_tf.py