code: def sample_once(self):
      for i in self.clean_index:
          compression = "ZLIB" if "zstd" in i else ""
  
          file = tf.data.TFRecordDataset(
              i, compression_type=compression
          ).map(self.parse_fn, num_parallel_calls=tf.data.AUTOTUNE)
          file = file.apply(
              tf.data.experimental.dense_to_ragged_batch(
                  np.prod(self.bs), drop_remainder=True
              )
          )
          file = file.prefetch(10)
  
          for file_idx, data in enumerate(file):
              data = jax.tree_map(lambda x: x.numpy(), data)
              data = self.map_fn(data)
  
              if not self.file_idx_init and file_idx <= self.file_idx:
                  if file_idx % 1000 == 0:
                      print(
                          f"skipping to batch {self.file_idx}, currently at {file_idx}"
                      )
                  continue
              self.file_idx_init = True
              self.file_idx = file_idx
              yield jax.tree_map(
                  lambda x: x.reshape(self.bs + x.shape[1:]), data
              )
          self.used.append(i)
          self.file_idx = 0
  
file_path: mesh-transformer-jax/tfrecord_loader.py
imports: import jax
  import tensorflow
  import numpy
  from transformers import GPT2TokenizerFast
  import itertools
label_desc: Use image node_id 3 for calling the TFRecordLoader obj's underlying member method code with example usage: obj.sample_once() without return types