def test_list_pop(self):
def test_fn():
l = [1, 2, 3]
utils.set_element_type(l, dtypes.int32, ())
s = l.pop()
return s, l
node = self.parse_and_analyze(
test_fn,
{
'utils': utils,
'dtypes': dtypes
},
include_type_analysis=True,
)
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
result.utils = utils
result.dtypes = dtypes
with self.test_session() as sess:
ts, tl = result.test_fn()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(sess.run(r), [1, 2])
self.assertAllEqual(sess.run(ts), 3)
def test_list_pop(self):
def test_fn():
l = [1, 2, 3]
utils.set_element_type(l, dtypes.int32, ())
s = l.pop()
return s, l
node = self.parse_and_analyze(
test_fn,
{
'utils': utils,
'dtypes': dtypes
},
include_type_analysis=True,
)
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
result.utils = utils
result.dtypes = dtypes
with self.test_session() as sess:
ts, tl = result.test_fn()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(sess.run(r), [1, 2])
self.assertAllEqual(sess.run(ts), 3)
def test_empty_list(self):
def test_fn():
return []
node = self.parse_and_analyze(test_fn, {})
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
tl = result.test_fn()
# Empty tensor lists cannot be evaluated or stacked.
self.assertTrue(isinstance(tl, ops.Tensor))
self.assertEqual(tl.dtype, dtypes.variant)
def test_initialized_list(self):
def test_fn():
return [1, 2, 3]
node = self.parse_and_analyze(test_fn, {})
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
with self.test_session() as sess:
tl = result.test_fn()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(sess.run(r), [1, 2, 3])
def test_initialized_list(self):
def test_fn():
return [1, 2, 3]
node = self.parse_and_analyze(test_fn, {})
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
with self.test_session() as sess:
tl = result.test_fn()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(sess.run(r), [1, 2, 3])
def test_empty_list(self):
def test_fn():
return []
node = self.parse_and_analyze(test_fn, {})
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
tl = result.test_fn()
# Empty tensor lists cannot be evaluated or stacked.
self.assertTrue(isinstance(tl, ops.Tensor))
self.assertEqual(tl.dtype, dtypes.variant)
def test_double_list_pop(self):
def test_fn(l):
s = l.pop().pop()
return s
node = self.parse_and_analyze(test_fn, {})
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
test_input = [1, 2, [1, 2, 3]]
# TODO(mdan): Pass a list of lists of tensor when we fully support that.
# For now, we just pass a regular Python list of lists just to verify that
# the two pop calls are sequenced properly.
self.assertAllEqual(result.test_fn(test_input), 3)
def test_list_append(self):
def test_fn():
l = [1]
l.append(2)
l.append(3)
return l
node = self.parse_and_analyze(test_fn, {})
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
with self.test_session() as sess:
tl = result.test_fn()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(sess.run(r), [1, 2, 3])
def test_double_list_pop(self):
def test_fn(l):
s = l.pop().pop()
return s
node = self.parse_and_analyze(test_fn, {})
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
test_input = [1, 2, [1, 2, 3]]
# TODO(mdan): Pass a list of lists of tensor when we fully support that.
# For now, we just pass a regular Python list of lists just to verify that
# the two pop calls are sequenced properly.
self.assertAllEqual(result.test_fn(test_input), 3)
def test_list_stack(self):
def test_fn():
l = [1, 2, 3]
return tf.stack(l)
node, ctx = self.prepare(test_fn, {})
def_, = anno.getanno(node.body[0].body[0].targets[0],
anno.Static.ORIG_DEFINITIONS)
def_.directives[directives.set_element_type] = {
'dtype': parser.parse_expression('tf.int32')
}
node = lists.transform(node, ctx)
with self.compiled(node, {}, array_ops.stack, dtypes.int32) as result:
with self.test_session() as sess:
self.assertAllEqual(sess.run(result.test_fn()), [1, 2, 3])
def test_list_append(self):
def test_fn():
l = [1]
l.append(2)
l.append(3)
return l
node = self.parse_and_analyze(test_fn, {})
node = lists.transform(node, self.ctx)
with self.compiled(node) as result:
with self.test_session() as sess:
tl = result.test_fn()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(sess.run(r), [1, 2, 3])
def test_list_stack(self):
def test_fn():
l = [1, 2, 3]
return tf.stack(l)
node, ctx = self.prepare(test_fn, {})
def_, = anno.getanno(node.body[0].targets[0],
anno.Static.ORIG_DEFINITIONS)
def_.directives[directives.set_element_type] = {
'dtype': parser.parse_expression('tf.int32')
}
node = lists.transform(node, ctx)
with self.compiled(node, {}, array_ops.stack, dtypes.int32) as result:
with self.test_session() as sess:
self.assertAllEqual(sess.run(result.test_fn()), [1, 2, 3])
def test_empty_annotated_list(self):
def test_fn():
l = []
utils.set_element_type(l, dtypes.int32)
l.append(1)
return l
node = self.parse_and_analyze(test_fn, {'dtypes': dtypes, 'utils': utils})
node = lists.transform(node, self.ctx)
with self.compiled(node, tensor_array_ops.TensorArray,
dtypes.int32) as result:
# TODO(mdan): Attach these additional modules automatically.
result.utils = utils
result.dtypes = dtypes
with self.test_session() as sess:
self.assertAllEqual([1], sess.run(result.test_fn().stack()))
def test_empty_annotated_list(self):
def test_fn():
l = []
utils.set_element_type(l, dtypes.int32)
l.append(1)
return l
node = self.parse_and_analyze(test_fn, {'dtypes': dtypes, 'utils': utils})
node = lists.transform(node, self.ctx)
with self.compiled(node, tensor_array_ops.TensorArray,
dtypes.int32) as result:
# TODO(mdan): Attach these additional modules automatically.
result.utils = utils
result.dtypes = dtypes
with self.test_session() as sess:
self.assertEqual(test_fn(), sess.run(result.test_fn().stack()))
def test_list_pop(self):
def test_fn():
l = [1, 2, 3]
s = l.pop()
return s, l
node, ctx = self.prepare(test_fn, {})
def_, = anno.getanno(node.body[0].body[0].targets[0],
anno.Static.ORIG_DEFINITIONS)
def_.directives[directives.set_element_type] = {
'dtype': parser.parse_expression('tf.int32'),
'shape': parser.parse_expression('()'),
}
node = lists.transform(node, ctx)
with self.compiled(node, {}, dtypes.int32) as result:
with self.test_session() as sess:
ts, tl = result.test_fn()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(sess.run(r), [1, 2])
self.assertAllEqual(sess.run(ts), 3)
def test_empty_annotated_lists_list_unpacked(self):
def test_fn():
[l, m] = [], []
utils.set_element_type(l, dtypes.int32)
utils.set_element_type(m, dtypes.int32)
l.append(1)
m.append(2)
return l, m
node = self.parse_and_analyze(test_fn, {'dtypes': dtypes, 'utils': utils})
node = lists.transform(node, self.ctx)
with self.compiled(node, tensor_array_ops.TensorArray,
dtypes.int32) as result:
result.utils = utils
result.dtypes = dtypes
with self.test_session() as sess:
res_l, res_m = result.test_fn()
self.assertEqual([1], sess.run(res_l.stack()))
self.assertEqual([2], sess.run(res_m.stack()))
def test_list_pop(self):
def test_fn():
l = [1, 2, 3]
s = l.pop()
return s, l
node, ctx = self.prepare(test_fn, {})
def_, = anno.getanno(node.body[0].body[0].targets[0],
anno.Static.ORIG_DEFINITIONS)
def_.directives[directives.set_element_type] = {
'dtype': parser.parse_expression('tf.int32'),
'shape': parser.parse_expression('()'),
}
node = lists.transform(node, ctx)
with self.compiled(node, {}, dtypes.int32) as result:
with self.test_session() as sess:
ts, tl = result.test_fn()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(sess.run(r), [1, 2])
self.assertAllEqual(sess.run(ts), 3)
def test_list_stack(self):
tf = None # Will be replaced with a mock.
def test_fn():
l = [1, 2, 3]
utils.set_element_type(l, dtypes.int32)
return tf.stack(l)
node = self.parse_and_analyze(
test_fn,
{
'utils': utils,
'dtypes': dtypes
},
include_type_analysis=True,
)
node = lists.transform(node, self.ctx)
with self.compiled(node, array_ops.stack, dtypes.int32) as result:
result.utils = utils
result.dtypes = dtypes
with self.test_session() as sess:
self.assertAllEqual(sess.run(result.test_fn()), [1, 2, 3])
def test_list_stack(self):
tf = None # Will be replaced with a mock.
def test_fn():
l = [1, 2, 3]
utils.set_element_type(l, dtypes.int32)
return tf.stack(l)
node = self.parse_and_analyze(
test_fn,
{
'utils': utils,
'dtypes': dtypes
},
include_type_analysis=True,
)
node = lists.transform(node, self.ctx)
with self.compiled(node, array_ops.stack, dtypes.int32) as result:
result.utils = utils
result.dtypes = dtypes
with self.test_session() as sess:
self.assertAllEqual(sess.run(result.test_fn()), [1, 2, 3])
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
