code: def input_fn(
      path: str,
      input_shape: Tuple[int, int, int],
      max_examples: Optional[int] = None,
      mode: str = "train",
  ) -> Optional[tf.data.Dataset]:
      """Load the dataset and shuffle for training."""
      assert mode in ["train", "validation"]
      is_training = mode == "train"
      dataset = None
      for pattern in path.split(","):
          one_dataset = tf.data.TFRecordDataset.list_files(
              sharded_file_utils.normalize_to_sharded_file_pattern(
                  pattern
              ),
              shuffle=is_training,
          )
          dataset = (
              dataset.concatenate(one_dataset)
              if dataset
              else one_dataset
          )
  
      if dataset is None:
          return None
  
      if is_training:
          dataset = dataset.shuffle(
              _DEFAULT_INITIAL_SHUFFLE_BUFFER_ELEMENTS,
              reshuffle_each_iteration=True,
          )
  
      dataset = dataset.interleave(
          load_dataset,
          cycle_length=_DEFAULT_INPUT_READ_THREADS,
          num_parallel_calls=tf.data.AUTOTUNE,
          deterministic=False,
      )
  
      if max_examples is not None:
          dataset = dataset.take(max_examples)
  
      dataset = dataset.repeat()
  
      if is_training:
          dataset = dataset.shuffle(
              _DEFAULT_INITIAL_SHUFFLE_BUFFER_ELEMENTS,
              reshuffle_each_iteration=True,
          )
      # Best practices suggest batching before mapping, but this fails so I swapped
      # the order.
      dataset = dataset.map(
          map_func=lambda example: parse_example(example, input_shape),
          num_parallel_calls=tf.data.AUTOTUNE,
          deterministic=False,
      )
      batch_size = (
          FLAGS.train_config.batch_size_train
          if mode == "train"
          else FLAGS.train_config.batch_size_eval
      )
      dataset = dataset.batch(
          batch_size=batch_size, drop_remainder=True
      )
  
      # Prefetch overlaps in-feed with training
      dataset = dataset.prefetch(tf.data.AUTOTUNE)
      return dataset
  
imports: import gc
  import os
  import re
  from typing import Any, Dict, Optional, Tuple
  import warnings
  from absl import app
  from absl import flags
  from absl import logging
  from ml_collections.config_flags import config_flags
  import tensorflow
  import tensorflow_addons
  from deepvariant import data_providers
  from deepvariant import dv_constants
  from deepvariant import dv_utils
  from deepvariant import keras_modeling
  from deepvariant.average_model_checkpoint_patched import AverageModelCheckpointPatched
  from third_party.nucleus.io import sharded_file_utils
  from official.modeling import optimization
label_desc: Use image node_id 3 for calling a global function with example usage: input_fn(path, input_shape, max_examples, mode) and returns: dataset, None