def test_for_variable_defined_in_body(self):
def bad_for_loop(n):
for i in range(n):
s = i
return s
node, ctx = self.prepare(bad_for_loop, {})
with self.assertRaises(transformer.AutographParseError):
control_flow.transform(node, ctx)
def test_if_imbalanced_outputs(self):
def test_fn(n):
if n > 0:
b = 4
return b
node, ctx = self.prepare(test_fn, {})
with self.assertRaises(transformer.AutographParseError):
control_flow.transform(node, ctx)
def test_if_imbalanced_outputs(self):
def test_fn(n):
if n > 0:
b = 4
return b
node, ctx = self.prepare(test_fn, {})
with self.assertRaises(transformer.AutographParseError):
control_flow.transform(node, ctx)
def test_while_variable_defined_in_body(self):
def bad_while_loop(n):
while n > 0:
n -= 1
s = n
return s
node, ctx = self.prepare(bad_while_loop, {})
with self.assertRaises(transformer.AutographParseError):
control_flow.transform(node, ctx)
def test_handle_temp_variable(self):
def test_fn_using_temp(x, y, w):
if x < y:
z = x + y
else:
w = 2
tmp = w
z = x - tmp
return z, w
node = self.parse_and_analyze(test_fn_using_temp, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond,
array_ops.ones) as result:
with self.test_session() as sess:
z, w = sess.run(
result.test_fn_using_temp(constant_op.constant(-3),
constant_op.constant(3),
constant_op.constant(3)))
self.assertEqual(0, z)
self.assertEqual(3, w)
z, w = sess.run(
result.test_fn_using_temp(constant_op.constant(3),
constant_op.constant(-3),
constant_op.constant(3)))
self.assertEqual(1, z)
self.assertEqual(2, w)
def test_fn_ignoring_temp(x, y, w):
if x < y:
z = x + y
else:
w = 2
tmp = w
z = x - tmp
return z
node = self.parse_and_analyze(test_fn_ignoring_temp, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond,
array_ops.ones) as result:
with self.test_session() as sess:
z = sess.run(
result.test_fn_ignoring_temp(constant_op.constant(-3),
constant_op.constant(3),
constant_op.constant(3)))
self.assertEqual(0, z)
z = sess.run(
result.test_fn_ignoring_temp(constant_op.constant(3),
constant_op.constant(-3),
constant_op.constant(3)))
self.assertEqual(1, z)
def test_handle_temp_variable(self):
def test_fn_using_temp(x, y, w):
if x < y:
z = x + y
else:
w = 2
tmp = w
z = x - tmp
return z, w
node = self.parse_and_analyze(test_fn_using_temp, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond, array_ops.ones) as result:
with self.test_session() as sess:
z, w = sess.run(
result.test_fn_using_temp(
constant_op.constant(-3), constant_op.constant(3),
constant_op.constant(3)))
self.assertEqual(0, z)
self.assertEqual(3, w)
z, w = sess.run(
result.test_fn_using_temp(
constant_op.constant(3), constant_op.constant(-3),
constant_op.constant(3)))
self.assertEqual(1, z)
self.assertEqual(2, w)
def test_fn_ignoring_temp(x, y, w):
if x < y:
z = x + y
else:
w = 2
tmp = w
z = x - tmp
return z
node = self.parse_and_analyze(test_fn_ignoring_temp, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond, array_ops.ones) as result:
with self.test_session() as sess:
z = sess.run(
result.test_fn_ignoring_temp(
constant_op.constant(-3), constant_op.constant(3),
constant_op.constant(3)))
self.assertEqual(0, z)
z = sess.run(
result.test_fn_ignoring_temp(
constant_op.constant(3), constant_op.constant(-3),
constant_op.constant(3)))
self.assertEqual(1, z)
def test_simple_for(self):
def test_fn(l):
s1 = 0
s2 = 0
for e in l:
s1 += e
s2 += e * e
return s1, s2
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node) as result:
with self.test_session() as sess:
l = [1, 2, 3]
self.assertEqual(
test_fn(l),
sess.run(result.test_fn(constant_op.constant(l))))
l = []
self.assertEqual(
test_fn(l),
sess.run(
result.test_fn(
constant_op.constant(l,
shape=(0, ),
dtype=dtypes.int32))))
def test_while_single_var(self):
def test_fn(n):
while n > 0:
n -= 1
return n
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.while_loop) as result:
with self.test_session() as sess:
self.assertEqual(
0, sess.run(result.test_fn(constant_op.constant(5))))
def test_if_single_var(self):
def test_fn(n):
if n > 0:
n = -n
return n
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond) as result:
with self.test_session() as sess:
self.assertEqual(
-1, sess.run(result.test_fn(constant_op.constant(1))))
def test_if_single_var(self):
def test_fn(n):
if n > 0:
n = -n
return n
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond) as result:
with self.test_session() as sess:
self.assertEqual(-1, sess.run(result.test_fn(constant_op.constant(1))))
def test_while_single_var(self):
def test_fn(n):
while n > 0:
n -= 1
return n
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.while_loop) as result:
with self.test_session() as sess:
self.assertEqual(0, sess.run(result.test_fn(constant_op.constant(5))))
def test_imbalanced_aliasing(self):
def test_fn(n):
if n > 0:
n = 3
return n
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond) as result:
with self.test_session() as sess:
self.assertEqual(3, sess.run(result.test_fn(constant_op.constant(2))))
self.assertEqual(-3, sess.run(result.test_fn(constant_op.constant(-3))))
def test_ignore_unread_variable(self):
def test_fn(n):
b = 3 # pylint: disable=unused-variable
if n > 0:
b = 4
return n
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond, array_ops.ones) as result:
with self.test_session() as sess:
self.assertEqual(3, sess.run(result.test_fn(constant_op.constant(3))))
self.assertEqual(-3, sess.run(result.test_fn(constant_op.constant(-3))))
def test_simple_while(self):
def test_fn(n):
i = 0
s = 0
while i < n:
s += i
i += 1
return s, i, n
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.while_loop) as result:
with self.test_session() as sess:
self.assertEqual((10, 5, 5),
sess.run(result.test_fn(constant_op.constant(5))))
def test_simple_while(self):
def test_fn(n):
i = 0
s = 0
while i < n:
s += i
i += 1
return s, i, n
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.while_loop) as result:
with self.test_session() as sess:
self.assertEqual(
(10, 5, 5),
sess.run(result.test_fn(constant_op.constant(5))))
def test_simple_if(self):
def test_fn(n):
a = 0
b = 0
if n > 0:
a = -n
else:
b = 2 * n
return a, b
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond) as result:
with self.test_session() as sess:
self.assertEqual((-1, 0),
sess.run(result.test_fn(constant_op.constant(1))))
self.assertEqual((0, -2),
sess.run(result.test_fn(constant_op.constant(-1))))
def test_simple_if(self):
def test_fn(n):
a = 0
b = 0
if n > 0:
a = -n
else:
b = 2 * n
return a, b
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node, control_flow_ops.cond) as result:
with self.test_session() as sess:
self.assertEqual(
(-1, 0), sess.run(result.test_fn(constant_op.constant(1))))
self.assertEqual(
(0, -2),
sess.run(result.test_fn(constant_op.constant(-1))))
def test_for_iterated_expression(self):
eval_count = [0]
def count_evals(x):
eval_count[0] += 1
return x
def test_fn(n):
s = 0
for e in count_evals(range(n)):
s += e
return s
ns = {'count_evals': count_evals}
node, ctx = self.prepare(test_fn, ns)
node = control_flow.transform(node, ctx)
with self.compiled(node, ns) as result:
self.assertEqual(result.test_fn(5), 10)
self.assertEqual(eval_count[0], 1)
def test_for_iterated_expression(self):
eval_count = [0]
def count_evals(x):
eval_count[0] += 1
return x
def test_fn(n):
s = 0
for e in count_evals(range(n)):
s += e
return s
ns = {'count_evals': count_evals}
node, ctx = self.prepare(test_fn, ns)
node = control_flow.transform(node, ctx)
with self.compiled(node, ns) as result:
self.assertEqual(result.test_fn(5), 10)
self.assertEqual(eval_count[0], 1)
def test_for_with_iterated_expression(self):
eval_count = [0]
def count_evals(x):
eval_count[0] += 1
return x
def test_fn(n):
s = 0
for e in count_evals(range(n)):
s += e
return s
node = self.parse_and_analyze(test_fn, {'count_evals': count_evals})
node = control_flow.transform(node, self.ctx)
with self.compiled(node) as result:
result.count_evals = count_evals
self.assertEqual(test_fn(5), result.test_fn(5))
# count_evals ran twice, once for test_fn and another for result.test_fn
self.assertEqual(eval_count[0], 2)
def test_for_with_iterated_expression(self):
eval_count = [0]
def count_evals(x):
eval_count[0] += 1
return x
def test_fn(n):
s = 0
for e in count_evals(range(n)):
s += e
return s
node = self.parse_and_analyze(test_fn, {'count_evals': count_evals})
node = control_flow.transform(node, self.ctx)
with self.compiled(node) as result:
result.count_evals = count_evals
self.assertEqual(test_fn(5), result.test_fn(5))
# count_evals ran twice, once for test_fn and another for result.test_fn
self.assertEqual(eval_count[0], 2)
def test_for_single_var(self):
def test_fn(l):
s = 0
for e in l:
s += e
return s
node = self.parse_and_analyze(test_fn, {})
node = control_flow.transform(node, self.ctx)
with self.compiled(node) as result:
with self.test_session() as sess:
l = [1, 2, 3]
self.assertEqual(
test_fn(l), sess.run(result.test_fn(constant_op.constant(l))))
l = []
self.assertEqual(
test_fn(l),
sess.run(
result.test_fn(
constant_op.constant(l, shape=(0,), dtype=dtypes.int32))))
def node_to_graph(node, ctx, nocompile_decorators):
"""Convert Python code to equivalent TF graph mode code.
Args:
node: A Python AST node representing the code to convert.
ctx: An EntityContext object.
nocompile_decorators: A tuple containing decorators to be stripped from
functions during conversion.
Returns:
A tuple (node, deps):
* node: A Python ast node, representing the converted code.
* deps: A set of strings, the fully qualified names of entity
dependencies that this node has.
"""
# TODO(mdan): Verify arguments for correctness.
# TODO(mdan): Factor out common elements.
# These include:
# * code move between blocks
# * visiting blocks in transformers
# Certain steps, especially canonicalization, insert new symbols into the
# tree, which must be accounted. Although less efficient, it is most robust
# to re-run the analysis.
node = _static_analysis_pass(node, ctx)
# TODO(mdan): Clean this up.
# Some intermediate analyses are not required, and some comments got orphaned.
# Past this point, line numbers are no longer accurate so we ignore the
# source.
# TODO(mdan): Is it feasible to reconstruct intermediate source code?
ctx.source_code = None
node = ifexp.transform(node, ctx)
node, deps = decorators.transform(node, nocompile_decorators)
node = break_statements.transform(node, ctx)
node = asserts.transform(node, ctx)
# Note: sequencing continue canonicalization before for loop one avoids
# dealing with the extra loop increment operation that the for
# canonicalization creates.
node = continue_statements.transform(node, ctx)
ctx.namespace['len'] = len
node = _static_analysis_pass(node, ctx)
node = single_return.transform(node, ctx)
node = _static_analysis_pass(node, ctx)
node = lists.transform(node, ctx)
node = builtin_functions.transform(node, ctx)
node = _static_analysis_pass(node, ctx)
node = call_trees.transform(node, ctx, config.DEFAULT_UNCOMPILED_MODULES,
nocompile_decorators)
node = control_flow.transform(node, ctx)
# control_flow may create new symbols and change scopes.
node = _static_analysis_pass(node, ctx)
node = logical_expressions.transform(node, ctx)
node = side_effect_guards.transform(node, ctx)
node = name_scopes.transform(node, ctx)
return node, deps
def node_to_graph(node, ctx, nocompile_decorators):
"""Convert Python code to equivalent TF graph mode code.
Args:
node: A Python AST node representing the code to convert.
ctx: An EntityContext object.
nocompile_decorators: A tuple containing decorators to be stripped from
functions during conversion.
Returns:
A tuple (node, deps):
* node: A Python ast node, representing the converted code.
* deps: A set of strings, the fully qualified names of entity
dependencies that this node has.
"""
# TODO(mdan): Verify arguments for correctness.
# TODO(mdan): Factor out common elements.
# These include:
# * code move between blocks
# * visiting blocks in transformers
# Certain steps, especially canonicalization, insert new symbols into the
# tree, which must be accounted. Although less efficient, it is most robust
# to re-run the analysis.
node = _static_analysis_pass(node, ctx)
# TODO(mdan): Clean this up.
# Some intermediate analyses are not required, and some comments got orphaned.
# Past this point, line numbers are no longer accurate so we ignore the
# source.
# TODO(mdan): Is it feasible to reconstruct intermediate source code?
ctx.source_code = None
node = ifexp.transform(node, ctx)
node, deps = decorators.transform(node, nocompile_decorators)
node = break_statements.transform(node, ctx)
node = asserts.transform(node, ctx)
# Note: sequencing continue canonicalization before for loop one avoids
# dealing with the extra loop increment operation that the for
# canonicalization creates.
node = continue_statements.transform(node, ctx)
ctx.namespace['len'] = len
node = _static_analysis_pass(node, ctx)
node = single_return.transform(node, ctx)
node = _static_analysis_pass(node, ctx)
node = lists.transform(node, ctx)
node = builtin_functions.transform(node, ctx)
node = _static_analysis_pass(node, ctx)
node = call_trees.transform(node, ctx, config.DEFAULT_UNCOMPILED_MODULES,
nocompile_decorators)
node = control_flow.transform(node, ctx)
# control_flow may create new symbols and change scopes.
node = _static_analysis_pass(node, ctx)
node = logical_expressions.transform(node, ctx)
node = side_effect_guards.transform(node, ctx)
node = name_scopes.transform(node, ctx)
return node, deps