code: def load_data(prefix, is_train=False, parallel_reads=4):
      """Loads the metaphase dataset.
  
      Args:
        prefix: str; Dataset path prefix to tfrecords.
        is_train: bool; is the dataset used for training?
        parallel_reads: int; Number of parallel readers (set to 1 for determinism).
  
      Returns:
        tf.data.Datasets for training, testing and validation.
      """
      num_hosts = jax.process_count()
      host_id = jax.process_index()
  
      # We shard the data between different hosts and create a Dataset that includes
      # only 1/num_shards of full dataset.
      filenames = tf.io.matching_files(prefix + "*")
      filenames_host_split = np.array_split(filenames, num_hosts)[
          host_id
      ]
      logging.info(
          "Host id=%d assigned %d out of %d dataset filenames (train=%r).",
          host_id,
          len(filenames_host_split),
          len(filenames),
          is_train,
      )
      if not list(filenames_host_split):
          raise ValueError(
              "Zero dataset filenames assigned to host %d for reading; %d available"
              " for all hosts" % (host_id, len(filenames))
          )
      files = tf.data.Dataset.list_files(filenames_host_split)
  
      data = files.interleave(
          tf.data.TFRecordDataset,
          cycle_length=1 if not is_train else parallel_reads,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      data = data.map(
          parse_example,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      return data
  
file_path: scenic/scenic/projects/tasseo/dataset_utils.py
imports: from absl import logging
  import jax
  import numpy
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: load_data(prefix, is_train, parallel_reads) and returns: data