Method main

tests/python/relax/test_frontend_from_exported_program.py:1361–1380

(input: R.Tensor((10, 10), dtype="float32"))

    class expected_triu:
        @R.function
        def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
            R.Tensor((10, 10), dtype="float32")
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((10,), dtype="int64") = R.arange(
                    R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
                )
                lv1: R.Tensor((1, 10), dtype="int64") = R.expand_dims(lv, axis=[-2])
                lv2: R.Tensor((10,), dtype="int64") = R.arange(
                    R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
                )
                lv3: R.Tensor((10, 1), dtype="int64") = R.expand_dims(lv2, axis=[-1])
                lv4: R.Tensor((10, 10), dtype="int64") = R.subtract(lv1, lv3)
                lv5: R.Tensor((10, 10), dtype="bool") = R.greater_equal(lv4, R.const(1, "int64"))
                lv6: R.Tensor((), dtype="float32") = R.const(0.0, "float32")
                lv7: R.Tensor((10, 10), dtype="float32") = R.where(lv5, input, lv6)
                gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv7,)
                R.output(gv)
            return gv

    verify_model(Triu(), example_args, {}, expected_triu)