Function parse

python/tvm/script/parser/core/entry.py:76–148

Register a method for a operand type, AST operator node and operand index. Parameters ---------- program : Union[doc.AST, Any, str] The TVMScript code to parse. extra_vars : Dict[str, Any] The extra variable table for parsing. check_well_formed : bool W

(
    program: doc.AST | Any | str,
    extra_vars: dict[str, Any] | None = None,
    check_well_formed: bool = True,
    s_tir: bool = False,
)

def parse(
    program: doc.AST | Any | str,
    extra_vars: dict[str, Any] | None = None,
    check_well_formed: bool = True,
    s_tir: bool = False,
) -> Any:
    """Register a method for a operand type, AST operator node and operand index.

    Parameters
    ----------
    program : Union[doc.AST, Any, str]
        The TVMScript code to parse.

    extra_vars : Dict[str, Any]
        The extra variable table for parsing.

    check_well_formed : bool
        Whether to check well-formedness after parsing.

    Returns
    -------
    func : Any
        The parsed TVMScript program.
    """
    if extra_vars is None:
        extra_vars = _default_globals()

    ann = {}
    all_pyfuncs = {}
    if inspect.isfunction(program):
        ann = {program.__name__: program.__annotations__}
    elif inspect.isclass(program):
        for name, func in program.__dict__.items():
            if inspect.isfunction(func):
                ann[name] = func.__annotations__
                all_pyfuncs[name] = func

    source = Source(program)
    parser = Parser(source, ann)
    with IRBuilder() as builder:
        try:
            parser.parse(extra_vars=extra_vars)
        except ParserError as err:
            parser.report_error(err.node, err.args[0])
    ret = builder.get()
    # Attach pyfuncs to the IRModule
    if inspect.isclass(program) and isinstance(ret, IRModule):
        _attach_pyfuncs_to_irmodule(ret, all_pyfuncs)

    # check well-formedness in both Relax and TIR
    if check_well_formed:
        check_ret = ret
        if not isinstance(check_ret, IRModule):
            check_ret = IRModule.from_expr(ret)

        source_ast = source.as_ast()

        if isinstance(