chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
License Check / license-check (push) Failing after 2s

This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
@@ -0,0 +1,402 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
py_library(
name = "slices",
srcs = ["slices.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/lang:directives",
"//tensorflow/python/autograph/pyct:templates",
"@pypi//gast",
],
)
py_library(
name = "continue_statements",
srcs = ["continue_statements.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:templates",
"//tensorflow/python/autograph/pyct/static_analysis:activity",
"//tensorflow/python/autograph/pyct/static_analysis:annos",
],
)
py_library(
name = "logical_expressions",
srcs = ["logical_expressions.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:templates",
"@pypi//gast",
],
)
py_library(
name = "break_statements",
srcs = ["break_statements.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:templates",
"//tensorflow/python/autograph/pyct/static_analysis:activity",
"//tensorflow/python/autograph/pyct/static_analysis:annos",
],
)
py_library(
name = "conditional_expressions",
srcs = ["conditional_expressions.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:templates",
"@pypi//gast",
],
)
py_library(
name = "call_trees",
srcs = ["call_trees.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:templates",
"//tensorflow/python/autograph/utils:ag_logging",
"@pypi//gast",
],
)
py_library(
name = "directives",
srcs = ["directives.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/lang:directives",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/util:tf_inspect",
"@pypi//gast",
],
)
py_library(
name = "variables",
srcs = ["variables.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:templates",
"@pypi//gast",
],
)
py_library(
name = "control_flow",
srcs = ["control_flow.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/lang:directives",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:origin_info",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:templates",
"//tensorflow/python/autograph/pyct/static_analysis:activity",
"//tensorflow/python/autograph/pyct/static_analysis:annos",
"//tensorflow/python/autograph/pyct/static_analysis:liveness",
"//tensorflow/python/autograph/pyct/static_analysis:reaching_definitions",
"//tensorflow/python/autograph/pyct/static_analysis:reaching_fndefs",
"@pypi//gast",
],
)
py_library(
name = "functions",
srcs = ["functions.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:templates",
"//tensorflow/python/autograph/pyct/static_analysis:activity",
"//tensorflow/python/autograph/pyct/static_analysis:annos",
"@pypi//gast",
],
)
py_library(
name = "lists",
srcs = ["lists.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/lang:directives",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:templates",
"//tensorflow/python/autograph/pyct/static_analysis:activity",
"//tensorflow/python/autograph/pyct/static_analysis:annos",
"@pypi//gast",
],
)
py_library(
name = "return_statements",
srcs = ["return_statements.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:parser",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:templates",
"//tensorflow/python/autograph/pyct/static_analysis:activity",
"//tensorflow/python/autograph/pyct/static_analysis:annos",
"@pypi//gast",
],
)
py_library(
name = "asserts",
srcs = ["asserts.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/pyct:templates",
"@pypi//gast",
],
)
py_test(
name = "asserts_test",
srcs = ["asserts_test.py"],
strict_deps = True,
deps = [
":asserts",
":functions",
":return_statements",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "break_statements_test",
srcs = ["break_statements_test.py"],
strict_deps = True,
deps = [
":break_statements",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "call_trees_test",
srcs = ["call_trees_test.py"],
strict_deps = True,
deps = [
":call_trees",
":functions",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "conditional_expressions_test",
srcs = ["conditional_expressions_test.py"],
strict_deps = True,
deps = [
":conditional_expressions",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "continue_statements_test",
srcs = ["continue_statements_test.py"],
strict_deps = True,
deps = [
":continue_statements",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "control_flow_test",
srcs = ["control_flow_test.py"],
strict_deps = True,
deps = [
":break_statements",
":continue_statements",
":control_flow",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
],
)
py_test(
name = "directives_test",
srcs = ["directives_test.py"],
strict_deps = True,
deps = [
":directives",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/autograph/lang:directives",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "functions_test",
srcs = ["functions_test.py"],
strict_deps = True,
deps = [
":functions",
":return_statements",
#internal proto upb dep
"//tensorflow/python/autograph/core:ag_ctx",
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/autograph/impl:api",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "lists_test",
srcs = ["lists_test.py"],
strict_deps = True,
deps = [
":directives",
":lists",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/autograph/lang:directives",
"//tensorflow/python/autograph/lang:special_functions",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "logical_expressions_test",
srcs = ["logical_expressions_test.py"],
strict_deps = True,
deps = [
":logical_expressions",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "return_statements_test",
srcs = ["return_statements_test.py"],
strict_deps = True,
deps = [
":functions",
":return_statements",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "slices_test",
srcs = ["slices_test.py"],
strict_deps = True,
deps = [
":directives",
":slices",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/autograph/lang:directives",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/platform:client_testlib",
],
)
py_test(
name = "variables_test",
srcs = ["variables_test.py"],
strict_deps = True,
deps = [
":variables",
#internal proto upb dep
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/platform:client_testlib",
],
)
@@ -0,0 +1,48 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Converts assert statements to their corresponding TF calls."""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import templates
class AssertTransformer(converter.Base):
"""Transforms Assert nodes to Call so they can be handled as functions."""
def visit_Assert(self, node):
self.generic_visit(node)
# Note: The lone tf.Assert call will be wrapped with control_dependencies
# by side_effect_guards.
template = """
ag__.assert_stmt(test, lambda: msg)
"""
if node.msg is None:
return templates.replace(
template,
test=node.test,
msg=gast.Constant('Assertion error', kind=None))
elif isinstance(node.msg, gast.Constant):
return templates.replace(template, test=node.test, msg=node.msg)
else:
raise NotImplementedError('can only convert string messages for now.')
def transform(node, ctx):
node = AssertTransformer(ctx).visit(node)
return node
@@ -0,0 +1,42 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for asserts module."""
from tensorflow.python.autograph.converters import asserts
from tensorflow.python.autograph.converters import functions
from tensorflow.python.autograph.converters import return_statements
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors_impl
from tensorflow.python.platform import test
class AssertsTest(converter_testing.TestCase):
def test_basic(self):
def f(a):
assert a, 'testmsg'
return a
tr = self.transform(f, (functions, asserts, return_statements))
op = tr(constant_op.constant(False))
with self.assertRaisesRegex(errors_impl.InvalidArgumentError, 'testmsg'):
self.evaluate(op)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,185 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Lowers break statements to conditionals."""
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct.static_analysis import activity
from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
class _Break(object):
def __init__(self):
self.used = False
self.control_var_name = None
def __repr__(self):
return 'used: %s, var: %s' % (self.used, self.control_var_name)
class BreakTransformer(converter.Base):
"""Canonicalizes break statements into additional conditionals."""
def visit_Break(self, node):
self.state[_Break].used = True
var_name = self.state[_Break].control_var_name
# TODO(mdan): This will fail when expanded inside a top-level else block.
template = """
var_name = True
continue
"""
return templates.replace(template, var_name=var_name)
def _guard_if_present(self, block, var_name):
"""Prevents the block from executing if var_name is set."""
if not block:
return block
template = """
if not var_name:
block
"""
node = templates.replace(
template,
var_name=var_name,
block=block)
return node
def _process_body(self, nodes, break_var):
self.state[_Break].enter()
self.state[_Break].control_var_name = break_var
nodes = self.visit_block(nodes)
break_used = self.state[_Break].used
self.state[_Break].exit()
return nodes, break_used
def visit_While(self, node):
original_node = node
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.ctx.namer.new_symbol('break_', scope.referenced)
node.test = self.visit(node.test)
node.body, break_used = self._process_body(node.body, break_var)
# A break in the else clause applies to the containing scope.
node.orelse = self.visit_block(node.orelse)
if not break_used:
template = """
while test:
body
orelse
"""
node = templates.replace(
template, test=node.test, body=node.body, orelse=node.orelse)
new_while_node = node[0]
anno.copyanno(original_node, new_while_node, anno.Basic.DIRECTIVES)
return node
# Python's else clause only triggers if the loop exited cleanly (e.g.
# break did not trigger).
guarded_orelse = self._guard_if_present(node.orelse, break_var)
template = """
var_name = False
while not var_name and test:
body
orelse
"""
node = templates.replace(
template,
var_name=break_var,
test=node.test,
body=node.body,
orelse=guarded_orelse)
new_while_node = node[1]
anno.copyanno(original_node, new_while_node, anno.Basic.DIRECTIVES)
return node
def visit_For(self, node):
original_node = node
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.ctx.namer.new_symbol('break_', scope.referenced)
node.target = self.visit(node.target)
node.iter = self.visit(node.iter)
node.body, break_used = self._process_body(node.body, break_var)
# A break in the else clause applies to the containing scope.
node.orelse = self.visit_block(node.orelse)
if not break_used:
template = """
for target in iter_:
body
orelse
"""
node = templates.replace(
template,
iter_=node.iter,
target=node.target,
body=node.body,
orelse=node.orelse)
new_for_node = node[0]
anno.copyanno(original_node, new_for_node, anno.Basic.EXTRA_LOOP_TEST)
anno.copyanno(original_node, new_for_node, anno.Basic.DIRECTIVES)
return node
# Python's else clause only triggers if the loop exited cleanly (e.g.
# break did not trigger).
guarded_orelse = self._guard_if_present(node.orelse, break_var)
extra_test = templates.replace_as_expression(
'not var_name', var_name=break_var)
# The extra test is hidden in the AST, which will confuse the static
# analysis. To mitigate that, we insert a no-op statement that ensures
# the control variable is marked as used.
# TODO(mdan): Use a marker instead, e.g. ag__.condition_loop_on(var_name)
template = """
var_name = False
for target in iter_:
(var_name,)
body
orelse
"""
node = templates.replace(
template,
var_name=break_var,
iter_=node.iter,
target=node.target,
body=node.body,
orelse=guarded_orelse)
new_for_node = node[1]
anno.setanno(new_for_node, anno.Basic.EXTRA_LOOP_TEST, extra_test)
anno.copyanno(original_node, new_for_node, anno.Basic.DIRECTIVES)
return node
def transform(node, ctx):
node = qual_names.resolve(node)
node = activity.resolve(node, ctx, None)
transformer = BreakTransformer(ctx)
node = transformer.visit(node)
return node
@@ -0,0 +1,170 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for break_statements module."""
from tensorflow.python.autograph.converters import break_statements
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.platform import test
class BreakCanonicalizationTest(converter_testing.TestCase):
def assertTransformedEquivalent(self, f, *inputs):
tr = self.transform(f, break_statements)
self.assertEqual(f(*inputs), tr(*inputs))
def test_while_loop(self):
def f(x):
v = []
while x > 0:
x -= 1
if x % 2 == 0:
break
v.append(x)
return v
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 4)
def test_while_loop_preserves_directives(self):
def f(x):
while x > 0:
x -= 1
if x % 2 == 0:
break
_, node, ctx = self.transform(f, (), include_ast=True)
fake_annotation = object()
anno.setanno(node.body[0], anno.Basic.DIRECTIVES, fake_annotation)
node = break_statements.transform(node, ctx)
self.assertIs(
anno.getanno(node.body[1], anno.Basic.DIRECTIVES), fake_annotation)
def test_for_loop(self):
def f(a):
v = []
for x in a:
x -= 1
if x % 2 == 0:
break
v.append(x)
return v
tr = self.transform(f, break_statements)
self.assertEqual([3], tr([5, 4]))
def test_for_loop_preserves_directives(self):
def f(a):
for x in a:
if x % 2 == 0:
break
_, node, ctx = self.transform(f, (), include_ast=True)
fake_annotation = object()
anno.setanno(node.body[0], anno.Basic.DIRECTIVES, fake_annotation)
node = break_statements.transform(node, ctx)
self.assertIs(
anno.getanno(node.body[1], anno.Basic.DIRECTIVES), fake_annotation)
def test_nested(self):
def f(x):
v = []
u = []
w = []
while x > 0:
x -= 1
if x % 2 == 0:
if x % 3 != 0:
u.append(x)
else:
w.append(x)
break
v.append(x)
return v, u, w
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 11)
def test_nested_loops(self):
def f(x):
v = []
u = []
while x > 0:
x -= 1
y = x
while y > 0:
y -= 1
if y % 2 == 0:
break
u.append(y)
if x == 0:
break
v.append(x)
return v, u
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 5)
def test_loop_orelse(self):
def f(x):
v = []
u = []
while x > 0:
x -= 1
y = x
while y > 1:
break
else:
u.append(y)
break
v.append(x)
return v, u
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, 3)
def test_multiple_correlated_breaks_with_side_effects(self):
def f(cond1):
lst = []
while True:
if cond1:
lst.append(1)
else:
break
if lst[-1] > 0: # lst always has an element here
break
return lst
self.assertTransformedEquivalent(f, True)
self.assertTransformedEquivalent(f, False)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,221 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Handles function calls, by generating compiled function names and calls.
Note: this transformer does not rename the top level object being converted;
that is the caller's responsibility.
Requires function_scopes.
"""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.utils import ag_logging
# TODO(mdan): Rename to FunctionCallsTransformer.
class _Function(object):
no_root = True
def __init__(self):
self.context_name = None
set_trace_warned = False
class _ArgTemplateBuilder(object):
"""Constructs a tuple representing the positional arguments in a call.
Example (yes, it's legal Python 3):
f(*args1, b, *args2, c, d) -> args1 + (b,) + args2 + (c, d)
"""
def __init__(self):
self._arg_accumulator = []
self._argspec = []
self._finalized = False
def _consume_args(self):
if self._arg_accumulator:
self._argspec.append(
gast.Tuple(elts=self._arg_accumulator, ctx=gast.Load()))
self._arg_accumulator = []
def add_arg(self, a):
self._arg_accumulator.append(a)
def add_stararg(self, a):
self._consume_args()
self._argspec.append(
gast.Call(
gast.Name(
'tuple', ctx=gast.Load(), annotation=None, type_comment=None),
args=[a],
keywords=()))
def finalize(self):
self._consume_args()
self._finalized = True
def to_ast(self):
assert self._finalized
if self._argspec:
result = self._argspec[0]
for i in range(1, len(self._argspec)):
result = gast.BinOp(result, gast.Add(), self._argspec[i])
return result
return gast.Tuple([], gast.Load())
class CallTreeTransformer(converter.Base):
"""Transforms the call tree by renaming transformed symbols."""
def visit_Lambda(self, node):
if not anno.hasanno(node, 'function_context_name'):
# Lambda functions created during the conversion process have no
# context manager.
return self.generic_visit(node)
with self.state[_Function] as fn_scope:
fn_scope.context_name = anno.getanno(node, 'function_context_name')
return self.generic_visit(node)
def visit_FunctionDef(self, node):
# Decorators and arg defaults are part of the outer scope.
node.decorator_list = self.visit_block(node.decorator_list)
node.args.defaults = self.visit_block(node.args.defaults)
for i, d in enumerate(node.args.kw_defaults):
if d is not None:
node.args.kw_defaults[i] = self.visit(d)
with self.state[_Function] as fn_scope:
# Note: if the conversion process ever creates helper functions, this
# assumption will no longer hold.
assert anno.hasanno(node, 'function_context_name'), (
'The function_scopes converter always creates a scope for functions.')
fn_scope.context_name = anno.getanno(node, 'function_context_name')
node.body = self.visit_block(node.body)
if node.returns:
node.returns = self.visit(node.returns)
return node
def visit_With(self, node):
# Context manager calls (in node.items) are not converted.
node.body = self.visit_block(node.body)
return node
def _args_to_tuple(self, node):
"""Ties together all positional and *arg arguments in a single tuple."""
# TODO(mdan): We could rewrite this to just a call to tuple(). Maybe better?
# For example for
# f(a, b, *args)
# instead of writing:
# (a, b) + args
# just write this?
# tuple(a, b, *args)
builder = _ArgTemplateBuilder()
for a in node.args:
if isinstance(a, gast.Starred):
builder.add_stararg(a.value)
else:
builder.add_arg(a)
builder.finalize()
return builder.to_ast()
def _kwargs_to_dict(self, node):
"""Ties together all keyword and **kwarg arguments in a single dict."""
if node.keywords:
return gast.Call(
gast.Name(
'dict', ctx=gast.Load(), annotation=None, type_comment=None),
args=(),
keywords=node.keywords)
else:
return parser.parse_expression('None')
def visit_Call(self, node):
full_name = str(anno.getanno(node.func, anno.Basic.QN, default=''))
function_context_name = self.state[_Function].context_name
node = self.generic_visit(node)
# TODO(mdan): Refactor converted_call as a 'Call' operator.
# Calls to the internal 'ag__' module are never converted (though their
# arguments might be).
if full_name.startswith('ag__.'):
return node
# Calls to the function context manager (inserted by function_scopes) are
# also safe.
if full_name.startswith(function_context_name + '.'):
return node
# Calls to pdb.set_trace or ipdb.set_trace are never converted. We don't use
# the normal mechanisms to bypass these literals because they are sensitive
# to the frame they are being called from.
# TODO(mdan): Generalize this to a "static allowlist" config.
if full_name in ('pdb.set_trace', 'ipdb.set_trace', 'breakpoint'):
global set_trace_warned
if not set_trace_warned:
# TODO(mdan): Update and shorten once available on tensorflow.org.
ag_logging.warning(
'Detected `pdb.set_trace()` in user code. The code'
' generated by AutoGraph is not optimized for step-by-step'
' debugging. See https://github.com/tensorflow/tensorflow/'
'blob/master/tensorflow/python/autograph/g3doc/reference/'
'debugging.md.')
set_trace_warned = True
return node
if (full_name == 'print' and
not self.ctx.user.options.uses(converter.Feature.BUILTIN_FUNCTIONS)):
return node
template = """
ag__.converted_call(func, args, kwargs, function_ctx)
"""
new_call = templates.replace_as_expression(
template,
func=node.func,
args=self._args_to_tuple(node),
kwargs=self._kwargs_to_dict(node),
function_ctx=function_context_name)
return new_call
def transform(node, ctx):
"""Transform function call to the compiled counterparts.
Args:
node: AST
ctx: EntityContext
Returns:
A tuple (node, new_names):
node: The transformed AST
new_names: set(string), containing any newly-generated names
"""
node = qual_names.resolve(node)
node = CallTreeTransformer(ctx).visit(node)
return node
@@ -0,0 +1,262 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for call_trees module."""
import types
from tensorflow.python.autograph.converters import call_trees
from tensorflow.python.autograph.converters import functions
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.platform import test
class MockConvertedCall(object):
def __init__(self):
self.calls = []
def __call__(self, f, args, kwargs, caller_fn_scope=None, options=None):
del caller_fn_scope, options
self.calls.append((args, kwargs))
kwargs = kwargs or {}
return f(*args, **kwargs)
class CallTreesTest(converter_testing.TestCase):
def _transform_with_mock(self, f):
mock = MockConvertedCall()
tr = self.transform(
f, (functions, call_trees),
ag_overrides={'converted_call': mock})
return tr, mock
def test_function_no_args(self):
def f(f):
return f() + 20
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(lambda: 1), 21)
self.assertListEqual(mock.calls, [((), None)])
def test_function_with_expression_in_argument(self):
def f(f, g):
return f(g() + 20) + 4000
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(lambda x: x + 300, lambda: 1), 4321)
self.assertListEqual(mock.calls, [
((), None),
((21,), None),
])
def test_function_with_call_in_argument(self):
def f(f, g):
return f(g()) + 300
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(lambda x: x + 20, lambda: 1), 321)
self.assertListEqual(mock.calls, [
((), None),
((1,), None),
])
def test_function_chaining(self):
def get_one():
return 1
def f():
return get_one().__add__(20)
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(), 21)
self.assertListEqual(mock.calls, [
((), None),
((20,), None),
])
def test_function_with_single_arg(self):
def f(f, a):
return f(a) + 20
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(lambda a: a, 1), 21)
self.assertListEqual(mock.calls, [((1,), None)])
def test_function_with_args_only(self):
def f(f, a, b):
return f(a, b) + 300
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(lambda a, b: a + b, 1, 20), 321)
self.assertListEqual(mock.calls, [((1, 20), None)])
def test_function_with_kwarg(self):
def f(f, a, b):
return f(a, c=b) + 300
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(lambda a, c: a + c, 1, 20), 321)
self.assertListEqual(mock.calls, [((1,), {'c': 20})])
def test_function_with_kwargs_starargs(self):
def f(f, a, *args, **kwargs):
return f(a, *args, **kwargs) + 5
tr, mock = self._transform_with_mock(f)
self.assertEqual(
tr(lambda *args, **kwargs: 7, 1, *[2, 3], **{
'b': 4,
'c': 5
}), 12)
self.assertListEqual(mock.calls, [((1, 2, 3), {'b': 4, 'c': 5})])
def test_function_with_starargs_only(self):
def g(*args):
return sum(args)
def f():
args = [1, 20, 300]
return g(*args) + 4000
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(), 4321)
self.assertListEqual(mock.calls, [((1, 20, 300), None)])
def test_function_with_starargs_mixed(self):
def g(a, b, c, d):
return a * 1000 + b * 100 + c * 10 + d
def f():
args1 = (1,)
args2 = [3]
return g(*args1, 2, *args2, 4)
tr, mock = self._transform_with_mock(f)
self.assertEqual(tr(), 1234)
self.assertListEqual(mock.calls, [((1, 2, 3, 4), None)])
def test_function_with_kwargs_keywords(self):
def f(f, a, b, **kwargs):
return f(a, b=b, **kwargs) + 5
tr, mock = self._transform_with_mock(f)
self.assertEqual(
tr(lambda *args, **kwargs: 7, 1, 2, **{'c': 3}), 12)
self.assertListEqual(mock.calls, [((1,), {'b': 2, 'c': 3})])
def test_function_with_multiple_kwargs(self):
def f(f, a, b, c, kwargs1, kwargs2):
return f(a, b=b, **kwargs1, c=c, **kwargs2) + 5
tr, mock = self._transform_with_mock(f)
self.assertEqual(
tr(lambda *args, **kwargs: 7, 1, 2, 3, {'d': 4}, {'e': 5}), 12)
self.assertListEqual(mock.calls, [((1,), {
'b': 2,
'c': 3,
'd': 4,
'e': 5
})])
def test_function_with_call_in_lambda_argument(self):
def h(l, a):
return l(a) + 4000
def g(a, *args):
return a + sum(args)
def f(h, g, a, *args):
return h(lambda x: g(x, *args), a)
tr, _ = self._transform_with_mock(f)
self.assertEqual(tr(h, g, 1, *(20, 300)), 4321)
def test_debugger_set_trace(self):
tracking_list = []
pdb = types.ModuleType('fake_pdb')
pdb.set_trace = lambda: tracking_list.append(1)
def f():
return pdb.set_trace()
tr, _ = self._transform_with_mock(f)
tr()
self.assertListEqual(tracking_list, [1])
def test_class_method(self):
class TestClass(object):
def other_method(self, x):
return x + 20
def test_method(self, a):
return self.other_method(a) + 300
tc = TestClass()
tr, mock = self._transform_with_mock(TestClass.test_method)
self.assertEqual(321, tr(tc, 1))
self.assertListEqual(mock.calls, [((1,), None)])
def test_object_method(self):
class TestClass(object):
def other_method(self, x):
return x + 20
def test_method(self, a):
return self.other_method(a) + 300
tc = TestClass()
tr, mock = self._transform_with_mock(tc.test_method)
self.assertEqual(321, tr(tc, 1))
self.assertListEqual(mock.calls, [((1,), None)])
if __name__ == '__main__':
test.main()
@@ -0,0 +1,46 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Converts the ternary conditional operator."""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import templates
class ConditionalExpressionTransformer(converter.Base):
"""Converts conditional expressions to functional form."""
def visit_IfExp(self, node):
template = '''
ag__.if_exp(
test,
lambda: true_expr,
lambda: false_expr,
expr_repr)
'''
expr_repr = parser.unparse(node.test, include_encoding_marker=False).strip()
return templates.replace_as_expression(
template,
test=node.test,
true_expr=node.body,
false_expr=node.orelse,
expr_repr=gast.Constant(expr_repr, kind=None))
def transform(node, ctx):
node = ConditionalExpressionTransformer(ctx).visit(node)
return node
@@ -0,0 +1,48 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for conditional_expressions module."""
from tensorflow.python.autograph.converters import conditional_expressions
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.platform import test
class ConditionalExpressionsTest(converter_testing.TestCase):
def assertTransformedEquivalent(self, f, *inputs):
tr = self.transform(f, conditional_expressions)
self.assertEqual(f(*inputs), tr(*inputs))
def test_basic(self):
def f(x):
return 1 if x else 0
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 3)
def test_nested_orelse(self):
def f(x):
y = x * x if x > 0 else x if x else 1
return y
self.assertTransformedEquivalent(f, -2)
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 2)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,164 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Canonicalizes continue statements by de-sugaring into a control boolean."""
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct.static_analysis import activity
from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
class _Continue(object):
def __init__(self):
self.used = False
self.control_var_name = None
def __repr__(self):
return '<_Continue(used: {}, var: {})>'.format(self.used,
self.control_var_name)
class _Block(object):
"""Tracks information about lexical blocks as they are visited in the AST.
Mainly, this object tracks the creation of block guards that replace
`continue` statements (e.g. `if not continue_:`).
Attributes:
create_guard_current: bool, whether to create a guard for the current
statement.
create_guard_next: bool, whether to create a guard for the next
statement.
is_loop_type: bool, whether this block is the body of a loop.
"""
def __init__(self):
self.is_loop_type = False
self.create_guard_current = False
self.create_guard_next = False
class ContinueCanonicalizationTransformer(converter.Base):
"""Canonicalizes continue statements into additional conditionals."""
def visit_Continue(self, node):
self.state[_Continue].used = True
for block in reversed(self.state[_Block].stack):
# See ContinueCanonicalizationTest.test_multiple_continues for an example
# it's necessary to create guards for all enclosing affected blocks, not
# just that of the current block.
block.create_guard_next = True
if block.is_loop_type:
# continue only affects the innermost loop
break
template = """
var_name = True
"""
return templates.replace(
template, var_name=self.state[_Continue].control_var_name)
def _postprocess_statement(self, node):
if self.state[_Continue].used:
block = self.state[_Block]
should_wrap_current = block.create_guard_current
# After processing propagate whether to guard the next statement
block.create_guard_current = block.create_guard_next
block.create_guard_next = False
if should_wrap_current:
template = """
if not var_name:
original_node
"""
cond, = templates.replace(
template,
var_name=self.state[_Continue].control_var_name,
original_node=node)
return cond, cond.body
return node, None
def _visit_loop_body(self, node, nodes):
self.state[_Continue].enter()
self.state[_Block].enter()
self.state[_Block].is_loop_type = True
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
continue_var = self.ctx.namer.new_symbol('continue_', scope.referenced)
self.state[_Continue].control_var_name = continue_var
nodes = self.visit_block(nodes, after_visit=self._postprocess_statement)
if self.state[_Continue].used:
template = """
var_name = False
"""
control_var_init = templates.replace(template, var_name=continue_var)
nodes = control_var_init + nodes
self.state[_Block].exit()
self.state[_Continue].exit()
return nodes
def _visit_non_loop_body(self, nodes):
self.state[_Block].enter()
nodes = self.visit_block(nodes, after_visit=self._postprocess_statement)
self.state[_Block].exit()
return nodes
def visit_While(self, node):
node.test = self.visit(node.test)
node.body = self._visit_loop_body(node, node.body)
# A continue in the else clause applies to the containing scope.
node.orelse = self._visit_non_loop_body(node.orelse)
return node
def visit_For(self, node):
node.target = self.generic_visit(node.target)
node.iter = self.generic_visit(node.iter)
node.body = self._visit_loop_body(node, node.body)
# A continue in the else clause applies to the containing scope.
node.orelse = self._visit_non_loop_body(node.orelse)
return node
def visit_If(self, node):
node.body = self._visit_non_loop_body(node.body)
node.orelse = self._visit_non_loop_body(node.orelse)
return node
def visit_With(self, node):
node.items = self.visit_block(node.items)
node.body = self._visit_non_loop_body(node.body)
return node
def visit_Try(self, node):
node.body = self._visit_non_loop_body(node.body)
node.orelse = self._visit_non_loop_body(node.orelse)
# In Python 3.8 and later continue is allowed in finally blocks
node.finalbody = self._visit_non_loop_body(node.finalbody)
node.handlers = self.visit_block(node.handlers)
return node
def visit_ExceptHandler(self, node):
node.body = self._visit_non_loop_body(node.body)
return node
def transform(node, ctx):
node = qual_names.resolve(node)
node = activity.resolve(node, ctx, None)
node = ContinueCanonicalizationTransformer(ctx).visit(node)
return node
@@ -0,0 +1,223 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for continue_statements module."""
from tensorflow.python.autograph.converters import continue_statements
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.framework import ops
from tensorflow.python.platform import test
class ContinueCanonicalizationTest(converter_testing.TestCase):
def assertTransformedEquivalent(self, f, *inputs):
tr = self.transform(f, continue_statements)
self.assertEqual(f(*inputs), tr(*inputs))
def test_basic(self):
def f(x):
v = []
while x > 0:
x -= 1
if x % 2 == 0:
continue
v.append(x)
return v
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 4)
def test_multiple_continues(self):
def f(x):
v = []
while x > 0:
x -= 1
if x > 1:
continue
if x > 2:
continue
v.append(x)
return v
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 4)
def test_multiple_continues_in_nested_scope(self):
def f(a):
v = []
for x in a:
x -= 1
if x > 100:
continue
try:
raise ValueError('intentional')
except ValueError:
continue
v.append(x)
return v
self.assertTransformedEquivalent(f, [])
self.assertTransformedEquivalent(f, [1])
self.assertTransformedEquivalent(f, [2])
self.assertTransformedEquivalent(f, [1, 2, 3])
def test_for_loop(self):
def f(a):
v = []
for x in a:
x -= 1
if x % 2 == 0:
continue
v.append(x)
return v
self.assertTransformedEquivalent(f, [])
self.assertTransformedEquivalent(f, [1])
self.assertTransformedEquivalent(f, [2])
self.assertTransformedEquivalent(f, [1, 2, 3])
def test_nested_with(self):
def f(x):
v = []
while x > 0:
x -= 1
with ops.name_scope(''):
if x % 2 == 0:
continue
v.append(x)
return v
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 4)
def test_nested_multiple_withs(self):
def f(x):
v = []
while x > 0:
x -= 1
with ops.name_scope(''):
if x % 2 == 0:
continue
with ops.name_scope(''):
v.append(x)
v.append(x)
return v
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 4)
def test_nested_multiple_withs_and_statements(self):
def f(x):
v = []
while x > 0:
x -= 1
with ops.name_scope(''):
if x % 2 == 0:
continue
v.append(x)
v.append(x)
with ops.name_scope(''):
v.append(x)
v.append(x)
return v
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 4)
def test_nested_multiple_withs_and_nested_withs(self):
def f(x):
v = []
while x > 0:
x -= 1
with ops.name_scope(''):
if x % 2 == 0:
continue
with ops.name_scope(''):
v.append(x)
v.append(x)
with ops.name_scope(''):
v.append(x)
v.append(x)
return v
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 4)
def test_nested(self):
def f(x):
v = []
u = []
w = []
while x > 0:
x -= 1
if x % 2 == 0:
if x % 3 != 0:
u.append(x)
else:
w.append(x)
continue
v.append(x)
return v, u, w
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 4)
def test_multiple_guarded_continues_with_side_effects(self):
def f(x):
def track(u, x):
u.append(x)
return x
u = []
v = []
while x > 0:
x -= 1
if track(u, x) > 1:
continue
if track(u, x) > 2:
continue
v.append(x)
return u, v
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 2)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,413 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Handles control flow statements: while, for, if."""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.lang import directives
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import cfg
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct.static_analysis import activity
from tensorflow.python.autograph.pyct.static_analysis import annos
from tensorflow.python.autograph.pyct.static_analysis import liveness
from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions
from tensorflow.python.autograph.pyct.static_analysis import reaching_fndefs
class _Function(object):
scope = None
class ControlFlowTransformer(converter.Base):
"""Transforms control flow structures like loops an conditionals."""
def visit_Lambda(self, node):
with self.state[_Function] as fn:
fn.scope = anno.getanno(node, anno.Static.SCOPE)
return self.generic_visit(node)
def visit_FunctionDef(self, node):
with self.state[_Function] as fn:
fn.scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
return self.generic_visit(node)
def _create_nonlocal_declarations(self, vars_):
vars_ = set(vars_)
results = []
global_vars = self.state[_Function].scope.globals & vars_
if global_vars:
results.append(gast.Global([str(v) for v in global_vars]))
nonlocal_vars = [
v for v in vars_ if not v.is_composite() and v not in global_vars]
if nonlocal_vars:
results.append(gast.Nonlocal([str(v) for v in nonlocal_vars]))
return results
def _create_state_functions(
self, block_vars, nonlocal_declarations, getter_name, setter_name):
if not block_vars:
template = """
def getter_name():
return ()
def setter_name(block_vars):
pass
"""
return templates.replace(
template, getter_name=getter_name, setter_name=setter_name)
guarded_block_vars = []
for v in block_vars:
if v.is_simple():
guarded_block_vars.append(v)
else:
guarded_block_vars.append(
templates.replace_as_expression(
'ag__.ldu(lambda: var_, name)',
var_=v,
name=gast.Constant(str(v), kind=None)))
template = """
def getter_name():
return guarded_state_vars,
def setter_name(vars_):
nonlocal_declarations
state_vars, = vars_
"""
return templates.replace(
template,
nonlocal_declarations=nonlocal_declarations,
getter_name=getter_name,
guarded_state_vars=guarded_block_vars,
setter_name=setter_name,
state_vars=tuple(block_vars))
def _create_loop_options(self, node):
if not anno.hasanno(node, anno.Basic.DIRECTIVES):
return gast.Dict([], [])
loop_directives = anno.getanno(node, anno.Basic.DIRECTIVES)
if directives.set_loop_options not in loop_directives:
return gast.Dict([], [])
opts_dict = loop_directives[directives.set_loop_options]
str_keys, values = zip(*opts_dict.items())
keys = [gast.Constant(s, kind=None) for s in str_keys]
values = list(values) # ast and gast don't play well with tuples.
return gast.Dict(keys, values)
def _create_undefined_assigns(self, undefined_symbols):
assignments = []
for s in undefined_symbols:
template = '''
var = ag__.Undefined(symbol_name)
'''
assignments += templates.replace(
template,
var=s,
symbol_name=gast.Constant(s.ssf(), kind=None))
return assignments
def _get_block_basic_vars(self, modified, live_in, live_out):
nonlocals = self.state[_Function].scope.nonlocals
basic_scope_vars = []
for s in modified:
if s.is_composite():
# TODO(mdan): Raise an error when this happens for a TF scope.
continue
# Variables not live into or out of the scope are considered local to the
# scope.
if s in live_in or s in live_out or s in nonlocals:
basic_scope_vars.append(s)
continue
return frozenset(basic_scope_vars)
def _get_block_composite_vars(self, modified, live_in):
# The scope variables corresponding to composite symbols (e.g. `self.x`).
composite_scope_vars = []
for s in modified:
if not s.is_composite():
continue
# Mutations made to objects created inside the scope will appear as writes
# to composite symbols. Because these mutations appear as modifications
# made to composite symbols, we check whether the composite's parent is
# actually live into the scope.
# Example:
# while cond:
# x = Foo()
# x.foo = 2 * x.foo # x.foo is live into the scope, but x is not.
#
# Note that some parents might not be symbols - for example, in x['foo'],
# 'foo' is a parent, but it's a literal, not a symbol. We don't check the
# liveness of literals.
support_set_symbols = tuple(
sss for sss in s.support_set if sss.is_symbol())
if not all(sss in live_in for sss in support_set_symbols):
continue
composite_scope_vars.append(s)
return frozenset(composite_scope_vars)
def _get_block_vars(self, node, modified):
"""Determines the variables affected inside a control flow statement."""
defined_in = anno.getanno(node, anno.Static.DEFINED_VARS_IN)
live_in = anno.getanno(node, anno.Static.LIVE_VARS_IN)
live_out = anno.getanno(node, anno.Static.LIVE_VARS_OUT)
fn_scope = self.state[_Function].scope
basic_scope_vars = self._get_block_basic_vars(
modified,
live_in,
live_out)
composite_scope_vars = self._get_block_composite_vars(modified, live_in)
scope_vars = tuple(basic_scope_vars | composite_scope_vars)
# Variables that are modified inside the scope, but not defined
# before entering it. Only simple variables must be defined. The
# composite ones will be implicitly checked at runtime.
possibly_undefined = (
modified - defined_in - fn_scope.globals - fn_scope.nonlocals)
undefined = tuple(v for v in possibly_undefined if not v.is_composite())
# Variables that are modified inside the scope, and depend on values outside
# it.
input_only = basic_scope_vars & live_in - live_out
# Place the outputs first, then sort lexicographically.
scope_vars = sorted(scope_vars, key=lambda v: (v in input_only, v))
nouts = len(scope_vars) - len(input_only)
return scope_vars, undefined, nouts
def visit_If(self, node):
node = self.generic_visit(node)
body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
orelse_scope = anno.getanno(node, annos.NodeAnno.ORELSE_SCOPE)
cond_vars, undefined, nouts = self._get_block_vars(
node, body_scope.bound | orelse_scope.bound)
undefined_assigns = self._create_undefined_assigns(undefined)
nonlocal_declarations = self._create_nonlocal_declarations(cond_vars)
reserved = body_scope.referenced | orelse_scope.referenced
state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
state_functions = self._create_state_functions(
cond_vars, nonlocal_declarations, state_getter_name, state_setter_name)
orelse_body = node.orelse
if not orelse_body:
orelse_body = [gast.Pass()]
template = """
state_functions
def body_name():
nonlocal_declarations
body
def orelse_name():
nonlocal_declarations
orelse
undefined_assigns
ag__.if_stmt(
test,
body_name,
orelse_name,
state_getter_name,
state_setter_name,
(symbol_names,),
nouts)
"""
new_nodes = templates.replace(
template,
body=node.body,
body_name=self.ctx.namer.new_symbol('if_body', reserved),
orelse=orelse_body,
orelse_name=self.ctx.namer.new_symbol('else_body', reserved),
nonlocal_declarations=nonlocal_declarations,
nouts=gast.Constant(nouts, kind=None),
state_functions=state_functions,
state_getter_name=state_getter_name,
state_setter_name=state_setter_name,
symbol_names=tuple(gast.Constant(str(s), kind=None) for s in cond_vars),
test=node.test,
undefined_assigns=undefined_assigns)
origin_info.copy_origin(node, new_nodes[-1])
return new_nodes
def visit_While(self, node):
node = self.generic_visit(node)
body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
loop_vars, undefined, _ = self._get_block_vars(node, body_scope.bound)
undefined_assigns = self._create_undefined_assigns(undefined)
nonlocal_declarations = self._create_nonlocal_declarations(loop_vars)
reserved = body_scope.referenced
state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
state_functions = self._create_state_functions(
loop_vars, nonlocal_declarations, state_getter_name, state_setter_name)
opts = self._create_loop_options(node)
template = """
state_functions
def body_name():
nonlocal_declarations
body
def test_name():
return test
undefined_assigns
ag__.while_stmt(
test_name,
body_name,
state_getter_name,
state_setter_name,
(symbol_names,),
opts)
"""
new_nodes = templates.replace(
template,
body=node.body,
body_name=self.ctx.namer.new_symbol('loop_body', reserved),
nonlocal_declarations=nonlocal_declarations,
opts=opts,
state_functions=state_functions,
state_getter_name=state_getter_name,
state_setter_name=state_setter_name,
symbol_names=tuple(gast.Constant(str(s), kind=None) for s in loop_vars),
test=node.test,
test_name=self.ctx.namer.new_symbol('loop_test', reserved),
undefined_assigns=undefined_assigns)
origin_info.copy_origin(node, new_nodes[-1])
return new_nodes
def visit_For(self, node):
node = self.generic_visit(node)
body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
iter_scope = anno.getanno(node, annos.NodeAnno.ITERATE_SCOPE)
loop_vars, undefined, _ = self._get_block_vars(
node, body_scope.bound | iter_scope.bound)
undefined_assigns = self._create_undefined_assigns(undefined)
nonlocal_declarations = self._create_nonlocal_declarations(loop_vars)
reserved = body_scope.referenced | iter_scope.referenced
state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
state_functions = self._create_state_functions(
loop_vars, nonlocal_declarations, state_getter_name, state_setter_name)
opts = self._create_loop_options(node)
opts.keys.append(gast.Constant('iterate_names', kind=None))
opts.values.append(gast.Constant(
parser.unparse(node.target, include_encoding_marker=False), kind=None))
if anno.hasanno(node, anno.Basic.EXTRA_LOOP_TEST):
extra_test = anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST)
extra_test_name = self.ctx.namer.new_symbol(
'extra_test', reserved)
template = """
def extra_test_name():
nonlocal_declarations
return extra_test_expr
"""
extra_test_function = templates.replace(
template,
extra_test_expr=extra_test,
extra_test_name=extra_test_name,
loop_vars=loop_vars,
nonlocal_declarations=nonlocal_declarations)
else:
extra_test_name = parser.parse_expression('None')
extra_test_function = []
# iterate_arg_name holds a single arg with the iterates, which may be a
# tuple.
iterate_arg_name = self.ctx.namer.new_symbol('itr', reserved)
template = """
iterates = iterate_arg_name
"""
iterate_expansion = templates.replace(
template, iterate_arg_name=iterate_arg_name, iterates=node.target)
origin_info.copy_origin(node, iterate_expansion)
template = """
state_functions
def body_name(iterate_arg_name):
nonlocal_declarations
iterate_expansion
body
extra_test_function
undefined_assigns
ag__.for_stmt(
iterated,
extra_test_name,
body_name,
state_getter_name,
state_setter_name,
(symbol_names,),
opts)
"""
new_nodes = templates.replace(
template,
body=node.body,
body_name=self.ctx.namer.new_symbol('loop_body', reserved),
extra_test_function=extra_test_function,
extra_test_name=extra_test_name,
iterate_arg_name=iterate_arg_name,
iterate_expansion=iterate_expansion,
iterated=node.iter,
nonlocal_declarations=nonlocal_declarations,
opts=opts,
symbol_names=tuple(gast.Constant(str(s), kind=None) for s in loop_vars),
state_functions=state_functions,
state_getter_name=state_getter_name,
state_setter_name=state_setter_name,
undefined_assigns=undefined_assigns)
origin_info.copy_origin(node, new_nodes[-1])
return new_nodes
class AnnotatedDef(reaching_definitions.Definition):
def __init__(self):
super(AnnotatedDef, self).__init__()
self.directives = {}
def transform(node, ctx):
graphs = cfg.build(node)
node = qual_names.resolve(node)
node = activity.resolve(node, ctx, None)
node = reaching_definitions.resolve(node, ctx, graphs)
node = reaching_fndefs.resolve(node, ctx, graphs)
node = liveness.resolve(node, ctx, graphs)
node = ControlFlowTransformer(ctx).visit(node)
return node
@@ -0,0 +1,786 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for control_flow module."""
import collections
import numpy as np
from tensorflow.python.autograph.converters import break_statements
from tensorflow.python.autograph.converters import continue_statements
from tensorflow.python.autograph.converters import control_flow
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor_util
from tensorflow.python.platform import test
from tensorflow.python.util import nest
for_unaffected_global = None
for_mixed_globals_nonglobals = None
for_test_global_local = None
class ControlFlowTestBase(converter_testing.TestCase):
def assertValuesEqual(self, actual, expected):
values = nest.map_structure(
lambda x: self.evaluate(x) if tensor_util.is_tf_type(x) else x,
actual)
self.assertAllEqual(values, expected)
def assertTransformedResult(self, f, inputs, expected):
if not isinstance(inputs, tuple):
inputs = (inputs,)
tr = self.transform(f, control_flow)
returns = tr(*inputs)
self.assertValuesEqual(returns, expected)
class NestedControlFlowTest(ControlFlowTestBase):
def test_basic(self):
def f(n):
i = 0
j = 0
s = 0
while i < n:
while j < i:
j += 3
u = i + j # 'u' is not defined within the inner loop
s += u
i += 1
j = 0
return s, i, j, n
self.assertTransformedResult(f, constant_op.constant(5),
(25, 5, 0, 5))
def test_mixed_globals_nonglobals(self):
def f(n):
global for_mixed_globals_nonglobals
i = 0
j = 0
for_mixed_globals_nonglobals = 0
while i < n:
while j < i:
j += 3
u = i + j # 'u' is not defined within the inner loop
for_mixed_globals_nonglobals += u
i += 1
j = 0
return for_mixed_globals_nonglobals, i, j, n
self.assertTransformedResult(f, constant_op.constant(5),
(25, 5, 0, 5))
def test_composite_state_complex(self):
class TestClassX(object):
def __init__(self, x):
self.x = x
class TestClassY(object):
def __init__(self, y):
self.y = y
def f(n):
tc = TestClassX(TestClassY({'z': TestClassX(n)}))
if n > 0:
while n > 0:
if n < 2:
tc.x.y['z'].x += 1
n -= 1
return n, tc
tr = self.transform(f, control_flow)
n, tc = tr(constant_op.constant(5))
self.assertValuesEqual((n, tc.x.y['z'].x), (0, 6))
class WhileStatementTest(ControlFlowTestBase):
def test_basic(self):
def f(n):
i = 0
s = 0
while i < n:
s += i
i += 1
return s, i, n
self.assertTransformedResult(f, constant_op.constant(5), (10, 5, 5))
def test_single_output(self):
def f(n):
while n > 0:
n -= 1
return n
self.assertTransformedResult(f, constant_op.constant(5), 0)
def test_composite_state_attr(self):
class TestClass(object):
def __init__(self):
self.x = constant_op.constant(3)
def f(n):
tc = TestClass()
while n > 0:
tc.x += 1
n -= 1
return n
self.assertTransformedResult(f, constant_op.constant(5), 0)
def test_composite_state_slice(self):
def f(n):
d = {'a': n}
k = 'a'
while n > 0:
d[k] += 1
n -= 1
return d[k], n
self.assertTransformedResult(f, constant_op.constant(5), (10, 0))
def test_composite_state_literal_slice(self):
def f(n):
d = {'a': n}
while n > 0:
d['a'] += 1
n -= 1
return d['a'], n
self.assertTransformedResult(f, constant_op.constant(5), (10, 0))
def test_composite_state_attr_initialized_in_loop(self):
class TestClass(object):
pass
def f(n, x):
tc = TestClass()
while n < 5:
if n == 0:
tc.subattr = x
else:
tc.subattr = tc.subattr + 1
n += 1
return tc.subattr
self.assertTransformedResult(f, (0, constant_op.constant(10)), 14)
tr = self.transform(f, control_flow)
with self.assertRaisesRegex(
ValueError, "'tc.subattr' must be defined before the loop"):
tr(constant_op.constant(0), 0)
def test_composite_state_slice_initialized_in_loop(self):
def f(n, x):
d = {}
k = 'subkey'
while n < 5:
if n == 0:
d[k] = x
else:
d[k] = d[k] + 1
n += 1
return d
self.assertTransformedResult(f, (0, constant_op.constant(10)),
{'subkey': 14})
tr = self.transform(f, control_flow)
with self.assertRaisesRegex(
ValueError, r"'d\[k\]' must be defined before the loop"):
tr(constant_op.constant(0), 0)
def test_composite_state_literal_slice_initialized_in_loop(self):
def f(n, x):
d = {}
while n < 5:
if n == 0:
d['subkey'] = x
else:
d['subkey'] = d['subkey'] + 1
n += 1
return d
self.assertTransformedResult(f, (0, constant_op.constant(10)),
{'subkey': 14})
tr = self.transform(f, control_flow)
with self.assertRaisesRegex(
ValueError, r"'d\['subkey'\]' must be defined before the loop"):
tr(constant_op.constant(0), 0)
def test_composite_state_slice_aliased_to_local(self):
def f(n, x):
d = {}
while n < 5:
k = 'subkey'
d[k] = x + 1
n += 1
return d
self.assertTransformedResult(f, (0, constant_op.constant(10)),
{'subkey': 11})
tr = self.transform(f, control_flow)
# TODO(b/136999953): Better error message.
# Note that this error happens at execution time.
with self.assertRaises(errors.InaccessibleTensorError):
graph_fn = def_function.function(tr, autograph=False)
self.evaluate(
graph_fn(constant_op.constant(0), constant_op.constant(5)))
def test_local_composite_attr(self):
class TestClass(object):
def __init__(self):
self.x = constant_op.constant(3)
def f(n):
while n > 0:
tc = TestClass()
tc.x = tc.x
n -= 1
return n
self.assertTransformedResult(f, constant_op.constant(5), 0)
def test_local_composite_slice(self):
def f(n):
while n > 0:
d = {'x': n}
k = 'x'
d[k] = d[k]
n -= 1
return n
self.assertTransformedResult(f, constant_op.constant(5), 0)
def test_local_composite_literal_slice(self):
def f(n):
while n > 0:
d = {'x': n}
d['x'] = d['x']
n -= 1
return n
self.assertTransformedResult(f, constant_op.constant(5), 0)
def test_non_tensor_state(self):
# This class is ok to be in a tf.while's state.
class TestClass(collections.namedtuple('TestClass', ('x'))):
pass
def f(n):
tc = TestClass([constant_op.constant(0)])
while n > 0:
tc = TestClass([constant_op.constant(3)])
tc.x[0] = tc.x[0] + 1
n -= 1
return tc.x[0]
self.assertTransformedResult(f, constant_op.constant(5), 4)
def test_non_tensor_state_illegal_type(self):
class TestClass(object):
def __init__(self):
self.x = [constant_op.constant(3)]
def f(n):
while n > 0:
tc = TestClass()
tc.x[0] = tc.x[0] + 1
n -= 1
return tc.x[0]
tr = self.transform(f, control_flow)
# The tested function would require `tc` to become part of the while loop
# state, but TensorFlow doesn't support classes at the moment.
with self.assertRaisesRegex(
ValueError, 'tc.*must be defined before the loop'):
tr(constant_op.constant(5))
def test_dispatches_by_cond_only(self):
class TensorIncompatibleNumeric(object):
"""Works in arithmetic expression, but errors out with TF ops."""
def __init__(self, val):
self.val = val
def __add__(self, other):
return TensorIncompatibleNumeric(self.val + other)
def f(n, s):
while n > 0:
n -= 1
s += n
return s
self.assertTransformedResult(f, (constant_op.constant(5), 0), 10)
tr = self.transform(f, control_flow)
# n alone controls the staging. When the loop is not staged, Python
# knows how to add the two objects. But when staged, tf.while will
# not know how to deal with the TensorIncompatibleNumeric object.
self.assertEqual(tr(5, TensorIncompatibleNumeric(0)).val, 10)
with self.assertRaises(TypeError):
tr(constant_op.constant(5), TensorIncompatibleNumeric(0))
class IfStatementTest(ControlFlowTestBase):
def test_basic(self):
def f(n):
a = 0
b = 0
if n > 0:
a = -n
else:
b = 2 * n
return a, b
self.assertTransformedResult(f, constant_op.constant(1), (-1, 0))
self.assertTransformedResult(f, constant_op.constant(-1), (0, -2))
def test_sparse_tensor(self):
def f(cond, a):
if cond:
a = -a
return a
st = sparse_tensor.SparseTensor(
indices=((0,),), values=(0,), dense_shape=(1,))
self.assertTransformedResult(f, (st, constant_op.constant(1)), -1)
self.assertTransformedResult(f, (None, constant_op.constant(1)), 1)
def test_complex_outputs(self):
class TestClass(object):
def __init__(self, a, b):
self.a = a
self.b = b
def f(n, obj):
obj.a = 0
obj.b = 0
if n > 0:
obj.a = -n
else:
obj.b = 2 * n
return obj
tr = self.transform(f, control_flow)
res_obj = tr(constant_op.constant(1), TestClass(0, 0))
self.assertValuesEqual((res_obj.a, res_obj.b), (-1, 0))
res_obj = tr(constant_op.constant(-1), TestClass(0, 0))
self.assertValuesEqual((res_obj.a, res_obj.b), (0, -2))
def test_single_output(self):
def f(n):
if n > 0:
n = -n
return n
self.assertTransformedResult(f, constant_op.constant(1), -1)
def test_unbalanced(self):
def f(n):
if n > 0:
n = 3
return n
self.assertTransformedResult(f, constant_op.constant(2), 3)
self.assertTransformedResult(f, constant_op.constant(-3), -3)
def test_unbalanced_raising(self):
def f(n):
if n > 0:
n = n + 1
raise ValueError()
return n
self.assertTransformedResult(f, -3, -3)
tr = self.transform(f, control_flow)
with self.assertRaises(ValueError):
tr(1)
def test_local_var(self):
def f(n):
if n > 0:
b = 4
n = b + 1
return n
self.assertTransformedResult(f, constant_op.constant(1), 5)
self.assertTransformedResult(f, constant_op.constant(-1), -1)
def test_local_remains_local(self):
def f(n):
if n > 0:
b = 4
n = b + 1
return n
self.assertTransformedResult(f, constant_op.constant(1), 5)
self.assertTransformedResult(f, constant_op.constant(-1), -1)
def test_global_local(self):
def f(n):
if n > 0:
global for_test_global_local
if for_test_global_local is None:
for_test_global_local = 1
else:
for_test_global_local += 1
n += for_test_global_local
return n
tr = self.transform(f, control_flow)
assert for_test_global_local is None
self.assertEqual(tr(1), 2)
self.assertEqual(for_test_global_local, 1)
def test_no_outputs(self):
def f(n):
if n > 0:
b = 4 # pylint:disable=unused-variable
return n
self.assertTransformedResult(f, constant_op.constant(1), 1)
self.assertTransformedResult(f, constant_op.constant(-1), -1)
def test_created_outputs(self):
def f(i):
if i == 0:
result = i - 1
else:
result = i + 1
return result
self.assertTransformedResult(f, 0, -1)
self.assertTransformedResult(f, 1, 2)
def test_created_loop_local_outputs(self):
def f(n, x):
for i in n:
if i == 0:
result = i - 1
else:
result = i + 1
if result > 0:
x += 1
return x
self.assertTransformedResult(f, (range(5), 10), 14)
def test_created_loop_variable(self):
def f(n, x):
for i in n:
if i == 0:
result = i - 1
if i > 0: # Using the result from previous iteration.
if result < 0:
x += 1
return x
self.assertTransformedResult(f, (range(5), 10), 14)
def test_unaffected_global(self):
global for_unaffected_global
for_unaffected_global = 3
def f(i):
global for_unaffected_global
if i == 0:
for_unaffected_global = i - 1
return for_unaffected_global
self.assertTransformedResult(f, 1, 3)
self.assertTransformedResult(f, 0, -1)
self.assertEqual(for_unaffected_global, -1)
def test_unaffected_nonlocal(self):
def f(i):
def inner_fn():
nonlocal n
if i == 0:
n = i - 1
n = 3
inner_fn()
return n
self.assertTransformedResult(f, 1, 3)
self.assertTransformedResult(f, 0, -1)
def test_output_defined_in_prior_except(self):
def f(i):
try:
raise ValueError()
except ValueError:
x = 1
if i == 0:
x = i - 1
return x
self.assertTransformedResult(f, 1, 1)
self.assertTransformedResult(f, 0, -1)
def test_unbalanced_multiple_composites(self):
class Foo(object):
def __init__(self):
self.b = 2
self.c = 3
def f(x, condition):
z = 5
if condition:
x.b = 7
x.c = 11
z = 13
return x.b, x.c, z
self.assertTransformedResult(f, (Foo(), constant_op.constant(True)),
(7, 11, 13))
self.assertTransformedResult(f, (Foo(), constant_op.constant(False)),
(2, 3, 5))
def test_unbalanced_composite(self):
class Foo(object):
def __init__(self):
self.b = 2
def f(x, condition):
z = 5
if condition:
x.b = 7
z = 13
return x.b, z
self.assertTransformedResult(f, (Foo(), constant_op.constant(True)),
(7, 13))
self.assertTransformedResult(f, (Foo(), constant_op.constant(False)),
(2, 5))
class ForStatementTest(ControlFlowTestBase):
def test_basic(self):
def f(l):
s1 = 0
s2 = 0
for e in l:
s1 += e
s2 += e * e
return s1, s2
self.assertTransformedResult(f, constant_op.constant([1, 3]), (4, 10))
empty_vector = constant_op.constant([], shape=(0,), dtype=dtypes.int32)
self.assertTransformedResult(f, empty_vector, (0, 0))
def test_single_output(self):
def f(l):
s = 0
for e in l:
s += e
return s
self.assertTransformedResult(f, constant_op.constant([1, 3]), 4)
empty_vector = constant_op.constant([], shape=(0,), dtype=dtypes.int32)
self.assertTransformedResult(f, empty_vector, 0)
def test_iterated_expression(self):
eval_count = [0]
def count_evals(x):
eval_count[0] += 1
return x
def f(n):
s = 0
for e in count_evals(range(n)):
s += e
return s
tr = self.transform(f, control_flow)
self.assertEqual(tr(5), 10)
self.assertEqual(eval_count[0], 1)
def test_composite_state_initialized_in_loop(self):
class TestClass(object):
pass
def f(n, x):
tc = TestClass()
for i in n:
if i == 0:
tc.x = x
else:
tc.x = tc.x + i
return tc.x
self.assertTransformedResult(f, (range(5), constant_op.constant(10)), 20)
tr = self.transform(f, control_flow)
with self.assertRaisesRegex(
ValueError, "'tc.x' must be defined before the loop"):
tr(constant_op.constant(list(range(5))), 0)
def test_tuple_unpacking(self):
def f(x_list):
z = constant_op.constant(0) # pylint:disable=undefined-variable
for i, x in enumerate(x_list):
z = z + x + i
return z
self.assertTransformedResult(f, [3, 3], 7)
def test_with_comprehension_in_body(self):
def f(l, n):
s = constant_op.constant(list(range(n)))
for _ in l:
s += constant_op.constant([a for a in range(n)])
return s
self.assertTransformedResult(f, (constant_op.constant([1, 2, 3]), 5),
np.array(range(5)) * 4)
class AdvancedControlFlowTest(ControlFlowTestBase):
def assertTransformedEquivalent(self, f, *inputs):
tr = self.transform(
f, (break_statements, continue_statements, control_flow))
self.assertEqual(f(*inputs), tr(*inputs))
def test_while_with_else(self):
def f(x):
while x > 2:
x /= 2
else:
x += 1
return x
self.assertTransformedEquivalent(f, 4)
self.assertTransformedEquivalent(f, 2)
def test_while_with_else_and_break(self):
def f(cond1):
x = 8
while x > 2:
x /= 2
if cond1:
break
else:
x += 1
return x
self.assertTransformedEquivalent(f, True)
self.assertTransformedEquivalent(f, False)
def test_for_with_else(self):
def f(l):
res = 0
for x in l:
res += x
else:
res += 1
return res
self.assertTransformedEquivalent(f, [])
self.assertTransformedEquivalent(f, [1, 2])
def test_for_with_else_and_break(self):
def f(flag):
l = [1, 2, 3]
res = 0
for x in l:
res += x
if flag:
break
else:
res += 1
return res
self.assertTransformedEquivalent(f, True)
self.assertTransformedEquivalent(f, False)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,177 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Handles directives.
This converter removes the directive functions from the code and moves the
information they specify into AST annotations. It is a specialized form of
static analysis, one that is specific to AutoGraph.
Note that this requires that the actual directive functions are static - that
is, they do not change at runtime. So if you do something like this:
tf.autograph.set_loop_options = <new function>
Then the directive will may no longer be recognized. Furthermore, if the
converted function is cached, such an action may be irreversible.
"""
import inspect
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.lang import directives
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.util import tf_inspect
STATIC_VALUE = 'static_value'
"""Used for AST annotations, see visit_Name."""
class _LoopScope(object):
def __init__(self):
self.ast_node = None
self.statements_visited = 0
def _map_args(call_node, function):
"""Maps AST call nodes to the actual function's arguments.
Args:
call_node: ast.Call
function: Callable[..., Any], the actual function matching call_node
Returns:
Dict[Text, ast.AST], mapping each of the function's argument names to
the respective AST node.
Raises:
ValueError: if the default arguments are not correctly set
"""
args = call_node.args
kwds = {kwd.arg: kwd.value for kwd in call_node.keywords}
call_args = tf_inspect.getcallargs(function, *args, **kwds)
# Keyword arguments not specified in kwds will be mapped to their defaults,
# which are Python values. Since we don't currently have a way to transform
# those into AST references, we simply remove them. By convention, directives
# use UNSPECIFIED as default value for optional arguments. No other
# defaults should be present.
unexpected_defaults = []
for k in call_args:
if (k not in kwds
and call_args[k] not in args
and call_args[k] is not directives.UNSPECIFIED):
unexpected_defaults.append(k)
if unexpected_defaults:
raise ValueError('Unexpected keyword argument values, %s, for function %s'
% (zip(unexpected_defaults,
[call_args[k] for k in unexpected_defaults]),
function))
return {k: v for k, v in call_args.items() if v is not directives.UNSPECIFIED}
class DirectivesTransformer(converter.Base):
"""Parses compiler directives and converts them into AST annotations."""
def _process_symbol_directive(self, call_node, directive):
if len(call_node.args) < 1:
raise ValueError('"%s" requires a positional first argument'
' as the target' % directive.__name__)
target = call_node.args[0]
defs = anno.getanno(target, anno.Static.ORIG_DEFINITIONS)
for def_ in defs:
def_.directives[directive] = _map_args(call_node, directive)
return call_node
def _process_statement_directive(self, call_node, directive):
if self.state[_LoopScope].statements_visited > 1:
raise ValueError(
'"%s" must be the first statement in the loop block' % (
directive.__name__))
if self.state[_LoopScope].level < 2:
raise ValueError(
'"%s" must be used inside a statement' % directive.__name__)
target = self.state[_LoopScope].ast_node
node_anno = anno.getanno(target, anno.Basic.DIRECTIVES, {})
node_anno[directive] = _map_args(call_node, directive)
anno.setanno(target, anno.Basic.DIRECTIVES, node_anno)
return call_node
def visit_Name(self, node):
node = self.generic_visit(node)
if isinstance(node.ctx, gast.Load):
defs = anno.getanno(node, anno.Static.DEFINITIONS, ())
is_defined = bool(defs)
if not is_defined and node.id in self.ctx.info.namespace:
anno.setanno(node, STATIC_VALUE, self.ctx.info.namespace[node.id])
return node
def visit_Attribute(self, node):
node = self.generic_visit(node)
parent_val = anno.getanno(node.value, STATIC_VALUE, default=None)
if parent_val is not None and inspect.ismodule(parent_val):
if hasattr(parent_val, node.attr):
anno.setanno(node, STATIC_VALUE, getattr(parent_val, node.attr))
return node
def visit_Assign(self, node):
self.state[_LoopScope].statements_visited += 1
return self.generic_visit(node)
def visit_AugAssign(self, node):
self.state[_LoopScope].statements_visited += 1
return self.generic_visit(node)
def visit_Expr(self, node):
self.state[_LoopScope].statements_visited += 1
node = self.generic_visit(node)
if isinstance(node.value, gast.Call):
call_node = node.value
static_val = anno.getanno(call_node.func, STATIC_VALUE, default=None)
if static_val is not None:
# Note: directive calls are not output in the generated code, hence
# the removal from the code by returning None.
if static_val is directives.set_element_type:
self._process_symbol_directive(call_node, static_val)
return None
elif static_val is directives.set_loop_options:
self._process_statement_directive(call_node, static_val)
return None
return node
# TODO(mdan): This will be insufficient for other control flow.
# That means that if we ever have a directive that affects things other than
# loops, we'll need support for parallel scopes, or have multiple converters.
def _track_and_visit_loop(self, node):
self.state[_LoopScope].enter()
self.state[_LoopScope].ast_node = node
node = self.generic_visit(node)
# Edge case: a loop with just one directive statement would become empty.
if not node.body:
node.body = [gast.Pass()]
self.state[_LoopScope].exit()
return node
def visit_While(self, node):
return self._track_and_visit_loop(node)
def visit_For(self, node):
return self._track_and_visit_loop(node)
def transform(node, ctx):
return DirectivesTransformer(ctx).visit(node)
@@ -0,0 +1,133 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for directives module."""
from tensorflow.python.autograph.converters import directives as directives_converter
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.autograph.lang import directives
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.platform import test
class DirectivesTest(converter_testing.TestCase):
def test_local_target(self):
def f():
l = []
string_var = 0
directives.set_element_type(l, 'a', string_var)
_, node, _ = self.transform(f, directives_converter, include_ast=True)
def_, = anno.getanno(node.body[0].targets[0],
anno.Static.DEFINITIONS)
d = def_.directives[directives.set_element_type]
self.assertEqual(d['dtype'].value, 'a')
self.assertEqual(d['shape'].id, 'string_var')
def test_argument_target(self):
def f(a):
directives.set_element_type(a, 1, shape=2)
pass
_, node, _ = self.transform(f, directives_converter, include_ast=True)
def_, = anno.getanno(node.args.args[0], anno.Static.DEFINITIONS)
d = def_.directives[directives.set_element_type]
self.assertEqual(d['dtype'].value, 1)
self.assertEqual(d['shape'].value, 2)
def test_loop_target(self):
def f():
a = True
while True:
directives.set_loop_options(parallel_iterations=10, back_prop=a) # pylint: disable=unexpected-keyword-arg
pass
_, node, _ = self.transform(f, directives_converter, include_ast=True)
d = anno.getanno(node.body[1], anno.Basic.DIRECTIVES)
d = d[directives.set_loop_options]
self.assertEqual(d['parallel_iterations'].value, 10)
self.assertEqual(d['back_prop'].id, 'a')
self.assertNotIn('swap_memory', d)
def test_loop_target_no_loop(self):
def f():
directives.set_loop_options()
pass
with self.assertRaisesRegex(ValueError, 'must be used inside a statement'):
self.transform(f, directives_converter, include_ast=True)
def test_loop_target_not_first(self):
def f():
a = 1
while True:
a = 2
directives.set_loop_options(parallel_iterations=10, back_prop=a) # pylint: disable=unexpected-keyword-arg
with self.assertRaisesRegex(ValueError, 'must be the first statement'):
self.transform(f, directives_converter, include_ast=True)
def test_value_verification_does_not_trigger_properties(self):
self_test = self
class TestClass(object):
@property
def b(self):
self_test.fail('This should never be evaluated')
tc = TestClass()
def f():
return tc.b + 1
_, node, _ = self.transform(f, directives_converter, include_ast=True)
self.assertIsNotNone(node)
def test_value_verification_does_not_trigger_getattr(self):
class TestClass(object):
def __init__(self):
self.getattr_called = False
def __getattr__(self, _):
# Note: seems that any exception raised here is absorbed by hasattr.
# So we can't call test.fail or raise.
self.getattr_called = True
tc = TestClass()
def f():
return tc.b + 1
_, node, _ = self.transform(f, directives_converter, include_ast=True)
self.assertIsNotNone(node)
self.assertFalse(tc.getattr_called)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,134 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Converts function definitions and lambdas by adding necessary boilerplate."""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct.static_analysis import activity
from tensorflow.python.autograph.pyct.static_analysis import annos
class _Function(object):
def __init__(self):
self.context_name = None
class FunctionTransformer(converter.Base):
"""Wraps function bodies around autograph-specific boilerplate."""
def _function_scope_options(self, fn_scope):
"""Returns the options with which to create function scopes."""
# Top-level function receive the options that were directly requested.
# All others receive the options corresponding to a recursive conversion.
# Note: this mainly controls the user_requested flag, which is important
# primarily because the FunctionScope context also creates a
# ControlStatusCtx(autograph=ENABLED) when user_requested is True. See
# function_wrappers.py.
if fn_scope.level == 2:
return self.ctx.user.options
return self.ctx.user.options.call_options()
def visit_Lambda(self, node):
with self.state[_Function] as fn_scope:
node = self.generic_visit(node)
# TODO(mdan): Fix the tests so that we can always add this decorator.
if fn_scope.level > 2:
return templates.replace_as_expression(
'ag__.autograph_artifact(l)', l=node)
scope = anno.getanno(node, anno.Static.SCOPE)
function_context_name = self.ctx.namer.new_symbol('lscope',
scope.referenced)
fn_scope.context_name = function_context_name
anno.setanno(node, 'function_context_name', function_context_name)
template = """
ag__.with_function_scope(
lambda function_context: body, function_context_name, options)
"""
node.body = templates.replace_as_expression(
template,
options=self._function_scope_options(fn_scope).to_ast(),
function_context=function_context_name,
function_context_name=gast.Constant(function_context_name, kind=None),
body=node.body)
return node
def visit_FunctionDef(self, node):
with self.state[_Function] as fn_scope:
scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
function_context_name = self.ctx.namer.new_symbol('fscope',
scope.referenced)
fn_scope.context_name = function_context_name
anno.setanno(node, 'function_context_name', function_context_name)
node = self.generic_visit(node)
if fn_scope.level <= 2:
# Top-level functions lose their decorator because the conversion is
# always just-in-time and by the time it happens the decorators are
# already set to be applied.
node.decorator_list = []
else:
# TODO(mdan): Fix the tests so that we can always add this decorator.
# Inner functions are converted already, so we insert a decorator to
# prevent double conversion. Double conversion would work too, but this
# saves the overhead.
node.decorator_list.append(
parser.parse_expression('ag__.autograph_artifact'))
docstring_node = None
if node.body:
first_statement = node.body[0]
if (isinstance(first_statement, gast.Expr) and
isinstance(first_statement.value, gast.Constant)):
docstring_node = first_statement
node.body = node.body[1:]
template = """
with ag__.FunctionScope(
function_name, context_name, options) as function_context:
body
"""
wrapped_body = templates.replace(
template,
function_name=gast.Constant(node.name, kind=None),
context_name=gast.Constant(function_context_name, kind=None),
options=self._function_scope_options(fn_scope).to_ast(),
function_context=function_context_name,
body=node.body)
if docstring_node is not None:
wrapped_body = [docstring_node] + wrapped_body
node.body = wrapped_body
return node
def transform(node, ctx):
node = qual_names.resolve(node)
node = activity.resolve(node, ctx, None)
return FunctionTransformer(ctx).visit(node)
@@ -0,0 +1,127 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for functions module."""
from tensorflow.python.autograph.converters import functions
from tensorflow.python.autograph.converters import return_statements
from tensorflow.python.autograph.core import ag_ctx
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.autograph.impl import api
from tensorflow.python.framework import constant_op
from tensorflow.python.platform import test
class FunctionTransformer(converter_testing.TestCase):
def test_basic(self):
def f(l):
"""Docstring."""
a = 1
l += a
return l
tr = self.transform(f, functions)
result_op = tr(constant_op.constant(1))
self.assertIn('f/', result_op.op.name)
self.assertEqual('Docstring.', tr.__doc__)
def test_multiline_docstring(self):
def f():
"""First sentence.
Second sentence.
Returns:
Something.
"""
return constant_op.constant(1)
tr = self.transform(f, functions)
result_op = tr()
self.assertIn('f/', result_op.op.name)
self.assertIn('First sentence.', tr.__doc__)
self.assertIn('Second sentence.', tr.__doc__)
def test_nested_functions(self):
def f(l):
def inner_fn(i):
return i + 1
l += 1
return l, inner_fn(l)
tr = self.transform(f, (functions, return_statements))
first, second = tr(constant_op.constant(1))
self.assertIn('f/', first.op.name)
self.assertNotIn('inner_fn', first.op.name)
self.assertIn('f/inner_fn/', second.op.inputs[0].name)
def test_conversion_context_preserves_in_inner_functions(self):
def inner_fn_callee():
self.assertEqual(
ag_ctx.control_status_ctx().status, ag_ctx.Status.DISABLED)
def f():
def inner_fn():
inner_fn_callee()
with ag_ctx.ControlStatusCtx(
ag_ctx.Status.DISABLED, converter.ConversionOptions(recursive=True)):
inner_fn()
tr = self.transform(f, functions)
tr()
def test_method(self):
class TestClass(object):
def f(self, l):
def inner_fn(i):
return i + 1
l += 1
return l, inner_fn(l)
tr = self.transform(TestClass.f, (functions, return_statements))
first, second = tr(TestClass(), constant_op.constant(1))
self.assertIn('f/', first.op.name)
self.assertNotIn('inner_fn', first.op.name)
self.assertIn('f/inner_fn/', second.op.inputs[0].name)
def test_lambda_in_return_value(self):
def f():
return lambda x: x + 1
tr = self.transform(f, functions)
result_l = tr()
self.assertTrue(api.is_autograph_artifact(result_l))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,239 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Converter for list operations.
This includes converting Python lists to TensorArray/TensorList.
"""
# TODO(mdan): Elaborate the logic here.
# TODO(mdan): Does it even make sense to attempt to try to use TAs?
# The current rule (always convert to TensorArray) is naive and insufficient.
# In general, a better mechanism could look like:
# * convert to TensorList by default
# * leave as Python list if the user explicitly forbids it
# * convert to TensorArray only when complete write once behavior can be
# guaranteed (e.g. list comprehensions)
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.lang import directives
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct.static_analysis import activity
from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
class _Statement(object):
def __init__(self):
self.pop_uses = None
class ListTransformer(converter.Base):
"""Converts lists and related operations to their TF counterpart."""
def visit_List(self, node):
node = self.generic_visit(node)
template = """
ag__.new_list(elements)
"""
return templates.replace_as_expression(template, elements=node)
def _replace_append_call(self, node):
assert len(node.args) == 1
assert isinstance(node.func, gast.Attribute)
template = """
target = ag__.list_append(target, element)
"""
return templates.replace(
template,
target=node.func.value,
element=node.args[0])
def _replace_pop_call(self, node):
# Expressions that use pop() are converted to a statement + expression.
#
# For example:
#
# print(target.pop())
#
# ... is converted to:
#
# target, target_pop = ag__.list_pop(target)
# print(target_pop)
#
# Here, we just generate the variable name and swap it in,
# and _generate_pop_operation will handle the rest.
#
# Multiple uses of pop() are allowed:
#
# print(tartget.pop(), target.pop())
# print(tartget.pop().pop())
#
assert isinstance(node.func, gast.Attribute)
scope = anno.getanno(node, NodeAnno.ARGS_SCOPE)
target_node = node.func.value
# Attempt to use a related name if one exists. Otherwise use something
# generic.
if anno.hasanno(target_node, anno.Basic.QN):
target_name = anno.getanno(target_node, anno.Basic.QN).ssf()
else:
target_name = 'list_'
pop_var_name = self.ctx.namer.new_symbol(target_name, scope.referenced)
stmt = self.state[_Statement]
if stmt.pop_uses is None:
stmt.pop_uses = []
stmt.pop_uses.append((node, pop_var_name))
return templates.replace_as_expression('var_name', var_name=pop_var_name)
def _replace_stack_call(self, node):
assert len(node.args) == 1
dtype = self.get_definition_directive(
node.args[0],
directives.set_element_type,
'dtype',
default=templates.replace_as_expression('None'))
template = """
ag__.list_stack(
target,
opts=ag__.ListStackOpts(
element_dtype=dtype,
original_call=orig_call))
"""
return templates.replace_as_expression(
template,
dtype=dtype,
target=node.args[0],
orig_call=node.func)
def visit_Call(self, node):
node = self.generic_visit(node)
# TODO(mdan): This is insufficient if target is a function argument.
# In the case of function arguments, we need to add the list to the
# function's return value, because it is being modified.
# TODO(mdan): Checking just the name is brittle, can it be improved?
if isinstance(node.func, gast.Attribute):
func_name = node.func.attr
if func_name == 'append' and (len(node.args) == 1):
node = self._replace_append_call(node)
elif func_name == 'pop' and (len(node.args) <= 1):
node = self._replace_pop_call(node)
elif (func_name == 'stack' and (len(node.args) == 1) and
(not node.keywords or node.keywords[0].arg == 'strict')):
# This avoids false positives with keyword args.
# TODO(mdan): handle kwargs properly.
node = self._replace_stack_call(node)
return node
def _generate_pop_operation(self, original_call_node, pop_var_name):
assert isinstance(original_call_node.func, gast.Attribute)
if original_call_node.args:
pop_element = original_call_node.args[0]
else:
pop_element = parser.parse_expression('None')
# The call will be something like "target.pop()", and the dtype is hooked to
# target, hence the func.value.
# TODO(mdan): For lists of lists, this won't work.
# The reason why it won't work is because it's unclear how to annotate
# the list as a "list of lists with a certain element type" when using
# operations like `l.pop().pop()`.
dtype = self.get_definition_directive(
original_call_node.func.value,
directives.set_element_type,
'dtype',
default=templates.replace_as_expression('None'))
shape = self.get_definition_directive(
original_call_node.func.value,
directives.set_element_type,
'shape',
default=templates.replace_as_expression('None'))
template = """
target, pop_var_name = ag__.list_pop(
target, element,
opts=ag__.ListPopOpts(element_dtype=dtype, element_shape=shape))
"""
return templates.replace(
template,
target=original_call_node.func.value,
pop_var_name=pop_var_name,
element=pop_element,
dtype=dtype,
shape=shape)
def _postprocess_statement(self, node):
"""Inserts any separate pop() calls that node may use."""
pop_uses = self.state[_Statement].pop_uses
if pop_uses:
replacements = []
for original_call_node, pop_var_name in pop_uses:
replacements.extend(
self._generate_pop_operation(original_call_node, pop_var_name))
replacements.append(node)
node = replacements
self.state[_Statement].exit()
return node, None
def _visit_and_process_block(self, block):
return self.visit_block(
block,
before_visit=self.state[_Statement].enter,
after_visit=self._postprocess_statement)
def visit_FunctionDef(self, node):
node.args = self.generic_visit(node.args)
node.decorator_list = self.visit_block(node.decorator_list)
node.body = self._visit_and_process_block(node.body)
return node
def visit_For(self, node):
node.target = self.visit(node.target)
node.body = self._visit_and_process_block(node.body)
node.orelse = self._visit_and_process_block(node.orelse)
return node
def visit_While(self, node):
node.test = self.visit(node.test)
node.body = self._visit_and_process_block(node.body)
node.orelse = self._visit_and_process_block(node.orelse)
return node
def visit_If(self, node):
node.test = self.visit(node.test)
node.body = self._visit_and_process_block(node.body)
node.orelse = self._visit_and_process_block(node.orelse)
return node
def visit_With(self, node):
node.items = self.visit_block(node.items)
node.body = self._visit_and_process_block(node.body)
return node
def transform(node, ctx):
node = qual_names.resolve(node)
node = activity.resolve(node, ctx, None)
return ListTransformer(ctx).visit(node)
@@ -0,0 +1,107 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for lists module."""
from tensorflow.python.autograph.converters import directives as directives_converter
from tensorflow.python.autograph.converters import lists
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.autograph.lang import directives
from tensorflow.python.autograph.lang import special_functions
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import list_ops
from tensorflow.python.platform import test
class ListTest(converter_testing.TestCase):
def test_empty_list(self):
def f():
return []
tr = self.transform(f, lists)
tl = tr()
# Empty tensor lists cannot be evaluated or stacked.
self.assertIsInstance(tl, tensor.Tensor)
self.assertEqual(tl.dtype, dtypes.variant)
def test_initialized_list(self):
def f():
return [1, 2, 3]
tr = self.transform(f, lists)
self.assertAllEqual(tr(), [1, 2, 3])
def test_list_append(self):
def f():
l = special_functions.tensor_list([1])
l.append(2)
l.append(3)
return l
tr = self.transform(f, lists)
tl = tr()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(self.evaluate(r), [1, 2, 3])
def test_list_pop(self):
def f():
l = special_functions.tensor_list([1, 2, 3])
directives.set_element_type(l, dtype=dtypes.int32, shape=())
s = l.pop()
return s, l
tr = self.transform(f, (directives_converter, lists))
ts, tl = tr()
r = list_ops.tensor_list_stack(tl, dtypes.int32)
self.assertAllEqual(self.evaluate(r), [1, 2])
self.assertAllEqual(self.evaluate(ts), 3)
def test_double_list_pop(self):
def f(l):
s = l.pop().pop()
return s
tr = self.transform(f, lists)
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(tr(test_input), 3)
def test_list_stack(self):
def f():
l = [1, 2, 3]
return array_ops_stack.stack(l)
tr = self.transform(f, lists)
self.assertAllEqual(self.evaluate(tr()), [1, 2, 3])
if __name__ == '__main__':
test.main()
@@ -0,0 +1,131 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Converter for logical expressions, e.g. `a and b -> tf.logical_and(a, b)`."""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import templates
# TODO(mdan): Properly extract boolean ops according to lazy eval rules.
# Note that this isn't completely safe either, because tensors may have control
# dependencies.
# Note that for loops that should be done after the loop was converted to
# tf.while_loop so that the expanded conditionals are properly scoped.
# Used to signal that an operand is safe for non-lazy evaluation.
SAFE_BOOLEAN_OPERAND = 'SAFE_BOOLEAN_OPERAND'
LOGICAL_OPERATORS = {
gast.And: 'ag__.and_',
gast.Not: 'ag__.not_',
gast.Or: 'ag__.or_',
}
EQUALITY_OPERATORS = {
gast.Eq: 'ag__.eq',
gast.NotEq: 'ag__.not_eq',
}
class LogicalExpressionTransformer(converter.Base):
"""Converts logical expressions to corresponding TF calls."""
def _overload_of(self, operator):
op_type = type(operator)
if op_type in LOGICAL_OPERATORS:
return LOGICAL_OPERATORS[op_type]
if self.ctx.user.options.uses(converter.Feature.EQUALITY_OPERATORS):
if op_type in EQUALITY_OPERATORS:
return EQUALITY_OPERATORS[op_type]
return None
def _as_lambda(self, expr):
return templates.replace_as_expression('lambda: expr', expr=expr)
def _as_binary_function(self, func_name, arg1, arg2):
return templates.replace_as_expression(
'func_name(arg1, arg2)',
func_name=parser.parse_expression(func_name),
arg1=arg1,
arg2=arg2)
def _as_binary_operation(self, op, arg1, arg2):
template = templates.replace_as_expression(
'arg1 is arg2', # Note: `is` will be replaced with `op` below.
arg1=arg1,
arg2=arg2)
template.ops[0] = op
return template
def _as_unary_function(self, func_name, arg):
return templates.replace_as_expression(
'func_name(arg)', func_name=parser.parse_expression(func_name), arg=arg)
def _process_binop(self, op, left, right):
overload = self._overload_of(op)
if overload is None:
return self._as_binary_operation(op, left, right)
return self._as_binary_function(overload, left, right)
def visit_Compare(self, node):
node = self.generic_visit(node)
ops_and_comps = list(zip(node.ops, node.comparators))
left = node.left
# Repeated comparisons are converted to conjunctions:
# a < b < c -> a < b and b < c
op_tree = None
while ops_and_comps:
op, right = ops_and_comps.pop(0)
binary_comparison = self._process_binop(op, left, right)
if op_tree is not None:
op_tree = self._as_binary_function('ag__.and_',
self._as_lambda(op_tree),
self._as_lambda(binary_comparison))
else:
op_tree = binary_comparison
left = right
assert op_tree is not None
return op_tree
def visit_UnaryOp(self, node):
node = self.generic_visit(node)
overload = self._overload_of(node.op)
if overload is None:
return node
return self._as_unary_function(overload, node.operand)
def visit_BoolOp(self, node):
node = self.generic_visit(node)
node_values = node.values
right = node.values.pop()
while node_values:
left = node_values.pop()
right = self._as_binary_function(
self._overload_of(node.op), self._as_lambda(left),
self._as_lambda(right))
return right
def transform(node, ctx):
transformer = LogicalExpressionTransformer(ctx)
return transformer.visit(node)
@@ -0,0 +1,82 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for logical_expressions module."""
from tensorflow.python.autograph.converters import logical_expressions
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
class LogicalExpressionTest(converter_testing.TestCase):
def test_equals(self):
def f(a, b):
return a == b
tr = self.transform(f, logical_expressions)
self.assertTrue(self.evaluate(tr(constant_op.constant(1), 1)))
self.assertFalse(self.evaluate(tr(constant_op.constant(1), 2)))
@test_util.run_deprecated_v1
def test_bool_ops(self):
def f(a, b, c):
return (a or b) and (a or b or c) and not c
tr = self.transform(f, logical_expressions)
self.assertTrue(self.evaluate(tr(constant_op.constant(True), False, False)))
self.assertFalse(self.evaluate(tr(constant_op.constant(True), False, True)))
def test_comparison(self):
def f(a, b, c, d):
return a < b == c > d
tr = self.transform(f, logical_expressions)
# Note: having just the first constant a tensor tests that the
# operations execute in the correct order. If anything other than
# a < b executed first, the result would be a Python scalar and not a
# Tensor. This is valid as long as the dispat is automatic based on
# type.
self.assertTrue(self.evaluate(tr(constant_op.constant(1), 2, 2, 1)))
self.assertFalse(self.evaluate(tr(constant_op.constant(1), 2, 2, 3)))
def test_default_ops(self):
def f(a, b):
return a in b
tr = self.transform(f, logical_expressions)
self.assertTrue(tr('a', ('a',)))
def test_unary_ops(self):
def f(a):
return ~a, -a, +a
tr = self.transform(f, logical_expressions)
self.assertEqual(tr(1), (-2, -1, 1))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,402 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Canonicalizes functions with multiple returns to use just one."""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct.static_analysis import activity
from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
BODY_DEFINITELY_RETURNS = 'BODY_DEFINITELY_RETURNS'
ORELSE_DEFINITELY_RETURNS = 'ORELSE_DEFINITELY_RETURNS'
STMT_DEFINITELY_RETURNS = 'STMT_DEFINITELY_RETURNS'
class _RewriteBlock(object):
def __init__(self):
self.definitely_returns = False
class ConditionalReturnRewriter(converter.Base):
"""Rewrites a pattern where it's unobvious that all paths return a value.
This rewrite allows avoiding intermediate None return values.
The following pattern:
if cond:
<block 1>
return
else:
<block 2>
<block 3>
is converted to:
if cond:
<block 1>
return
else:
<block 2>
<block 3>
and vice-versa (if the else returns, subsequent statements are moved under the
if branch).
"""
def visit_Return(self, node):
self.state[_RewriteBlock].definitely_returns = True
return node
def _postprocess_statement(self, node):
# If the node definitely returns (e.g. it's a with statement with a
# return statement in it), then the current block also definitely returns.
if anno.getanno(node, STMT_DEFINITELY_RETURNS, default=False):
self.state[_RewriteBlock].definitely_returns = True
# The special case: collapse a typical conditional return pattern into
# a single conditional with possibly returns on both branches. This
# reduces the use of None return values, which don't work with TF
# conditionals.
if (isinstance(node, gast.If)
and anno.getanno(node, BODY_DEFINITELY_RETURNS, default=False)):
return node, node.orelse
elif (isinstance(node, gast.If)
and anno.getanno(node, ORELSE_DEFINITELY_RETURNS, default=False)):
return node, node.body
return node, None
def _visit_statement_block(self, node, nodes):
self.state[_RewriteBlock].enter()
new_nodes = self.visit_block(nodes, after_visit=self._postprocess_statement)
block_definitely_returns = self.state[_RewriteBlock].definitely_returns
self.state[_RewriteBlock].exit()
return new_nodes, block_definitely_returns
def visit_While(self, node):
node.test = self.visit(node.test)
node.body, _ = self._visit_statement_block(node, node.body)
node.orelse, _ = self._visit_statement_block(node, node.orelse)
return node
def visit_For(self, node):
node.iter = self.visit(node.iter)
node.target = self.visit(node.target)
node.body, _ = self._visit_statement_block(node, node.body)
node.orelse, _ = self._visit_statement_block(node, node.orelse)
return node
def visit_With(self, node):
node.items = self.visit_block(node.items)
node.body, definitely_returns = self._visit_statement_block(node, node.body)
if definitely_returns:
anno.setanno(node, STMT_DEFINITELY_RETURNS, True)
return node
def visit_Try(self, node):
# We could decide whether a 'try' DEFINITELY_RETURNS based on its components
# It is not clear whether we want to do anything with this given
# a 'try' is likely to throw an exception in some circumstances.
node.body, _ = self._visit_statement_block(node, node.body)
node.orelse, _ = self._visit_statement_block(node, node.orelse)
node.finalbody, _ = self._visit_statement_block(node, node.finalbody)
node.handlers = self.visit_block(node.handlers)
return node
def visit_ExceptHandler(self, node):
# To determine whether `try` DEFINITELY_RETURNS we need to revisit this.
node.body, _ = self._visit_statement_block(node, node.body)
return node
def visit_If(self, node):
node.test = self.visit(node.test)
node.body, body_definitely_returns = self._visit_statement_block(
node, node.body)
if body_definitely_returns:
anno.setanno(node, BODY_DEFINITELY_RETURNS, True)
node.orelse, orelse_definitely_returns = self._visit_statement_block(
node, node.orelse)
if orelse_definitely_returns:
anno.setanno(node, ORELSE_DEFINITELY_RETURNS, True)
if body_definitely_returns and orelse_definitely_returns:
self.state[_RewriteBlock].definitely_returns = True
return node
def visit_FunctionDef(self, node):
node.args = self.visit(node.args)
node.body, _ = self._visit_statement_block(node, node.body)
return node
class _Block(object):
def __init__(self):
self.is_function = False
self.return_used = False
self.create_guard_next = False
self.create_guard_now = False
def __repr__(self):
return 'used: {}'.format(
self.return_used)
class _Function(object):
def __init__(self):
self.do_return_var_name = None
self.retval_var_name = None
def __repr__(self):
return 'return control: {}, return value: {}'.format(
self.do_return_var_name, self.retval_var_name)
class ReturnStatementsTransformer(converter.Base):
"""Lowers return statements into variables and conditionals.
Specifically, the following pattern:
<block 1>
return val
<block 2>
is converted to:
do_return = False
retval = None
<block 1>
do_return = True
retval = val
if not do_return:
<block 2>
return retval
The conversion adjusts loops as well:
<block 1>
while cond:
<block 2>
return retval
is converted to:
<block 1>
while not do_return and cond:
<block 2>
do_return = True
retval = val
"""
def __init__(self, ctx, allow_missing_return):
super(ReturnStatementsTransformer, self).__init__(ctx)
self.allow_missing_return = allow_missing_return
def visit_Return(self, node):
for block in reversed(self.state[_Block].stack):
block.return_used = True
block.create_guard_next = True
if block.is_function:
break
retval = node.value if node.value else parser.parse_expression('None')
# Note: If `return <expr> raises, then the return is aborted.
# The try-catch below ensures the variables remain consistent in that case.
template = """
try:
do_return_var_name = True
retval_var_name = retval
except:
do_return_var_name = False
raise
"""
node = templates.replace(
template,
do_return_var_name=self.state[_Function].do_return_var_name,
retval_var_name=self.state[_Function].retval_var_name,
retval=retval)
return node
def _postprocess_statement(self, node):
if not self.state[_Block].return_used:
return node, None
state = self.state[_Block]
if state.create_guard_now:
template = """
if not do_return_var_name:
original_node
"""
cond, = templates.replace(
template,
do_return_var_name=self.state[_Function].do_return_var_name,
original_node=node)
node, block = cond, cond.body
else:
node, block = node, None
state.create_guard_now = state.create_guard_next
state.create_guard_next = False
return node, block
def _visit_statement_block(self, node, nodes):
self.state[_Block].enter()
nodes = self.visit_block(nodes, after_visit=self._postprocess_statement)
self.state[_Block].exit()
return nodes
def visit_While(self, node):
node.test = self.visit(node.test)
# Add the check for return to the loop condition.
node.body = self._visit_statement_block(node, node.body)
if self.state[_Block].return_used:
node.test = templates.replace_as_expression(
'not control_var and test',
test=node.test,
control_var=self.state[_Function].do_return_var_name)
node.orelse = self._visit_statement_block(node, node.orelse)
return node
def visit_For(self, node):
node.iter = self.visit(node.iter)
node.target = self.visit(node.target)
# Add the check for return to the loop condition.
node.body = self._visit_statement_block(node, node.body)
if self.state[_Block].return_used:
extra_test = anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST, default=None)
if extra_test is not None:
extra_test = templates.replace_as_expression(
'not control_var and extra_test',
extra_test=extra_test,
control_var=self.state[_Function].do_return_var_name)
else:
extra_test = templates.replace_as_expression(
'not control_var',
control_var=self.state[_Function].do_return_var_name)
anno.setanno(node, anno.Basic.EXTRA_LOOP_TEST, extra_test)
node.orelse = self._visit_statement_block(node, node.orelse)
return node
def visit_With(self, node):
node.items = self.visit_block(node.items)
node.body = self._visit_statement_block(node, node.body)
return node
def visit_Try(self, node):
node.body = self._visit_statement_block(node, node.body)
node.orelse = self._visit_statement_block(node, node.orelse)
node.finalbody = self._visit_statement_block(node, node.finalbody)
node.handlers = self.visit_block(node.handlers)
return node
def visit_ExceptHandler(self, node):
node.body = self._visit_statement_block(node, node.body)
return node
def visit_If(self, node):
node.test = self.visit(node.test)
node.body = self._visit_statement_block(node, node.body)
node.orelse = self._visit_statement_block(node, node.orelse)
return node
def visit_FunctionDef(self, node):
with self.state[_Function] as fn:
with self.state[_Block] as block:
block.is_function = True
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
do_return_var_name = self.ctx.namer.new_symbol('do_return',
scope.referenced)
retval_var_name = self.ctx.namer.new_symbol('retval_', scope.referenced)
fn.do_return_var_name = do_return_var_name
fn.retval_var_name = retval_var_name
node.body = self._visit_statement_block(node, node.body)
if block.return_used:
if self.allow_missing_return:
# The function would have a single `with` node that wraps the
# entire body. If the function had a docstring, the body has two
# nodes, with the `with` as the second node.
wrapper_node = node.body[-1]
assert isinstance(wrapper_node, gast.With), (
'This transformer requires the functions converter.')
template = """
do_return_var_name = False
retval_var_name = ag__.UndefinedReturnValue()
body
return function_context.ret(retval_var_name, do_return_var_name)
"""
wrapper_node.body = templates.replace(
template,
body=wrapper_node.body,
do_return_var_name=do_return_var_name,
function_context=anno.getanno(node, 'function_context_name'),
retval_var_name=retval_var_name)
else:
template = """
body
return retval_var_name
"""
node.body = templates.replace(
template,
body=node.body,
do_return_var_name=do_return_var_name,
retval_var_name=retval_var_name)
return node
def transform(node, ctx, default_to_null_return=True):
"""Ensure a function has only a single return, at the end."""
node = qual_names.resolve(node)
node = activity.resolve(node, ctx, None)
# Note: Technically, these two could be merged into a single walk, but
# keeping them separate helps with readability.
node = ConditionalReturnRewriter(ctx).visit(node)
node = qual_names.resolve(node)
node = activity.resolve(node, ctx, None)
transformer = ReturnStatementsTransformer(
ctx, allow_missing_return=default_to_null_return)
node = transformer.visit(node)
return node
@@ -0,0 +1,254 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for return_statements module."""
from tensorflow.python.autograph.converters import functions
from tensorflow.python.autograph.converters import return_statements
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.framework import ops
from tensorflow.python.platform import test
class SingleReturnTest(converter_testing.TestCase):
def assertTransformedEquivalent(self, f, *inputs):
tr = self.transform(f, (functions, return_statements))
self.assertEqual(f(*inputs), tr(*inputs))
def test_straightline(self):
def f(x):
return x * x
self.assertTransformedEquivalent(f, 2)
def test_superfluous_returns(self):
def f():
retval = 1
return retval
retval = 2 # pylint:disable=unreachable
return retval
self.assertTransformedEquivalent(f)
def test_superfluous_returns_adjacent(self):
def f():
return 1
return 2 # pylint:disable=unreachable
self.assertTransformedEquivalent(f)
def test_conditional(self):
def f(x):
if x > 0:
return x
else:
return x * x
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_conditional_missing_else(self):
def f(x):
if x > 0:
return x
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_conditional_missing_else_then_default(self):
def f(x):
if x > 0:
return x
return x * x
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_conditional_else_only_then_default(self):
def f(x):
if x < 0:
x *= x
else:
return x
return x
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_conditional_nested(self):
def f(x):
if x > 0:
if x < 5:
return x
else:
return x * x
else:
return x * x * x
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
self.assertTransformedEquivalent(f, 5)
def test_context_manager(self):
def f(x):
with ops.name_scope(''):
return x * x
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_context_manager_in_conditional(self):
def f(x):
if x > 0:
with ops.name_scope(''):
return x * x
else:
return x
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_conditional_in_context_manager(self):
def f(x):
with ops.name_scope(''):
if x > 0:
return x * x
else:
return x
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_no_return(self):
def f(x):
x *= x
self.assertTransformedEquivalent(f, 2)
def test_nested_function(self):
def f(x):
def inner_fn(y):
if y > 0:
return y * y
else:
return y
return inner_fn(x)
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_nested_function_in_control_flow(self):
def f(x):
if x:
def inner_fn(y):
return y
inner_fn(x)
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, -2)
def test_for_loop(self):
def f(n):
for _ in range(n):
return 1
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, 0)
def test_while_loop(self):
def f(n):
i = 0
s = 0
while i < n:
i += 1
s += i
if s > 4:
return s
return -1
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 2)
self.assertTransformedEquivalent(f, 4)
def test_null_return(self):
def f(n):
if n > 4:
return
return
self.assertTransformedEquivalent(f, 4)
self.assertTransformedEquivalent(f, 5)
def test_nested_multiple_withs(self):
def f(x):
v = []
while x > 0:
x -= 1
with ops.name_scope(''):
if x % 2 == 0:
return v
with ops.name_scope(''):
v.append(x)
v.append(x)
return v
self.assertTransformedEquivalent(f, 0)
self.assertTransformedEquivalent(f, 1)
self.assertTransformedEquivalent(f, 3)
self.assertTransformedEquivalent(f, 4)
def test_multiple_returns_in_nested_scope(self):
def f(a):
v = []
for x in a:
x -= 1
if x > 100:
return v
try:
raise ValueError('intentional')
except ValueError: # pylint:disable=bare-except
return v
v.append(x)
return v
self.assertTransformedEquivalent(f, [])
self.assertTransformedEquivalent(f, [1])
self.assertTransformedEquivalent(f, [2])
self.assertTransformedEquivalent(f, [1, 2, 3])
if __name__ == '__main__':
test.main()
@@ -0,0 +1,83 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Converter for slice operations."""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.lang import directives
from tensorflow.python.autograph.pyct import templates
class SliceTransformer(converter.Base):
"""Converts slicing operations to their TF counterpart.
Currently, relying on the default slice operator that Tensor uses is
insufficient, because TensorArray and tensor lists use dedicated index read
and write functions.
"""
def _process_single_assignment(self, target, value):
if not isinstance(target, gast.Subscript):
return None
s = target.slice
if isinstance(s, (gast.Tuple, gast.Slice)):
return None
template = """
target = ag__.set_item(target, key, item)
"""
return templates.replace(
template, target=target.value, key=target.slice, item=value)
def visit_Assign(self, node):
node = self.generic_visit(node)
# TODO(mdan): Support unpackings and multiple assignments.
if len(node.targets) != 1:
raise NotImplementedError('multiple assignment')
replacement = self._process_single_assignment(node.targets[0], node.value)
if replacement is not None:
return replacement
return node
def visit_Subscript(self, node):
node = self.generic_visit(node)
s = node.slice
if isinstance(s, (gast.Tuple, gast.Slice)):
return node
if not isinstance(node.ctx, gast.Load):
# Index writes are handled at a higher level, one at which the rvalue is
# also available.
return node
dtype = self.get_definition_directive(
node.value,
directives.set_element_type,
'dtype',
default=templates.replace_as_expression('None'))
template = """
ag__.get_item(
target,
key,
opts=ag__.GetItemOpts(element_dtype=dtype))
"""
return templates.replace_as_expression(
template, target=node.value, key=s, dtype=dtype)
def transform(node, ctx):
return SliceTransformer(ctx).visit(node)
@@ -0,0 +1,54 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for slices module."""
from tensorflow.python.autograph.converters import directives as directives_converter
from tensorflow.python.autograph.converters import slices
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.autograph.lang import directives
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import list_ops
from tensorflow.python.platform import test
class SliceTest(converter_testing.TestCase):
def test_index_access(self):
def f(l):
directives.set_element_type(l, dtypes.int32)
return l[1]
tr = self.transform(f, (directives_converter, slices))
tl = list_ops.tensor_list_from_tensor(
[1, 2], element_shape=constant_op.constant([], dtype=dtypes.int32))
y = tr(tl)
self.assertEqual(2, self.evaluate(y))
def test_index_access_multiple_definitions(self):
def f(l):
directives.set_element_type(l, dtypes.int32)
if l:
l = []
return l[1]
self.transform(f, (directives_converter, slices))
if __name__ == '__main__':
test.main()
@@ -0,0 +1,97 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Overloads all variable read operations."""
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import templates
class VariableAccessTransformer(converter.Base):
"""Rewrites basic symbol reads.
This transformer rewrites variable reads with a "read" operator which allows
tracking activity.
Example:
For a basic statement:
a = b + c
This is translated to:
a = ld(b) + ld(c)
Augmented assignment operations also introduce a `ld` operator:
a += b
The assignment target also receives an operator to properly represent the
read:
a = ld(a)
a += ld(b)
"""
def visit_Name(self, node):
# Only the loads which existed in the original code are overloaded.
if not anno.hasanno(node, anno.Static.ORIG_DEFINITIONS):
return node
if isinstance(node.ctx, gast.Load):
node = templates.replace_as_expression('ag__.ld(var_)', var_=node)
return node
def visit_Delete(self, node):
node = self.generic_visit(node)
rewrite_targets = []
for tgt in node.targets:
# Don't rewrite composites like `del a[0]`.
if isinstance(tgt, gast.Name):
rewrite_targets.append(tgt)
if not rewrite_targets:
return node
results = []
for tgt in rewrite_targets:
template = """
var_ = ag__.Undefined(var_name)
"""
results.extend(templates.replace(
template, var_=tgt, var_name=gast.Constant(tgt.id, kind=None)))
remaining_targets = [n for n in node.targets if n not in rewrite_targets]
if remaining_targets:
results.append(gast.Delete(targets=remaining_targets))
return results
def visit_AugAssign(self, node):
if isinstance(node.target, gast.Name):
template = """
var_ = ag__.ld(var_)
original
"""
node = templates.replace(template, var_=node.target, original=node)
else:
node = self.generic_visit(node)
return node
def transform(node, ctx):
return VariableAccessTransformer(ctx).visit(node)
@@ -0,0 +1,210 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for variables module."""
from tensorflow.python.autograph.converters import variables
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.platform import test
class VariablesTest(converter_testing.TestCase):
def transform_with_test_ld(self, f):
"""Generates code which adds 1 to all variable reads."""
return self.transform(f, variables, ag_overrides={'ld': lambda x: x + 1})
def test_read(self):
def f(l):
return l
tr = self.transform_with_test_ld(f)
self.assertEqual(tr(1), 2)
def test_aug_assign(self):
def f(l):
l *= 10
return l
tr = self.transform_with_test_ld(f)
self.assertEqual(tr(1), (1 + 1) * 10 + 1) # two reads
def test_del(self):
def f(l):
del l
return l
tr = self.transform(f, variables)
with self.assertRaisesRegex(NameError, "'l' is used before assignment"):
tr(1)
def test_del_getitem_ignored_basic_slice(self):
def f(l):
del l[0]
return l
tr = self.transform(f, variables)
self.assertListEqual([2], tr([1, 2]))
def test_del_getitem_ignored_range_slice(self):
def f(l):
del l[0:2]
return l
tr = self.transform(f, variables)
self.assertListEqual([], tr([1, 2]))
def test_del_getattr_ignored(self):
def f(l):
del l.a
return l
class TestClass(object):
def __init__(self):
self.a = 1
self.b = 2
tr = self.transform(f, variables)
self.assertFalse(hasattr(tr(TestClass()), 'a'))
self.assertEqual(tr(TestClass()).b, 2)
def test_del_packing_ignored_list(self):
# Note: testing for UnboundLocalError, not NameError because in this case we
# don't rewrite the del.
def f(a, b):
del [a, b]
return a
tr = self.transform(f, variables)
with self.assertRaises(UnboundLocalError):
tr(1, 2)
def test_del_packing_ignored_nested(self):
# Note: testing for UnboundLocalError, not NameError because in this case we
# don't rewrite the del.
def f(a, b, c):
del [a, (b, c)]
return c
tr = self.transform(f, variables)
with self.assertRaises(UnboundLocalError):
tr(1, 2, 3)
def test_del_item_multiple_mixed_used_after(self):
def f(a, b, c):
del a, b, c[0]
a = 1
return a, b, c
tr = self.transform(f, variables)
with self.assertRaisesRegex(NameError, "'b' is used before assignment"):
tr(1, 2, [1, 2])
def test_del_item_multiple_mixed_unused_after(self):
def f(a, b, c):
del a, b, c[0]
a = 1
b = 2
return c
tr = self.transform(f, variables)
self.assertListEqual([2], tr(1, 2, [1, 2]))
def test_attribute(self):
class TestClass(object):
def __init__(self):
self.v = 1
def __add__(self, other):
self.v += other
return self
def f(l):
return l.v
tc = TestClass()
tr = self.transform_with_test_ld(f)
self.assertEqual(tr(tc), 2)
def test_subscript(self):
class TestClass(object):
def __init__(self):
self.v = 1
def __add__(self, other):
self.v += other
return self
def __getitem__(self, _):
return self.v
def f(l):
return l[0]
tc = TestClass()
tr = self.transform_with_test_ld(f)
self.assertEqual(tr(tc), 2)
def test_call(self):
class TestClass(object):
def __init__(self):
self.v = 1
def __add__(self, other):
self.v += other
return self
def __call__(self):
return self.v
def f(l):
return l()
tc = TestClass()
tr = self.transform_with_test_ld(f)
self.assertEqual(tr(tc), 2)
if __name__ == '__main__':
test.main()