1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
|
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for wrap_py_func module."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensorflow.python.autograph.utils import py_func
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import test
class PyFuncTest(test.TestCase):
def test_wrap_py_func_simple(self):
def test_fn(a, b, c):
return a + b + c
with self.cached_session() as sess:
result = py_func.wrap_py_func(test_fn, dtypes.int64,
(1, constant_op.constant(1), 1))
self.assertEqual(3, sess.run(result))
result = py_func.wrap_py_func(test_fn, dtypes.int64, (1, 1, 1))
self.assertEqual(3, sess.run(result))
result = py_func.wrap_py_func(
test_fn, dtypes.int64,
(constant_op.constant(1), 1, constant_op.constant(1)))
self.assertEqual(3, sess.run(result))
def test_wrap_py_func_complex_args(self):
class TestClass(object):
def __init__(self):
self.foo = 5
def test_fn(a, b):
return a * b.foo
with self.cached_session() as sess:
result = py_func.wrap_py_func(test_fn, dtypes.int64, (7, TestClass()))
self.assertEqual(35, sess.run(result))
result = py_func.wrap_py_func(test_fn, dtypes.int64,
(constant_op.constant(7), TestClass()))
self.assertEqual(35, sess.run(result))
def test_wrap_py_func_kwargs(self):
class TestClass(object):
def __init__(self, foo):
self.foo = foo
def test_fn(a, b, c, d):
return a * b.foo + c * d.foo
with self.cached_session() as sess:
result = py_func.wrap_py_func(test_fn, dtypes.int64, (7, TestClass(5)), {
'c': 11,
'd': TestClass(13)
})
self.assertEqual(178, sess.run(result))
result = py_func.wrap_py_func(test_fn, dtypes.int64,
(constant_op.constant(7), TestClass(5)), {
'c': constant_op.constant(11),
'd': TestClass(13)
})
self.assertEqual(178, sess.run(result))
def test_wrap_py_func_dummy_return(self):
side_counter = [0]
def test_fn(_):
side_counter[0] += 1
with self.cached_session() as sess:
result = py_func.wrap_py_func(test_fn, None, (5,), use_dummy_return=True)
self.assertEqual(1, sess.run(result))
self.assertEqual([1], side_counter)
result = py_func.wrap_py_func(
test_fn, None, (constant_op.constant(5),), use_dummy_return=True)
self.assertEqual(1, sess.run(result))
self.assertEqual([2], side_counter)
if __name__ == '__main__':
test.main()
|
