code: def create_tf_record_dataset(
      filenames: MutableSequence[Text],
      batch_size: int,
      shuffle_buffer_size_per_record: int = 100,
      shuffle_buffer_size: int = 100,
      load_buffer_size: int = 100000000,
      num_shards: int = 50,
      cycle_length: int = tf.data.experimental.AUTOTUNE,
      block_length: int = 10,
      num_parallel_reads: Optional[int] = None,
      num_parallel_calls: int = tf.data.experimental.AUTOTUNE,
      num_prefetch: int = 10,
      strategy: Optional[tf.distribute.Strategy] = None,
      reward_shift: float = 0.0,
      action_clipping: Optional[Tuple[float, float]] = None,
      use_trajectories: bool = True,
  ):
      """Create a TF dataset from a list of filenames.
  
      A dataset is created for each record file and these are interleaved together
      to create the final dataset.
  
      Args:
        filenames: List of filenames of a TFRecord dataset containing TF Examples.
        batch_size: The batch size of tensors in the returned dataset.
        shuffle_buffer_size_per_record: The buffer size used for shuffling within a
          Record file.
        shuffle_buffer_size: The shuffle buffer size for the interleaved dataset.
        load_buffer_size: Number of bytes in the read buffer. 0 means no buffering.
        num_shards: The number of shards, each consisting of 1 or more record file
          datasets, that are then interleaved together.
        cycle_length: The number of input elements processed concurrently while
          interleaving.
        block_length: The number of consecutive elements to produce from each input
          element before cycling to another input element.
        num_parallel_reads: Optional, number of parallel reads in the TFRecord
          dataset. If not specified, len(filenames) will be used.
        num_parallel_calls: Number of parallel calls for interleave.
        num_prefetch: Number of batches to prefetch.
        strategy: Optional, `tf.distribute.Strategy` being used in training.
        reward_shift: Value to add to the reward.
        action_clipping: Optional, (minimum, maximum) values to clip actions.
        use_trajectories: Whether to use trajectories. If false, use transitions.
  
      Returns:
        A TF.Dataset containing a batch of nested Tensors.
      """
      initial_len = len(filenames)
      remainder = initial_len % num_shards
      for _ in range(num_shards - remainder):
          filenames.append(
              filenames[np.random.randint(low=0, high=initial_len)]
          )
      filenames = np.array(filenames)
      np.random.shuffle(filenames)
      filenames = np.array_split(filenames, num_shards)
  
      record_file_ds = tf.data.Dataset.from_tensor_slices(filenames)
      record_file_ds = record_file_ds.repeat().shuffle(len(filenames))
  
      spec_path = filenames[0][0] + ".spec"
      record_spec = (
          example_encoding_dataset.parse_encoded_spec_from_file(
              spec_path
          )
      )
      decoder = example_encoding.get_example_decoder(record_spec)
  
      example_ds = record_file_ds.interleave(
          functools.partial(
              create_single_tf_record_dataset,
              load_buffer_size=load_buffer_size,
              shuffle_buffer_size=shuffle_buffer_size_per_record,
              num_parallel_reads=num_parallel_reads,
              decoder=decoder,
              reward_shift=reward_shift,
              action_clipping=action_clipping,
              use_trajectories=use_trajectories,
          ),
          cycle_length=cycle_length,
          block_length=block_length,
          num_parallel_calls=num_parallel_calls,
      )
      example_ds = example_ds.shuffle(shuffle_buffer_size)
  
      use_tpu = isinstance(
          strategy,
          (
              tf.distribute.experimental.TPUStrategy,
              tf.distribute.TPUStrategy,
          ),
      )
      example_ds = example_ds.batch(
          batch_size, drop_remainder=use_tpu
      ).prefetch(num_prefetch)
      return example_ds
  
file_path: agents/tf_agents/examples/cql_sac/kumar20/data_utils.py
imports: import functools
  from typing import Any, Callable, MutableSequence, Optional, Text, Tuple
  import numpy
  import tensorflow
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import example_encoding
  from tf_agents.utils import example_encoding_dataset
label_desc: Use image node_id 3 for calling a global function with example usage: create_tf_record_dataset(filenames, batch_size, shuffle_buffer_size_per_record, shuffle_buffer_size, load_buffer_size, num_shards, cycle_length, block_length, num_parallel_reads, num_parallel_calls, num_prefetch, strategy, reward_shift, action_clipping, use_trajectories) and returns: example_ds