def test_basic_non_tff_function_as_nondecorator_succeeds(self): def f(): return 5 f_wrapped = computation_wrapper_instances.check_returns_type( f, tf.int32) self.assertEqual(f_wrapped(), 5)
def test_basic_non_tff_function_as_nondecorator_fails(self): def f(): return [5] f_wrapped = computation_wrapper_instances.check_returns_type( f, tf.int32) with self.assertRaises(TypeError): f_wrapped()