def test_embed_tensorflow_computation_fails_with_bogus_device(self):
@computations.tf_computation(tf.int32)
def comp(x):
return tf.add(x, 1)
comp_proto = computation_impl.ComputationImpl.get_proto(comp)
with self.assertRaises(ValueError):
eager_tf_executor.embed_tensorflow_computation(
comp_proto, comp.type_signature, device='/there_is_no_such_device')
def test_embed_tensorflow_computation_with_no_arg_and_int_result(self):
@computations.tf_computation
def comp():
return 1000
fn = eager_tf_executor.embed_tensorflow_computation(
computation_impl.ComputationImpl.get_proto(comp))
result = fn()
self.assertIsInstance(result, tf.Tensor)
self.assertEqual(result, 1000)
def test_embed_tensorflow_computation_with_float(self):
@computations.tf_computation(tf.float32)
def comp(x):
return x + 0.5
fn = eager_tf_executor.embed_tensorflow_computation(
computation_impl.ComputationImpl.get_proto(comp))
result = fn(tf.constant(10.0))
self.assertIsInstance(result, tf.Tensor)
self.assertEqual(result, 10.5)
def test_embed_tensorflow_computation_with_int_arg_and_result(self):
@computations.tf_computation(tf.int32)
def comp(x):
return x + 1
fn = eager_tf_executor.embed_tensorflow_computation(
computation_impl.ComputationImpl.get_proto(comp))
result = fn(10)
self.assertIsInstance(result, tf.Tensor)
self.assertEqual(result.numpy(), 11)
def test_embed_tensorflow_computation_with_dataset_arg_and_int_result(
self):
@computations.tf_computation(computation_types.SequenceType(tf.int32))
def comp(ds):
return ds.reduce(np.int32(0), lambda p, q: p + q)
fn = eager_tf_executor.embed_tensorflow_computation(
computation_impl.ComputationImpl.get_proto(comp))
result = fn(tf.data.Dataset.from_tensor_slices([10, 20]))
self.assertIsInstance(result, tf.Tensor)
self.assertEqual(result, 30)
def _get_embed_tensorflow_computation_succeeds_with_device(self, device):
@computations.tf_computation(tf.int32)
def comp(x):
return tf.add(x, 1)
comp_proto = computation_impl.ComputationImpl.get_proto(comp)
fn = eager_tf_executor.embed_tensorflow_computation(
comp_proto, comp.type_signature, device=device)
result = fn(tf.constant(20))
return result
def test_embed_tensorflow_computation_with_variable_v1(self):
@computations.tf_computation
def comp():
x = tf.Variable(10)
with tf.control_dependencies([x.initializer]):
return tf.add(x, 20)
fn = eager_tf_executor.embed_tensorflow_computation(
computation_impl.ComputationImpl.get_proto(comp))
result = fn()
self.assertIsInstance(result, tf.Tensor)
self.assertEqual(result, 30)
def test_embed_tensorflow_computation_succeeds_with_cpu_or_gpu(
self, device):
@computations.tf_computation(tf.int32)
def comp(x):
return tf.add(x, 1)
comp_proto = computation_impl.ComputationImpl.get_proto(comp)
fn = eager_tf_executor.embed_tensorflow_computation(
comp_proto, comp.type_signature, device='/{}'.format(device))
result = fn(tf.constant(20))
self.assertTrue(result.device.endswith(device))
def test_embed_tensorflow_computation_with_variable_v2(self):
@computations.tf_computation(tf.int32)
def comp(x):
v = tf.Variable(10)
with tf.control_dependencies([v.initializer]):
with tf.control_dependencies([v.assign_add(20)]):
return tf.add(x, v)
fn = eager_tf_executor.embed_tensorflow_computation(
computation_impl.ComputationImpl.get_proto(comp))
result = fn(tf.constant(30))
self.assertIsInstance(result, tf.Tensor)
self.assertEqual(result, 60)
def test_embed_tensorflow_computation_with_tuple_arg_and_result(self):
@computations.tf_computation([('a', tf.int32), ('b', tf.int32)])
def comp(a, b):
return {'sum': a + b}
fn = eager_tf_executor.embed_tensorflow_computation(
computation_impl.ComputationImpl.get_proto(comp))
p = tf.constant(10)
q = tf.constant(20)
result = fn(structure.Struct([('a', p), ('b', q)]))
self.assertIsInstance(result, structure.Struct)
self.assertCountEqual(dir(result), ['sum'])
self.assertIsInstance(result.sum, tf.Tensor)
self.assertEqual(result.sum, 30)
def test_embed_tensorflow_computation_with_float_variables_same_name(self):
@computations.tf_computation
def comp1():
x = tf.Variable(0.5, name='bob')
with tf.control_dependencies([x.initializer]):
return tf.add(x, 0.6)
@computations.tf_computation
def comp2():
x = tf.Variable(0.5, name='bob')
with tf.control_dependencies([x.initializer]):
return tf.add(x, 0.7)
fns = [
eager_tf_executor.embed_tensorflow_computation(
computation_impl.ComputationImpl.get_proto(x))
for x in [comp1, comp2]
]
results = [f() for f in fns]
for res in results:
self.assertIsInstance(res, tf.Tensor)
self.assertAlmostEqual(results[0], 1.1)
self.assertAlmostEqual(results[1], 1.2)
