code: def load_dataset(filename: str) -> tf.data.Dataset:
      return tf.data.TFRecordDataset(
          filename,
          buffer_size=_DEFAULT_PREFETCH_BUFFER_BYTES,
          compression_type="GZIP",
      )
  
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 2 for calling a global function with example usage: load_dataset(filename) and returns: tf