Method while_loop

tests/python/relax/test_transform_dead_code_elimination.py:594–607

(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.call_pure_packed(
                    "test.vm.less", i, R.const(10), sinfo_args=R.Tensor((), dtype="bool")
                )
                c = R.const(1, dtype="int32")
                if cond:
                    new_i = R.add(i, c)
                    new_s = R.add(s, x)
                    r = while_loop(new_i, new_s)
                else:
                    r = s
                return r

            gv = while_loop(R.const(0), x)
            return gv