Function visit_tvm_declare_function

python/tvm/relax/script/parser/parser.py:268–288

(self: Parser, node: doc.FunctionDef)

@dispatch.register(token="relax", type_name="tvm_declare_function")
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
    with self.var_table.with_frame():
        collect_symbolic_var_from_params(self, node)

        if node.returns is None:
            # Use ObjectStructInfo as unknown return type
            # NOTE: Cannot use VoidStructInfo here because the return type can be refined later.
            ret_sinfo = relax.ObjectStructInfo()
        else:
            ret_sinfo = eval_struct_info(self, node.returns, eval_str=True)
        params = []
        for arg in node.args.args:
            if arg.annotation is None:
                self.report_error(arg, "Type annotation is required for function parameters.")
            param_sinfo = eval_struct_info(self, arg.annotation, eval_str=True)
            params.append(relax.Var(arg.arg, param_sinfo))

    is_pure = find_decorator_annotation(node, "pure")

    func_signature = relax.Function.create_empty(params, ret_sinfo, is_pure=is_pure)
    return I.decl_function(node.name, func_signature)


@dispatch.register(token="relax", type_name="pre_visit_local_function")