code: def reduce(r, elems):
      """Reduces `elems` along the first dimension with `reducer`."""
      elems = tf.convert_to_tensor(elems)
      state = r.initialize(elems[0])
  
      def body(i, state):
          return i + 1, r.one_step(elems[i], state)
  
      _, state = tf.while_loop(
          lambda i, _: i < elems.shape[0], body, (0, state)
      )
      return r.finalize(state)
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import reducer
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.random import random_ops
label_desc: Use image node_id 1 for calling a global function with example usage: reduce(r, elems) and returns: r, unknown,  r