code: def input_fn(
      path: str,
      config: ml_collections.ConfigDict,
      mode: str = "train",
      strategy: tf.distribute.Strategy = tf.distribute.get_strategy(),
      n_epochs: int = -1,
      limit: Optional[int] = None,
  ) -> tf.data.Dataset:
      """tf.data.Dataset loading function.
  
      Args:
        path: the input filename for a tfrecord[.gz] file containing examples. Can
          contain sharding designators.
        config: A configuration file.
        mode: One of ['train', 'tune', 'predict']
        strategy: A tf.distribute.Strategy.
        n_epochs: Number of epochs.
        limit: Limit the number of batches for testing purposes.
  
      Returns:
        tf.data.Dataset
      """
  
      if mode not in ["train", "tune", "predict"]:
          raise ValueError(
              'Mode must be set to one of "train", "tune", or "predict"'
          )
      is_training = mode in ["train", "tune"]
  
      # Get input shape from input path.
      input_shape = dv_utils.get_shape_from_examples_path(path)
  
      def load_dataset(filename: str) -> tf.data.Dataset:
          return tf.data.TFRecordDataset(
              filename,
              buffer_size=config.prefetch_buffer_bytes,
              compression_type="GZIP",
          )
  
      file_list = [
          tf.io.gfile.glob(
              sharded_file_utils.normalize_to_sharded_file_pattern(x)
          )
          for x in path.split(",")
      ]
      file_list = list(itertools.chain(*file_list))
      ds = tf.data.Dataset.from_tensor_slices(file_list)
  
      if is_training:
          ds = ds.shuffle(
              ds.cardinality(), reshuffle_each_iteration=True
          )
  
      ds = ds.interleave(
          load_dataset,
          cycle_length=config.input_read_threads,
          num_parallel_calls=tf.data.AUTOTUNE,
          deterministic=False,
      )
  
      if is_training and n_epochs > 0:
          ds = ds.repeat(n_epochs)
  
      if is_training:
          ds = ds.shuffle(
              config.shuffle_buffer_elements,
              reshuffle_each_iteration=True,
          )
  
      # Retrieve preprocess function
      parse_example = create_parse_example_fn(config)
  
      ds = ds.map(
          map_func=lambda example: parse_example(example, input_shape),
          num_parallel_calls=tf.data.AUTOTUNE,
          deterministic=False,
      )
      ds = ds.batch(batch_size=config.batch_size, drop_remainder=True)
  
      # Limit the number of batches.
      if limit:
          ds = ds.take(limit)
  
      if n_epochs > 0:
          ds = ds.repeat(n_epochs)
      elif mode == "tune":
          ds = ds.repeat()
  
      # Prefetch overlaps in-feed with training
      ds = ds.prefetch(tf.data.AUTOTUNE)
  
      # Distribute the dataset
      ds = strategy.experimental_distribute_dataset(ds)
  
      return ds
  
imports: import itertools
  from typing import Callable, Optional, Tuple
  from absl import logging
  import ml_collections
  import tensorflow
  from tensorflow import estimator
  from deepvariant import dv_constants
  from deepvariant import dv_utils
  from deepvariant import keras_modeling
  from deepvariant.protos import deepvariant_pb2
  from google.protobuf import text_format
  from third_party.nucleus.io import sharded_file_utils
label_desc: Use image node_id 2 for calling a global function with example usage: input_fn(path, config, mode, strategy, n_epochs, limit) and returns: ds, tf