Function forward

tests/python/relax/test_frontend_nn_modules.py:35–44

(
        x: R.Tensor((3, 3), dtype="float32"),
        _io: R.Object,
    )

def test_relu():
    @R.function
    def forward(
        x: R.Tensor((3, 3), dtype="float32"),
        _io: R.Object,
    ) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Object)):
        R.func_attr({"num_input": 2})
        with R.dataflow():
            relu: R.Tensor((3, 3), dtype="float32") = R.nn.relu(x)
            gv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Object)) = relu, (_io,)
            R.output(gv1)
        return gv1

    mod = modules.ReLU()
    tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((3, 3), "float32")}}, debug=True)