text: # Copyright (c) 2017-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import sys
  
  import tensorflow as tf
  from tensorflow.python.data.util import nest
  from tensorflow.python.framework import tensor_shape
  from nvidia.dali import types
  from nvidia.dali import internal as _internal
  
  from nvidia.dali.external_source import _is_external_source, _is_external_source_with_callback, _has_external_source
  
  from nvidia.dali._utils.external_source_impl import _get_generator_from_source_desc, _cycle_enabled
  
  from collections import Iterable
  from distutils.version import LooseVersion
  import warnings
  
  from nvidia.dali_tf_plugin import dali_tf_plugin
  
  from collections.abc import Mapping
  
  _dali_tf_module = dali_tf_plugin.load_dali_tf_plugin()
  _dali_tf = _dali_tf_module.dali
  _dali_tf.__doc__ = _dali_tf.__doc__ + """
  
      Please keep in mind that TensorFlow allocates almost all available device memory by default. This might cause errors in
      DALI due to insufficient memory. On how to change this behaviour please look into the TensorFlow documentation, as it may
      differ based on your use case.
  """
  
  _experimental_dataset_docstring = """Experimental variant of
  :class:`~nvidia.dali.plugin.tf.DALIDataset`. This dataset adds support for input tf.data.Datasets.
  Support for input tf.data.Datasets is available only for TensorFlow 2.4.1 and newer.
  
  **Input dataset specification**
  
  Each of the input datasets must be mapped to a :meth:`~nvidia.dali.fn.external_source` operator
  that will represent the input to the DALI pipeline. In the pipeline the input is represented as
  the ``name`` parameter of :meth:`~nvidia.dali.fn.external_source`. Input datasets must be provided
  as a mapping from that ``name`` to the dataset object via the ``input_datasets`` dictionary argument
  of DALIDatasetWithInputs.
  
  **Per-sample and batch mode**
  
  The input datasets can operate in per-sample mode or in batch mode.
  
  In per-sample mode, the values produced by the source dataset are interpreted
  as individual samples. The batch dimension is absent. For example, a 640x480 RGB image would
  have a shape ``[480, 640, 3]``.
  
  In batch mode, the tensors produced by the source dataset are interpreted as batches,
  with an additional outer dimension denoting the samples in the batch. For example, a batch of
  ten 640x480 RGB images would have a shape ``[10, 480, 640, 3]``.
  
  In both cases (per-sample and batch mode), the layout of those inputs should be denoted as "HWC".
  
  In per-sample mode DALIDataset will query the inputs dataset ``batch_size``-times to build a batch
  that would be fed into the DALI Pipeline.
  In per-sample mode, each sample produced by the input dataset can have a different shape,
  but the number of dimension and the layout must remain constant.
  
  **External Source with** ``source`` **parameter**
  
  This experimental DALIDataset accepts pipelines with :meth:`~nvidia.dali.fn.external_source`
  nodes that have ``source`` parameter specified.
  In that case, the ``source`` will be converted automatically into appropriate
  ``tf.data.Dataset.from_generator`` dataset with correct placement and
  ``tf.data.experimental.copy_to_device`` directives.
  
  Those nodes can also work in per-sample or in batch mode. The data in batch mode must be
  a dense, uniform tensor (each sample has the same dimensions). Only CPU data is accepted.
  
  This allows TensorFlow DALIDataset to work with most Pipelines that have External Source
  ``source`` already specified.
  
  .. warning::
      This class is experimental and its API might change without notice.
  
  .. note::
      Setting ``no_copy`` on the external source nodes when defining the pipeline is considered
      a no-op when used with DALI Dataset. The ``no_copy`` option is handled internally
      and enabled automatically if possible.
  
  The operator adds additional parameters to the ones supported by the
  :class:`~nvidia.dali.plugin.tf.DALIDataset`:
  
  Parameters
  ----------
      input_datasets : dict[str, tf.data.Dataset] or dict[str, nvidia.dali.plugin.tf.experimental.Input]
          input datasets to the DALI Pipeline. It must be provided as a dictionary mapping from
          the names of the ``External Source`` nodes to the datasets objects or to the
          :meth:`~nvidia.dali.plugin.tf.experimental.Input` wrapper.
  
          For example::
  
              {
                  'tensor_input': tf.data.Dataset.from_tensors(tensor).repeat(),
                  'generator_input': tf.data.Dataset.from_generator(some_generator)
              }
  
          can be passed as ``input_datasets`` for Pipeline like::
  
              @pipeline_def
              def external_source_pipe():
                  input_0 = fn.external_source(name='tensor_input')
                  input_1 = fn.external_source(name='generator_input')
                  return fn.resize(input_1, resize_x=input_0)
  
          Entries that use ``tf.data.Dataset`` directly, like::
  
              {
                  'input': tf.data.Dataset.from_tensors(tensor)
              }
  
          are equivalent to following specification using ``nvidia.dali.plugin.tf.experimental.Input``::
  
              {
                  'input' : nvidia.dali.plugin.tf.experimental.Input(
                                dataset=tf.data.Dataset.from_tensors(tensor),
                                layout=None,
                                batch=False)
              }
  
          This means that inputs, specified as ``tf.data.Dataset`` directly, are considered
          sample inputs.
  
          .. warning::
              Input dataset must be placed on the same device as ``DALIDatasetWithInputs``.
              If the input has different placement (for instance, input is placed on CPU, while
              ``DALIDatasetWithInputs`` is placed on GPU) the ``tf.data.experimental.copy_to_device``
              with GPU argument must be first applied to input.
  """
  
  
  _experimental_input_docstring = """Wrapper for an input passed to DALIDataset.
  Allows to pass additional options that can override some of the ones specified
  in the External Source node in the Python Pipeline object.
  Passing None indicates, that the value should be looked up in the pipeline definition.
  
  Parameters
  ----------
  dataset : tf.data.Dataset
      The dataset used as an input
  layout : str, optional, default = None
      Layout of the input. If None, the layout will be taken from the corresponding
      External Source node in the Python Pipeline object. If it is not provided there,
      empty layout will be used.
  batch: bool, optional, default = False
      Batch mode of a given input. If None, the batch mode will be taken from the
      corresponding External Source node in the Python Pipeline object.
  
      If the ``batch = False``, the input dataset is considered sample input.
  
      If the ``batch = True``, the input dataset is expected to return batches.
  """
  
  def serialize_pipeline(pipeline):
      try:
          return pipeline.serialize()
      except RuntimeError as e:
          raise RuntimeError("Error during pipeline initialization. Note that some operators "
                             "(e.g. Python Operators) cannot be used with "
                             "tensorflow data set API and DALIIterator.") from e
  
  
  def DALIIteratorWrapper(pipeline=None,
                          serialized_pipeline=None,
                          sparse=[],
                          shapes=[],
                          dtypes=[],
                          batch_size=-1,
                          prefetch_queue_depth=2,
                          **kwargs):
      """
    TF Plugin Wrapper
  
    This operator works in the same way as DALI TensorFlow plugin, with the exception that it also
    accepts Pipeline objects as an input, which are serialized internally. For more information,
    see :meth:`nvidia.dali.plugin.tf.DALIRawIterator`.
    """
      if type(prefetch_queue_depth) is dict:
          exec_separated = True
          cpu_prefetch_queue_depth = prefetch_queue_depth["cpu_size"]
          gpu_prefetch_queue_depth = prefetch_queue_depth["gpu_size"]
      elif type(prefetch_queue_depth) is int:
          exec_separated = False
          cpu_prefetch_queue_depth = -1  # dummy: wont' be used
          gpu_prefetch_queue_depth = prefetch_queue_depth
  
      if serialized_pipeline is None:
          serialized_pipeline = serialize_pipeline(pipeline)
  
      # if batch_size is not provided we need to extract if from the shape arg
      if (not isinstance(shapes, Iterable) or len(shapes) == 0) and batch_size == -1:
          raise Exception(
              'shapes and batch_size arguments cannot be empty, '
              'please provide at leas one shape argument element with the BATCH size or set batch_size'
          )
  
      if len(sparse) > 0 and sparse[0] and batch_size == -1:
          if isinstance(shapes[0], Iterable) and len(shapes[0]) == 1:
              shapes[0] = (shapes[0][0], 1)
          else:
              shapes[0] = (shapes[0], 1)
  
      # shapes and dtypes need to take into account that sparse tensor will produce 3 output tensors
      new_dtypes = []
      new_shapes = []
      for i in range(len(dtypes)):
          if i < len(sparse) and sparse[i]:
              # indices type of sparse tensor is tf.int64
              new_dtypes.append(tf.int64)
              new_dtypes.append(dtypes[i])
              # dense shape type of sparse tensor is tf.int64
              new_dtypes.append(tf.int64)
              if len(shapes) > i and len(shapes[i]) > 0:
                  new_shapes.append((shapes[i][0], 1))
                  new_shapes.append((shapes[i][0]))
              else:
                  new_shapes.append(())
                  new_shapes.append(())
              new_shapes.append(())
          else:
              new_dtypes.append(dtypes[i])
              if len(shapes) > i:
                  new_shapes.append(shapes[i])
  
      # gpu_prefetch_queue_depth correspond to the global queue depth in the uniform case
      out = _dali_tf(serialized_pipeline=serialized_pipeline,
                     shapes=new_shapes,
                     dtypes=new_dtypes,
                     sparse=sparse,
                     batch_size=batch_size,
                     exec_separated=exec_separated,
                     gpu_prefetch_queue_depth=gpu_prefetch_queue_depth,
                     cpu_prefetch_queue_depth=cpu_prefetch_queue_depth,
                     **kwargs)
      new_out = []
      j = 0
      for i in range(len(dtypes)):
          if i < len(sparse) and sparse[i]:
              new_out.append(
                  tf.SparseTensor(indices=out[j], values=out[j + 1], dense_shape=out[j + 2]))
              j += 3
          else:
              new_out.append(out[j])
              j += 1
      return new_out
  
  
  def DALIIterator():
      return DALIIteratorWrapper
  
  
  # Vanilla raw operator legacy
  def DALIRawIterator():
      return _dali_tf
  
  
  def _get_tf_version():
      return LooseVersion(tf.__version__)
  
  
  MIN_TENSORFLOW_VERSION = LooseVersion('1.15')
  
  
  def dataset_compatible_tensorflow():
      """Returns ``True`` if current TensorFlow version is compatible with DALIDataset."""
      return LooseVersion(tf.__version__) >= MIN_TENSORFLOW_VERSION
  
  
  def dataset_inputs_compatible_tensorflow():
      """Returns ``True`` if the current TensorFlow version is compatible with
      experimental.DALIDatasetWithInputs and input Datasets can be used with DALI.
      """
      return LooseVersion(tf.__version__) >= LooseVersion('2.4.1')
  
  
  def dataset_distributed_compatible_tensorflow():
      """Returns ``True`` if the tf.distribute APIs for current TensorFlow version are compatible
      with DALIDataset.
      """
      return LooseVersion(tf.__version__) >= LooseVersion('2.5.0')
  
  
  def _get_experimental():
      # TODO(klecki): this is WAR only for experimental module
      current_module = sys.modules[__name__]
      experimental = _internal.get_submodule(current_module, "experimental")
      return experimental
  
  
  def _insert_experimental_member(member, name):
      experimental_module = _get_experimental()
      member.__module__ = experimental_module
      setattr(experimental_module, name, member)
  
  def _get_external_source_param(input_name, input_value, name_es_map, param_name):
      """Get value of the parameter `param_name` specified for the External Source node
         named `input_name`. It can be specified either via `input_value` or in the op instance
         passed in `name_es_map`.
         Not `None` value in `input_value` overwrites the one specified in the Operator instances.
         Otherwise, the one from pipeline definition (the op instance) is used.
  
      Parameters
      ----------
      input_name : str
          Name of the input
      input_value : Input, optional
          Description of the input
      name_es_map : dict[str, ExternalSource]
          Mapping from the External Source names to operator nodes.
      param_name : str
          name of the parameter we want to access
      """
      def get_param_from_pipe(input_name, name_es_map, param_name):
          es_op = name_es_map[input_name]
          # Check the OpInstance and the `_op`
          try:
              return getattr(es_op, "_" + param_name)
          except AttributeError:
              return getattr(es_op._op, "_" + param_name, None)
  
      # We didn't get input through input_datasets
      if input_value is None or getattr(input_value, param_name) is None:
          return get_param_from_pipe(input_name, name_es_map, param_name)
      else:
          return getattr(input_value, param_name)
  
  
  def _get_signature(dtype, shape):
      # TODO(klecki): Find out how we can use ragged tensors for non-uniform batches
      return tf.TensorSpec(shape=shape, dtype=dtype)
  
  
  def _get_current_device_spec():
      """Best guess at checking the current device string in eager and graph mode.
  
      Using callable in `with tf.device(...)` for Graph mode will probably break it.
      The graph in use is assumed to be current default graph.
      """
      if tf.executing_eagerly():
          # We are not using this `tf.device` with `with ...`,
          # so we do not change the context, it returns _EagerDeviceContext
          dummy_context_manager = tf.device(None)
          # Get the eager.context singleton instance for this thread
          context = dummy_context_manager._ctx
          # DeviceSpec
          return context.device_spec
      else:
          # Get the default graf, we assume that it's the one in use
          g = tf.compat.v1.get_default_graph()
          # Get the top element of _UserDeviceSpec stack - `with tf.device()` pushes to the stack
          # in graph mode.
          spec = g._device_function_stack.peek_top_obj()
          # Try to normalize to DeviceSpec
          return tf.DeviceSpec.from_string(spec.display_name)
  
  
  
  if dataset_compatible_tensorflow():
      from tensorflow.python.framework import ops
      from tensorflow.python.data.ops import dataset_ops
      from tensorflow.python.data.util import structure
      import functools
  
      def dataset_options():
          options = tf.data.Options()
          options.experimental_optimization.apply_default_optimizations = False
          options.experimental_optimization.autotune = False
  
          return options
  
      class _DALIDatasetV2(dataset_ops.DatasetV2):
          @classmethod
          def experimental_constructor(cls):
              return cls(randint(0, 33))
  
          def __init__(
                  self,
                  pipeline,
                  output_dtypes=None,
                  output_shapes=None,
                  fail_on_device_mismatch=True,
                  *,
                  input_datasets=None,
                  batch_size=1,
                  num_threads=4,
                  device_id=0,
                  exec_separated=False,
                  prefetch_queue_depth=2,
                  cpu_prefetch_queue_depth=2,
                  gpu_prefetch_queue_depth=2,
                  dtypes=None,
                  shapes=None):
  
              output_shapes = self._handle_deprecation(output_shapes, shapes, "shapes")
              output_dtypes = self._handle_deprecation(output_dtypes, dtypes, "dtypes")
  
              if not self._check_dtypes(output_dtypes, tf.DType):
                  raise TypeError(("`output_dtypes` should be provided as single tf.DType value "
                      "or a tuple of tf.DType values. Got value `{}` of type `{}`.") \
                          .format(output_dtypes, type(output_dtypes)))
  
              if output_shapes is None:
                  output_shapes = nest.map_structure(lambda _: tensor_shape.TensorShape(None),
                                                     output_dtypes)
              else:
                  output_shapes = nest.map_structure_up_to(output_dtypes, tensor_shape.as_shape,
                                                           output_shapes)
  
              if not isinstance(output_dtypes, tuple):
                  output_dtypes = (output_dtypes, )
                  output_shapes = (output_shapes, )
  
              output_classes = nest.map_structure(lambda _: ops.Tensor, output_dtypes)
  
              self._pipeline_instance = pipeline  # keep the live Pipeline object
              self._pipeline_serialized = serialize_pipeline(pipeline)
              self._batch_size = batch_size
              self._num_threads = num_threads
              if device_id is None:
                  device_id = types.CPU_ONLY_DEVICE_ID
              self._device_id = device_id
              self._exec_separated = exec_separated
              self._prefetch_queue_depth = prefetch_queue_depth
              self._cpu_prefetch_queue_depth = cpu_prefetch_queue_depth
              self._gpu_prefetch_queue_depth = gpu_prefetch_queue_depth
              self._output_shapes = output_shapes
              self._output_dtypes = output_dtypes
              self._fail_on_device_mismatch = fail_on_device_mismatch
  
              self._setup_inputs(input_datasets)
  
              self._structure = structure.convert_legacy_structure(self._output_dtypes,
                                                                   self._output_shapes,
                                                                   output_classes)
  
              super(_DALIDatasetV2, self).__init__(self._as_variant_tensor())
  
          def _input_lists_from_input_datasets(self, input_datasets, name_es_map):
              """Extract the input specification from the input_datasets dictionary.
  
              Validate if the inputs exist in the pipeline and the types are correct
  
              Returns
              -------
              list, list, list, list
                  input_datasets, input_names, input_layouts, input_batched
              """
  
              if input_datasets is None:
                  return [], [], [], []
  
              def _get_dataset(value):
                  if isinstance(value, dataset_ops.DatasetV2):
                      return value
                  else:
                      return value.dataset
  
              in_datasets_list = []
              in_names_list = []
              in_layouts_list = []
              in_batched_list = []
  
              error_str = (
                  "`input_datasets` must be a dictionary that maps input names (the `name` "
                  "specified for External Source node in DALI pipeline) to input datasets "
                  "objects (`tf.data.Dataset`) or `nvidia.dali.plugin.tf.experimental.Input` wrapper "
                  "objects")
  
              if not isinstance(input_datasets, Mapping):
                  raise TypeError(error_str +
                                  ", got: `{}` of type: {} instead.".format(input_datasets, type(input_datasets)))
  
              for input_name, input_value in input_datasets.items():
                  # keys are str
                  if not isinstance(input_name, str):
                      raise TypeError(error_str + (". Expected the keys (representing the input "
                                                   "names) to be of type `str`, got: `{}` of type: "
                                                   "{} instead.").format(input_name, type(input_name)))
  
                  # values are tf.data.Dataset or Input
                  is_dataset_only = isinstance(input_value, dataset_ops.DatasetV2)
                  experimental = _get_experimental()
                  if not is_dataset_only and not isinstance(input_value, experimental.Input):
                      raise TypeError(error_str + (". Expected the values of the dictionary "
                                                   "(representing the inputs) "
                                                   " to be of type `tf.data.Dataset` or "
                                                   "`nvidia.dali.plugin.tf.Input` got: `{}` of "
                                                   "type: {} instead.").format(input_value, type(input_value)))
  
                  # there is External Source with name equal to `input_name`
                  if input_name not in name_es_map.keys():
                      raise ValueError(("Did not find an External Source placeholder node with "
                                        "name='{}' in the provided pipeline - required by the name "
                                        "specified in the `input_datasets`. Names of available  "
                                        "placeholder External Source nodes are: {}. Placeholder "
                                        "nodes cannot have `source` argument specified.").format(input_name,
                                                                        list(name_es_map.keys())))
  
                  in_names_list.append(input_name)
                  in_datasets_list.append(_get_dataset(input_value))
  
                  if is_dataset_only:
                      # Set the defaults used in lookup
                      as_input = experimental.Input(input_value, layout=None, batch=False)
                  else:
                      as_input = input_value
  
                  # TODO(klecki): Do we want all Python-only ES parameters to be overridable here?
                  layout = _get_external_source_param(input_name, as_input, name_es_map, 'layout')
                  in_layouts_list.append(layout or "")
  
                  # Batched mode is supported by default
                  batched = _get_external_source_param(input_name, as_input, name_es_map, 'batch')
                  in_batched_list.append(batched if batched is not None else True)
  
              return in_datasets_list, in_names_list, in_layouts_list, in_batched_list
  
  
          def _input_lists_from_source(self, callbacked_es_map):
  
              # TODO(klecki): Warn about this in the doc.
              # We do it only when the users wants to use ExternalSource with `source` specified,
              # as it has some additional limitations.
  
              # Capture the device that DALI was placed, as we may need to copy the CPU callbacks
              # to that device.
              dali_device_spec = _get_current_device_spec()
              is_dali_on_gpu = dali_device_spec.device_type == "GPU"
  
              in_datasets_list = []
              in_names_list = []
              in_layouts_list = []
              in_batched_list = []
  
              for input_name, external_source in callbacked_es_map.items():
                  in_names_list.append(input_name)
                  layout = _get_external_source_param(input_name, None, callbacked_es_map, 'layout')
                  in_layouts_list.append(layout or "")
  
                  # Batched mode is supported by default
                  batched = _get_external_source_param(input_name, None, callbacked_es_map, 'batch')
                  in_batched_list.append(batched if batched is not None else True)
  
                  source_desc = external_source._op._source_desc
                  if source_desc.cycle == 'raise':
                      raise NotImplementedError(("External Source node: '{}' got argument "
                                                 "cycle='raise' which is not supported.").format(input_name))
  
                  # All generator datasets must be placed on CPU.
                  with tf.device('/cpu:0'):
                      tf_gen, dtype, shape = _get_generator_from_source_desc(
                          source_desc, self._batch_size, external_source._batch)
                      # dataset = tf.data.Dataset.from_generator(tf_gen, output_types=dtype)
                      signature = _get_signature(dtype, shape)
                      dataset = tf.data.Dataset.from_generator(tf_gen, output_signature=signature)
                      if _cycle_enabled(source_desc.cycle):
                          dataset = dataset.repeat()
                      # if DALIDataset was placed on GPU, we need to add the copy targetting
                      # that device (with proper id).
                      if is_dali_on_gpu:
                          dataset = dataset.apply(tf.data.experimental.copy_to_device(dali_device_spec.to_string()))
                      in_datasets_list.append(dataset)
  
              return in_datasets_list, in_names_list, in_layouts_list, in_batched_list
  
  
          def _setup_inputs(self, input_datasets):
              """Verify the input specification and assign it to private members in
              normalized form.
              """
  
              has_es = _has_external_source(self._pipeline_instance)
  
              # If no inputs are specified, input handling is no-op
              if input_datasets is None and not has_es:
                  self._input_datasets = ()
                  self._input_names = ()
                  self._input_layouts = ()
                  self._input_batched = ()
                  return
  
              self._assert_pipeline_instance()
  
              # To not check everywhere for None
              if input_datasets is None:
                  input_datasets = {}
  
              name_es_map, callbacked_es_map = self._get_name_es_instance_map()
  
              inputs_from_dict = self._input_lists_from_input_datasets(input_datasets, name_es_map)
  
              inputs_from_source = self._input_lists_from_source(callbacked_es_map)
  
              # Check if someone passed an entry in `input_datasets` for the ES with callback
              if not input_datasets.keys().isdisjoint(callbacked_es_map.keys()):
                  overlapped = input_datasets.keys().intersection(callbacked_es_map.keys())
                  raise ValueError(("Double specification of External Source input is not allowed. "
                                    "External Source nodes named: `{}` got inputs specified via "
                                    "`input_datasets` DALIDataset argument and ExternalSource "
                                    "`source` argument at the same time.").format(overlapped))
  
              # We covered all inputs
              non_matched = set(name_es_map.keys()) - set(input_datasets.keys()) - set(callbacked_es_map.keys())
              if len(non_matched) != 0:
                  raise ValueError(("Found External Source nodes in the Pipeline, that were not "
                                    "assigned any inputs. Nodes without inputs: \n{}.\nNodes that "
                                    "were assigned inputs:\n{}.").format(list(non_matched), list(input_datasets.keys())))
  
  
              self._input_datasets = tuple(inputs_from_dict[0] + inputs_from_source[0])
              self._input_names = tuple(inputs_from_dict[1] + inputs_from_source[1])
              self._input_layouts = tuple(inputs_from_dict[2] + inputs_from_source[2])
              # Map it to integers, to pass as vector<int> instead of vector<bool> to C++
              self._input_batched = tuple(int(b) for b in inputs_from_dict[3] + inputs_from_source[3])
  
          def _assert_pipeline_instance(self):
              """Ensure that the pipeline is built, and check if the Python part is available.
              """
              self._pipeline_instance.build()
              if not self._pipeline_instance._py_graph_built and self._pipeline_instance._built:
                  raise ValueError("Deserialized pipelines cannot be used with `input_datasets`. "
                                   "Please provide a pipeline that was created directly in Python "
                                   "and not recreated from serialized one.")
  
          def _assert_correct_external_sources(self, external_source):
              """Validate that the external source nodes used are properly configured"""
              if external_source._op._num_outputs is not None:
                  raise ValueError("Found placeholder External Source node (without `source` "
                                   "argument) in the Pipeline that was created with `num_outputs` "
                                   "`num_outputs` parameter. Only single-output "
                                   "(with `num_outputs=None`), named (with `name` argument "
                                   "specified) External Source nodes are supported as inputs "
                                   "placeholders for DALIDataset integration. "
                                   "Alternatively, External Source can be used with `source` "
                                   "parameter specified.")
              if external_source._op._name is None:
                  raise ValueError("Found placeholder External Source node (without `source` "
                                   "argument) in the Pipeline that was not named "
                                   "(no `name` argument set). Only single-output "
                                   "(with `num_outputs=None`), named (with `name` argument "
                                   "specified) External Source nodes are supported as inputs "
                                   "placeholders for DALIDataset integration. "
                                   "Alternatively, External Source can be used with `source` "
                                   "parameter specified.")
  
          def _get_name_es_instance_map(self):
              """Return mappings between name of External Source and the op.
  
              Returns
              -------
              mapping for placeholders nodes, mapping for nodes with Python source
                  Two mappings are returned, separating the placeholder nodes without a `source`
                  and nodes that got a `source` parameter.
              """
              name_es = {}
              name_es_with_callback = {}
              for op in self._pipeline_instance._ops:
                  if _is_external_source_with_callback(op):
                      # use the internal op name (generated automatically in most cases)
                      name_es_with_callback[op.name] = op
                  elif _is_external_source(op):
                      self._assert_correct_external_sources(op)
                      # use the user provided name
                      name_es[op._op._name] = op
              return name_es, name_es_with_callback
  
          def _check_dtypes(self, values, expected_elem_type):
              """Check whether `values` is instance of `expected_elem_type`
              or tuple of `expected_elem_type`. TF doesn't treat list as a nesting type, but as a Tensor.
              """
              if isinstance(values, expected_elem_type):
                  return True
              elif isinstance(values, tuple) \
                  and all(isinstance(elem, expected_elem_type) for elem in values):
                  return True
              else:
                  return False
  
          def _handle_deprecation(self, supported_arg, deprecated_arg, name):
              if deprecated_arg is not None:
                  if supported_arg is not None:
                      raise ValueError((
                          "Usage of `{name}` is deprecated in favor of `output_{name}`. "
                          "Both arguments were provided, but only `output_{name}` should be provided."
                      ).format(name=name))
                  # show only this warning
                  warnings.warn(("Use of argument `{name}` is deprecated. Please use `output_{name}` instead. " \
                      + "`output_{name}` should be provided as a tuple or a single value.").format(name=name),
                      Warning, stacklevel=2)
                  if isinstance(deprecated_arg, list):
                      return tuple(deprecated_arg)
                  return deprecated_arg
              else:
                  return supported_arg
  
          @property
          def element_spec(self):
              return self._structure
  
          @property
          def _element_structure(self):
              return self._structure
  
          def _inputs(self):
              # Apparently here TF is happy with a list
              return nest.flatten(self._input_datasets)
  
          def _as_variant_tensor(self):
              return _dali_tf_module.dali_dataset(
                  # Experimental dataset inputs
                  nest.map_structure(lambda d: d._variant_tensor, self._input_datasets),
                  # Description of inputs
                  input_names=self._input_names,
                  input_layouts=self._input_layouts,
                  input_batched=self._input_batched,
                  # End of experimental inputs
                  pipeline=self._pipeline_serialized,
                  batch_size=self._batch_size,
                  num_threads=self._num_threads,
                  device_id=self._device_id,
                  exec_separated=self._exec_separated,
                  prefetch_queue_depth=self._prefetch_queue_depth,
                  cpu_prefetch_queue_depth=self._cpu_prefetch_queue_depth,
                  gpu_prefetch_queue_depth=self._gpu_prefetch_queue_depth,
                  output_shapes=self._output_shapes,
                  output_dtypes=self._output_dtypes,
                  fail_on_device_mismatch=self._fail_on_device_mismatch)
  
      if _get_tf_version() < LooseVersion('2.0'):
  
          class _DALIDatasetImpl(dataset_ops.DatasetV1Adapter):
              @functools.wraps(_DALIDatasetV2.__init__)
              def __init__(self, pipeline, **kwargs):
                  self._wrapped = _DALIDatasetV2(pipeline, **kwargs)
                  super(_DALIDatasetImpl, self).__init__(self._wrapped)
      else:
          _DALIDatasetImpl = _DALIDatasetV2
  
      _experimental_kwargs = ['input_datasets']
  
      class DALIDataset(dataset_ops._OptionsDataset):
          @functools.wraps(_DALIDatasetV2.__init__)
          def __init__(self, pipeline, **kwargs):
  
              # TODO(klecki): Remove this when we move support for inputs from experimental.
              for disallowed_kwarg in _experimental_kwargs:
                  if disallowed_kwarg in kwargs.keys():
                      raise TypeError((
                          "__init__() got an unexpected keyword argument '{}'. "
                          "Dataset inputs are allowed only in 'experimental.DALIDatasetWithInputs'.").format(
                          disallowed_kwarg))
              # We detected External Source nodes in the Pipeline
              if _has_external_source(pipeline):
                  raise ValueError(("DALIDataset got a DALI pipeline containing External Source "
                      "operator nodes. External Source nodes can be used to express placeholders "
                      "for tf.data.Dataset inputs to DALI or to run user-provided Python code "
                      "via `source` parameter. Support for Dataset inputs and External Source's "
                      "`source` is allowed only in 'experimental.DALIDatasetWithInputs'."))
  
              dataset_impl = _DALIDatasetImpl(pipeline, **kwargs)
              super(DALIDataset, self).__init__(dataset_impl, dataset_options())
  
  else:
  
      class DALIDataset:
          def __init__(self,
                       pipeline,
                       output_dtypes=None,
                       output_shapes=None,
                       fail_on_device_mismatch=True,
                       *,
                       batch_size=1,
                       num_threads=4,
                       device_id=0,
                       exec_separated=False,
                       prefetch_queue_depth=2,
                       cpu_prefetch_queue_depth=2,
                       gpu_prefetch_queue_depth=2,
                       dtypes=None,
                       shapes=None):
              raise RuntimeError(
                  'DALIDataset is not supported for detected version of TensorFlow.  DALIDataset supports versions: 1.15, 2.x family'
              )
  
  
  if dataset_inputs_compatible_tensorflow():
      def _load_experimental_dataset():
          class DALIDatasetWithInputs(dataset_ops._OptionsDataset):
              @functools.wraps(_DALIDatasetV2.__init__)
              def __init__(self, pipeline, **kwargs):
                  dataset_impl = _DALIDatasetImpl(pipeline, **kwargs)
                  super(DALIDatasetWithInputs, self).__init__(dataset_impl, dataset_options())
  
          DALIDatasetWithInputs.__doc__ = _experimental_dataset_docstring
          _insert_experimental_member(DALIDatasetWithInputs, "DALIDatasetWithInputs")
  
          class Input:
              def __init__(self, dataset, *, layout=None, batch=False):
                  if not isinstance(dataset, dataset_ops.DatasetV2):
                      raise TypeError(
                          ("The inputs specified to DALIDataset must be instances of "
                           "type `tf.data.Dataset` got: `{}` of type: {} instead.").format(
                               dataset, type(dataset)))
                  self.dataset = dataset
                  self.layout = layout
                  self.batch = batch
  
          Input.__doc__ = _experimental_input_docstring
  
          _insert_experimental_member(Input, "Input")
  
      _load_experimental_dataset()
  
  else:
      def _load_experimental_dataset():
          class DALIDatasetWithInputs:
              def __init__(self, *args, **kwargs):
                  raise RuntimeError(
                      'experimental.DALIDatasetWithInputs is not supported for detected version of TensorFlow. '
                      + 'DALIDataset supports versions: 2.4.1 and above.')
  
          DALIDatasetWithInputs.__doc__ = _experimental_dataset_docstring
          _insert_experimental_member(DALIDatasetWithInputs, "DALIDatasetWithInputs")
  
          class Input:
              def __init__(self, *args, **kwargs):
                  pass
  
          Input.__doc__ = _experimental_input_docstring
  
          _insert_experimental_member(Input, "Input")
  
      _load_experimental_dataset()
  
  DALIDataset.__doc__ = """Creates a ``DALIDataset`` compatible with
      `tf.data.Dataset <https://www.tensorflow.org/api_docs/python/tf/data/Dataset>`_ from a DALI
      pipeline. It supports TensorFlow 1.15 and 2.x family.
  
      ``DALIDataset`` can be placed on CPU and GPU.
  
      Please keep in mind that TensorFlow allocates almost all available device memory by default.
      This might cause errors in DALI due to insufficient memory. On how to change this behaviour
      please look into the TensorFlow documentation, as it may differ based on your use case.
  
      .. warning::
         Most TensorFlow Datasets have only CPU variant. To process GPU-placed ``DALIDataset`` by
         other TensorFlow dataset you need to first copy it back to CPU using explicit
         ``tf.data.experimental.copy_to_device`` - roundtrip from CPU to GPU back to CPU would
         probably degrade performance a lot and is thus discouraged.
  
         Additionally, it is advised to not use datasets like ``repeat()`` or similar after
         ``DALIDataset``, which may interfere with DALI memory allocations and prefetching.
  
      Parameters
      ----------
      pipeline : :class:`nvidia.dali.Pipeline`
          defining the data processing to be performed.
      output_dtypes: tf.DType or tuple of tf.DType, default = None
          expected output types
      output_shapes: tuple of shapes, optional, default = None
          expected output shapes. If provided, must match arity of the ``output_dtypes``.
          When set to None, DALI will infer the shapes on its own.
          Individual shapes can be also set to None or contain None to indicate unknown dimensions.
          If specified must be compatible with shape returned from DALI Pipeline
          and with ``batch_size`` argument which will be the outermost dimension of returned tensors.
          In case of ``batch_size = 1`` it can be omitted in the shape.
          DALI Dataset will try to match requested shape by squeezing 1-sized dimensions
          from shape obtained from Pipeline.
      fail_on_device_mismatch : bool, optional, default = True
          When set to ``True`` runtime check will be performed to ensure DALI device and TF device are
          both CPU or both GPU. In some contexts this check might be inaccurate. When set to ``False``
          will skip the check but print additional logs to check the devices. Keep in mind that this
          may allow hidden GPU to CPU copies in the workflow and impact performance.
      batch_size : int, optional, default = 1
          batch size of the pipeline.
      num_threads : int, optional, default = 4
          number of CPU threads used by the pipeline.
      device_id : int, optional, default = 0
          id of GPU used by the pipeline.
          A None value for this parameter means that DALI should not use GPU nor CUDA runtime.
          This limits the pipeline to only CPU operators but allows it to run on any CPU capable machine.
      exec_separated : bool, optional, default = False
          Whether to execute the pipeline in a way that enables
          overlapping CPU and GPU computation, typically resulting
          in faster execution speed, but larger memory consumption.
      prefetch_queue_depth : int, optional, default = 2
          depth of the executor queue. Deeper queue makes DALI more
          resistant to uneven execution time of each batch, but it also
          consumes more memory for internal buffers.
          Value will be used with ``exec_separated`` set to ``False``.
      cpu_prefetch_queue_depth : int, optional, default = 2
          depth of the executor cpu queue. Deeper queue makes DALI more
          resistant to uneven execution time of each batch, but it also
          consumes more memory for internal buffers.
          Value will be used with ``exec_separated`` set to ``True``.
      gpu_prefetch_queue_depth : int, optional, default = 2
          depth of the executor gpu queue. Deeper queue makes DALI more
          resistant to uneven execution time of each batch, but it also
          consumes more memory for internal buffers.
          Value will be used with ``exec_separated`` set to ``True``.
  
      Returns
      -------
      ``DALIDataset`` object based on DALI pipeline and compatible with ``tf.data.Dataset`` API.
  
      """
  
  DALIIterator.__doc__ = DALIIteratorWrapper.__doc__
  DALIRawIterator.__doc__ = _dali_tf.__doc__
  