code: def consume(running_stat, elems, chunk_axis=None):
      def body(running_stat, elem):
          if chunk_axis is None:
              return running_stat.update(elem)
          else:
              return running_stat.update(elem, axis=chunk_axis)
  
      return tf.foldl(body, elems, running_stat)
  
imports: from absl.testing import parameterized
  import numpy
  import scipy.stats
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import diagnostic
label_desc: Use image node_id 1 for calling a global function with example usage: consume(running_stat, elems, chunk_axis) and returns: tf, running_stat, running_stat