code: def map_on_gpu(map_func):
      """Maps `map_func` across the elements of this dataset.
  
      NOTE: This is a highly experimental version of `tf.data.Dataset.map` that runs
      `map_func` on GPU. It must be used after applying the
      `tf.data.experimental.copy_to_device` transformation with a GPU device
      argument.
  
      Args:
        map_func: A function mapping a nested structure of tensors (having shapes
          and types defined by `self.output_shapes` and `self.output_types`) to
          another nested structure of tensors.
  
      Returns:
        A `Dataset` transformation function, which can be passed to
        `tf.data.Dataset.apply`.
      """
  
      def _apply_fn(dataset):
          return _MapOnGpuDataset(dataset, map_func)
  
      return _apply_fn
  
imports: from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import iterator_ops
  from tensorflow.python.data.ops import structured_function
  from tensorflow.python.data.util import structure
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import device
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import functional_ops
  from tensorflow.python.ops import gen_dataset_ops
  from tensorflow.python.ops import gen_experimental_dataset_ops
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 3 for calling a global function with example usage: map_on_gpu(map_func) and returns: _apply_fn, _MapOnGpuDataset