code: def make_reverb_dataset(
      server_address: str,
      table: str,
      data_spec,
      max_in_flight_samples_per_worker=10,
      batch_size=None,
      prefetch_size=None,
      sequence_length=None,
      cycle_length=None,
      num_parallel_calls=None,
      per_sequence_fn=None,
      dataset_transformation=None,
      num_workers_per_iterator=-1,
      max_samples_per_stream=-1,
      rate_limiter_timeout_ms=-1,
  ) -> tf.data.Dataset:
      """Makes a TensorFlow dataset.
  
      Args:
        server_address: The server address of the replay server.
        table: The name of the table to sample from replay.
        data_spec: The data's spec.
        max_in_flight_samples_per_worker: Optional, dataset buffer capacity.
        batch_size: Optional. If specified the dataset returned will combine
          consecutive elements into batches. This argument is also used to determine
          the cycle_length for `tf.data.Dataset.interleave` -- if unspecified the
          cycle length is set to `tf.data.experimental.AUTOTUNE`.
        prefetch_size: How many batches to prefectch in the pipeline.
        sequence_length: Optional. If specified consecutive elements of each
          interleaved dataset will be combined into sequences.
        cycle_length: Optional. When equal to batch_size it would make take a sample
          from a different sequence. For reducing memory usage use a smaller number.
        num_parallel_calls: Optional. If specified number of parallel calls in
          iterleave. By default use `tf.data.experimental.AUTOTUNE`.
        per_sequence_fn: Optional, per sequence function.
        dataset_transformation: Optional, per dataset interleave transformation.
        num_workers_per_iterator: (Defaults to -1, i.e auto selected) The number of
          worker threads to create per dataset iterator. When the selected table
          uses a FIFO sampler (i.e a queue) then exactly 1 worker must be used to
          avoid races causing invalid ordering of items. For all other samplers,
          this value should be roughly equal to the number of threads available on
          the CPU.
        max_samples_per_stream: (Defaults to -1, i.e auto selected) The maximum
          number of samples to fetch from a stream before a new call is made.
          Keeping this number low ensures that the data is fetched uniformly from
          all server.
        rate_limiter_timeout_ms: Timeout (in milliseconds) to wait on the rate
          limiter when sampling from the table. If `rate_limiter_timeout_ms >= 0`,
          this is the timeout passed to `Table::Sample` describing how long to wait
          for the rate limiter to allow sampling.
  
      Returns:
        A tf.data.Dataset that streams data from the replay server.
      """
      # Extract the shapes and dtypes from these specs.
      get_dtype = lambda x: tf.as_dtype(x.dtype)
      get_shape = lambda x: (sequence_length,) + x.shape
      shapes = tf.nest.map_structure(get_shape, data_spec)
      dtypes = tf.nest.map_structure(get_dtype, data_spec)
      # TODO(b/144858901): Validate Tableinfo when it's available from reverb.
  
      def generate_reverb_dataset(_):
          dataset = reverb.TrajectoryDataset(
              server_address,
              table,
              dtypes,
              shapes,
              max_in_flight_samples_per_worker=max_in_flight_samples_per_worker,
              num_workers_per_iterator=num_workers_per_iterator,
              max_samples_per_stream=max_samples_per_stream,
              rate_limiter_timeout_ms=rate_limiter_timeout_ms,
          )
  
          def broadcast_info(
              info_traj: types.ReverbReplaySample,
          ) -> types.ReverbReplaySample:
              # Assumes that the first element of traj is shaped
              # (sequence_length, ...); and we extract this length.
              info, traj = info_traj
              first_elem = tf.nest.flatten(traj)[0]
              length = first_elem.shape[0] or tf.shape(first_elem)[0]
              info = tf.nest.map_structure(
                  lambda t: tf.repeat(t, [length]), info
              )
              return reverb.ReplaySample(info, traj)
  
          dataset = dataset.map(broadcast_info)
  
          if per_sequence_fn:
              dataset = dataset.map(per_sequence_fn)
          if dataset_transformation:
              dataset = dataset_transformation(dataset)
          return dataset
  
      cycle_length = cycle_length or batch_size or 1
      num_parallel_calls = (
          num_parallel_calls or tf.data.experimental.AUTOTUNE
      )
  
      if cycle_length == 1:
          dataset = generate_reverb_dataset(0)
      else:
          dataset = tf.data.Dataset.range(cycle_length).interleave(
              generate_reverb_dataset,
              cycle_length=cycle_length,
              num_parallel_calls=num_parallel_calls,
          )
  
      # Allows interleave to retrieve data from the first `reverb.TrajectoryDataset`
      # available.
      options = tf.data.Options()
      # reverb replay buffers are not considered deterministic for tf.data.
      options.experimental_deterministic = False
      dataset = dataset.with_options(options)
  
      if batch_size:
          dataset = dataset.batch(batch_size, drop_remainder=True)
      if prefetch_size:
          dataset = dataset.prefetch(prefetch_size)
      return dataset
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import tensorflow.compat.v2
  from tf_agents.replay_buffers import replay_buffer
  from tf_agents.typing import types
  from tf_agents.utils import lazy_loader
label_desc: Use image node_id 1 for calling a global function with example usage: make_reverb_dataset(server_address, table, data_spec, max_in_flight_samples_per_worker, batch_size, prefetch_size, sequence_length, cycle_length, num_parallel_calls, per_sequence_fn, dataset_transformation, num_workers_per_iterator, max_samples_per_stream, rate_limiter_timeout_ms) and returns: dataset, dataset, reverb