|
16 | 16 |
|
17 | 17 | from tensorflow.python.eager import def_function |
18 | 18 | from tensorflow.python.framework import dtypes |
| 19 | +from tensorflow.python.framework import errors |
19 | 20 | from tensorflow.python.framework import test_util |
20 | 21 | from tensorflow.python.framework import constant_op |
| 22 | +from tensorflow.python.ops import gen_script_ops |
| 23 | +from tensorflow.python.ops import resource_variable_ops |
21 | 24 | from tensorflow.python.ops import script_ops |
22 | 25 | from tensorflow.python.ops.script_ops import numpy_function |
23 | 26 | from tensorflow.python.platform import test |
@@ -87,5 +90,30 @@ def func_stateful(a, b): |
87 | 90 | 2) # as stateful, func is guaranteed to execute twice |
88 | 91 |
|
89 | 92 |
|
| 93 | +class PyFunctionTest(test.TestCase): |
| 94 | + |
| 95 | + @test_util.run_in_graph_and_eager_modes |
| 96 | + def test_variable_arguments(self): |
| 97 | + |
| 98 | + def plus(a, b): |
| 99 | + return a + b |
| 100 | + |
| 101 | + v1 = resource_variable_ops.ResourceVariable(1) |
| 102 | + self.evaluate(v1.initializer) |
| 103 | + |
| 104 | + actual_result = script_ops.eager_py_func(plus, [v1, 2], dtypes.int32) |
| 105 | + expect_result = constant_op.constant(3, dtypes.int32) |
| 106 | + self.assertAllEqual(actual_result, expect_result) |
| 107 | + |
| 108 | + @test_util.run_in_graph_and_eager_modes |
| 109 | + def test_fail_on_non_utf8_token(self): |
| 110 | + value = constant_op.constant(value=[1, 2]) |
| 111 | + token = b"\xb0" |
| 112 | + data_type = [dtypes.int32] |
| 113 | + with self.assertRaises((errors.InternalError, UnicodeDecodeError)): |
| 114 | + self.evaluate( |
| 115 | + gen_script_ops.py_func(input=[value], token=token, Tout=data_type)) |
| 116 | + |
| 117 | + |
90 | 118 | if __name__ == "__main__": |
91 | 119 | test.main() |
0 commit comments