code: def testSampleSingleCorrectProbabilityAsDataset(
      self, buffer_batch_size
  ):
      max_length = 3
      spec = specs.TensorSpec([], tf.int32, "action")
      replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
          spec, batch_size=buffer_batch_size, max_length=max_length
      )
  
      actions = tf.stack(
          [tf.Variable(0).count_up_to(9)] * buffer_batch_size
      )
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
  
      ds = replay_buffer.as_dataset()
      # Turn off `inject_prefetch` optimization
      options = tf.data.Options()
      options.experimental_optimization.inject_prefetch = False
      ds = ds.with_options(options)
      if tf.executing_eagerly():
          add_op = lambda: replay_buffer.add_batch(actions)
          itr = iter(ds)
          sample = lambda: next(itr)
      else:
          add_op = replay_buffer.add_batch(actions)
          itr = tf.compat.v1.data.make_initializable_iterator(ds)
          self.evaluate(itr.initializer)
          sample = itr.get_next()
  
      num_adds = 5
      for i in range(1, num_adds):
          self.evaluate(add_op)
          probabilities_ = self.evaluate(sample)[1].probabilities
          expected_probability = 1.0 / min(
              i * buffer_batch_size, max_length * buffer_batch_size
          )
          self.assertAllClose(expected_probability, probabilities_)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents import specs
  from tf_agents.replay_buffers import tf_uniform_replay_buffer
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
label_desc: Use image node_id 18 for calling the TFUniformReplayBufferTest obj's underlying member method code with example usage: obj.testSampleSingleCorrectProbabilityAsDataset(buffer_batch_size) without return types