Add while loop test

Author: Ch0ronomato

tests/link/pytorch/test_basic.py

@@ -318,3 +318,16 @@ def test_ScalarLoop():
     np.testing.assert_allclose(fn(5, 0, 1), 5)
     np.testing.assert_allclose(fn(5, 0, 2), 10)
     np.testing.assert_allclose(fn(4, 3, -1), -1)
+
+
+def test_ScalarLoop_while():
+    n_steps = int64("n_steps")
+    x0 = float64("x0")
+    x = x0 + 1
+    until = x >= 10
+
+    op = ScalarLoop(init=[x0], update=[x], until=until)
+    fn = function([n_steps, x0], op(n_steps, x0), mode=pytorch_mode)
+    np.testing.assert_allclose(fn(n_steps=20, x0=0), [10, True])
+    np.testing.assert_allclose(fn(n_steps=20, x0=1), [10, True])
+    np.testing.assert_allclose(fn(n_steps=5, x0=1), [6, False])