def test_get_no_arg_wrapped_function_check_dataset_reduce_in_multi_gpu(self):
@computations.tf_computation
def comp():
return tf.data.Dataset.range(10).reduce(np.int64(0), lambda p, q: p + q)
with self.assertRaisesRegex(
ValueError, 'Detected dataset reduce op in multi-GPU TFF simulation.*'):
eager_tf_executor._get_wrapped_function_from_comp(
computation_impl.ComputationImpl.get_proto(comp),
must_pin_function_to_cpu=False,
param_type=None,
device=None)
def test_get_no_arg_wrapped_function_with_composed_fn_and_variables(
self, device_type):
self._skip_in_multi_gpus()
@tf.function
def reduce_fn(x, y):
return x + y
@tf.function
def dataset_reduce_fn(ds, initial_val):
return ds.reduce(initial_val, reduce_fn)
@computations.tf_computation(computation_types.SequenceType(tf.int64),
computation_types.TensorType(tf.int64))
def dataset_reduce_fn_wrapper(ds, dummy_val):
initial_val = tf.Variable(np.int64(1.0)) + dummy_val
return dataset_reduce_fn(ds, initial_val)
@computations.tf_computation
def comp():
ds = tf.data.Dataset.range(10).map(lambda x: x + 1)
dummy_val = tf.Variable(np.int64(1.0))
return dataset_reduce_fn_wrapper(ds, dummy_val)
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
computation_impl.ComputationImpl.get_proto(comp),
must_pin_function_to_cpu=False,
param_type=None,
device=_get_first_logical_device(device_type))
self.assertEqual(wrapped_fn(), np.int64(57))
def test_get_no_arg_wrapped_function_from_comp_with_dataset_reduce(self):
@computations.tf_computation
def comp():
return tf.data.Dataset.range(10).reduce(np.int64(0),
lambda p, q: p + q)
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
computation_impl.ComputationImpl.get_proto(comp),
must_pin_function_to_cpu=False,
param_type=None,
device=None)
self.assertEqual(wrapped_fn(), np.int64(45))
def test_get_wrapped_function_from_comp_raises_with_incorrect_binding(
self, device_type):
self._skip_in_multi_gpus()
with tf.Graph().as_default() as graph:
var = tf.Variable(initial_value=0.0, name='var1', import_scope='')
assign_op = var.assign_add(tf.constant(1.0))
tf.add(1.0, assign_op)
result_binding = pb.TensorFlow.Binding(
tensor=pb.TensorFlow.TensorBinding(tensor_name='Invalid'))
comp = pb.Computation(tensorflow=pb.TensorFlow(
graph_def=serialization_utils.pack_graph_def(graph.as_graph_def()),
result=result_binding))
with self.assertRaisesRegex(
TypeError, 'Caught exception trying to prune graph.*'):
eager_tf_executor._get_wrapped_function_from_comp(
comp,
must_pin_function_to_cpu=False,
param_type=None,
device=_get_first_logical_device(device_type))
def test_get_no_arg_wrapped_function_multi_gpu_no_reduce(self):
@computations.tf_computation
@tf.function
def comp():
value = tf.constant(0, dtype=tf.int64)
for d in iter(tf.data.Dataset.range(10)):
value += d
return value
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
computation_impl.ComputationImpl.get_proto(comp),
must_pin_function_to_cpu=False,
param_type=None,
device=None)
self.assertEqual(wrapped_fn(), np.int64(45))
def test_get_no_arg_wrapped_function_multi_gpu_tf_device(self):
logical_gpus = tf.config.list_logical_devices('GPU')
@computations.tf_computation
def comp():
with tf.device(logical_gpus[0].name):
return tf.data.Dataset.range(10).reduce(np.int64(0), lambda p, q: p + q)
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
computation_impl.ComputationImpl.get_proto(comp),
must_pin_function_to_cpu=False,
param_type=None,
device=None)
self.assertEqual(wrapped_fn(), np.int64(45))
def test_get_no_arg_wrapped_function_with_variables(self, device_type):
self._skip_in_multi_gpus()
@computations.tf_computation
def comp():
initial_val = tf.Variable(np.int64(1.0))
return tf.data.Dataset.range(10).map(lambda x: x + 1).reduce(
initial_val, lambda p, q: p + q)
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
computation_impl.ComputationImpl.get_proto(comp),
must_pin_function_to_cpu=False,
param_type=None,
device=_get_first_logical_device(device_type))
self.assertEqual(wrapped_fn(), np.int64(56))
def test_get_no_arg_wrapped_function_from_comp_with_dataset_reduce(
self, device_type):
self._skip_in_multi_gpus()
@tensorflow_computation.tf_computation
def comp():
return tf.data.Dataset.range(10).reduce(np.int64(0),
lambda p, q: p + q)
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
computation_impl.ConcreteComputation.get_proto(comp),
must_pin_function_to_cpu=False,
param_type=None,
device=_get_first_logical_device(device_type))
self.assertEqual(wrapped_fn(), np.int64(45))
def test_get_wrapped_function_from_comp_raises_with_incorrect_binding(self):
with tf.Graph().as_default() as graph:
var = tf.Variable(initial_value=0.0, name='var1', import_scope='')
assign_op = var.assign_add(tf.constant(1.0))
tf.add(1.0, assign_op)
result_binding = pb.TensorFlow.Binding(
tensor=pb.TensorFlow.TensorBinding(tensor_name='Invalid'))
comp = pb.Computation(
tensorflow=pb.TensorFlow(
graph_def=serialization_utils.pack_graph_def(graph.as_graph_def()),
result=result_binding))
with self.assertRaises(TypeError):
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
comp, must_pin_function_to_cpu=False, param_type=None, device=None)
wrapped_fn()
def test_get_no_arg_wrapped_function_from_comp_with_iter_dataset(
self, device_type):
self._skip_in_multi_gpus()
@computations.tf_computation
@tf.function
def comp():
value = tf.constant(0, dtype=tf.int64)
for d in iter(tf.data.Dataset.range(10)):
value += d
return value
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
computation_impl.ComputationImpl.get_proto(comp),
must_pin_function_to_cpu=False,
param_type=None,
device=_get_first_logical_device(device_type))
self.assertEqual(wrapped_fn(), np.int64(45))
