[Dy2st]Fix while to static bug (#70287)

zhangbo9674 · web-flow · commit 46195827a1ac · 2024-12-19T13:19:27.000+08:00
* fix

* add ut
diff --git a/python/paddle/static/nn/control_flow.py b/python/paddle/static/nn/control_flow.py
@@ -911,6 +911,23 @@ def infer_loop_var_type_with_next_var(loop_var, next_var):
                 loop_vars, next_vars
             )
 
+            from paddle.jit.dy2static.convert_operators import (
+                to_static_variable,
+            )
+
+            def check_next_var(loop_var):
+                if not loop_var.is_variable_curr_var:
+                    return
+                if not isinstance(
+                    loop_var.next_var, paddle.pir.Value
+                ) and not isinstance(loop_var.next_var, (bool, float, int)):
+                    raise ValueError(
+                        "The loop var in the while op is variable, but the corresponding yielded var is not variable, and it is not a constant of type bool, int, or float."
+                    )
+                loop_var.next_var = to_static_variable(loop_var.next_var)
+
+            paddle.utils.map_structure(check_next_var, loop_vars)
+
             next_cond = cond(
                 *map_structure(lambda var: var.next_var, loop_vars)
             )
diff --git a/test/dygraph_to_static/test_loop.py b/test/dygraph_to_static/test_loop.py
@@ -485,5 +485,23 @@ def test_loop_with_inner_mutate_list(self):
         np.testing.assert_allclose(dygraph_res.numpy(), static_res.numpy())
 
 
+def loop_change_value_to_int():
+    x = paddle.to_tensor(1, dtype='float32')
+    y = paddle.to_tensor(False, dtype='bool')
+    while y:
+        x = 2
+    return x
+
+
+class TestLoopChangeValueToInt(Dy2StTestBase):
+    def test_loop_change_value_to_int(self):
+        static_fn = paddle.jit.to_static(
+            loop_change_value_to_int, full_graph=True
+        )
+        static_res = static_fn()
+        dygraph_res = loop_change_value_to_int()
+        np.testing.assert_allclose(dygraph_res.numpy(), static_res.numpy())
+
+
 if __name__ == '__main__':
     unittest.main()
