code: def next_dataset_element(test_case, dataset):
      """Utility function to iterate over tf.data.Datasets in both TF 1.x and 2.x.
  
      TensorFlow 1.x and 2.x have different mechanisms for iterating over elements
      of a tf.data.Dataset. TensorFlow 1.x would require something like:
  
      itr = tf.data.Dataset.range(10).make_one_shot_iterator()
      get_next = itr.get_next()
      with tf.Session() as sess:
        for _ in range(10):
          item = sess.run(get_next)
          process(item)
  
      While TensorFlow 2.x enables something simpler like:
  
      for item in tf.data.Dataset.range(10):
        process(item)
  
      That simpler latter form is also available in TensorFlow 1.x when running
      with eager execution enabled.
  
      This function accomodates for the differing styles using:
  
      next_element = next_dataset_element(self, tf.data.Dataset.range(10))
      for _ in range 10:
        process(next_element())
  
      Args:
        test_case: The tf.test.TestCase object of the test calling this function.
        dataset: A tf.data.Dataset object.
  
      Returns:
        A Python function that returns successive elements from dataset on each call
        (using test_case.evaluate() in TensorFlow 1.x).
      """
      if tf.executing_eagerly():
          itr = iter(dataset)
          return lambda: next(itr)
      get_next = tf.compat.v1.data.make_one_shot_iterator(
          dataset
      ).get_next()
      return lambda: test_case.evaluate(get_next)
  
imports: from __future__ import division
  from __future__ import unicode_literals
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.replay_buffers import py_hashed_replay_buffer
  from tf_agents.replay_buffers import py_uniform_replay_buffer
  from tf_agents.specs import array_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 for calling a global function with example usage: next_dataset_element(test_case, dataset) and returns: function, function