code: def _load(  # pylint: disable=unused-private-name
      path: str,
      element_spec: Any,
      compression: Optional[str],
      reader_func: Optional[
          Callable[[dataset_ops.Dataset], dataset_ops.Dataset]
      ],
  ) -> dataset_ops.Dataset:
      """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,
          )
  
      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
          )
  
      if element_spec is None:
          element_spec = _load_element_spec(path)
      return _LoadDataset(path, element_spec, compression, reader_func)
  
imports: import multiprocessing
  import os
  import time
  from typing import Any, Callable, Optional, Union
  from absl import logging
  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