def test_tensor(self):
def body():
nonlocal i, s
s = s * 10 + i
i += 1
def set_state(loop_vars):
nonlocal i, s
i, s = loop_vars
i = 0
n = constant_op.constant(5)
s = 0
control_flow.while_stmt(
test=lambda: i < n,
body=body,
get_state=lambda: (i, s),
set_state=set_state,
symbol_names=('i', 's'),
opts={})
self.assertEqual(i, 5)
self.assertEqual(s, 1234)
self.assertOpCreated('StatelessWhile')
def test_tensor_with_side_effecting_condition(self):
v = self.variable('v', 0, dtypes.int32)
def cond():
v.assign(v.read_value() * 10 + i)
return i < n
def body():
nonlocal i
i += 1
def set_state(loop_vars):
nonlocal i
i, = loop_vars
i = 0
n = constant_op.constant(5)
control_flow.while_stmt(test=cond,
body=body,
get_state=lambda: (i, ),
set_state=set_state,
symbol_names=('i', ),
opts={})
self.assertEqual(i, (5, ))
self.assertEqual(v, (12345, ))
self.assertOpCreated('While')
def test_python_for_large_unroll_warning(self):
if not __debug__:
self.skipTest('Feature disabled in optimized mode.')
with test.mock.patch.object(control_flow,
'INEFFICIENT_UNROLL_MIN_ITERATIONS', 10):
with ops.Graph().as_default():
out_capturer = six.StringIO()
with test.mock.patch.object(sys, 'stdout', out_capturer):
with test.mock.patch.object(ag_logging,
'echo_log_to_stdout', True):
def body():
nonlocal i
gen_math_ops.add(i, 1)
i += 1
i = 0
control_flow.while_stmt(test=lambda: i < 100,
body=body,
get_state=None,
set_state=None,
symbol_names=('i', ),
opts={})
self.assertTrue(
re.match(r'.* Large unrolled loop.*Add.*',
out_capturer.getvalue()))
def test_tensor_with_python_state(self):
class MutableObject(object):
field = constant_op.constant(0, dtype=dtypes.int32)
state = MutableObject()
def body():
nonlocal i
state.field = state.field * 10 + i
i += 1
def set_state(loop_vars):
nonlocal i
i, state.field = loop_vars
i = 0
n = constant_op.constant(5)
control_flow.while_stmt(test=lambda: i < n,
body=body,
get_state=lambda: (i, state.field),
set_state=set_state,
symbol_names=('i', 'state.field'),
opts={})
self.assertEqual(i, 5)
self.assertEqual(state.field, 1234)
self.assertOpCreated('StatelessWhile')
def test_tensor_failing_to_determine_placeholder(self):
class UserType:
pass
def body():
nonlocal v
v = UserType()
def set_state(loop_vars):
nonlocal v
v, = loop_vars
v = variable_operators.Undefined('v')
with self.assertRaisesRegex(
ValueError,
re.compile(
'must be defined.*tried to define.*unsupported type',
re.DOTALL)):
control_flow.while_stmt(test=lambda: constant_op.constant(True),
body=body,
get_state=lambda: (v, ),
set_state=set_state,
symbol_names=('v', ),
opts={})
def test_tensor_creating_variable_of_dynamic_shape(self):
def body():
nonlocal i, s
i = array_ops.ones(
[random_ops.random_uniform(minval=1, maxval=4, shape=()), 7])
s = math_ops.reduce_sum(i)
def set_state(loop_vars):
nonlocal i, s
i, s = loop_vars
i = variable_operators.Undefined('i')
s = constant_op.constant(0.0)
control_flow.while_stmt(
test=lambda: math_ops.equal(s, 0),
body=body,
get_state=lambda: (i, s),
set_state=set_state,
symbol_names=('i', 's'),
opts={})
self.assertEqual(i[0][0], 1)
self.assertGreaterEqual(s, 7)
self.assertOpCreated('While') # Not stateless because of the random op.
def test_tensor_creating_variable(self):
def body():
nonlocal i, s
i = constant_op.constant(2)
s = i**5
def set_state(loop_vars):
nonlocal i, s
i, s = loop_vars
i = variable_operators.Undefined('i')
s = constant_op.constant(0)
control_flow.while_stmt(test=lambda: math_ops.equal(s, 0),
body=body,
get_state=lambda: (i, s),
set_state=set_state,
symbol_names=('i', 's'),
opts={})
self.assertEqual(i, 2)
self.assertEqual(s, 32)
self.assertOpCreated('StatelessWhile')
# Check that the temporary staging of the body did not create extra ops.
# Node naming is inconsistent between V1 and V2.
self.assertGraphContains(r'(while/)?pow$', 1)
def test_python_infinite_loop(self):
if __debug__:
with test.mock.patch.object(control_flow, 'PYTHON_MAX_ITERATIONS', 100):
with self.assertRaisesRegexp(errors.ExecutionError, 'iteration limit'):
control_flow.while_stmt(
test=lambda _: True,
body=lambda i: (i + 1,),
init_state=(0,))
def test_python_infinite_loop(self):
if __debug__:
with test.mock.patch.object(control_flow, 'PYTHON_MAX_ITERATIONS',
1000):
with self.assertRaisesRegexp(errors.ExecutionError,
'iteration limit'):
control_flow.while_stmt(test=lambda _: True,
body=lambda i: (i + 1, ),
init_state=(0, ))
def test_python_due_to_hidden_cond_type(self):
n = 5
# TODO(b/124002646): Improve the error message.
with self.assertRaises(Exception):
control_flow.while_stmt(test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
init_state=(constant_op.constant(0),
constant_op.constant(0)),
extra_deps=())
def test_python_while_infinite(self):
if not __debug__:
self.skipTest('Feature disabled in optimized mode.')
with test.mock.patch.object(control_flow, 'PYTHON_MAX_ITERATIONS', 100):
with self.assertRaisesRegexp(ValueError, 'iteration limit'):
control_flow.while_stmt(
test=lambda: True,
body=lambda: None,
get_state=None,
set_state=None,
symbol_names=(),
opts={})
def test_python_long_loop_unroll_warning(self):
if __debug__:
with ops.Graph().as_default():
out_capturer = six.StringIO()
with test.mock.patch.object(sys, 'stdout', out_capturer):
ag_logging.echo_log_to_stdout = True
sys.stdout = out_capturer
control_flow.while_stmt(
test=lambda i, _: i < 10000,
body=lambda i, _: (i + 1, gen_math_ops.add(i, 1),),
init_state=(0, None))
self.assertTrue(re.match(
r'.*ops.*loop.*large.*iterations.*Add.*', out_capturer.getvalue()))
def test_python_infinite_loop(self):
if __debug__:
with test.mock.patch.object(control_flow, 'PYTHON_MAX_ITERATIONS', 100):
with self.assertRaisesRegexp(ValueError, 'iteration limit'):
control_flow.while_stmt(
test=lambda _: True,
body=lambda i: (i + 1,),
get_state=None,
set_state=None,
init_vars=(0,),
basic_symbol_names=('i',),
composite_symbol_names=(),
opts={})
def test_python(self):
n = 5
results = control_flow.while_stmt(
test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
init_state=(0, 0))
self.assertEqual((5, 10), results)
def test_python(self):
n = 5
results = control_flow.while_stmt(test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
init_state=(0, 0),
extra_deps=(n, ))
self.assertEqual((5, 10), results)
def test_tensor_with_python_state(self):
n = constant_op.constant(5)
class MutableObject(object):
field = constant_op.constant(0, dtype=dtypes.int32)
state = MutableObject()
def get_state():
return (state.field, )
def set_state(new_state):
state.field, = new_state
def body(i, s):
state.field += i
return (i + 1, s + i)
s = control_flow.while_stmt(test=lambda i, s: i < n,
body=body,
get_state=get_state,
set_state=set_state,
init_vars=(0, 0))
self.assertEqual(self.evaluate(s), (5, 10))
self.assertEqual(self.evaluate(state.field), 10)
def test_fn():
return control_flow.while_stmt(
test=lambda i: get_and_increment(v) < n,
body=lambda i: (i + 1, ),
get_state=lambda: (),
set_state=lambda _: None,
init_vars=(0, ))
def test_python_long_loop_unroll_warning(self):
if __debug__:
with ops.Graph().as_default():
out_capturer = six.StringIO()
with test.mock.patch.object(sys, 'stdout', out_capturer):
ag_logging.echo_log_to_stdout = True
sys.stdout = out_capturer
control_flow.while_stmt(test=lambda i, _: i < 10000,
body=lambda i, _: (
i + 1,
gen_math_ops.add(i, 1),
),
init_state=(0, None))
self.assertTrue(
re.match(r'.*ops.*loop.*large.*iterations.*Add.*',
out_capturer.getvalue()))
def test_tensor(self):
n = constant_op.constant(5)
results = control_flow.while_stmt(
test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
init_state=(0, 0))
self.assertEqual((5, 10), self.evaluate(results))
def test_python(self):
def body():
nonlocal i, s
s = s * 10 + i
i += 1
i = 0
s = 0
n = 5
control_flow.while_stmt(test=lambda: i < n,
body=body,
get_state=None,
set_state=None,
symbol_names=('i', 's'),
opts={})
self.assertEqual(s, 1234)
def test_tensor(self):
n = constant_op.constant(5)
results = control_flow.while_stmt(test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
get_state=lambda: (),
set_state=lambda _: None,
init_vars=(0, 0))
self.assertEqual((5, 10), self.evaluate(results))
def test_python_with_tensor_state(self):
def body():
nonlocal i, s
s = s * 10 + i
i += 1
i = 0
s = constant_op.constant(0)
n = 5
control_flow.while_stmt(test=lambda: i < n,
body=body,
get_state=None,
set_state=None,
symbol_names=('i', 's'),
opts={})
self.assertEqual(i, 5)
self.assertEqual(s, 1234)
def test_python(self):
n = 5
results = control_flow.while_stmt(test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
get_state=None,
set_state=None,
init_vars=(0, 0))
self.assertEqual((5, 10), results)
def test_tensor(self):
n = constant_op.constant(5)
results = control_flow.while_stmt(
test=lambda i, sum: i < n,
body=lambda i, sum: (i + 1, sum + i,),
init_state=(0, 0),
extra_deps=(n,))
with self.cached_session():
self.assertEqual((5, 10), self.evaluate(results))
def test_tensor(self):
n = constant_op.constant(5)
results = control_flow.while_stmt(
test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i,),
init_state=(0, 0),
extra_deps=(n,))
with self.cached_session() as sess:
self.assertEqual((5, 10), self.evaluate(results))
def test_python_with_tensor_state(self):
n = 5
results = control_flow.while_stmt(test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
init_state=(0,
constant_op.constant(0)))
result_i, result_s = results
self.assertEqual(5, result_i)
self.assertEqual(10, self.evaluate(result_s))
def test_python_with_tensor_state(self):
n = 5
results = control_flow.while_stmt(
test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
init_state=(0, constant_op.constant(0)))
result_i, result_s = results
self.assertEqual(5, result_i)
self.assertEqual(10, self.evaluate(result_s))
def test_tensor_dict_state(self):
n = 5
init_state = {'i': constant_op.constant(0), 'sum': constant_op.constant(0)}
results = control_flow.while_stmt(
test=lambda s: s['i'] < n,
body=lambda s: ({'i': s['i'] + 1, 'sum': s['sum'] + s['i']},),
init_state=(init_state,),
extra_deps=())
with self.cached_session():
self.assertEqual(({'i': 5, 'sum': 10},), self.evaluate(results))
def test_fn():
return control_flow.while_stmt(
test=lambda i: get_and_increment(v) < n,
body=lambda i: (i + 1,),
get_state=lambda: (),
set_state=lambda _: None,
init_vars=(0,),
basic_symbol_names=('i',),
composite_symbol_names=(),
opts={})
def test_tensor_dict_state(self):
n = 5
init_state = {'i': constant_op.constant(0), 'sum': constant_op.constant(0)}
results = control_flow.while_stmt(
test=lambda s: s['i'] < n,
body=lambda s: ({'i': s['i'] + 1, 'sum': s['sum'] + s['i']},),
init_state=(init_state,),
extra_deps=())
with self.cached_session():
self.assertEqual(({'i': 5, 'sum': 10},), self.evaluate(results))
def _fixed_while_loop(self, cond_fn):
def test_():
return cond_fn(s)
def body():
nonlocal s
s += 1
def set_state(loop_vars):
nonlocal s
s, = loop_vars
s = constant_op.constant(0)
control_flow.while_stmt(test=test_,
body=body,
get_state=lambda: (s, ),
set_state=set_state,
symbol_names=('s', ),
opts={})
return s
def test_python(self):
n = 5
results = control_flow.while_stmt(test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
get_state=None,
set_state=None,
init_vars=(0, 0),
basic_symbol_names=('i', 's'),
composite_symbol_names=(),
opts={})
self.assertEqual((5, 10), results)
def test_tensor(self):
n = constant_op.constant(5)
results = control_flow.while_stmt(test=lambda i, s: i < n,
body=lambda i, s: (i + 1, s + i),
get_state=lambda: (),
set_state=lambda _: None,
init_vars=(0, 0),
basic_symbol_names=('i', 's'),
composite_symbol_names=(),
opts={})
self.assertEqual((5, 10), self.evaluate(results))
def _basic_loop(self, init_value, body_fn):
def body():
nonlocal i, s
s = body_fn(i, s)
i += 1
def set_state(loop_vars):
nonlocal i, s
i, s = loop_vars
i = 0
n = constant_op.constant(5)
s = init_value
control_flow.while_stmt(test=lambda: i < n,
body=body,
get_state=lambda: (i, s),
set_state=set_state,
symbol_names=('i', 's'),
opts={})
return s
def test_python_long_loop_unroll_warning(self):
if __debug__:
with test.mock.patch.object(
control_flow, 'INEFFICIENT_UNROLL_MIN_ITERATIONS', 10):
with ops.Graph().as_default():
out_capturer = six.StringIO()
with test.mock.patch.object(sys, 'stdout', out_capturer):
ag_logging.echo_log_to_stdout = True
sys.stdout = out_capturer
control_flow.while_stmt(
test=lambda i, _: i < 100,
body=lambda i, _: (i + 1, gen_math_ops.add(i, 1),),
get_state=None,
set_state=None,
init_vars=(0, None),
basic_symbol_names=('i',),
composite_symbol_names=(),
opts={})
self.assertTrue(re.match(
r'.*ops.*loop.*large.*iterations.*Add.*',
out_capturer.getvalue()))
def test_tensor_creating_dynamic_shape_variable(self):
def body():
nonlocal i, y
i += 1
y = random_ops.random_uniform([i])
def set_state(loop_vars):
nonlocal i, y
i, y = loop_vars
i = constant_op.constant(0)
y = variable_operators.Undefined('y')
control_flow.while_stmt(test=lambda: math_ops.less(i, 3),
body=body,
get_state=lambda: (i, y),
set_state=set_state,
symbol_names=('i', 'y'),
opts={})
self.assertEqual(i, 3)
self.assertLess(y[0], 3)
def test_tensor_creating_dynamic_shape_variable_preserves_shape_invar(self):
def body():
nonlocal i, y
i += 1
y = array_ops.zeros([1])
def set_state(loop_vars):
nonlocal i, y
i, y = loop_vars
i = constant_op.constant(0)
y = variable_operators.Undefined('y')
control_flow.while_stmt(
test=lambda: math_ops.less(i, 3),
body=body,
get_state=lambda: (i, y),
set_state=set_state,
symbol_names=('i', 'y'),
opts={'shape_invariants': ((y, tensor_shape.TensorShape([1])),)})
self.evaluate(y)
def test_fn():
return control_flow.while_stmt(
test=lambda i: get_and_increment(v) < n,
body=lambda i: (i + 1,),
init_state=(0,))
