code: def repeat(x, repeats, total_repeat_length=None):
      """Like jnp.repeat."""
      res = tf.repeat(x, repeats)
      if total_repeat_length is not None:
          res.set_shape(
              [total_repeat_length] + [None] * (len(res.shape) - 1)
          )
      return res
  
imports: import functools
  import numpy
  import six
  import tensorflow.compat.v2
  from tensorflow.python.util import nest
label_desc: Use image node_id 23 for calling a global function with example usage: repeat(x, repeats, total_repeat_length) and returns: res