chore: import upstream snapshot with attribution

tensorflow/python/autograph/tests/BUILD

@@ -0,0 +1,152 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load(":reference_test.bzl", "reference_test")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+licenses(["notice"])
+
+py_library(
+    name = "reference_tests",
+    srcs = ["reference_test_base.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//third_party/py/numpy",
+    ],
+)
+
+# # Misc
+
+# ## Boolean expressions, equality testing
+
+reference_test(name = "logical_expression_test")
+
+reference_test(name = "ext_slice_test")
+
+# ## Type annotations
+
+reference_test(name = "type_annotations_test")
+
+# ## Scoping and modularity
+
+reference_test(
+    name = "loop_scoping_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+# ## Composite names
+
+reference_test(
+    name = "composite_names_in_control_flow_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+# ## Function Calls
+
+reference_test(name = "call_to_builtin_function_test")
+
+reference_test(name = "call_to_lambda_function_test")
+
+reference_test(name = "call_to_named_tuple_test")
+
+reference_test(
+    name = "call_to_numpy_function_test",
+    additional_deps = ["//third_party/py/numpy"],
+)
+
+reference_test(
+    name = "call_to_print_function_test",
+    additional_deps = ["//third_party/py/numpy"],
+)
+
+reference_test(name = "call_to_tf_api_test")
+
+reference_test(name = "call_to_user_function_test")
+
+# # Control Flow
+
+reference_test(name = "basic_ifexp_test")
+
+reference_test(
+    name = "cond_basic_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+reference_test(
+    name = "early_return_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+reference_test(
+    name = "loop_basic_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+reference_test(
+    name = "loop_control_flow_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+    shard_count = 20,
+)
+
+reference_test(
+    name = "loop_control_flow_illegal_cases_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+reference_test(
+    name = "loop_created_variables_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+reference_test(
+    name = "loop_distributed_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+    tags = [
+        "no_oss",  # TODO(mdan): Investigate timeout reason.
+    ],
+)
+
+reference_test(
+    name = "loop_with_variable_type_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+reference_test(
+    name = "loop_with_variable_type_illegal_cases_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+reference_test(
+    name = "loop_with_function_call_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+reference_test(
+    name = "nested_control_flow_test",
+    additional_deps = ["@absl_py//absl/testing:parameterized"],
+)
+
+# ## Assert
+
+reference_test(name = "assertion_test")
+
+# # Data Structures
+
+# ## Generators
+
+reference_test(name = "generator_test")
+
+# ## Iterators
+
+reference_test(name = "datasets_test")
+
+# ## Lists
+
+reference_test(
+    name = "basic_list_test",
+    additional_deps = ["//tensorflow/python/autograph"],
+    tags = [
+        "no_oss",
+        "notap",
+    ],
+)

tensorflow/python/autograph/tests/README.md

@@ -0,0 +1,23 @@
+# AutoGraph reference tests
+
+This directory contains tests for various Python idioms under AutoGraph.
+Since they are easy to read, they also double as small code samples.
+For constructs which are not supported, these tests indicate the error being
+raised.
+
+The BUILD file contains the full list of tests.
+
+## Locating the samples inside tests
+
+Each test is structured as:
+
+    <imports>
+
+    <sample functions>
+
+    <test class>
+
+The sample functions are what demonstrate how code is authored for AutoGraph.
+
+The test in generale ensure that the sample code produces the same results when
+run in a TF graph as it would when executed as regular Python.

tensorflow/python/autograph/tests/assertion_test.py

@@ -0,0 +1,43 @@
+# 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.
+# ==============================================================================
+"""Basic assertions."""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def simple_assertion(x):
+  assert x > 0
+  return x
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def setUp(self):
+    super(ReferenceTest, self).setUp()
+    self.autograph_opts = tf.autograph.experimental.Feature.ASSERT_STATEMENTS
+
+  def test_basic(self):
+    self.assertFunctionMatchesEager(simple_assertion, 1)
+    self.assertFunctionMatchesEager(simple_assertion, tf.constant(1))
+    with self.assertRaises(AssertionError):
+      self.function(simple_assertion)(0)
+    with self.assertRaises(tf.errors.InvalidArgumentError):
+      self.function(simple_assertion)(tf.constant(0))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/basic_ifexp_test.py

@@ -0,0 +1,58 @@
+# 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.
+# ==============================================================================
+"""Basic if conditional.
+
+The loop is converted to tf.cond.
+"""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def consecutive_conds(x):
+  if x > 0:
+    x = -x if x < 5 else x
+  else:
+    x = -2 * x if x < 5 else 1
+  if x > 0:
+    x = -x if x < 5 else x
+  else:
+    x = (3 * x) if x < 5 else x
+  return x
+
+
+def cond_with_multiple_values(x):
+  if x > 0:
+    x = -x if x < 5 else x
+    y = 2 * x if x < 5 else x
+    z = -y if y < 5 else y
+  else:
+    x = 2 * x if x < 5 else x
+    y = -x if x < 5 else x
+    z = -y if y < 5 else y
+  return x, y, z
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_basic(self):
+    for x in [-1, 1, 5, tf.constant(-1), tf.constant(1), tf.constant(5)]:
+      self.assertFunctionMatchesEager(consecutive_conds, x)
+      self.assertFunctionMatchesEager(cond_with_multiple_values, x)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/basic_list_test.py

@@ -0,0 +1,142 @@
+# 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.
+# ==============================================================================
+"""Basic list operations."""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.lang import directives
+from tensorflow.python.autograph.lang import special_functions
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def type_not_annotated(n):
+  l = []
+  # TODO(mdan): Here, we ought to infer the dtype and shape when i is staged.
+  for i in range(n):
+    l.append(i)
+  return special_functions.stack(l, strict=False)
+
+
+def element_access():
+  l = []
+  l.append(1)
+  l.append(2)
+  l.append(3)
+  directives.set_element_type(l, tf.int32)
+  return 2 * l[1]
+
+
+def element_update():
+  l = []
+  l.append(1)
+  l.append(2)
+  l.append(3)
+  directives.set_element_type(l, tf.int32)
+  l[1] = 5
+  return special_functions.stack(l, strict=False)
+
+
+def simple_fill(n):
+  l = []
+  directives.set_element_type(l, tf.int32)
+  for i in range(n):
+    l.append(i)
+  return special_functions.stack(l, strict=False)
+
+
+def nested_fill(m, n):
+  mat = []
+  directives.set_element_type(mat, tf.int32)
+  for _ in range(m):
+    l = []
+    directives.set_element_type(l, tf.int32)
+    for j in range(n):
+      l.append(j)
+    mat.append(special_functions.stack(l, strict=False))
+  return special_functions.stack(mat, strict=False)
+
+
+def read_write_loop(n):
+  l = []
+  l.append(1)
+  l.append(1)
+  directives.set_element_type(l, tf.int32)
+  for i in range(2, n):
+    l.append(l[i-1] + l[i-2])
+    l[i-2] = -l[i-2]
+  return special_functions.stack(l, strict=False)
+
+
+def simple_empty(n):
+  l = []
+  l.append(1)
+  l.append(2)
+  l.append(3)
+  l.append(4)
+  directives.set_element_type(l, tf.int32, ())
+  s = 0
+  for _ in range(n):
+    s += l.pop()
+  return special_functions.stack(l, strict=False), s
+
+
+def mutation(t, n):
+  for i in range(n):
+    t[i] = i
+  return t
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def setUp(self):
+    super(ReferenceTest, self).setUp()
+    self.autograph_opts = tf.autograph.experimental.Feature.LISTS
+
+  def test_tensor_mutation(self):
+    self.assertConvertedMatchesNative(mutation, [0] * 10, 10)
+
+  def test_basic(self):
+    self.all_inputs_tensors = True
+    self.assertFunctionMatchesEager(element_access)
+    self.assertFunctionMatchesEager(element_update)
+
+    # TODO(mdan): This should raise a compilation, not runtime, error.
+    with self.assertRaisesRegex(
+        ValueError,
+        'cannot stack a list without knowing its element type; '
+        'use set_element_type to annotate it'):
+      self.function(type_not_annotated)(3)
+
+    self.assertFunctionMatchesEager(simple_fill, 5)
+    self.assertFunctionMatchesEager(nested_fill, 5, 3)
+    self.assertFunctionMatchesEager(read_write_loop, 4)
+    self.assertFunctionMatchesEager(simple_empty, 0)
+    self.assertFunctionMatchesEager(simple_empty, 2)
+    self.assertFunctionMatchesEager(simple_empty, 4)
+
+    # TODO(mdan): Allow explicitly setting the element shape to mitigate these.
+    # TODO(mdan): This should raise a friendlier runtime error.
+    # The error should spell out that empty lists cannot be stacked.
+    # Alternatively, we can also insert conditionals that construct a zero-sized
+    # Tensor of the appropriate type and shape, but we first want to make sure
+    # that doesn't degrade performance.
+    with self.assertRaises(ValueError):
+      self.function(simple_fill)(0)
+    with self.assertRaises(ValueError):
+      self.function(nested_fill)(0, 3)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/call_to_builtin_function_test.py

@@ -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.
+# ==============================================================================
+"""Simple call to a builtin function."""
+
+import unittest
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+# TODO(mdan): Add tests for all builtins.
+
+
+def dict_call(x):
+  return dict(foo=x)
+
+
+def dict_call_aliased(x):
+  def fake_dict(x):
+    return x
+
+  dict = fake_dict  # pylint:disable=redefined-builtin
+  return dict(x)
+
+
+def dict_call_dynamic(x):
+  def gen_dict():
+    return dict
+
+  d = gen_dict()
+  return d(foo=x)
+
+
+def len_call(x):
+  return len(x)
+
+
+def nested_call(x):
+  return list(range(len(x)))
+
+
+def nested_cast(x):
+  return float(int(x))
+
+
+def len_call_aliased(x):
+
+  def fake_len(x):
+    return x
+
+  len = fake_len  # pylint:disable=redefined-builtin
+  return len(x)
+
+
+def len_call_dynamic(x):
+
+  def gen_len():
+    return len
+
+  l = gen_len()
+  return l(x)
+
+
+def len_call_on_mock():
+  x = unittest.mock.MagicMock()
+  return len(x)
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_basic(self):
+    self.assertFunctionMatchesEager(dict_call, 1)
+    self.assertFunctionMatchesEager(len_call, [1, 2])
+    self.assertFunctionMatchesEager(dict_call_aliased, 1)
+    self.assertFunctionMatchesEager(len_call_aliased, [1, 2])
+    self.assertFunctionMatchesEager(dict_call_dynamic, 1)
+    self.assertFunctionMatchesEager(len_call_dynamic, [1, 2])
+    self.assertFunctionMatchesEager(nested_call, [])
+    self.assertFunctionMatchesEager(nested_call, [1, 2, 3])
+
+  def test_basic_tensor(self):
+    self.all_inputs_tensors = True
+    self.assertFunctionMatchesEager(dict_call, 1)
+    self.assertFunctionMatchesEager(len_call, [1, 2])
+    self.assertFunctionMatchesEager(dict_call_aliased, 1)
+    self.assertFunctionMatchesEager(len_call_aliased, [1, 2])
+    self.assertFunctionMatchesEager(dict_call_dynamic, 1)
+    self.assertFunctionMatchesEager(len_call_dynamic, [1, 2])
+    self.assertFunctionMatchesEager(nested_call, [])
+    self.assertFunctionMatchesEager(nested_call, [1, 2, 3])
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/call_to_lambda_function_test.py

@@ -0,0 +1,43 @@
+# 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.
+# ==============================================================================
+"""Simple call to lambda functions."""
+
+import tensorflow.compat.v1 as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def inline_lambda(x):
+  l = lambda x: x * x if x > 0 else -x
+  return l(x)
+
+
+def external_lambda(x, l):
+  return l(x)
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_inline(self):
+    self.assertFunctionMatchesEager(inline_lambda, 1)
+    self.assertFunctionMatchesEager(inline_lambda, tf.constant(1))
+
+  def test_external(self):
+    self.assertFunctionMatchesEager(external_lambda, 1, lambda x: x == 0)
+    self.assertFunctionMatchesEager(
+        external_lambda, tf.constant(1), lambda x: x == 0)
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/call_to_named_tuple_test.py

@@ -0,0 +1,62 @@
+# 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.
+# ==============================================================================
+"""Simple call to construct a namedtuple."""
+
+import collections
+
+import tensorflow.compat.v1 as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def inline_namedtuple(x):
+  nt = collections.namedtuple('TestNamedTuple', ('a', 'b'))
+  n = nt(a=1, b=x)
+  return n
+
+
+def external_namedtuple(x, nt):
+  return nt(a=1, b=x)
+
+
+class NamedTupleSubclass(collections.namedtuple('TestNamedTuple', ('a',))):
+
+  def foo(self):
+    return self.a + 1
+
+
+def namedtuple_subclass(x):
+  nt = NamedTupleSubclass(x)
+  return nt.foo()
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_inline(self):
+    self.assertFunctionMatchesEager(inline_namedtuple, 1)
+    self.assertFunctionMatchesEager(inline_namedtuple, tf.constant(1))
+
+  def test_external(self):
+    nt = collections.namedtuple('TestNamedTuple', ('a', 'b'))
+    self.assertFunctionMatchesEager(external_namedtuple, 1, nt)
+    self.assertFunctionMatchesEager(external_namedtuple, tf.constant(1), nt)
+
+  def test_subclass(self):
+    self.assertFunctionMatchesEager(namedtuple_subclass, 1)
+    self.assertFunctionMatchesEager(namedtuple_subclass, tf.constant(1))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/call_to_numpy_function_test.py

@@ -0,0 +1,38 @@
+# 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.
+# ==============================================================================
+"""Simple call to a whitelisted Numpy function.
+
+The call should be wrapped in py_func.
+"""
+
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def f():
+  np.random.seed(1)
+  return 2 * np.random.binomial(1, 0.5, size=(10,)) - 1
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_basic(self):
+    self.assertFunctionMatchesEager(f)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/call_to_print_function_test.py

@@ -0,0 +1,86 @@
+# 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.
+# ==============================================================================
+"""Simple call to a print function preceding other computations.
+
+The call may be wrapped inside a py_func, but tf.Print should be used if
+possible. The subsequent computations will be gated by the print function
+execution.
+"""
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def lone_print(x):
+  print(x)
+
+
+def print_multiple_values(x):
+  print('x is', x)
+
+
+def multiple_prints(x, y):
+  tf.print('x is', x)
+  tf.print('y is', y)
+
+
+def print_with_nontf_values(x):
+  print('x is', x, {'foo': 'bar'})
+
+
+def print_in_cond(x):
+  if x == 0:
+    print(x)
+
+
+def tf_print(x):
+  tf.print(x)
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def setUp(self):
+    super(ReferenceTest, self).setUp()
+    self.autograph_opts = tf.autograph.experimental.Feature.BUILTIN_FUNCTIONS
+
+  def test_lone_print(self):
+    self.assertFunctionMatchesEager(lone_print, 1)
+    self.assertFunctionMatchesEager(lone_print, np.array([1, 2, 3]))
+
+  def test_print_multiple_values(self):
+    self.assertFunctionMatchesEager(print_multiple_values, 1)
+    self.assertFunctionMatchesEager(print_multiple_values, np.array([1, 2, 3]))
+
+  def test_multiple_prints(self):
+    self.assertFunctionMatchesEager(multiple_prints, 1, 2)
+    self.assertFunctionMatchesEager(multiple_prints, np.array([1, 2, 3]), 4)
+
+  def test_print_with_nontf_values(self):
+    self.assertFunctionMatchesEager(print_with_nontf_values, 1)
+    self.assertFunctionMatchesEager(print_with_nontf_values, np.array([1, 2,
+                                                                       3]))
+
+  def test_print_in_cond(self):
+    self.assertFunctionMatchesEager(print_in_cond, 0)
+    self.assertFunctionMatchesEager(print_in_cond, 1)
+
+  def test_tf_print(self):
+    self.assertFunctionMatchesEager(tf_print, 0)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/call_to_tf_api_test.py

@@ -0,0 +1,37 @@
+# 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.
+# ==============================================================================
+"""Simple call to a TF API function.
+
+The call will remain unchanged.
+"""
+
+import tensorflow.compat.v1 as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def core_tf_call(x):
+  return x * tf.constant(2)
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_basic(self):
+    self.assertFunctionMatchesEager(core_tf_call, 1)
+    self.assertFunctionMatchesEager(core_tf_call, tf.constant(1))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/call_to_user_function_test.py

@@ -0,0 +1,98 @@
+# 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.
+# ==============================================================================
+"""Calls to dynamic (i.e. nonglobal) functions.
+
+Examples:
+ * function variables
+ * function parameters
+ * factories
+"""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def function_1(x):
+  return x * x * x
+
+
+def function_2(x):
+  return -1 * x + 11
+
+
+def factory(n):
+  if n == 1:
+    return function_1
+  return function_2
+
+
+def static_fn(x):
+  a = function_1(x)
+  b = function_2(x)
+  return a + b
+
+
+def factory_dynamic_fn(x):
+  f = factory(1)
+  a = f(x)
+  f = factory(2)
+  b = f(x)
+  return a + b
+
+
+def param_dynamic_fn(f, x):
+  return f(x)
+
+
+def variable_dynamic_fn(x):
+  f = function_1
+  a = f(x)
+  f = function_2
+  b = f(x)
+  return a + b
+
+
+def variable_dynamic_whitelisted_fn(x):
+  f = tf.identity
+  return f(x)
+
+
+def dynamic_fn_with_kwargs(f, x):
+  return f(x=x)
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_basic(self):
+    self.assertFunctionMatchesEager(static_fn, 1)
+    self.assertFunctionMatchesEager(factory_dynamic_fn, 1)
+    self.assertFunctionMatchesEager(param_dynamic_fn, function_1, 1)
+    self.assertFunctionMatchesEager(variable_dynamic_fn, 1)
+    self.assertFunctionMatchesEager(variable_dynamic_whitelisted_fn, 1)
+    self.assertFunctionMatchesEager(dynamic_fn_with_kwargs, function_1, 1)
+
+  def test_basic_tensor(self):
+    self.all_inputs_tensors = True
+    self.assertFunctionMatchesEager(static_fn, 1)
+    self.assertFunctionMatchesEager(factory_dynamic_fn, 1)
+    self.assertFunctionMatchesEager(param_dynamic_fn, function_1, 1)
+    self.assertFunctionMatchesEager(variable_dynamic_fn, 1)
+    self.assertFunctionMatchesEager(variable_dynamic_whitelisted_fn, 1)
+    self.assertFunctionMatchesEager(dynamic_fn_with_kwargs, function_1, 1)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/composite_names_in_control_flow_test.py

@@ -0,0 +1,380 @@
+# 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.
+# ==============================================================================
+"""Composite names (attributes) in control flow.
+
+Generally, composite symbols should be treated like regular ones.
+"""
+
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def if_basic_dict(a):
+  x = {'a': a}
+
+  if x['a'] > 0:
+    x['b'] = 1
+  else:
+    x['b'] = -1
+  return x
+
+
+def if_basic_list(a):
+  x = [a, 0]
+
+  if x[0] > 0:
+    x[1] = 1
+  else:
+    x[1] = -1
+  return x
+
+
+def if_imbalanced(a):
+  x = {'a': a}
+
+  if x['a'] > 0:
+    x['b'] = 1
+  return x
+
+
+def else_imbalanced(a):
+  x = {'a': a}
+
+  if x['a'] > 0:
+    pass
+  else:
+    x['b'] = 1
+  return x
+
+
+def if_imbalanced_nested(a):
+  x = {'a': {'a': a}}
+
+  if x['a']['a'] > 0:
+    x['a']['b'] = 1
+  return x
+
+
+def else_imbalanced_nested(a):
+  x = {'a': {'a': a}}
+
+  if x['a']['a'] > 0:
+    pass
+  else:
+    x['a']['b'] = 1
+  return x
+
+
+# Note: this is incorrect code. We test that the error message is clear about
+# that.
+def if_buggy(a):
+  x = {'a': {'a': a}}
+
+  if x['a']['a'] > 0:
+    x['b']['a'] = 1
+  else:
+    x['b']['a'] = -1
+  return x
+
+
+# Note: this is incorrect code. We test that the error message is clear about
+# that.
+def if_imbalanced_buggy(a):
+  x = {'a': {'a': a}}
+
+  if x['a']['a'] > 0:
+    x['b']['a'] = 1
+  return x
+
+
+def while_basic_dict(x, a, b):
+  y = {'a': a, 'b': b}
+  while x > 0:
+    x -= 1
+    y['a'] += 1
+  return y
+
+
+def while_basic_list(x, a, b):
+  y = [a, b]
+  while x > 0:
+    x -= 1
+    y[0] += 1
+  return y
+
+
+def while_state_only_dict(a, b):
+  y = {'a': a, 'b': b}
+  while y['b'] <= 10:
+    y['a'] += 1
+    y['b'] *= 2
+  return y
+
+
+def while_state_only_list(a, b):
+  y = [a, b]
+  while y[1] <= 10:
+    y[0] += 1
+    y[1] *= 2
+  return y
+
+
+def while_imbalanced(b):
+  y = {'b': b}
+  while y['b'] <= 10:
+    y['a'] = y['b'] + 1
+    y['b'] *= 2
+  return y
+
+
+def for_basic_dict(n, x, a, b):
+  y = {'a': a, 'b': b}
+  for i in range(n):
+    x -= 1
+    y['a'] += i
+  return y
+
+
+def for_basic_list(n, x, a, b):
+  y = [a, b]
+  for i in range(n):
+    x -= 1
+    y[0] += i
+  return y
+
+
+def for_state_only_dict(n, a, b):
+  y = {'a': a, 'b': b}
+  for _ in range(n):
+    y['a'] += 1
+  return y
+
+
+def for_state_only_list(n, a, b):
+  y = [a, b]
+  for _ in range(n):
+    y[0] += 1
+  return y
+
+
+def for_imbalanced(n, x):
+  y = {}
+  for i in range(n):
+    x -= i
+    y['a'] = x
+  return y
+
+
+# TODO(mdan): More tests needed. Many pitfalls around mutating objects this way.
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (
+          if_basic_dict,
+          if_basic_list,
+      ),
+      (
+          0,
+          1,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+  ))
+  def test_if_basic(self, target, a, type_):
+    a = type_(a)
+    self.assertFunctionMatchesEager(target, a)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          if_imbalanced,
+          else_imbalanced,
+          if_imbalanced_nested,
+          else_imbalanced_nested,
+      ),
+      (
+          0,
+          1,
+      )
+  ))
+  def test_if_imbalanced_legal(self, target, a):
+    self.assertFunctionMatchesEager(target, a)
+
+  @parameterized.parameters(
+      (if_imbalanced, 0, tf.constant, ValueError,
+       r"'x\['b'\]' must also be initialized in the else"),
+      (if_imbalanced, 1, tf.constant, ValueError,
+       r"'x\['b'\]' must also be initialized in the else"),
+      (else_imbalanced, 0, tf.constant, ValueError,
+       r"'x\['b'\]' must also be initialized in the main"),
+      (else_imbalanced, 1, tf.constant, ValueError,
+       r"'x\['b'\]' must also be initialized in the main"),
+      (if_imbalanced_nested, 0, tf.constant, ValueError,
+       r"'x\['a'\]\['b'\]' must also be initialized in the else"),
+      (if_imbalanced_nested, 1, tf.constant, ValueError,
+       r"'x\['a'\]\['b'\]' must also be initialized in the else"),
+      (else_imbalanced_nested, 0, tf.constant, ValueError,
+       r"'x\['a'\]\['b'\]' must also be initialized in the main"),
+      (else_imbalanced_nested, 1, tf.constant, ValueError,
+       r"'x\['a'\]\['b'\]' must also be initialized in the main"),
+      (if_buggy, 1, int, KeyError, "'b'"),
+      (if_buggy, 1, tf.constant, KeyError, "'b'"),
+      (if_buggy, 0, tf.constant, KeyError, "'b'"),
+      (if_imbalanced_buggy, 1, int, KeyError, "'b'"),
+      (if_imbalanced_buggy, 1, tf.constant, KeyError, "'b'"),
+      (if_imbalanced_buggy, 0, tf.constant, KeyError, "'b'"),
+  )
+  def test_if_imbalanced_illegal(self, target, a, type_, exc_type, exc_regex):
+    a = type_(a)
+    with self.assertRaisesRegex(exc_type, exc_regex):
+      tf.function(target)(a)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          while_basic_dict,
+          while_basic_list,
+      ),
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+      (
+          3,
+          7,
+      ),
+  ))
+  def test_while_basic(self, target, x, type_, a):
+    x = type_(x)
+    self.assertFunctionMatchesEager(target, x, a, 0)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          while_state_only_dict,
+          while_state_only_list,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+      (
+          3,
+          4,
+          5,
+          6,
+      ),
+  ))
+  def test_while_state_only(self, target, type_, b):
+    b = type_(b)
+    self.assertFunctionMatchesEager(target, 0, b)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          5,
+          10,
+          11,
+      ),
+      (
+          int,
+          tf.constant,
+      ),
+  ))
+  def test_while_imbalanced_legal(self, b, type_):
+    if b == 11 and type_ is tf.constant:
+      self.skipTest("TF loop must initialize y['a']")
+
+    b = type_(b)
+    self.assertFunctionMatchesEager(while_imbalanced, b)
+
+  def test_while_imbalanced_illegal(self):
+    with self.assertRaisesRegex(
+        tf.errors.InvalidArgumentError,
+        r"loop must iterate at least once to initialize y\[\\'a\\'\]"):
+      tf.function(while_imbalanced)(tf.constant(11))
+
+  @parameterized.parameters(*itertools.product(
+      (
+          for_basic_dict,
+          for_basic_list,
+      ),
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+  ))
+  def test_for_basic(self, target, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(target, n, 1, 1, 1)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          for_state_only_dict,
+          for_state_only_list,
+      ),
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+  ))
+  def test_for_state_only(self, target, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(target, n, 1, 1)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 1, 2,
+      ),
+      (
+          int,
+          tf.constant,
+      ),
+  ))
+  def test_for_imbalanced_legal(self, n, type_):
+    if n == 0 and type_ is tf.constant:
+      self.skipTest("TF loop must initialize y['a']")
+
+    n = type_(n)
+    self.assertFunctionMatchesEager(for_imbalanced, n, 0)
+
+  def test_for_imbalanced_illegal(self):
+    with self.assertRaisesRegex(
+        tf.errors.InvalidArgumentError,
+        r"loop must iterate at least once to initialize y\[\\'a\\'\]"):
+      tf.function(for_imbalanced)(tf.constant(0), 0)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/cond_basic_test.py

@@ -0,0 +1,437 @@
+# 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.
+# ==============================================================================
+"""Basic conditionals."""
+
+import itertools
+import re
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def if_no_vars(c, v):
+  v.assign(0)
+  if c:
+    v.assign_add(1)
+  return v.read_value()
+
+
+def if_else_no_vars(c, v):
+  v.assign(0)
+  if c:
+    v.assign_add(1)
+  else:
+    v.assign_add(2)
+  return v.read_value()
+
+
+def if_one_var(n):
+  i = 0
+  if i < n:
+    i += 1
+  return i
+
+
+def if_else_one_var(n):
+  i = 0
+  if i < n:
+    i += 1
+  else:
+    i += 2
+  return i
+
+
+def if_two_vars(n):
+  i = 0
+  j = 1
+  if i < n:
+    i += 1
+    j *= 10
+  return i, j
+
+
+def if_else_two_vars(n):
+  i = 0
+  j = 1
+  if i < n:
+    i += 1
+    j *= 10
+  else:
+    i += 2
+    j *= 20
+  return i, j
+
+
+def if_creates_var(c):
+  if c:
+    i = 1
+  return i
+
+
+def if_else_creates_var(c):
+  if c:
+    i = 1
+  else:
+    i = 2
+  return i
+
+
+def else_creates_var(c):
+  if c:
+    pass
+  else:
+    i = 2
+  return i
+
+
+def if_destroys_var(c):
+  i = 1
+  if c:
+    del i
+  return i
+
+
+def if_else_destroys_var(c):
+  i = 1
+  if c:
+    del i
+  else:
+    del i
+  return i
+
+
+def else_destroys_var(c):
+  i = 2
+  if c:
+    pass
+  else:
+    del i
+  return i
+
+
+def if_returns_none(c):
+  i = 0
+  j = 1
+  if c:
+    i = None
+    j = 2
+  return i, j
+
+
+def if_else_returns_none(c):
+  if c:
+    i = None
+    j = 1
+  else:
+    i = None
+    j = 2
+  return i, j
+
+
+def else_returns_none(c):
+  i = 1
+  j = 1
+  if c:
+    pass
+  else:
+    i = None
+    j = 2
+  return i, j
+
+
+def if_local_var(c):
+  i = 0
+  if c:
+    j = 1
+    i = j + 1
+  return i
+
+
+def if_else_local_var(c):
+  i = 0
+  if c:
+    j = 1
+  else:
+    j = 2
+    i = j + 1
+  return i
+
+
+def if_locally_modified_var(c):
+  i = 0
+  j = 2
+  if c:
+    j = j + 1
+    i = j + 1
+  return i
+
+
+def successive_ifs(n1, n2):
+  s = 0
+  i = 0
+  if i < n1:
+    s = s * 10 + i
+    i += 1
+  i = 0
+  if i < n2:
+    s = s * 10 + i
+    i += 1
+  return s
+
+
+def successive_if_elses(n1, n2):
+  s = 0
+  i = 0
+  if i < n1:
+    s = s * 10 + i
+    i += 1
+  else:
+    s = s * 11 + i
+    i += 2
+  i = 0
+  if i < n2:
+    s = s * 10 + i
+    i += 1
+  else:
+    s = s * 11 + i
+    i += 2
+  return s
+
+
+def nested_ifs(n1, n2):
+  i = 0
+  l = 0
+  if i < n1:
+    j = 0
+    s = 0
+    if j < n2:
+      s = s * 10 + i * j
+      j += 1
+    l = l * 1000 + s
+    i += 1
+  return l
+
+
+def nested_if_temporarily_undefined_return(c1, c2):
+  if c1:
+    if c2:
+      return 1
+  return 2
+
+
+def nested_if_elses(n1, n2):
+  i = 0
+  l = 0
+  if i < n1:
+    j = 0
+    s = 0
+    if j < n2:
+      s = s * 10 + i * j
+      j += 1
+    else:
+      s = s * 11 + i * j
+      j += 2
+    l = l * 1000 + s
+    i += 1
+  else:
+    j = 0
+    s = 0
+    if j < n2:
+      s = s * 12 + i * j
+      j += 3
+    else:
+      s = s * 13 + i * j
+      j += 4
+    l = l * 2000 + s
+    i += 1
+  return l
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (
+          if_no_vars,
+          if_else_no_vars,
+      ),
+      (
+          True,
+          False,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+  ))
+  def test_no_vars(self, target, c, type_):
+    c = type_(c)
+    self.assertFunctionMatchesEager(target, c, tf.Variable(0))
+
+  @parameterized.parameters(*itertools.product(
+      (
+          if_one_var,
+          if_else_one_var,
+          if_two_vars,
+          if_else_two_vars,
+      ),
+      (
+          0,
+          1,
+      ),
+      (
+          int,
+          tf.constant,
+      ),
+  ))
+  def test_several_vars(self, target, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(target, n)
+
+  def test_var_lifetime_imbalanced_legal(self):
+    self.assertFunctionMatchesEager(if_creates_var, True)
+    self.assertFunctionMatchesEager(else_creates_var, False)
+    self.assertFunctionMatchesEager(if_destroys_var, False)
+    self.assertFunctionMatchesEager(else_destroys_var, True)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          True,
+          False,
+      ),
+      (
+          int,
+          tf.constant,
+      ),
+  ))
+  def test_if_else_var_lifetime(self, c, type_):
+    c = type_(c)
+    self.assertFunctionMatchesEager(if_else_creates_var, c)
+    if type_ is int:
+      with self.assertRaisesRegex(UnboundLocalError, "'i'"):
+        tf.function(if_else_destroys_var)(c)
+    else:
+      with self.assertRaisesRegex(ValueError, "'i' must also be initialized"):
+        tf.function(if_else_destroys_var)(c)
+
+  @parameterized.parameters(
+      (if_creates_var, False, bool, UnboundLocalError,
+       "'i' is used before assignment"),
+      (if_creates_var, True, tf.constant, ValueError,
+       "'i' must also be initialized in the else branch"),
+      (if_creates_var, False, tf.constant, ValueError,
+       "'i' must also be initialized in the else branch"),
+      (else_creates_var, True, bool, UnboundLocalError,
+       "'i' is used before assignment"),
+      (else_creates_var, True, tf.constant, ValueError,
+       "'i' must also be initialized in the main branch"),
+      (else_creates_var, False, tf.constant, ValueError,
+       "'i' must also be initialized in the main branch"),
+  )
+  def test_creates_var_imbalanced_illegal(self, target, c, type_, exc_type,
+                                          exc_regex):
+    c = type_(c)
+    with self.assertRaisesRegex(exc_type, exc_regex):
+      tf.function(target)(c)
+
+  def test_returns_none_legal(self):
+    self.assertFunctionMatchesEager(if_returns_none, True)
+    self.assertFunctionMatchesEager(if_else_returns_none, False)
+    self.assertFunctionMatchesEager(else_returns_none, False)
+
+  @parameterized.parameters(
+      (if_returns_none, True),
+      (if_returns_none, False),
+      (else_returns_none, True),
+      (else_returns_none, False),
+      (if_else_returns_none, True),
+      (if_else_returns_none, False),
+  )
+  def test_returns_none_illegal(self, target, c):
+    c = tf.constant(c)
+    with self.assertRaisesRegex(ValueError, re.compile("'i' is None",
+                                                       re.DOTALL)):
+      tf.function(target)(c)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          if_local_var,
+          if_else_local_var,
+          if_locally_modified_var,
+      ),
+      (
+          True,
+          False,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+  ))
+  def test_local_vars(self, target, c, type_):
+    c = type_(c)
+    self.assertFunctionMatchesEager(target, c)
+
+  @parameterized.parameters(*itertools.product(
+      (True, False),
+      (True, False),
+  ))
+  def test_nested_if_temporarily_undefined_return_legal(self, c1, c2):
+    self.assertFunctionMatchesEager(
+        nested_if_temporarily_undefined_return, c1, c2)
+
+  @parameterized.parameters(*itertools.product(
+      (True, False),
+      (True, False),
+  ))
+  def test_nested_if_temporarily_undefined_return_illegal(self, c1, c2):
+    c1 = tf.constant(c1)
+    c2 = tf.constant(c2)
+    with self.assertRaisesRegex(ValueError, "must also have a return"):
+      tf.function(nested_if_temporarily_undefined_return)(c1, c2)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          successive_ifs,
+          successive_if_elses,
+          nested_ifs,
+          nested_if_elses,
+      ),
+      (
+          0,
+          1,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+      (
+          0,
+          1,
+      ),
+      (
+          bool,
+          tf.constant,
+      ),
+  ))
+  def test_composition(self, target, n1, n1_type, n2, n2_type):
+    n1 = n1_type(n1)
+    n2 = n2_type(n2)
+    self.assertFunctionMatchesEager(target, n1, n2)
+
+
+if __name__ == "__main__":
+  tf.test.main()

tensorflow/python/autograph/tests/datasets_test.py

@@ -0,0 +1,219 @@
+# 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 involving the tf.data.Datasets API."""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def dataset_no_vars_loop(ds):
+  for e in ds:
+    tf.print(e)
+
+
+def iterator_no_vars_loop(ds):
+  for e in iter(ds):
+    tf.print(e)
+
+
+def dataset_single_var_loop(ds):
+  s = tf.constant(0, dtype=tf.int64)
+  for e in ds:
+    s = s * 10 + e
+  return s
+
+
+def iterator_single_var_loop(ds):
+  s = tf.constant(0, dtype=tf.int64)
+  for e in iter(ds):
+    s = s * 10 + e
+  return s
+
+
+def dataset_two_vars_loop(ds):
+  s = tf.constant(0, dtype=tf.int64)
+  p = tf.constant(1, dtype=tf.int64)
+  for e in ds:
+    s += e
+    p *= e
+  return s, p
+
+
+def iterator_two_vars_loop(ds):
+  s = tf.constant(0, dtype=tf.int64)
+  p = tf.constant(1, dtype=tf.int64)
+  for e in iter(ds):
+    s += e
+    p *= e
+  return s, p
+
+
+def dataset_loop_with_break(ds):
+  s = tf.constant(0, dtype=tf.int64)
+  for e in ds:
+    s = s * 10 + e
+    if s > 100:
+      break
+  return s
+
+
+def iterator_loop_with_break(ds):
+  s = tf.constant(0, dtype=tf.int64)
+  for e in iter(ds):
+    s = s + e
+    if s > 100:
+      break
+  return s
+
+
+def iterator_resuming_loop(ds):
+  s = tf.constant(0, dtype=tf.int64)
+  itr = iter(ds)
+  for e in itr:
+    s = s * 10 + e
+    break
+  for e in itr:
+    s = s * 10 + e
+    break
+  for e in itr:
+    s = s * 10 + e
+  return s
+
+
+def dataset_loop_with_return(ds):
+  y = tf.constant(0, dtype=tf.int64)
+  for e in ds:
+    y = e
+    return y
+  return y
+
+
+def iterator_loop_with_return(ds):
+  y = tf.constant(0, dtype=tf.int64)
+  for e in iter(ds):
+    y = e
+    return y
+  return y
+
+
+def iterator_next(ds):
+  itr = iter(ds)
+  return next(itr)
+
+
+def iterator_next_multiple_calls(ds):
+  itr = iter(ds)
+  return 10 * next(itr) + next(itr)
+
+
+def iterator_next_in_loop(ds, n):
+  itr = iter(ds)
+  s = tf.constant(0, dtype=tf.int64)
+  for _ in range(n):
+    s = s * 10 + next(itr)
+  return s
+
+
+def iterator_next_stopping(ds, cond):
+  # This case will raise, but not the expected StopIteration error.
+  itr = iter(ds)
+  while cond:
+    next(itr)
+
+
+def iterator_next_with_catching_stop_iteration(ds, cond):
+  # This is the only instance when the use of TF iterators does not work as
+  # intended. In graph mode, the `except` below will never catch, and the
+  # tf.function will raise the error instead.
+  # TODO(b/132311724): The error should be friendlier here.
+  # Note: b/132298783 covers actually supporting this pattern.
+  itr = iter(ds)
+  try:
+    while cond:
+      next(itr)
+  except StopIteration:
+    pass
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def setUp(self):
+    super(ReferenceTest, self).setUp()
+    self.ds = tf.data.Dataset.range(7)
+
+  def test_dataset_no_vars_loop(self):
+    self.assertFunctionMatchesEager(dataset_no_vars_loop, self.ds)
+
+  def test_iterator_no_vars_loop(self):
+    self.assertFunctionMatchesEager(iterator_no_vars_loop, self.ds)
+
+  def test_dataset_single_var_loop(self):
+    self.assertFunctionMatchesEager(dataset_single_var_loop, self.ds)
+
+  def test_iterator_single_var_loop(self):
+    self.assertFunctionMatchesEager(iterator_single_var_loop, self.ds)
+
+  def test_dataset_two_vars_loop(self):
+    self.assertFunctionMatchesEager(dataset_two_vars_loop, self.ds)
+
+  def test_iterator_two_vars_loop(self):
+    self.assertFunctionMatchesEager(iterator_two_vars_loop, self.ds)
+
+  def test_dataset_loop_with_break(self):
+    self.assertFunctionMatchesEager(dataset_loop_with_break, self.ds)
+
+  def test_iterator_loop_with_break(self):
+    self.assertFunctionMatchesEager(iterator_loop_with_break, self.ds)
+
+  def test_dataset_loop_with_return_raises(self):
+    # This is for the same reason why returns in loops aren't allowed.
+    # TODO(mdan): This might be resolved by unrolling the loop once.
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        'a return statement cannot be placed inside this TensorFlow loop'):
+      tf.function(dataset_loop_with_return)(self.ds)
+
+  def test_iterator_loop_with_return_raises(self):
+    # This is for the same reason why returns in loops aren't allowed.
+    # TODO(mdan): This might be resolved by unrolling the loop once.
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        'a return statement cannot be placed inside this TensorFlow loop'):
+      tf.function(iterator_loop_with_return)(self.ds)
+
+  def test_iterator_next(self):
+    self.assertFunctionMatchesEager(iterator_next, self.ds)
+
+  def test_iterator_next_multiple_calls(self):
+    self.assertFunctionMatchesEager(iterator_next_multiple_calls, self.ds)
+
+  def test_iterator_next_in_loop(self):
+    self.assertFunctionMatchesEager(iterator_next_in_loop, self.ds, 7)
+
+  def test_iterator_next_stopping(self):
+    # Graph ops raise OutOfRangeError, but eager ops raise StopIteration
+    with self.assertRaises(tf.errors.OutOfRangeError):
+      tf.function(iterator_next_stopping)(self.ds, tf.constant(True))
+
+  def test_iterator_next_with_catching_stop_iteration(self):
+    # Graph ops raise OutOfRangeError, but eager ops raise StopIteration
+    with self.assertRaises(tf.errors.OutOfRangeError):
+      tf.function(iterator_next_with_catching_stop_iteration)(
+          self.ds, tf.constant(True))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/early_return_test.py

@@ -0,0 +1,280 @@
+# 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.
+# ==============================================================================
+"""Multiple returns, some in conditionals."""
+
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def return_with_default(x):
+  if x > 0:
+    tf.print('x', x)
+    return x
+  return x * x
+
+
+def return_dependent_on_local(c):
+  t = tf.constant(1)
+  if c:
+    return t
+  t = tf.stack([t, t])
+  return tf.reduce_sum(t)
+
+
+def return_possibly_undefined(x):
+  if x > 0:
+    if x < 5:
+      return x
+  else:
+    return x * x * x
+
+
+def nested_ifs(x):
+  if x > 0:
+    if x < 5:
+      return x
+    else:
+      return x * x
+  else:
+    return x * x * x
+
+
+def possible_return_before_loop(c1, c2, n):
+  if c1:
+    if c2:
+      return 1
+  for _ in range(n):
+    pass
+  return 2
+
+
+def nested_ifs_and_context_managers(x):
+  with tf.name_scope(''):
+    if x > 0:
+      if x < 5:
+        with tf.name_scope(''):
+          return x
+      else:
+        return x * x
+    else:
+      return x * x * x
+
+
+def unreachable_return(x):
+  with tf.name_scope(''):
+    if x > 0:
+      if x < 5:
+        with tf.name_scope(''):
+          return x
+      else:
+        return x * x
+    else:
+      return x * x * x
+  return x * x * x * x
+
+
+def return_with_default_in_contexmanager(x):
+  with tf.name_scope(''):
+    if x > 0:
+      return 1
+    return 0
+
+
+def return_in_try_with_finally(x):
+  try:
+    if x > 0:
+      return 1
+    else:
+      return 0
+  finally:
+    x = x + 1
+
+
+def return_with_default_in_try_with_finally(x):
+  try:
+    if x > 0:
+      return 1
+    return 0
+  finally:
+    x = x + 1
+
+
+def return_in_finally(x):
+  try:
+    return 2
+  finally:
+    if x > 0:
+      return 1  # pylint: disable=lost-exception
+    else:
+      return 0  # pylint: disable=lost-exception
+
+
+def return_with_default_in_finally(x):
+  try:
+    return 2
+  finally:
+    if x > 0:
+      return 1  # pylint: disable=lost-exception
+    return 0  # pylint: disable=lost-exception
+
+
+def return_in_finally_default_in_try(x):
+  try:
+    if x > 0:
+      return 0
+  finally:
+    return 1  # pylint: disable=lost-exception
+
+
+def _raising_helper():
+  raise ValueError()
+
+
+def raise_during_return_caught():
+  try:
+    return _raising_helper()
+  except ValueError:
+    pass
+  return 1
+
+
+def raise_during_return_caught_in_tail_branch(c):
+  if c:
+    return 2
+  try:
+    return _raising_helper()
+  except ValueError:
+    pass
+  return 1
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+  """Base class for the reference tests."""
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1),
+      (int, tf.constant),
+  ))
+  def test_return_with_default(self, n, type_):
+    self.assertFunctionMatchesEager(return_with_default, type_(n))
+
+  @parameterized.parameters(*itertools.product(
+      (True, False),
+      (int, tf.constant),
+  ))
+  def test_return_dependent_on_local(self, c, type_):
+    self.assertFunctionMatchesEager(return_dependent_on_local, type_(c))
+
+  @parameterized.parameters((0,), (3,), (5,))
+  def test_return_possibly_undefined_legal(self, n):
+    self.assertFunctionMatchesEager(return_possibly_undefined, n)
+
+  @parameterized.parameters((0,), (3,), (5,))
+  def test_return_possibly_undefined_illegal(self, n):
+    with self.assertRaisesRegex(
+        ValueError, 'else branch must also have a return'):
+      tf.function(return_possibly_undefined)(tf.constant(n))
+
+  @parameterized.parameters(*itertools.product(
+      (-1, 3, 6),
+      (int, tf.constant),
+  ))
+  def test_nested_ifs(self, n, type_):
+    self.assertFunctionMatchesEager(nested_ifs, type_(n))
+
+  @parameterized.parameters(*itertools.product(
+      (True, False),
+      (True, False),
+      (0, 1, 2),
+  ))
+  def test_possible_return_before_loop(self, c1, c2, n):
+    self.assertFunctionMatchesEager(possible_return_before_loop, c1, c2, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 3, 5),
+      (int, tf.constant),
+  ))
+  def test_nested_ifs_and_context_managers(self, x, type_):
+    self.assertFunctionMatchesEager(nested_ifs_and_context_managers, type_(x))
+
+  @parameterized.parameters(*itertools.product(
+      (0, 3, 5),
+      (int, tf.constant),
+  ))
+  def test_unreachable_return(self, x, type_):
+    self.assertFunctionMatchesEager(unreachable_return, type_(x))
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1),
+      (int, tf.constant),
+  ))
+  def test_return_with_default_in_contexmanager(self, x, type_):
+    self.assertFunctionMatchesEager(
+        return_with_default_in_contexmanager, type_(x))
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1),
+      (int, tf.constant),
+  ))
+  def test_return_in_try_finally(self, x, type_):
+    self.assertFunctionMatchesEager(return_in_try_with_finally, type_(x))
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1),
+      (int, tf.constant),
+  ))
+  def test_return_with_default_try_finally(self, x, type_):
+    self.assertFunctionMatchesEager(
+        return_with_default_in_try_with_finally, type_(x))
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1),
+      (int, tf.constant),
+  ))
+  def test_return_in_finally(self, x, type_):
+    self.assertFunctionMatchesEager(return_in_finally, type_(x))
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1),
+      (int, tf.constant),
+  ))
+  def test_return_with_default_in_finally(self, x, type_):
+    self.assertFunctionMatchesEager(return_with_default_in_finally, type_(x))
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1),
+      (int, tf.constant),
+  ))
+  def test_return_in_finally_default_in_try(self, x, type_):
+    self.assertFunctionMatchesEager(return_in_finally_default_in_try, type_(x))
+
+  def test_raise_during_return_caught(self):
+    self.assertFunctionMatchesEager(raise_during_return_caught)
+
+  @parameterized.parameters(*itertools.product(
+      (True, False),
+      (int, tf.constant),
+  ))
+  def test_raise_during_return_caught_in_tail_branch(self, c, type_):
+    self.assertFunctionMatchesEager(
+        raise_during_return_caught_in_tail_branch, type_(c))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/ext_slice_test.py

@@ -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.
+# ==============================================================================
+"""Extended slice operations."""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def basic_ext_slice(n):
+  return n[:, :], n[0, :], n[:, 0]
+
+
+def basic_expand_dims(n):
+  return n[:, tf.newaxis] - n[tf.newaxis, :]
+
+
+def slice_of_application(n, x):
+  return n(x)[:, tf.newaxis] - n(x)[tf.newaxis, :]
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_basic_ext_slice(self):
+    self.assertFunctionMatchesEager(basic_ext_slice, tf.eye(3))
+
+  def test_basic_expand_dims(self):
+    self.assertFunctionMatchesEager(basic_expand_dims, tf.eye(3))
+
+  def test_slice_of_application(self):
+    self.assertFunctionMatchesEager(slice_of_application, lambda x: x,
+                                    tf.eye(3))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/generator_test.py

@@ -0,0 +1,55 @@
+# 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.
+# ==============================================================================
+"""Generators."""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def basic_generator():
+  yield 1
+
+
+def generator_in_for(n):
+  for i in range(n):
+    yield i
+
+
+def generator_in_while(n):
+  i = 0
+  while i < n:
+    i += 1
+    yield i
+
+
+class LoopControlFlowTest(reference_test_base.TestCase):
+
+  def test_basic_generator(self):
+    with self.assertRaisesRegex(NotImplementedError, 'generators'):
+      tf.function(basic_generator)()
+
+  def test_generator_in_for(self):
+    with self.assertRaisesRegex(NotImplementedError, 'generators'):
+      tf.function(generator_in_for)([])
+
+  def test_generator_in_while(self):
+    with self.assertRaisesRegex(NotImplementedError, 'generators'):
+      tf.function(generator_in_while)(0)
+
+
+if __name__ == '__main__':
+  tf.test.main()
+

tensorflow/python/autograph/tests/logical_expression_test.py

@@ -0,0 +1,119 @@
+# 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.
+# ==============================================================================
+"""Basic logical expressions that are not autoboxed to TF."""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def composite_ors_with_callable(x, y, z):
+  z1 = lambda: z
+  return x or y or z1()
+
+
+def composite_ors(x, y, z):
+  return x or y or z
+
+
+def composite_ands(x, y, z):
+  return x and y and z
+
+
+def composite_mixed(x, y, z):
+  return x or y or z and y and z
+
+
+def equality(x, y):
+  return x == y
+
+
+def inequality(x, y):
+  return x != y
+
+
+def multiple_equality(x, y, z):
+  return x == y == z
+
+
+def comparison(x, y, z):
+  return x < y and y < z
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_basic(self):
+    self.assertFunctionMatchesEager(composite_ors, False, True, False)
+    self.assertFunctionMatchesEager(composite_ors, False, False, False)
+    self.assertFunctionMatchesEager(composite_ands, True, True, True)
+    self.assertFunctionMatchesEager(composite_ands, True, False, True)
+    self.assertFunctionMatchesEager(composite_mixed, False, True, True)
+    self.assertFunctionMatchesEager(composite_ors_with_callable, False, True,
+                                    False)
+    self.assertFunctionMatchesEager(composite_ors_with_callable, False, False,
+                                    True)
+    self.assertFunctionMatchesEager(composite_ors_with_callable, False, False,
+                                    False)
+
+    self.assertFunctionMatchesEager(comparison, 1, 2, 3)
+    self.assertFunctionMatchesEager(comparison, 2, 1, 3)
+    self.assertFunctionMatchesEager(comparison, 3, 2, 1)
+    self.assertFunctionMatchesEager(comparison, 3, 1, 2)
+    self.assertFunctionMatchesEager(comparison, 1, 3, 2)
+    self.assertFunctionMatchesEager(comparison, 2, 3, 1)
+
+  def test_basic_tensor(self):
+    self.all_inputs_tensors = True
+    self.assertFunctionMatchesEager(composite_ors, False, True, False)
+    self.assertFunctionMatchesEager(composite_ors, False, False, False)
+    self.assertFunctionMatchesEager(composite_ands, True, True, True)
+    self.assertFunctionMatchesEager(composite_ands, True, False, True)
+    self.assertFunctionMatchesEager(composite_mixed, False, True, True)
+    self.assertFunctionMatchesEager(composite_ors_with_callable, False, True,
+                                    False)
+    self.assertFunctionMatchesEager(composite_ors_with_callable, False, False,
+                                    True)
+    self.assertFunctionMatchesEager(composite_ors_with_callable, False, False,
+                                    False)
+
+    self.assertFunctionMatchesEager(comparison, 1, 2, 3)
+    self.assertFunctionMatchesEager(comparison, 2, 1, 3)
+    self.assertFunctionMatchesEager(comparison, 3, 2, 1)
+    self.assertFunctionMatchesEager(comparison, 3, 1, 2)
+    self.assertFunctionMatchesEager(comparison, 1, 3, 2)
+    self.assertFunctionMatchesEager(comparison, 2, 3, 1)
+
+  def test_equality(self):
+    self.assertFunctionMatchesEager(equality, 1, 1)
+    self.assertFunctionMatchesEager(equality, 1, 2)
+    self.assertFunctionMatchesEager(inequality, 1, 1)
+    self.assertFunctionMatchesEager(inequality, 1, 2)
+    self.assertFunctionMatchesEager(multiple_equality, 1, 1, 2)
+    self.assertFunctionMatchesEager(multiple_equality, 1, 1, 1)
+
+  def test_equality_tensor(self):
+    self.autograph_opts = tf.autograph.experimental.Feature.EQUALITY_OPERATORS
+    self.all_inputs_tensors = True
+
+    self.assertFunctionMatchesEager(equality, 1, 1)
+    self.assertFunctionMatchesEager(equality, 1, 2)
+    self.assertFunctionMatchesEager(inequality, 1, 1)
+    self.assertFunctionMatchesEager(inequality, 1, 2)
+    self.assertFunctionMatchesEager(multiple_equality, 1, 1, 2)
+    self.assertFunctionMatchesEager(multiple_equality, 1, 1, 1)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_basic_test.py

@@ -0,0 +1,490 @@
+# 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.
+# ==============================================================================
+"""Basic loops iterating over various types."""
+
+import functools
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def while_no_vars(n, v):
+  v.assign(0)
+  while v.read_value() < n:
+    v.assign_add(1)
+  return v.read_value()
+
+
+def for_no_vars(l, v):
+  v.assign(0)
+  for _ in l:
+    v.assign_add(1)
+  return v.read_value()
+
+
+def while_one_var(n):
+  i = 0
+  while i < n:
+    i = i * 10 + 1
+  return i
+
+
+def for_one_var(l):
+  i = 0
+  for i in l:
+    pass
+  return i
+
+
+def while_two_vars(n):
+  s = 0
+  i = 0
+  while i < n:
+    s = s * 10 + i
+    i += 1
+  return s
+
+
+def for_two_vars(l):
+  s = 0
+  for i in l:
+    s = s * 10 + i
+  return s
+
+
+def while_else(x, y):
+  s = 0
+  while x > 0:
+    x -= 1
+    if x > y:
+      break
+    s += x
+  else:
+    s = -100
+  return s
+
+
+def for_else(l1, l2):
+  s = 0
+  for c in l1:
+    if c in l2:
+      break
+    s = s * 10 + c
+  else:
+    s = -1000
+  return s
+
+
+def while_creates_var(n):
+  i = 0
+  while i < n:
+    i = n + 1
+    j = i + 1
+  return i, j
+
+
+def for_creates_var(l):
+  for i in l:
+    pass
+  return i  # pylint:disable=undefined-loop-variable
+
+
+def successive_while_loops(n1, n2):
+  s = 0
+  i = 0
+  while i < n1:
+    s = s * 10 + i
+    i += 1
+  i = 0
+  while i < n2:
+    s = s * 10 + i
+    i += 1
+  return s
+
+
+def successive_for_loops(l1, l2):
+  s = 0
+  for i in l1:
+    s = s * 10 + i
+  for i in l2:
+    s = s * 10 + i
+  return s
+
+
+def nested_while_loops(n1, n2):
+  i = 0
+  l = tf.TensorArray(tf.int32, size=0, dynamic_size=True, element_shape=())
+  while i < n1:
+    j = 0
+    s = 0
+    while j < n2:
+      s = s * 10 + i * j
+      j += 1
+    l = l.write(i, s)
+    i += 1
+  return l.stack()
+
+
+def nested_for_loops(m):
+  l = tf.TensorArray(tf.int32, size=0, dynamic_size=True, element_shape=())
+  for i in m:
+    s = 0
+    for j in i:
+      s = s * 10 + j
+    l = l.write(l.size(), s)
+  return l.stack()
+
+
+def _int_tensor(x):
+  return tf.constant(x, dtype=tf.int32)
+
+
+def _int_dataset(l):
+  return tf.data.Dataset.from_tensor_slices(tf.constant(l, dtype=tf.int32))
+
+
+def double_product(l1, l2):
+  for i in l1:
+    for j in l2:
+      for k in l1:
+        for l in l2:
+          yield i, j, k, l
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_while_no_vars(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_no_vars, n, tf.Variable(0))
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [],
+          [1],
+          [1, 2],
+      ),
+      (
+          list,
+          _int_tensor,
+          # TODO(mdan): Enable this once #35335 is fixed.
+          # _int_dataset,
+      ),
+      (
+          True,
+          False,
+      ),
+  ))
+  def test_for_no_vars(self, l, type_, xla):
+    if type_ is _int_tensor and xla and not l:
+      self.skipTest('Empty loops not supported in XLA')
+
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_no_vars, l, tf.Variable(0), xla=xla)
+
+  @parameterized.parameters(
+      ([],),
+      ([1],),
+      ([1, 2],),
+  )
+  def test_for_no_vars_ds_iterator(self, l):
+    inputs_ = lambda: (iter(_int_dataset(l)), tf.Variable(0))
+    self.assertFunctionMatchesEagerStatefulInput(for_no_vars, inputs_)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_while_one_var(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_one_var, n)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [],
+          [1],
+          [1, 2],
+      ),
+      (
+          list,
+          _int_tensor,
+          _int_dataset,
+      ),
+      (
+          True,
+          False,
+      ),
+  ))
+  def test_for_one_var(self, l, type_, xla):
+    if type_ is _int_dataset and xla:
+      self.skipTest('Datasets not supported in XLA')
+    if type_ is _int_tensor and xla and not l:
+      self.skipTest('Empty loops not supported in XLA')
+
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_one_var, l, xla=xla)
+
+  @parameterized.parameters(
+      ([],),
+      ([1],),
+      ([1, 2],),
+  )
+  def test_for_one_var_ds_iterator(self, l):
+    inputs_ = lambda: (iter(_int_dataset(l)), tf.Variable(0))
+    self.assertFunctionMatchesEagerStatefulInput(for_one_var, inputs_)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_while_two_vars(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_two_vars, n)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [],
+          [1],
+          [1, 2],
+      ),
+      (
+          list,
+          _int_tensor,
+          _int_dataset,
+      ),
+      (
+          True,
+          False,
+      ),
+  ))
+  def test_for_two_vars(self, l, type_, xla):
+    if type_ is _int_dataset and xla:
+      self.skipTest('Datasets not supported in XLA')
+    if type_ is _int_tensor and xla and not l:
+      self.skipTest('Empty loops not supported in XLA')
+
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_two_vars, l, xla=xla)
+
+  @parameterized.parameters(
+      ([],),
+      ([1],),
+      ([1, 2],),
+  )
+  def test_for_two_vars_ds_iterator(self, l):
+    inputs_ = lambda: (iter(_int_dataset(l)), tf.Variable(0))
+    self.assertFunctionMatchesEagerStatefulInput(for_two_vars, inputs_)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [],
+          [1, 2],
+      ),
+      (
+          [],
+          [1, 2],
+      ),
+      (list, _int_tensor),
+  ))
+  def test_for_else(self, l1, l2, type_):
+    l1 = type_(l1)
+    l2 = type_(l2)
+    with self.assertRaisesRegex(NotImplementedError, 'for/else'):
+      tf.function(for_else)(l1, l2)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 2),
+      (0, 2),
+      (int, _int_tensor),
+  ))
+  def test_while_else(self, n1, n2, type_):
+    n1 = type_(n1)
+    n2 = type_(n2)
+    with self.assertRaisesRegex(NotImplementedError, 'while/else'):
+      tf.function(while_else)(n1, n2)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          1,
+          2,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_while_creates_var_legal(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_creates_var, n)
+
+  def test_while_creates_var_illegal(self):
+    with self.assertRaisesRegex(UnboundLocalError, 'used before assignment'):
+      tf.function(while_creates_var)(0)
+    with self.assertRaisesRegex(
+        tf.errors.InvalidArgumentError, 'loop must iterate at least once'):
+      tf.function(while_creates_var)(tf.constant(0))
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [1],
+          [1, 2],
+      ),
+      (
+          list,
+          _int_tensor,
+      ),
+  ))
+  def test_for_creates_var_legal(self, l, type_):
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_creates_var, l)
+
+  def test_for_creates_var_illegal(self):
+    with self.assertRaisesRegex(UnboundLocalError, 'used before assignment'):
+      tf.function(for_creates_var)([])
+    with self.assertRaisesRegex(
+        tf.errors.InvalidArgumentError, 'loop must iterate at least once'):
+      tf.function(for_creates_var)(tf.constant([]))
+
+  @parameterized.parameters(*double_product(
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_successive_while_loops(self, n1, type1, n2, type2):
+    n1 = type1(n1)
+    n2 = type1(n2)
+    self.assertFunctionMatchesEager(successive_while_loops, n1, n2)
+
+  @parameterized.parameters(*double_product(
+      (
+          [],
+          [1],
+          [1, 2],
+      ),
+      (
+          list,
+          _int_tensor,
+          _int_dataset,
+      ),
+  ))
+  def test_successive_for_loops(self, l1, type1, l2, type2):
+    l1 = type1(l1)
+    l2 = type1(l2)
+    self.assertFunctionMatchesEager(successive_for_loops, l1, l2)
+
+  @parameterized.parameters(*double_product(
+      (
+          [],
+          [1],
+          [1, 2],
+      ),
+      (
+          list,
+          _int_dataset,
+      ),
+  ))
+  def test_successive_for_loops_iterators(self, l1, type1, l2, type2):
+    inputs_ = lambda: (iter(type1(l1)), iter(type2(l2)))
+    self.assertFunctionMatchesEagerStatefulInput(successive_for_loops, inputs_)
+
+  @parameterized.parameters(*double_product(
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_nested_while_loops(self, n1, type1, n2, type2):
+    n1 = type1(n1)
+    n2 = type1(n2)
+    self.assertFunctionMatchesEager(nested_while_loops, n1, n2)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [[]],
+          [[], []],
+          [[1]],
+          [[1], [2]],
+          [[1, 2]],
+          [[1, 2], [3, 4]],
+      ),
+      (
+          _int_tensor,
+          _int_dataset,
+      ),
+  ))
+  def test_nested_for_loops_dense(self, m, type_):
+    m = type_(m)
+    self.assertFunctionMatchesEager(nested_for_loops, m)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [[]],
+          [[], [1]],
+          [[], [1], [1, 2]],
+      ),
+      (
+          list,
+          functools.partial(tf.ragged.constant, dtype=tf.int32),
+      ),
+  ))
+  def test_nested_for_loops_ragged(self, m, type_):
+    m = type_(m)
+    self.assertFunctionMatchesEager(nested_for_loops, m)
+
+  def test_nested_for_loops_mixed_list(self):
+    m = [[], _int_tensor([]), [1], _int_tensor([1]), [1, 2]]
+    self.assertFunctionMatchesEager(nested_for_loops, m)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_control_flow_illegal_cases_test.py

@@ -0,0 +1,112 @@
+# 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.
+# ==============================================================================
+"""Loop control statements (e.g. break, return) in illegal patterns.
+
+Meant to verify that:
+  * break/return on a dynamic condition raises error inside Python loop
+"""
+
+import itertools
+import re
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def tf_break_in_py_for(l):
+  s = 0
+  for c in l:
+    if tf.greater(c % 2, 0):
+      break
+    s += c
+  return s
+
+
+def tf_return_in_py_for(l):
+  s = 0
+  for c in l:
+    if tf.greater(c % 2, 0):
+      return s
+    else:
+      return s
+    s += c
+  return s
+
+
+def tf_break_in_py_while(x):
+  s = 0
+  while x > 0:
+    x -= 1
+    if tf.greater(x % 2, 0):
+      break
+    s += x
+  return s
+
+
+def tf_return_in_py_while(x):
+  s = 0
+  while x > 0:
+    x -= 1
+    if tf.greater(x % 2, 0):
+      return s
+    else:
+      return s
+    s += x
+  return s
+
+
+class LoopControlFlowIllegalCasesTest(reference_test_base.TestCase,
+                                      parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [1],
+          [1, 2],
+          [1, 2, 3],
+      ),
+      (
+          tf_break_in_py_for,
+          tf_return_in_py_for,
+      ),
+  ))
+  def test_tf_control_flow_in_py_for(self, l, target):
+    with self.assertRaisesRegex(NotImplementedError,
+                                'not supported in Python for'):
+      tf.function(target)(l)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          1,
+          2,
+          3,
+      ),
+      (
+          tf_break_in_py_while,
+          tf_return_in_py_while,
+      ),
+  ))
+  def test_tf_control_flow_in_py_while(self, n, target):
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        re.compile(
+            r'.*condition of while loop started as non\-Tensor,'
+            r' then changed to Tensor.*', re.DOTALL)):
+      tf.function(target)(n)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_control_flow_test.py

@@ -0,0 +1,359 @@
+# 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.
+# ==============================================================================
+"""Nested loops and loop control statements (e.g. break and continue).
+
+Meant to verify that:
+  * break/continue in the inner loop does not affect outer loop
+  * break/continue inside nested conditionals still works
+"""
+
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def continue_in_single_for(l):
+  s = 0
+  for c in l:
+    if c % 2 > 0:
+      continue
+    s += c
+  return s
+
+
+def continue_in_single_while(x):
+  s = 0
+  while x > 0:
+    x -= 1
+    if x % 2 > 0:
+      continue
+    s += x
+  return s
+
+
+def continue_in_inner_for(m):
+  s = 0
+  for l in m:
+    for c in l:
+      if c % 2 > 0:
+        continue
+      s += c
+  return s
+
+
+def continue_in_inner_while(x, y):
+  s = 0
+  while x > 0:
+    x -= 1
+    while y > 0:
+      y -= 1
+      if (x + y) % 2 > 0:
+        continue
+      s += x + y
+  return s
+
+
+def break_in_single_for(l):
+  s = 0
+  for c in l:
+    if c % 2 > 0:
+      break
+    s += c
+  return s
+
+
+def unconditional_return_in_single_for(l):
+  s = 0
+  for c in l:
+    s += c
+    return s
+  return s
+
+
+def effectively_unconditional_return_in_single_for(l):
+  s = 0
+  for c in l:
+    s += c
+    if c % 2 > 0:
+      return s
+    else:
+      return -s
+  return s
+
+
+def unconditional_return_in_single_while(x):
+  s = 0
+  while x > 0:
+    x -= 1
+    s += x
+    return s
+  return s
+
+
+def effectively_unconditional_return_in_single_while(x):
+  s = 0
+  while x > 0:
+    x -= 1
+    s += x
+    if x % 2 > 0:
+      return s
+    else:
+      return -s
+  return s
+
+
+def break_in_single_while(x):
+  s = 0
+  while x > 0:
+    x -= 1
+    if x % 2 > 0:
+      break
+    s += x
+  return s
+
+
+def break_in_inner_for(m):
+  s = 0
+  for l in m:
+    for c in l:
+      if c % 2 > 0:
+        break
+      s += c
+  return s
+
+
+def break_in_inner_while(x, y):
+  s = 0
+  while x > 0:
+    x -= 1
+    while y > 0:
+      y -= 1
+      if ((x + y) % 2) == 0:
+        break
+      s += x + y
+  return s
+
+
+def break_continue_in_inner_for(m):
+  s = 0
+  for l in m:
+    for c in l:
+      if c % 2 > 0:
+        break
+      else:
+        continue
+      s += c
+  return s
+
+
+def break_continue_in_inner_while(x, y):
+  s = 0
+  while x > 0:
+    x -= 1
+    while y > 0:
+      y -= 1
+      if (x + y) % 2 > 0:
+        break
+      else:
+        continue
+      s += x + y
+  return s
+
+
+def break_followed_by_cond_in_single_for(x, y):
+  for i in range(y):
+    if i == 2:
+      break
+    if x > 0:
+      x -= 1
+  return x
+
+
+def break_followed_by_cond_in_single_while(x):
+  while x > 0:
+    if x == 2:
+      break
+    if x > 0:
+      x -= 1
+  return x
+
+
+def multiple_breaks_in_single_while(n):
+  s = 1
+  i = 0
+  while i < n:
+    i += 1
+    if i > 10 * n:
+      break
+    if i == n:
+      break
+    s = s * 10 + i
+  return i, s
+
+
+def _int_tensor(x):
+  return tf.constant(x, dtype=tf.int32)
+
+
+def _list_of_int_tensor(l):
+  return [_int_tensor(x) for x in l]
+
+
+def _int_dataset(l):
+  return tf.data.Dataset.from_tensor_slices(tf.constant(l, dtype=tf.int32))
+
+
+class LoopControlFlowTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [],
+          [1],
+          [1, 2],
+          [1, 2, 3],
+          [1, 2, 3, 4],
+      ),
+      (
+          list,
+          _int_tensor,
+          _int_dataset,
+      ),
+      (
+          continue_in_single_for,
+          break_in_single_for,
+          unconditional_return_in_single_for,
+          effectively_unconditional_return_in_single_for,
+      ),
+  ))
+  def test_single_for(self, l, type_, target):
+    if ((type_ is _int_dataset) and
+        (target in (unconditional_return_in_single_for,
+                    effectively_unconditional_return_in_single_for))):
+      # TODO(mdan): Enable in a separate improvement.
+      self.skipTest('Creating symbols in dataset loops.')
+
+    if ((not l) and
+        ((target in (unconditional_return_in_single_for,
+                     effectively_unconditional_return_in_single_for)))):
+      self.skipTest('Undefined symbols require at least one iteration.')
+
+    l = type_(l)
+    self.assertFunctionMatchesEager(target, l)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0,
+          1,
+          2,
+          3,
+          4,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+      (
+          continue_in_single_while,
+          break_in_single_while,
+          multiple_breaks_in_single_while,
+          break_followed_by_cond_in_single_while,
+          unconditional_return_in_single_while,
+          effectively_unconditional_return_in_single_while,
+      ),
+  ))
+  def test_single_while(self, n, type_, target):
+    if ((not n) and
+        ((target in (unconditional_return_in_single_while,
+                     effectively_unconditional_return_in_single_while)))):
+      self.skipTest('Undefined symbols require at least one iteration.')
+
+    n = type_(n)
+    self.assertFunctionMatchesEager(target, n)
+
+  @parameterized.parameters(
+      (unconditional_return_in_single_for, _int_tensor, []),
+      (effectively_unconditional_return_in_single_for, _int_tensor, []),
+      (unconditional_return_in_single_while, _int_tensor, 0),
+      (effectively_unconditional_return_in_single_while, _int_tensor, 0),
+  )
+  def test_single_loop_illegal_return(self, target, type_, l):
+    with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                'must iterate at least once to initialize'):
+      tf.function(target)(type_(l))
+
+  @parameterized.parameters(*itertools.product(
+      (
+          [[], []],
+          [[1], [2]],
+          [[1, 2], [3, 4]],
+          [[1, 2, 3], [4, 5, 6]],
+          # TODO(mdan): Add ragged tensors / variable-shape datasets.
+      ),
+      (
+          list,
+          _int_tensor,
+          _list_of_int_tensor,
+          _int_dataset,
+      ),
+      (
+          continue_in_inner_for,
+          break_in_inner_for,
+          break_continue_in_inner_for,
+      ),
+  ))
+  def test_nested_for(self, a, type_, target):
+    a = type_(a)
+    self.assertFunctionMatchesEager(target, a)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0,
+          1,
+          2,
+          3,
+          4,
+      ),
+      (
+          0,
+          1,
+          2,
+          3,
+          4,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+      (
+          continue_in_inner_while,
+          break_in_inner_while,
+          break_continue_in_inner_while,
+      ),
+  ))
+  def test_nested_while(self, m, n, m_type, n_type, target):
+    m = m_type(m)
+    n = m_type(n)
+    self.assertFunctionMatchesEager(target, m, n)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_created_variables_test.py

@@ -0,0 +1,137 @@
+# 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.
+# ==============================================================================
+"""Loops which create variables, with or without shape invariants."""
+
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+# TODO(mdan): When syntax allows shape_invariants on created vars, add tests.
+
+
+def while_creates_var_static_shape(n):
+  i = 0
+  while i < n:
+    v = tf.zeros([1, 2, 3])
+    i += 1
+  return v
+
+
+def while_creates_var_dynamic_shape(n):
+  i = 0
+  while i < n:
+    v = tf.zeros([1, tf.random.uniform((), i, i + 1, tf.int32), 2])
+    i += 1
+  return v
+
+
+def while_creates_var_dynamic_rank(n):
+  i = 0
+  while i < n:
+    v = tf.zeros(tf.range(tf.random.uniform((), i, i + 1, tf.int32)))
+    i += 1
+  return v
+
+
+def while_creates_var_dynamic_shape_py_init_var(n):
+  i = 0
+  while i < n:
+    v = tf.range(i)
+    i += 1
+  return v
+
+
+def while_creates_nested_var_static_shape(n):
+  i = 0
+  while i < n:
+    v = {'a': tf.zeros([1, 2, 3]), 'b': tf.ones([1, 2, 3])}
+    i += 1
+  return v['a'], v['b']
+
+
+def while_creates_nested_var_dynamic_shape(n):
+  i = 0
+  while i < n:
+    v = {
+        'a': tf.zeros([1, tf.random.uniform((), i, i + 1, tf.int32)]),
+        'b': tf.ones([tf.random.uniform((), i, i + 1, tf.int32), 2])
+    }
+    i += 1
+  return v['a'], v['b']
+
+
+def while_creates_nested_var_dynamic_rank(n):
+  i = 0
+  while i < n:
+    v = {
+        'a': tf.ones(tf.range(tf.random.uniform((), i, i + 1, tf.int32))),
+        'b': tf.ones([1, 2, 3])
+    }
+    i += 1
+  return v['a'], v['b']
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      while_creates_var_static_shape,
+      while_creates_var_dynamic_shape,
+      while_creates_var_dynamic_rank,
+      while_creates_var_dynamic_shape_py_init_var,
+      while_creates_nested_var_static_shape,
+      while_creates_nested_var_dynamic_shape,
+      while_creates_nested_var_dynamic_rank,
+  )
+  def test_while_creates_var_illegal_tf(self, target):
+    with self.assertRaises(tf.errors.InvalidArgumentError):
+      tf.function(target)(tf.constant(0))
+
+  @parameterized.parameters(
+      while_creates_var_static_shape,
+      while_creates_var_dynamic_shape,
+      while_creates_var_dynamic_rank,
+      while_creates_var_dynamic_shape_py_init_var,
+      while_creates_nested_var_static_shape,
+      while_creates_nested_var_dynamic_shape,
+      while_creates_nested_var_dynamic_rank,
+  )
+  def test_while_creates_var_illegal_py(self, target):
+    with self.assertRaises(UnboundLocalError):
+      tf.function(target)(0)
+
+  @parameterized.parameters(*itertools.product(
+      (1, 2),
+      (int, tf.constant),
+      (
+          while_creates_var_static_shape,
+          while_creates_var_dynamic_shape,
+          while_creates_var_dynamic_rank,
+          while_creates_var_dynamic_shape_py_init_var,
+          while_creates_nested_var_static_shape,
+          while_creates_nested_var_dynamic_shape,
+          while_creates_nested_var_dynamic_rank,
+      ),
+  ))
+  def test_while_creates_var(self, n, type_, target):
+    n = type_(n)
+    self.assertFunctionMatchesEager(target, n)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_distributed_test.py

@@ -0,0 +1,190 @@
+# 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 involving the tf.distributed datasets."""
+
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def no_vars_loop(strat, iterable):
+  for pr in iterable:
+    tf.print(strat.reduce('SUM', pr, axis=0))
+
+
+def single_var_loop(strat, iterable):
+  s = 0
+  for pr in iterable:
+    # TODO(mdan): It would be nice to be able to write s = s * 10 + pr.
+    s = s * 10 + strat.reduce('SUM', pr, axis=0)
+  return s
+
+
+def loop_with_break(strat, iterable):
+  s = 0
+  for pr in iterable:
+    if strat.reduce('SUM', pr, axis=0) % 5 == 0:
+      break
+    s = s * 10 + strat.reduce('SUM', pr, axis=0)
+  return s
+
+
+def loop_with_continue(strat, iterable):
+  s = 0
+  for pr in iterable:
+    if strat.reduce('SUM', pr, axis=0) % 2 == 0:
+      continue
+    s = s * 10 + strat.reduce('SUM', pr, axis=0)
+  return s
+
+
+def two_vars_loop(strat, iterable):
+  s = 0
+  p = 1
+  for pr in iterable:
+    e = strat.reduce('SUM', pr, axis=0)
+    s += e
+    p *= e
+  return s, p
+
+
+def enumeration(strat, iterable):
+  s = 0
+  p = 1
+  for i, pr in enumerate(iterable):
+    e = strat.reduce('SUM', pr, axis=0)
+    s = s * 10 + e
+    p *= i
+  return s, p
+
+
+def iterator_next(strat, iterable):
+  itr = iter(iterable)
+  return strat.reduce('SUM', next(itr), axis=0)
+
+
+def iterator_next_multiple_calls(strat, iterable):
+  itr = iter(iterable)
+  a = strat.reduce('SUM', next(itr), axis=0)
+  b = strat.reduce('SUM', next(itr), axis=0)
+  return a * 10 + b
+
+
+def iterator_next_in_limited_loop(strat, iterable, l):
+  itr = iter(iterable)
+  s = 0
+  for _ in l:
+    s = s * 10 + strat.reduce('SUM', next(itr), axis=0)
+  return s
+
+
+def iterator_next_stopping(strat, iterable, cond):
+  # This case will raise, but not the expected StopIteration error.
+  itr = iter(iterable)
+  while cond:
+    strat.reduce('SUM', next(itr), axis=0)
+
+
+def iterator_next_with_catching_stop_iteration(strat, iterable, cond):
+  # This is the one instance when the use of TF iterators does not work as
+  # intended. In graph mode, the `except` below will never catch, and the
+  # tf.function will raise the error instead.
+  # TODO(b/132311724): The error should be friendlier here.
+  # Note: b/132298783 covers actually supporting this pattern.
+  itr = iter(iterable)
+  try:
+    while cond:
+      strat.reduce('SUM', next(itr), axis=0)
+  except StopIteration:
+    pass
+
+
+def _distributed_dataset():
+  cpus = tf.config.experimental.list_physical_devices('CPU')
+  tf.config.experimental.set_virtual_device_configuration(
+      cpus[0], [tf.config.experimental.VirtualDeviceConfiguration()] * 2)
+
+  strat = tf.distribute.MirroredStrategy()
+  ds = tf.data.Dataset.from_tensor_slices(tf.reshape(tf.range(12), (3, 4)))
+
+  return strat, strat.experimental_distribute_dataset(ds)
+
+
+def _distributed_iterator():
+  strat, ds = _distributed_dataset()
+  return strat, iter(ds)
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (
+          no_vars_loop,
+          single_var_loop,
+          two_vars_loop,
+          loop_with_break,
+          loop_with_continue,
+      ),
+      (
+          _distributed_dataset,
+          _distributed_iterator,
+      ),
+  ))
+  def test_basic(self, test_fn, target):
+    if (test_fn in (loop_with_break, loop_with_continue) and
+        target is _distributed_dataset):
+      self.skipTest('b/162250181')
+    self.assertFunctionMatchesEagerStatefulInput(test_fn, target)
+
+  def test_iterator_next(self):
+    strat, ds = _distributed_dataset()
+    self.assertFunctionMatchesEager(iterator_next, strat, ds)
+
+  def test_iterator_next_multiple_calls(self):
+    strat, ds = _distributed_dataset()
+    self.assertFunctionMatchesEager(iterator_next_multiple_calls, strat, ds)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0,
+          1,
+          2,
+      ),
+      (
+          range,
+          tf.range,
+      ),
+  ))
+  def test_iterator_next_in_limited_loop(self, n, type_):
+    n = type_(n)
+    strat, ds = _distributed_dataset()
+    self.assertFunctionMatchesEager(iterator_next_in_limited_loop, strat, ds, n)
+
+  @parameterized.parameters(
+      (iterator_next_stopping,),
+      # Note that `except` has no effect in graph mode.
+      (iterator_next_with_catching_stop_iteration,),
+  )
+  def test_iterator_next_stopping(self, test_fn):
+    strat, ds = _distributed_dataset()
+    with self.assertRaises(tf.errors.OutOfRangeError):
+      tf.function(test_fn)(strat, ds, tf.constant(True))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_scoping_test.py

@@ -0,0 +1,322 @@
+# 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 that verify scoping around loops."""
+
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def for_with_local_var(l):
+  s = 0
+  for i in l:
+    x = i + 2
+    s = s * 10 + x
+  return s
+
+
+def while_with_local_var(x):
+  s = 0
+  while x > 0:
+    y = x + 2
+    s = s * 10 + y
+    x -= 1
+  return s
+
+
+def for_with_lambda_iter(l):
+  fns = []
+  results = []
+  for i in l:
+    fns.append(lambda: i)
+  for f in fns:
+    results.append(f())
+  return results
+
+
+def for_with_lambda_object():
+
+  class SomeRandomObject:
+
+    def bar(self, n):
+      return n + 1
+
+  def foo_init():
+    return tf.constant(0)
+
+  fns = []
+  results = []
+  foo = foo_init()
+  for i in tf.range(3):
+    foo = SomeRandomObject()
+    fns.append(lambda i=i: foo.bar(i))
+  for f in fns:
+    results.append(f())
+  return results
+
+
+def for_with_lambda_iter_local_var(l):
+  fns = []
+  results = []
+  for i in l:
+    fns.append(lambda i=i: i)
+  for f in fns:
+    results.append(f())
+  return results
+
+def for_initializes_local_var(l):
+  s = 0
+  for i in l:
+    if i == l[0]:
+      x = 0
+    else:
+      x += 1
+    s = s * 10 + x
+  return s
+
+
+def while_initializes_local_var(x):
+  s = 0
+  while x > 0:
+    if x > 0:
+      y = 0
+    else:
+      y += 1
+    s = s * 10 + y
+    x -= 1
+  return s
+
+
+def for_defines_var(l):
+  for i in l:
+    x = i + 2
+  return x
+
+
+def while_defines_var(x):
+  while x > 0:
+    y = x + 2
+    x -= 1
+  return y
+
+
+def for_defines_iterate(n, fn):
+  s = 0
+  for i in fn(n):
+    s = s * 10 + i
+  return i, s  # pylint:disable=undefined-loop-variable
+
+
+def for_reuses_iterate(n, fn):
+  i = 7
+  s = 0
+  for i in fn(n):
+    s = s * 10 + i
+  return i, s
+
+
+def for_alters_iterate(n, fn):
+  i = 7
+  s = 0
+  for i in fn(n):
+    i = 3 * i + 1
+    s = s * 10 + i
+  return i, s
+
+
+def _int_tensor(x):
+  return tf.constant(x, dtype=tf.int32)
+
+
+class LoopScopingTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+  ))
+  def test_for_with_local_var(self, l, type_):
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_local_var, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (range, tf.range),
+  ))
+  def test_for_with_local_var_range(self, l, type_):
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_local_var, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, _int_tensor),
+  ))
+  def test_while_with_local_var(self, x, type_):
+    x = type_(x)
+    self.assertFunctionMatchesEager(while_with_local_var, x)
+
+  @parameterized.parameters(
+      ([],),
+      ([1],),
+      ([1, 2],),
+  )
+  def test_for_initializes_local_var_legal_cases(self, l):
+    self.assertFunctionMatchesEager(for_initializes_local_var, l)
+
+  @parameterized.parameters(
+      ([],),
+      ([1],),
+      ([1, 2],),
+  )
+  def test_for_initializes_local_var_illegal_cases(self, l):
+    self.skipTest('TODO(mdanatg): Check')
+    l = tf.constant(l)
+    with self.assertRaisesRegex(ValueError, '"x" must be defined'):
+      tf.function(for_initializes_local_var)(l)
+
+  @parameterized.parameters(
+      0,
+      1,
+      2,
+  )
+  def test_while_initializes_local_var_legal_cases(self, x):
+    self.assertFunctionMatchesEager(while_initializes_local_var, x)
+
+  @parameterized.parameters(
+      0,
+      1,
+      2,
+  )
+  def test_while_initializes_local_var_illegal_cases(self, x):
+    self.skipTest('TODO(mdanatg): check')
+    x = tf.constant(x)
+    with self.assertRaisesRegex(ValueError, '"y" must be defined'):
+      tf.function(while_initializes_local_var)(x)
+
+  @parameterized.parameters(
+      # TODO(b/155171694): Enable once the error message here is corrected.
+      # ([],),
+      ([1],),
+      ([1, 2],),
+  )
+  def test_for_defines_var_legal_cases(self, l):
+    self.assertFunctionMatchesEager(for_defines_var, l)
+
+  @parameterized.parameters(
+      ([],),
+      ([1],),
+      ([1, 2],),
+  )
+  def test_for_defines_var_illegal_cases(self, l):
+    self.skipTest('TODO(mdanatg): check')
+    l = tf.constant(l)
+    with self.assertRaisesRegex(ValueError, '"x" must be defined'):
+      tf.function(for_defines_var)(l)
+
+  @parameterized.parameters(
+      # TODO(b/155171694): Enable once the error message here is corrected.
+      # (0,),
+      (1,),
+      (2,),
+  )
+  def test_while_defines_var_legal_cases(self, x):
+    self.assertFunctionMatchesEager(while_defines_var, x)
+
+  @parameterized.parameters(
+      (0,),
+      (1,),
+      (2,),
+  )
+  def test_while_defines_var_illegal_cases(self, x):
+    self.skipTest('TODO(mdanatg): check')
+    x = tf.constant(x)
+    with self.assertRaisesRegex(ValueError, '"y" must be defined'):
+      tf.function(while_defines_var)(x)
+
+  @parameterized.parameters(*itertools.product(
+      (1, 2),
+      (range, tf.range),
+  ))
+  def test_for_defines_iterate_legal_cases(self, n, fn):
+    self.assertFunctionMatchesEager(for_defines_iterate, n, fn)
+
+  def test_for_defines_iterate_range(self):
+    self.skipTest('b/155171694')
+
+  def test_for_defines_iterate_tf_range(self):
+    # Deviating from the normal Python semantics here to avoid inserting
+    # an extra assert op. If needed, we can insert it and raise an error
+    # to mimic the eager behavior, but this is an exceptionally uncummon
+    # use case.
+    self.assertAllEqual(tf.function(for_defines_iterate)(0, tf.range), (0, 0))
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+  ))
+  def test_for_reuses_iterate(self, l, fn):
+    self.assertFunctionMatchesEager(for_reuses_iterate, l, fn)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (range, tf.range),
+  ))
+  def test_for_reuses_iterate_range(self, n, fn):
+    self.assertFunctionMatchesEager(for_reuses_iterate, n, fn)
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+  ))
+  def test_for_alters_iterate(self, l, fn):
+    self.assertFunctionMatchesEager(for_alters_iterate, l, fn)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (range, tf.range),
+  ))
+  def test_for_alters_iterate_range(self, n, fn):
+    self.assertFunctionMatchesEager(for_alters_iterate, n, fn)
+
+
+class LoopLambdaScopingTest(reference_test_base.TestCase,
+                            parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2], [(1, 2), (3, 4)]),
+      (list, list),
+  ))
+  def test_for_with_lambda_iter(self, l, type_):
+    self.skipTest('https://github.com/tensorflow/tensorflow/issues/56089')
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_lambda_iter, l)
+
+  def test_for_with_lambda_object(self):
+    self.skipTest('https://github.com/tensorflow/tensorflow/issues/56089')
+    self.assertFunctionMatchesEager(for_with_lambda_object)
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2], [(1, 2), (3, 4)]),
+      (list, list),
+  ))
+  def test_for_with_lambda_iter_local_var(self, l, type_):
+    self.skipTest('https://github.com/tensorflow/tensorflow/issues/56089')
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_lambda_iter_local_var, l)
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_with_function_call_test.py

@@ -0,0 +1,317 @@
+# 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.
+# ==============================================================================
+"""Function calls inside the while loop body."""
+
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def while_with_call_in_cond(n, fn):
+  i = 0
+  s = 0
+  while i < fn(n):
+    s = s * 10 + i
+    i += 1
+  return s
+
+
+def for_with_call_in_target(l, fn):
+  s = 0
+  for i in fn(l):
+    s = s * 10 + i
+  return s
+
+
+def while_with_local_call_in_cond(n):
+
+  def local_fn(x):
+    return x * 3
+
+  i = 0
+  s = 0
+  while i < local_fn(n):
+    s = s * 10 + i
+    i += 1
+  return s
+
+
+def for_with_local_call_in_target(l):
+
+  def local_fn(l):
+    return l * 1
+
+  s = 0
+  for i in local_fn(l):
+    s = s * 10 + i
+  return s
+
+
+def while_with_call(n, fn):
+  i = 0
+  s = 0
+  while i < n:
+    s = s * 10 + fn(i)
+    i += 1
+  return s
+
+
+def for_with_call(l, fn):
+  s = 0
+  for i in l:
+    s = s * 10 + fn(i)
+  return s
+
+
+def while_with_local_call(n):
+
+  def local_fn(x):
+    return x * 3
+
+  i = 0
+  s = 0
+  while i < n:
+    s = s * 10 + local_fn(i)
+    i += 1
+  return s
+
+
+def for_with_local_call(l):
+
+  def local_fn(x):
+    return x * 3
+
+  s = 0
+  for i in l:
+    s = s * 10 + local_fn(i)
+  return s
+
+
+def while_with_closure_call(n):
+  i = 0
+
+  def i_via_closure():
+    return i + 2
+
+  i = 0
+  s = 0
+  while i < n:
+    s = s * 10 + i_via_closure()
+    i += 1
+  return s
+
+
+def for_with_closure_call(l):
+  i = 0
+
+  def i_via_closure():
+    return i + 2
+
+  s = 0
+  for i in l:
+    s = s * 10 + i_via_closure()
+  # TODO(b/134822197): Remove i from return values.
+  return s, i
+
+
+def while_with_lambda_closure_call(n):
+  i = 0
+  s = 0
+  i_via_closure = lambda: i + 2
+  while i < n:
+    s = s * 10 + i_via_closure()
+    i += 1
+  return s
+
+
+def for_with_lambda_closure_call(l):
+  i = 0
+  s = 0
+  i_via_closure = lambda: i + 2
+  for i in l:
+    s = s * 10 + i_via_closure()
+  # TODO(b/134822197): Remove i from return values.
+  return s, i
+
+
+def while_with_method_closure_call(n):
+  i = 0
+
+  class Callable(object):
+
+    def __call__(self):
+      return i
+
+  i_via_closure = Callable()
+  i = 0
+  s = 0
+  while i < n:
+    s = s * 10 + i_via_closure()
+    i += 1
+  return s
+
+
+def for_with_method_closure_call(l):
+  i = 0
+
+  class Callable(object):
+
+    def __call__(self):
+      return i
+
+  i_via_closure = Callable()
+  i = 0
+  s = 0
+  for i in l:
+    s = s * 10 + i_via_closure()
+  # TODO(b/134822197): Remove i from return values.
+  return s, i
+
+
+def global_fn(x):
+  return x * 2
+
+
+class TestClass(object):
+
+  def method(self, x):
+    return x * 4
+
+
+def _int_tensor(x):
+  return tf.constant(x, dtype=tf.int32)
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+      (global_fn, lambda x: x * 1, TestClass().method, abs),
+  ))
+  def test_while_with_call_in_cond(self, n, type_, fn):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_call_in_cond, n, fn)
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+      (global_fn, lambda x: x * 1, TestClass().method, tf.abs),
+  ))
+  def test_for_with_call_in_target(self, l, type_, fn):
+    if fn is tf.abs and type_ is list:
+      self.skipTest('tf.abs([]) defaults to float32')
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_call_in_target, l, fn)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, _int_tensor),
+      (range, tf.range),
+  ))
+  def test_for_with_range_call_in_target(self, l, type_, fn):
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_call_in_target, l, fn)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+      (global_fn, lambda x: x * 1, TestClass().method, abs),
+  ))
+  def test_while_with_call(self, n, type_, fn):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_call, n, fn)
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+      (global_fn, lambda x: x * 1, TestClass().method, abs),
+  ))
+  def test_for_with_call(self, l, type_, fn):
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_call, l, fn)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_local_call(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_local_call, n)
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+  ))
+  def test_for_with_local_call(self, l, type_):
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_local_call, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_closure_call(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_closure_call, n)
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+  ))
+  def test_for_with_closure_call(self, l, type_):
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_closure_call, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_lambda_closure_call(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_lambda_closure_call, n)
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+  ))
+  def test_for_with_lambda_closure_call(self, l, type_):
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_lambda_closure_call, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_method_closure_call(self, n, type_):
+    self.skipTest('fix static analysis for nested classes')
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_method_closure_call, n)
+
+  @parameterized.parameters(*itertools.product(
+      ([], [1], [1, 2]),
+      (list, _int_tensor),
+  ))
+  def test_for_with_method_closure_call(self, l, type_):
+    self.skipTest('fix static analysis for nested classes')
+    l = type_(l)
+    self.assertFunctionMatchesEager(for_with_method_closure_call, l)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_with_variable_type_illegal_cases_test.py

@@ -0,0 +1,272 @@
+# 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.
+# ==============================================================================
+"""Loops with type changing variables."""
+
+import re
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def while_with_variable_py_type():
+  n = tf.constant(0, dtype=tf.int32)
+  c = True
+  while c:
+    c = tf.constant(True)
+  return n
+
+
+def while_with_variable_dtype():
+  n = tf.constant(0, dtype=tf.int32)
+  while tf.constant(True):
+    n = tf.constant(0, dtype=tf.float32)
+  return n
+
+
+def while_with_variable_dtype_and_early_stopping():
+  n = tf.constant(0, dtype=tf.int32)
+  while tf.constant(True):
+    n = tf.constant(0, dtype=tf.float32)
+    break
+  return n
+
+
+def for_with_variable_dtype(l):
+  n = tf.constant(0, dtype=tf.int32)
+  for _ in l:
+    n = tf.constant(0, dtype=tf.float32)
+  return n
+
+
+def for_with_variable_dtype_and_early_stopping(l):
+  n = tf.constant(0, dtype=tf.int32)
+  for _ in l:
+    n = tf.constant(0, dtype=tf.float32)
+    break
+  return n
+
+
+def while_with_variable_shape():
+  t = tf.constant([1])
+  while tf.constant(True):
+    t = tf.constant([1, 1])
+  return t
+
+
+def for_with_variable_shape(l):
+  t = tf.constant([1])
+  for _ in l:
+    t = tf.constant([1, 1])
+  return t
+
+
+def while_with_shape_erasure():
+  t = tf.constant([1])
+  while tf.constant(True):
+    t = tf.range(tf.random.uniform((), 2, 3, dtype=tf.int32))
+  return t
+
+
+def for_with_shape_erasure(l):
+  t = tf.constant([1])
+  for _ in l:
+    t = tf.range(tf.random.uniform((), 2, 3, dtype=tf.int32))
+  return t
+
+
+def while_with_shape_invariant_violation():
+  t = tf.constant([1])
+  while tf.constant(True):
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=((t, tf.TensorShape([1])),))
+    t = tf.range(tf.random.uniform((), 2, 3, dtype=tf.int32))
+  return t
+
+
+def for_with_shape_invariant_violation(l):
+  t = tf.constant([1])
+  for _ in l:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=((t, tf.TensorShape([1])),))
+    t = tf.range(tf.random.uniform((), 2, 3, dtype=tf.int32))
+  return t
+
+
+def while_with_variable_structure():
+  s = {'a': tf.constant(0)}
+  while tf.constant(True):
+    s = tf.constant(7.0)
+  return s
+
+
+def for_with_variable_structure(l):
+  s = [tf.constant(0)]
+  for _ in l:
+    s = s + [tf.constant(0)]
+  return s
+
+
+def _tf_range(l):
+  return tf.range(len(l))
+
+
+def _dataset(l):
+  return tf.data.Dataset.from_tensor_slices(l)
+
+
+def _dataset_iterator(l):
+  return iter(tf.data.Dataset.from_tensor_slices(l))
+
+
+def _distributed_dataset(l):
+  ds = tf.data.Dataset.from_tensor_slices([l] * 2)
+  return tf.distribute.MirroredStrategy().experimental_distribute_dataset(ds)
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  def test_while_with_variable_py_type(self):
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        re.compile(
+            r'.*condition of while loop started as non\-Tensor,'
+            r' then changed to Tensor.*', re.DOTALL)):
+      tf.function(while_with_variable_py_type)()
+
+  def test_while_with_variable_dtype(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        "'n' has dtype int32 before the loop, but dtype float32 after"):
+      tf.function(while_with_variable_dtype)()
+
+  def test_while_with_variable_dtype_and_early_stopping(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        "'n' has dtype int32 before the loop, but dtype float32 after"):
+      tf.function(while_with_variable_dtype_and_early_stopping)()
+
+  @parameterized.parameters(
+      (tf.constant,),
+      (_tf_range,),
+      (_dataset,),
+      (_dataset_iterator,),
+      (_distributed_dataset,),
+  )
+  def test_for_with_variable_dtype(self, type_):
+    l = type_([1, 2, 3])
+    with self.assertRaisesRegex(
+        TypeError,
+        "'n' has dtype int32 before the loop, but dtype float32 after"):
+      tf.function(for_with_variable_dtype)(l)
+
+  # Note: distributed datasets don't allow early stopping.
+  @parameterized.parameters(
+      (tf.constant,),
+      (_tf_range,),
+      (_dataset,),
+      (_dataset_iterator,),
+  )
+  def test_for_with_variable_dtype_and_early_stopping(self, type_):
+    l = type_([1, 2, 3])
+    with self.assertRaisesRegex(
+        TypeError,
+        "'n' has dtype int32 before the loop, but dtype float32 after"):
+      tf.function(for_with_variable_dtype_and_early_stopping)(l)
+
+  def test_while_with_variable_shape(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        r"'t' has shape \(1,\) before the loop, but shape \(2,\) after"):
+      tf.function(while_with_variable_shape)()
+
+  # Note: datasets do allow variable shape.
+  @parameterized.parameters(
+      (tf.constant,),
+      (_tf_range,),
+      (_dataset_iterator,),
+      (_distributed_dataset,),
+  )
+  def test_for_with_variable_shape(self, type_):
+    l = type_([1, 2, 3])
+    with self.assertRaisesRegex(
+        ValueError,
+        r"'t' has shape \(1,\) before the loop, but shape \(2,\) after"):
+      tf.function(for_with_variable_shape)(l)
+
+  def test_while_with_shape_erasure(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        r"'t' has shape \(1,\) before the loop, but shape \(None,\) after"):
+      tf.function(while_with_shape_erasure)()
+
+  # Note: datasets do allow variable shape.
+  @parameterized.parameters(
+      (tf.constant,),
+      (_tf_range,),
+      (_dataset_iterator,),
+      (_distributed_dataset,),
+  )
+  def test_for_with_shape_erasure(self, type_):
+    l = type_([1, 2, 3])
+    with self.assertRaisesRegex(
+        ValueError,
+        r"'t' has shape \(1,\) before the loop, but shape \(None,\) after"):
+      tf.function(for_with_shape_erasure)(l)
+
+  def test_while_with_shape_invariant_violation(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        r"'t' has shape \(None,\) after one iteration, which does not conform"):
+      tf.function(while_with_shape_invariant_violation)()
+
+  # Note: dataset loops ignore shape invariants.
+  @parameterized.parameters(
+      (tf.constant,),
+      (_tf_range,),
+      (_dataset_iterator,),
+      (_distributed_dataset,),
+  )
+  def test_for_with_shape_invariant_violation(self, type_):
+    l = type_([1, 2, 3])
+    with self.assertRaisesRegex(
+        ValueError,
+        r"'t' has shape \(None,\) after one iteration, which does not conform"):
+      tf.function(for_with_shape_invariant_violation)(l)
+
+  def test_while_with_variable_structure(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        "'s' does not have the same nested structure"):
+      tf.function(while_with_variable_structure)()
+
+  @parameterized.parameters(
+      (tf.constant,),
+      (_tf_range,),
+      (_dataset,),
+      (_dataset_iterator,),
+      (_distributed_dataset,),
+  )
+  def test_for_with_variable_structure(self, type_):
+    l = type_([1, 2, 3])
+    with self.assertRaisesRegex(
+        TypeError,
+        "'s' does not have the same nested structure"):
+      tf.function(for_with_variable_structure)(l)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/loop_with_variable_type_test.py

@@ -0,0 +1,337 @@
+# 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.
+# ==============================================================================
+"""Loops with type changing variables."""
+
+import collections
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def while_with_variable_shape_growing_vector(n):
+  v = tf.constant([0, 0])
+  i = 0
+  while i < n:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(v, tf.TensorShape([None]))])
+    v = tf.concat((v, [i]), 0)
+    i += 1
+  return v
+
+
+def for_with_variable_shape_growing_vector(l):
+  v = tf.constant([0, 0])
+  for i in l:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(v, tf.TensorShape([None]))])
+    v = tf.concat((v, [i]), 0)
+  return v
+
+
+def while_with_variable_shape_growing_matrix_rows(n):
+  m = tf.constant([[0]])
+  i = 0
+  while i < n:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(m, tf.TensorShape([None, 1]))])
+    m = tf.concat((m, [[i]]), 0)
+    i += 1
+  return m
+
+
+def for_with_variable_shape_growing_matrix_rows(l):
+  m = tf.constant([[0]])
+  for i in l:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(m, tf.TensorShape([None, 1]))])
+    m = tf.concat((m, [[i]]), 0)
+  return m
+
+
+def while_with_variable_shape_growing_matrix_cols(n):
+  m = tf.constant([[0, 0]])
+  i = 0
+  while i < n:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(m, tf.TensorShape([1, None]))])
+    m = tf.concat((m, [[i]]), 1)
+    i += 1
+  return m
+
+
+def for_with_variable_shape_growing_matrix_cols(l):
+  m = tf.constant([[0, 0]])
+  for i in l:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(m, tf.TensorShape([1, None]))])
+    m = tf.concat((m, [[i]]), 1)
+  return m
+
+
+def while_with_variable_shape_growing_matrix(n):
+  m = tf.constant([[0, 0], [0, 0]])
+  i = 0
+  while i < n:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(m, tf.TensorShape(None))])
+    m = tf.pad(m, [[1, 1], [1, 1]], constant_values=i)
+    i += 1
+  return m
+
+
+def for_with_variable_shape_growing_matrix(l):
+  m = tf.constant([[0, 0], [0, 0]])
+  for i in l:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(m, tf.TensorShape(None))])
+    m = tf.pad(m, [[1, 1], [1, 1]], constant_values=i)
+  return m
+
+
+def while_with_variable_shape_inside_if(n):
+  v = tf.constant([0, 0])
+  i = 0
+  if n > 1:
+    while i < n:
+      tf.autograph.experimental.set_loop_options(
+          shape_invariants=[(v, tf.TensorShape([None]))])
+      v = tf.concat((v, [i]), 0)
+      i += 1
+  else:
+    v = tf.constant([1, 2, 3])
+  return v
+
+
+def for_with_variable_shape_inside_if(n):
+  v = tf.constant([0, 0])
+  if n > 1:
+    for i in range(n):
+      tf.autograph.experimental.set_loop_options(
+          shape_invariants=[(v, tf.TensorShape([None]))])
+      v = tf.concat((v, [i]), 0)
+      i += 1
+  else:
+    v = tf.constant([1, 2, 3])
+  return v
+
+
+def for_with_nested_variable_shape_inside_if(n):
+  Test = collections.namedtuple('Test', ['var'])
+  t = Test(var=tf.constant([0]))
+  v = tf.constant([0, 0])
+  if n > 1:
+    for i in range(n):
+      tf.autograph.experimental.set_loop_options(
+          shape_invariants=[(v, tf.TensorShape([None]))])
+      v = tf.concat((v, [i]), 0)
+      t = Test(var=t.var + 1)
+      i += 1
+  else:
+    v = tf.constant([1, 2, 3])
+    t = Test(var=tf.constant([3]))
+  return v
+
+
+def while_with_variable_shape_and_break(n):
+  v = tf.constant([0, 0])
+  i = 0
+  if n > 1:
+    while i < n:
+      tf.autograph.experimental.set_loop_options(
+          shape_invariants=[(v, tf.TensorShape([None]))])
+      v = tf.concat((v, [i]), 0)
+      i += 1
+      if i > 3:
+        break
+  else:
+    v = tf.constant([1, 2, 3])
+  return v
+
+
+def for_with_variable_shape_and_break(n):
+  v = tf.constant([0, 0])
+  if n > 1:
+    for i in range(n):
+      tf.autograph.experimental.set_loop_options(
+          shape_invariants=[(v, tf.TensorShape([None]))])
+      v = tf.concat((v, [i]), 0)
+      i += 1
+      if i > 3:
+        break
+  else:
+    v = tf.constant([1, 2, 3])
+  return v
+
+
+def while_with_composite_tensor_shape_invariant(n):
+  v = tf.SparseTensor(
+      indices=[[0, 0], [1, 1]], values=[1, 2], dense_shape=[3, 3])
+  i = 0
+  while i < n:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(v, tf.TensorShape(None))])
+    v = tf.sparse.expand_dims(v)
+    i += 1
+  return v
+
+
+def for_with_composite_tensor_shape_invariant(l):
+  v = tf.SparseTensor(
+      indices=[[0, 0], [1, 1]], values=[1, 2], dense_shape=[3, 3])
+  for _ in l:
+    tf.autograph.experimental.set_loop_options(
+        shape_invariants=[(v, tf.TensorShape(None))])
+    v = tf.sparse.expand_dims(v)
+  return v
+
+
+def _int_dataset_range(n):
+  return tf.data.Dataset.range(n).map(lambda x: tf.cast(x, tf.int32))
+
+
+class ReferenceTest(reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_variable_shape_growing_vector(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_variable_shape_growing_vector, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (range, tf.range, tf.data.Dataset.range),
+  ))
+  def test_for_with_variable_shape_growing_vector(self, n, list_type):
+    l = list_type(n)
+    self.assertFunctionMatchesEager(for_with_variable_shape_growing_vector, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_variable_shape_growing_matrix_rows(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(
+        while_with_variable_shape_growing_matrix_rows, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (range, tf.range, _int_dataset_range),
+  ))
+  def test_for_with_variable_shape_growing_matrix_rows(self, l, type_):
+    l = type_(l)
+    self.assertFunctionMatchesEager(
+        for_with_variable_shape_growing_matrix_rows, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_variable_shape_growing_matrix_cols(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(
+        while_with_variable_shape_growing_matrix_cols, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (range, tf.range, tf.data.Dataset.range),
+  ))
+  def test_for_with_variable_shape_growing_matrix_cols(self, l, type_):
+    l = type_(l)
+    self.assertFunctionMatchesEager(
+        for_with_variable_shape_growing_matrix_cols, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_variable_shape_growing_matrix(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_variable_shape_growing_matrix, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (range, tf.range, _int_dataset_range),
+  ))
+  def test_for_with_variable_shape_growing_matrix(self, n, type_):
+    l = type_(n)
+    self.assertFunctionMatchesEager(for_with_variable_shape_growing_matrix, l)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_variable_shape_inside_if(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_variable_shape_inside_if, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_for_with_variable_shape_inside_if(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(for_with_variable_shape_inside_if, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_for_with_nested_variable_shape_inside_if(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(for_with_nested_variable_shape_inside_if, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (int, tf.constant),
+  ))
+  def test_while_with_variable_shape_and_break(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(while_with_variable_shape_and_break, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2, 5),
+      (int, tf.constant),
+  ))
+  def test_for_with_variable_shape_and_break(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(for_with_variable_shape_and_break, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2, 5),
+      (int, tf.constant),
+  ))
+  def test_while_with_composite_tensor_shape_invariant(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(
+        while_with_composite_tensor_shape_invariant, n)
+
+  @parameterized.parameters(*itertools.product(
+      (0, 1, 2),
+      (range, tf.range, _int_dataset_range),
+  ))
+  def test_for_with_composite_tensor_shape_invariant(self, n, type_):
+    l = type_(n)
+    self.assertFunctionMatchesEager(
+        for_with_composite_tensor_shape_invariant, l)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/nested_control_flow_test.py

@@ -0,0 +1,474 @@
+# 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.
+# ==============================================================================
+"""Nested loops and conditional statements (e.g. while, for, if).
+
+Meant to verify that arbitrarily nested statements are processed correctly.
+"""
+
+import itertools
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def independent_ifs(x, y):
+  z = 0
+  if x > 0:
+    if y > 0:
+      z = x + y
+  return z
+
+
+def dependent_inner_if(x):
+  y = 0
+  if x > 0:
+    y = -2 * x
+    if y > 0:
+      x = -3 * x
+  else:
+    y = 4 * x
+  return x, y
+
+
+def dependent_imbalanced_inner_if(x):
+  y = 0
+  if x > 0:
+    if x < 3:
+      y = -2 * x
+      x = -3 * x
+  return x, y
+
+
+def _hidden_raise():
+  raise ValueError('exception used for control flow')
+
+
+def if_with_local_modification_masked_by_exception(x):
+  y = 0
+  if x > 0:
+    try:
+      if x > 1:
+        _hidden_raise()
+      y = 1
+    except ValueError:
+      pass
+    if y == 0:
+      y = 2
+  return y
+
+
+_test_global = None
+
+
+def if_nested_with_modification_of_global(x):
+  y = 0
+  if x > 0:
+    if x > 0:
+      global _test_global
+      if _test_global is None:
+        _test_global = 1
+      else:
+        _test_global += 1
+      y += _test_global
+  return y
+
+
+def independent_inner_for(a, b):
+  p = 0
+  for _ in a:
+    tmp = b
+    for j in tmp:
+      p += j
+  return p
+
+
+def independent_inner_while(a, b):
+  p = 0
+  while a > 0:
+    tmp = b
+    while tmp > 0:
+      p += 1
+      tmp -= 1
+    a -= 1
+  return p
+
+
+def dependent_inner_for(a, b):
+  r = 1
+  s = 0
+  for _ in a:
+    r += s
+    tmp = b
+    for j in tmp:
+      s += j
+  return r
+
+
+def dependent_inner_while(a, b):
+  r = 1
+  while a > 0:
+    r += 1
+    tmp = b
+    while tmp > 0:
+      a -= 1
+      tmp -= 1
+    a -= 1
+  return r
+
+
+def if_in_for(a):
+  k = 0
+  for i in a:
+    if i % 2 > 0:
+      j = i // 2
+      k += j
+  return k
+
+
+def while_with_continue_in_context_manager(x):
+  z = 0
+  while x > 0:
+    with tf.name_scope(''):
+      x = x - 1
+      if x < 5:
+        continue
+      z = z + 1
+  return z
+
+
+def while_continue_in_try(x):
+  z = 0
+  while x > 0:
+    x = x - 1
+    try:
+      if x < 5:
+        continue
+      z = z + 1
+    finally:
+      z = z + 10
+  return z
+
+
+def while_break_in_context_manager(x):
+  z = 0
+  while x > 0:
+    with tf.name_scope(''):
+      x = x - 1
+      if x < 5:
+        break
+      z = z + 1
+  return z
+
+
+def while_break_in_try(x):
+  z = 0
+  while x > 0:
+    x = x - 1
+    try:
+      if x < 5:
+        break
+      z = z + 1
+    finally:
+      z = z + 10
+  return z
+
+
+def loop_initializing_invariant_variable(n):
+  for i in range(n):
+    if i == 0:
+      a = 1
+    else:
+      a = 2
+  return a
+
+
+def loop_initializing_variant_variable(n):
+  for i in range(n):
+    if i == 0:
+      a = 1
+    else:
+      a = a + 1
+  return a
+
+
+def _int_tensor(x):
+  return tf.constant(x, dtype=tf.int32)
+
+
+class NestedControlFlowTest(
+    reference_test_base.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(*itertools.product(
+      (
+          -1, 1,
+      ),
+      (
+          -1, 1,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_independent_ifs(self, x, y, type_x, type_y):
+    x = type_x(x)
+    y = type_x(y)
+    self.assertFunctionMatchesEager(independent_ifs, x, y)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          -1, 1,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_dependent_inner_if(self, x, type_):
+    x = type_(x)
+    self.assertFunctionMatchesEager(dependent_inner_if, x)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          -1, 1,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_dependent_imbalanced_inner_if(self, x, type_):
+    x = type_(x)
+    self.assertFunctionMatchesEager(dependent_imbalanced_inner_if, x)
+
+  @parameterized.parameters(
+      (-1,),
+      (0,),
+      (1,),
+      (2,),
+  )
+  def test_if_with_local_modification_masked_by_exception(self, x):
+    # Note: If the input is a Tensor, the behavior is undefined.
+    self.assertFunctionMatchesEager(
+        if_with_local_modification_masked_by_exception, x)
+
+  def test_if_nested_with_modification_of_global(self):
+    global _test_global
+    _test_global = None
+    self.assertEqual(tf.function(if_nested_with_modification_of_global)(1), 1)
+    self.assertEqual(_test_global, 1)
+
+  def test_if_nested_with_modification_of_global_not_executed(self):
+    global _test_global
+    _test_global = None
+    self.assertEqual(tf.function(if_nested_with_modification_of_global)(0), 0)
+    self.assertIsNone(_test_global)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 1, 2,
+      ),
+      (
+          0, 1, 2,
+      ),
+      (
+          range,
+          tf.range,
+      ),
+      (
+          range,
+          tf.range,
+      ),
+  ))
+  def test_independent_inner_for(self, a, b, type_a, type_b):
+    a = type_a(a)
+    b = type_b(b)
+    self.assertFunctionMatchesEager(independent_inner_for, a, b)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 1, 2,
+      ),
+      (
+          0, 1, 2,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_independent_inner_while(self, a, b, type_a, type_b):
+    a = type_a(a)
+    b = type_b(b)
+    self.assertFunctionMatchesEager(independent_inner_while, a, b)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 1, 2,
+      ),
+      (
+          0, 1, 2,
+      ),
+      (
+          range,
+          tf.range,
+      ),
+      (
+          range,
+          tf.range,
+      ),
+  ))
+  def test_dependent_inner_for(self, a, b, type_a, type_b):
+    a = type_a(a)
+    b = type_b(b)
+    self.assertFunctionMatchesEager(dependent_inner_for, a, b)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 1, 2, 3, 4,
+      ),
+      (
+          0, 1, 2, 3, 4,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_dependent_inner_while(self, a, b, type_a, type_b):
+    if (type_a is int) and (type_b is _int_tensor):
+      self.skipTest('b/124378596')
+    a = type_a(a)
+    b = type_b(b)
+    self.assertFunctionMatchesEager(dependent_inner_while, a, b)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 1, 2,
+      ),
+      (
+          range,
+          tf.range,
+      ),
+  ))
+  def test_if_in_for(self, a, type_):
+    a = type_(a)
+    self.assertFunctionMatchesEager(if_in_for, a)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 4, 10,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_while_continue_in_context_manager(self, x, type_):
+    x = type_(x)
+    self.assertFunctionMatchesEager(while_with_continue_in_context_manager, x)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 4, 10,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_while_continue_in_try(self, x, type_):
+    x = type_(x)
+    self.assertFunctionMatchesEager(while_continue_in_try, x)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 4, 10,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_while_break_in_context_manager(self, x, type_):
+    x = type_(x)
+    self.assertFunctionMatchesEager(while_break_in_context_manager, x)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          0, 4, 10,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_while_break_in_try(self, x, type_):
+    x = type_(x)
+    self.assertFunctionMatchesEager(while_break_in_try, x)
+
+  @parameterized.parameters(*itertools.product(
+      (
+          1, 2,
+      ),
+      (
+          int,
+          _int_tensor,
+      ),
+  ))
+  def test_loop_initializing_invariant_variable_legal(self, n, type_):
+    n = type_(n)
+    self.assertFunctionMatchesEager(loop_initializing_invariant_variable, n)
+
+  def test_loop_initializing_invariant_variable_illegal(self):
+    with self.assertRaises(UnboundLocalError):
+      tf.function(loop_initializing_invariant_variable)(0)
+    with self.assertRaisesRegex(
+        tf.errors.InvalidArgumentError, 'loop must iterate at least once'):
+      tf.function(loop_initializing_invariant_variable)(tf.constant(0))
+
+  @parameterized.parameters(
+      (1,),
+      (2,),
+  )
+  def test_loop_initializing_variant_variable_legal(self, n):
+    tf.function(loop_initializing_variant_variable)(n)
+
+  @parameterized.parameters(
+      (0,),
+      (1,),
+      (2,),
+  )
+  def test_loop_initializing_variant_variable_illegal(self, n):
+    with self.assertRaisesRegex(ValueError, 'must be defined before the loop'):
+      tf.function(loop_initializing_variant_variable)(tf.constant(n))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/reference_test.bzl

@@ -0,0 +1,30 @@
+# 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.
+# ==============================================================================
+"""BUILD target helpers."""
+
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+def reference_test(name, additional_deps = [], tags = [], shard_count = 1):
+    py_test(
+        name = name,
+        srcs = [name + ".py"],
+        deps = [
+            ":reference_tests",
+            "//tensorflow:tensorflow_py_no_contrib",
+        ] + additional_deps,
+        python_version = "PY3",
+        shard_count = shard_count,
+        tags = tags + ["no_windows", "no_pip"],
+    )

tensorflow/python/autograph/tests/reference_test_base.py

@@ -0,0 +1,178 @@
+# 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.
+# ==============================================================================
+# 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.
+# ==============================================================================
+"""Reference tests check that a function is compiled correctly."""
+
+import io
+import numbers
+import os
+import sys
+import traceback
+
+import numpy as np
+import tensorflow as tf
+
+
+class TestCase(tf.test.TestCase):
+  """Base class for the reference tests."""
+
+  def setUp(self):
+    super(TestCase, self).setUp()
+    os.environ['AUTOGRAPH_STRICT_CONVERSION'] = '1'
+    self.autograph_opts = None
+    self.all_inputs_tensors = False
+    self.allow_exceptions = False
+
+  # TODO(mdan): Consider rewriting as a context manager.
+  def _run_with_output_capture(self, func):
+    """Executes `func`, capturing stdout."""
+    out_capturer = io.StringIO()
+    results = None
+    captured_out = None
+    captured_err = None
+    try:
+      sys.stdout = out_capturer
+      results = func()
+      captured_out = out_capturer.getvalue()
+    except Exception as e:  # pylint:disable=broad-except
+      sys.stdout = sys.__stdout__
+      captured_err = e
+      print('*** Capturing exception:\n{}\n'.format(traceback.format_exc()))
+    finally:
+      sys.stdout = sys.__stdout__
+      out_capturer.close()
+    return results, captured_out, captured_err
+
+  def _as_tensors(self, args):
+    """Converts args to tensors."""
+    tensor_args = []
+    for a in args:
+      if isinstance(a, (numbers.Number, list, np.ndarray)):
+        tensor_arg = tf.constant(a)
+      elif isinstance(a, dict):
+        keys = tuple(a.keys())
+        tensor_arg = dict(zip(keys, self._as_tensors([a[k] for k in keys])))
+      else:
+        tensor_arg = a
+      tensor_args.append(tensor_arg)
+    return tensor_args
+
+  def run_native(self, f, *args):
+    return self._run_with_output_capture(lambda: f(*args))
+
+  def _deep_equal(self, left, right):
+    """Compares two possibly-nested structures."""
+    if isinstance(left, tf.Tensor):
+      return self._deep_equal(left.numpy(), right)
+    if isinstance(right, tf.Tensor):
+      return self._deep_equal(left, right.numpy())
+    if isinstance(left, tf.SparseTensor) and isinstance(right, tf.SparseTensor):
+      return (self._deep_equal(left.indices, right.indices)
+              and self._deep_equal(left.values, right.values)
+              and self._deep_equal(left.shape, right.shape))
+    if isinstance(left, np.ndarray) or isinstance(right, np.ndarray):
+      return np.array_equal(left, right)
+    if isinstance(left, (list, tuple)) and isinstance(right, (list, tuple)):
+      return all(self._deep_equal(l, r) for l, r in zip(left, right))
+    return left == right
+
+  def assertResultsMatch(self,
+                         f,
+                         args,
+                         native_data,
+                         compiled_data):
+    """Asserts that native_data matches compiled_data."""
+    native_results, native_out, native_err = native_data
+    compiled_results, compiled_out, compiled_err = compiled_data
+    str_args = '(%s)' % ', '.join(str(a) for a in args)
+    # Using a manual verification to avoid a second compilation on success.
+    # For exceptions, we don't enforce that they are the same, only that
+    # both paths raised.
+    # TODO(mdan): Add an API that returns both object and source code instead.
+    outputs_equal = (
+        self._deep_equal(native_results, compiled_results) and
+        native_out == compiled_out)
+    errors_equivalent = type(native_err) == type(compiled_err)  # pylint:disable=unidiomatic-typecheck
+    if (not outputs_equal or not errors_equivalent):
+      self.fail('Native and compiled functions are not equivalent.\n\n'
+                'Native results: %s\n'
+                'Compiled results: %s\n'
+                'Native out: %s\n'
+                'Compiled out: %s\n'
+                'Native error: %s: %s\n'
+                'Compiled error: %s: %s\n'
+                'Native call: %s%s\n'
+                'Check the logs for the generated code.'
+                '' % (
+                    native_results,
+                    compiled_results,
+                    native_out,
+                    compiled_out,
+                    type(native_err).__name__,
+                    native_err,
+                    type(compiled_err).__name__,
+                    compiled_err,
+                    f.__name__,
+                    str_args,
+                ))
+
+  def function(self, f, xla=False):
+    return tf.function(
+        f,
+        experimental_autograph_options=self.autograph_opts,
+        experimental_compile=xla)
+
+  def convert(self, f):
+    return tf.autograph.to_graph(
+        f, experimental_optional_features=self.autograph_opts)
+
+  def assertFunctionMatchesEagerStatefulInput(self, f, args):
+    """Like assertFunctionMatchesEager but creates new inputs each time."""
+    compiled_data = self.run_native(self.function(f), *args())
+    native_data = self.run_native(f, *args())
+    self.assertResultsMatch(f, args(), native_data, compiled_data)
+
+  def assertFunctionMatchesEager(self, f, *args, xla=False):
+    if self.all_inputs_tensors:
+      args = self._as_tensors(args)
+    compiled_data = self.run_native(self.function(f, xla=xla), *args)
+    if not self.allow_exceptions:
+      _, _, compiled_err = compiled_data
+      if compiled_err is not None:
+        self.fail(str(compiled_err))
+    native_data = self.run_native(f, *args)
+    self.assertResultsMatch(f, args, native_data, compiled_data)
+
+  def assertConvertedMatchesNative(self, f, *args):
+    compiled_data = self.run_native(self.convert(f), *args)
+    native_data = self.run_native(f, *args)
+    self.assertResultsMatch(f, args, native_data, compiled_data)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/autograph/tests/type_annotations_test.py

@@ -0,0 +1,34 @@
+# 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.
+# ==============================================================================
+"""Type annotations."""
+
+import tensorflow as tf
+
+from tensorflow.python.autograph.tests import reference_test_base
+
+
+def pure_declaration():
+  n: int  # pylint:disable=unused-variable
+  return 1
+
+
+class ReferenceTest(reference_test_base.TestCase):
+
+  def test_pure_declaration(self):
+    self.assertFunctionMatchesEager(pure_declaration)
+
+
+if __name__ == '__main__':
+  tf.test.main()