def test_wrap_py_func_simple(self):
def test_fn(a, b, c):
return a + b + c
with self.test_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.test_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_dummy_return(self):
side_counter = [0]
def test_fn(_):
side_counter[0] += 1
with self.test_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)
def test_wrap_py_func_simple(self):
def test_fn(a, b, c):
return a + b + c
with self.test_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 py_func_wrapper(*args, **kwargs):
if kwargs:
raise NotImplementedError(
'RunMode.PY_FUNC does not yet support kwargs')
# TODO(mdan): Add support for kwargs.
return py_func.wrap_py_func(
f, return_dtypes, args, use_dummy_return=not return_dtypes)
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.test_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 py_func_wrapper(*args, **kwargs):
if kwargs:
raise NotImplementedError(
'RunMode.PY_FUNC does not yet support kwargs')
# TODO(mdan): Add support for kwargs.
return py_func.wrap_py_func(
f, return_dtypes, args, kwargs, use_dummy_return=not return_dtypes)
def test_wrap_py_func_dummy_return(self):
side_counter = [0]
def test_fn(_):
side_counter[0] += 1
with self.test_session() as sess:
self.assertEqual(
1, sess.run(py_func.wrap_py_func(test_fn, None, (5, ), True)))
self.assertEqual([1], side_counter)
self.assertEqual(
1,
sess.run(
py_func.wrap_py_func(test_fn, None,
(constant_op.constant(5), ), True)))
self.assertEqual([2], side_counter)
def test_wrap_py_func_simple(self):
def test_fn(a, b, c):
return a + b + c
with self.test_session() as sess:
tensor_1 = constant_op.constant(1)
self.assertEqual(3,
sess.run(
py_func.wrap_py_func(test_fn, dtypes.int64,
(1, tensor_1, 1))))
self.assertEqual(3,
sess.run(
py_func.wrap_py_func(test_fn, dtypes.int64,
(1, 1, 1))))
self.assertEqual(3,
sess.run(
py_func.wrap_py_func(test_fn, dtypes.int64,
(tensor_1, 1, tensor_1))))
def test_wrap_py_func_simple(self):
def test_fn(a, b, c):
return a + b + c
with self.test_session() as sess:
tensor_1 = constant_op.constant(1)
self.assertEqual(
3,
sess.run(
py_func.wrap_py_func(test_fn, dtypes.int64,
(1, tensor_1, 1))))
self.assertEqual(
3,
sess.run(py_func.wrap_py_func(test_fn, dtypes.int64,
(1, 1, 1))))
self.assertEqual(
3,
sess.run(
py_func.wrap_py_func(test_fn, dtypes.int64,
(tensor_1, 1, tensor_1))))
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.test_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_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.test_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 dynamic_print(*values):
"""Implementartion of print using dynamic dispatch.
The function attempts to use tf.Print if all the values are compatible.
Otherwise, it will fall back to py_func.
Args:
*values: values to print
Returns:
A dummy value indicating the print completed. If tf.
"""
if all(map(is_tf_print_compatible, values)):
return logging_ops.Print(1, values)
return py_func.wrap_py_func(print, None, values, use_dummy_return=True)
def call_print(*values):
"""Compiled counterpart of the print builtin.
The function attempts to use tf.Print if all the values are compatible.
Otherwise, it will fall back to py_func.
Args:
*values: values to print
Returns:
A dummy value indicating the print completed. If tf.
"""
if all(map(is_tf_print_compatible, values)):
return logging_ops.Print(1, values)
return py_func.wrap_py_func(print, None, values, use_dummy_return=True)
def call_print(*values):
"""Compiled counterpart of the print builtin.
The function attempts to use tf.Print if all the values are compatible.
Otherwise, it will fall back to py_func.
Args:
*values: values to print
Returns:
A dummy value indicating the print completed. If tf.
"""
if all(map(is_tf_print_compatible, values)):
return logging_ops.Print(1, values)
return py_func.wrap_py_func(print, None, values, use_dummy_return=True)
def dynamic_print(*values):
"""Implementartion of print using dynamic dispatch.
The function attempts to use tf.Print if all the values are compatible.
Otherwise, it will fall back to py_func.
Args:
*values: values to print
Returns:
A dummy value indicating the print completed. If tf.
"""
if all(map(is_tf_print_compatible, values)):
return logging_ops.Print(1, values)
return py_func.wrap_py_func(print, None, values, use_dummy_return=True)