Add wrap_function tests with explicit device placement.

PiperOrigin-RevId: 334215093
Change-Id: I0f9e2a7819a480d10ec1bf41db91b8511bac2318

tensorflow/python/eager/BUILD

@@ -982,6 +982,26 @@ tf_py_test(
     ],
 )
 
+cuda_py_test(
+    name = "wrap_function_device_test",
+    srcs = ["wrap_function_device_test.py"],
+    python_version = "PY3",
+    xla_enable_strict_auto_jit = False,
+    xla_enabled = False,
+    deps = [
+        ":def_function",
+        ":wrap_function",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:config",
+        "//tensorflow/python:constant_op",
+        "//tensorflow/python:dtypes",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:framework_ops",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
 py_library(
     name = "remote",
     srcs = ["remote.py"],

tensorflow/python/eager/wrap_function_device_test.py

@@ -0,0 +1,86 @@
+# Copyright 2018 The TensorFlow Authors. 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.
+# ==============================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+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 as graph_def_importer
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+def _dataset_reduce_sum(dataset):
+  return dataset.reduce(
+      constant_op.constant(0, dtype=dtypes.int64), lambda x, y: x + y)
+
+
+def _loop_dataset_sum(dataset):
+  value = constant_op.constant(0, dtype=dtypes.int64)
+  for d in dataset:
+    value += d
+  return value
+
+
+def _iter_dataset_sum(dataset):
+  value = constant_op.constant(0, dtype=dtypes.int64)
+  for d in iter(dataset):
+    value += d
+  return value
+
+
+class WrappedGraphTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ('cpu_reduce', 'CPU', _dataset_reduce_sum),
+      ('gpu_reduce', 'GPU', _dataset_reduce_sum),
+      ('cpu_loop', 'CPU', _loop_dataset_sum),
+      ('gpu_loop', 'GPU', _loop_dataset_sum),
+      ('cpu_iter', 'CPU', _iter_dataset_sum),
+      ('gpu_iter', 'GPU', _iter_dataset_sum),
+  )
+  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()
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()