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_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_executor.embed_tensorflow_computation( computation_impl.ComputationImpl.get_proto(comp)) result = fn() self.assertIsInstance(result, tf.Tensor) self.assertEqual(result.numpy(), 1000)
def test_embed_tensorflow_computation_with_int_arg_and_result(self): @computations.tf_computation(tf.int32) def comp(x): return x + 1 fn = eager_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_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.numpy(), 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_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_v1(self): @computations.tf_computation def comp(): x = tf.Variable(10) with tf.control_dependencies([x.initializer]): return tf.add(x, 20) fn = eager_executor.embed_tensorflow_computation( computation_impl.ComputationImpl.get_proto(comp)) result = fn() self.assertIsInstance(result, tf.Tensor) self.assertEqual(result.numpy(), 30)
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_executor.embed_tensorflow_computation( computation_impl.ComputationImpl.get_proto(comp)) result = fn(30) self.assertIsInstance(result, tf.Tensor) self.assertEqual(result.numpy(), 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_executor.embed_tensorflow_computation( computation_impl.ComputationImpl.get_proto(comp)) p = tf.constant(10) q = tf.constant(20) result = fn(anonymous_tuple.AnonymousTuple([('a', p), ('b', q)])) self.assertIsInstance(result, anonymous_tuple.AnonymousTuple) self.assertCountEqual(dir(result), ['sum']) self.assertIsInstance(result.sum, tf.Tensor) self.assertEqual(result.sum.numpy(), 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_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].numpy(), 1.1) self.assertAlmostEqual(results[1].numpy(), 1.2)