Method while_loop

tests/python/relax/test_transform_lambda_lift.py:222–235

(i: R.Tensor((), "int32"), s: R.Tensor((2, 3), "float32"))

        def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor:
            @R.function
            def while_loop(i: R.Tensor((), "int32"), s: R.Tensor((2, 3), "float32")) -> R.Tensor(
                (2, 3), "float32"
            ):
                cond: R.Tensor((), "bool") = R.call_pure_packed(
                    "test.vm.less", i, R.const(10), sinfo_args=(R.Tensor((), dtype="bool"))
                )
                c: R.Tensor((), "int32") = R.const(1, dtype="int32")
                if cond:
                    new_i: R.Tensor((), "int32") = R.add(i, c)
                    new_s: R.Tensor((2, 3), "float32") = R.add(s, x)
                    r: R.Tensor((2, 3), "float32") = while_loop(new_i, new_s)
                else:
                    r: R.Tensor((2, 3), "float32") = s
                return r

            gv: R.Tensor((2, 3), "float32") = while_loop(R.const(0), x)
            return gv