Adds a unit test of wrap_function for all physical devices.
@@ -18,6 +18,7 @@ import asyncio
import collections
from absl.testing import absltest
+from absl.testing import parameterized
import numpy as np
import tensorflow as tf
@@ -29,7 +30,17 @@ from tensorflow_federated.python.core.impl import eager_executor
from tensorflow_federated.python.core.impl import executor_test_utils
-class EagerExecutorTest(absltest.TestCase):
+def _get_physical_devices_for_testing():
+ result = []
+ for dev in tf.config.experimental.list_physical_devices():
+ parts = dev.name.split(':')
+ if ((len(parts) == 3) and (parts[0] == '/physical_device') and
+ (parts[1] in ['CPU', 'GPU'])):
+ result.append(':'.join(parts[1:]))
+ return result
+
+
+class EagerExecutorTest(parameterized.TestCase):
# TODO(b/134764569): Potentially take advantage of something similar to the
# `tf.test.TestCase.evaluate()` to avoid having to call `.numpy()` everywhere.
@@ -416,6 +427,32 @@ class EagerExecutorTest(absltest.TestCase):
executor_test_utils.test_mnist_training(self,
eager_executor.EagerExecutor())
+ @parameterized.parameters(
+ *[(dev,) for dev in _get_physical_devices_for_testing()])
+ def test_wrap_function_on_all_available_physical_devices(self, device):
+ with tf.Graph().as_default() as graph:
+ x = tf.placeholder(tf.int32, shape=[])
+ y = tf.add(x, tf.constant(1))
+
+ arg_for_tf_device = '/{}'.format(device)
+
+ def _function_to_wrap(arg):
+ with tf.device(arg_for_tf_device):
+ return tf.import_graph_def(
+ graph.as_graph_def(),
+ input_map={x.name: arg},
+ return_elements=[y.name])[0]
+
+ signature = [tf.TensorSpec([], tf.int32)]
+ wrapped_fn = tf.compat.v1.wrap_function(_function_to_wrap, signature)
+
+ def fn(arg):
+ with tf.device(arg_for_tf_device):
+ return wrapped_fn(arg)
+
+ result = fn(tf.constant(10))
+ self.assertTrue(result.device.endswith(device))
+
if __name__ == '__main__':
tf.compat.v1.enable_v2_behavior()
