code: def testWrapFuncDatasetDevice(self, device_type, dataset_reduce_fn):
      devices = config.list_logical_devices(device_type=device_type)
      if not devices:
          self.skipTest(
              "Skip when {} is not detected by TF".format(device_type)
          )
  
      @def_function.function
      def comp():
          return dataset_reduce_fn(dataset_ops.Dataset.range(10))
  
      graph = comp.get_concrete_function().graph
  
      def function_to_wrap():
          with ops.device(devices[0].name):
              return graph_def_importer.import_graph_def(
                  graph.as_graph_def()
              )
  
      with ops.device(devices[0].name):
          wrapped_noarg_fn = wrap_function.wrap_function(
              function_to_wrap, signature=[]
          )
  
      wrapped_noarg_fn()
  
file_path: tensorflow/tensorflow/python/eager/wrap_function_device_test.py
imports: from absl.testing import parameterized
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import wrap_function
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import importer
  from tensorflow.python.framework import ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 1 for calling the WrappedGraphTest obj's underlying member method code with example usage: obj.testWrapFuncDatasetDevice(device_type, dataset_reduce_fn) and returns: dataset_reduce_fn, graph_def_importer