code: def _load(path, element_spec, compression, reader_func):
      """Loads dataset from tf.data snapshot."""
  
      if reader_func is None:
          reader_func = lambda datasets: datasets.interleave(  # pylint:disable=g-long-lambda
              lambda x: x,
              cycle_length=multiprocessing.cpu_count(),
              num_parallel_calls=dataset_ops.AUTOTUNE,
          )
  
      if element_spec is None:
          with gfile.GFile(
              os.path.join(path, dataset_ops.DATASET_SPEC_FILENAME),
              "rb",
          ) as f:
              encoded_spec = f.read()
          element_spec = _parse_element_spec(encoded_spec)
  
      distributed_snapshot_metadata = (
          _load_distributed_snapshot_metadata(path)
      )
      if distributed_snapshot_metadata:
          _validate_snapshot(
              path,
              distributed_snapshot_metadata,
              element_spec,
              compression,
          )
          return _load_distributed_snapshot(
              path, distributed_snapshot_metadata, reader_func
          )
      return _LoadDataset(path, element_spec, compression, reader_func)
  
imports: import multiprocessing
  import os
  import time
  from typing import Optional
  from google.protobuf import message
  from google.protobuf import text_format
  from tensorflow.core.protobuf import snapshot_pb2
  from tensorflow.python.data.experimental.service import _pywrap_snapshot_utils
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import structured_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import gen_experimental_dataset_ops
  from tensorflow.python.platform import gfile
  from tensorflow.python.saved_model import nested_structure_coder
label_desc: Use image node_id 1 for calling a global function with example usage: _load(path, element_spec, compression, reader_func) and returns: _LoadDataset, _load_distributed_snapshot