Merge pull request #37813 from zhuzilin:autograph-loop-else-dev

PiperOrigin-RevId: 304383345
Change-Id: I9f5574e8c07edda607b8b43c662f32e045f0fba9

tensorflow/python/autograph/converters/break_statements.py

@@ -80,28 +80,40 @@ class BreakTransformer(converter.Base):
     # A break in the else clause applies to the containing scope.
     node.orelse = self.visit_block(node.orelse)
 
-    if break_used:
-      # Python's else clause only triggers if the loop exited cleanly (e.g.
-      # break did not trigger).
-      guarded_orelse = self._guard_if_present(node.orelse, break_var)
-
+    if not break_used:
       template = """
-        var_name = False
-        while ag__.and_(lambda: test, lambda: ag__.not_(var_name)):
+        while test:
           body
-        else:
-          orelse
+        orelse
       """
       node = templates.replace(
-          template,
-          var_name=break_var,
-          test=node.test,
-          body=node.body,
-          orelse=guarded_orelse)
+          template, test=node.test, body=node.body, orelse=node.orelse)
 
-      new_while_node = node[1]
+      new_while_node = node[0]
       anno.copyanno(original_node, new_while_node, anno.Basic.DIRECTIVES)
 
+      return node
+
+    # Python's else clause only triggers if the loop exited cleanly (e.g.
+    # break did not trigger).
+    guarded_orelse = self._guard_if_present(node.orelse, break_var)
+
+    template = """
+      var_name = False
+      while ag__.and_(lambda: test, lambda: ag__.not_(var_name)):
+        body
+      orelse
+    """
+    node = templates.replace(
+        template,
+        var_name=break_var,
+        test=node.test,
+        body=node.body,
+        orelse=guarded_orelse)
+
+    new_while_node = node[1]
+    anno.copyanno(original_node, new_while_node, anno.Basic.DIRECTIVES)
+
     return node
 
   def visit_For(self, node):
@@ -115,37 +127,54 @@ class BreakTransformer(converter.Base):
     # A break in the else clause applies to the containing scope.
     node.orelse = self.visit_block(node.orelse)
 
-    if break_used:
-      # Python's else clause only triggers if the loop exited cleanly (e.g.
-      # break did not trigger).
-      guarded_orelse = self._guard_if_present(node.orelse, break_var)
-      extra_test = templates.replace_as_expression(
-          'ag__.not_(var_name)', var_name=break_var)
-
-      # The extra test is hidden in the AST, which will confuse the static
-      # analysis. To mitigate that, we insert a no-op statement that ensures
-      # the control variable is marked as used.
-      # TODO(mdan): Use a marker instead, e.g. ag__.condition_loop_on(var_name)
+    if not break_used:
       template = """
-        var_name = False
         for target in iter_:
-          (var_name,)
           body
-        else:
-          orelse
+        orelse
       """
       node = templates.replace(
           template,
-          var_name=break_var,
           iter_=node.iter,
           target=node.target,
           body=node.body,
-          orelse=guarded_orelse)
+          orelse=node.orelse)
 
-      new_for_node = node[1]
-      anno.setanno(new_for_node, anno.Basic.EXTRA_LOOP_TEST, extra_test)
+      new_for_node = node[0]
+      anno.copyanno(original_node, new_for_node, anno.Basic.EXTRA_LOOP_TEST)
       anno.copyanno(original_node, new_for_node, anno.Basic.DIRECTIVES)
 
+      return node
+
+    # Python's else clause only triggers if the loop exited cleanly (e.g.
+    # break did not trigger).
+    guarded_orelse = self._guard_if_present(node.orelse, break_var)
+    extra_test = templates.replace_as_expression(
+        'ag__.not_(var_name)', var_name=break_var)
+
+    # The extra test is hidden in the AST, which will confuse the static
+    # analysis. To mitigate that, we insert a no-op statement that ensures
+    # the control variable is marked as used.
+    # TODO(mdan): Use a marker instead, e.g. ag__.condition_loop_on(var_name)
+    template = """
+      var_name = False
+      for target in iter_:
+        (var_name,)
+        body
+      orelse
+    """
+    node = templates.replace(
+        template,
+        var_name=break_var,
+        iter_=node.iter,
+        target=node.target,
+        body=node.body,
+        orelse=guarded_orelse)
+
+    new_for_node = node[1]
+    anno.setanno(new_for_node, anno.Basic.EXTRA_LOOP_TEST, extra_test)
+    anno.copyanno(original_node, new_for_node, anno.Basic.DIRECTIVES)
+
     return node
 
 

tensorflow/python/autograph/converters/loop_integration_test.py

@@ -0,0 +1,95 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Integration Tests for loop."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.autograph.converters import break_statements
+from tensorflow.python.autograph.converters import continue_statements
+from tensorflow.python.autograph.converters import control_flow
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.framework import constant_op
+from tensorflow.python.platform import test
+
+
+class LoopIntegrationTest(converter_testing.TestCase):
+
+  def assertTransformedEquivalent(self, test_fn, *inputs):
+    with self.converted(test_fn,
+                        [break_statements, continue_statements, control_flow],
+                        {}, (constant_op.constant,)) as result:
+      self.assertEqual(test_fn(*inputs), result.test_fn(*inputs))
+
+  def test_while_loop_with_else(self):
+
+    def test_fn(x):
+      while x > 2:
+        x /= 2
+      else:
+        x += 1
+      return x
+
+    self.assertTransformedEquivalent(test_fn, 4)
+    self.assertTransformedEquivalent(test_fn, 2)
+
+  def test_while_loop_with_else_and_break(self):
+
+    def test_fn(cond1):
+      x = 8
+      while x > 2:
+        x /= 2
+        if cond1:
+          break
+      else:
+        x += 1
+      return x
+
+    self.assertTransformedEquivalent(test_fn, True)
+    self.assertTransformedEquivalent(test_fn, False)
+
+  def test_for_loop_with_else(self):
+
+    def test_fn(l):
+      res = 0
+      for x in l:
+        res += x
+      else:
+        res += 1
+      return res
+
+    self.assertTransformedEquivalent(test_fn, [])
+    self.assertTransformedEquivalent(test_fn, [1, 2])
+
+  def test_for_loop_with_else_and_break(self):
+
+    def test_fn(flag):
+      l = [1, 2, 3]
+      res = 0
+      for x in l:
+        res += x
+        if flag:
+          break
+      else:
+        res += 1
+      return res
+
+    self.assertTransformedEquivalent(test_fn, True)
+    self.assertTransformedEquivalent(test_fn, False)
+
+
+if __name__ == '__main__':
+  test.main()