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