Function collect_symbolic_var_from_params

python/tvm/relax/script/parser/parser.py:179–196

(self: Parser, node: doc.FunctionDef)

def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> None:
    # Collect symbolic vars from parameters
    symbolic_vars = {}
    for arg in node.args.args:
        if arg.annotation is None:
            self.report_error(arg, "Type annotation is required for function parameters.")
        param_sinfo_proxy = eval_struct_info_proxy(self, arg.annotation)

        for var_name in param_sinfo_proxy.get_symbolic_vars():
            if var_name not in symbolic_vars:
                symbolic_vars[var_name] = tirx.Var(var_name, "int64")

    # Update symbolic vars based on
    symbolic_vars = collect_symbolic_var_from_prelude(self, node, symbolic_vars)

    # Define symbolic vars to the current var_table frame
    for var_name, var in symbolic_vars.items():
        self.var_table.add(var_name, var, allow_shadowing=False)


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