Function find_decorator_annotation

python/tvm/tirx/script/parser/parser.py:240–252

Check the value of given annotation (argument name) in the prim_func decorator. Returns the value of the annotation if present, otherwise giving the default value.

(node: doc.FunctionDef, annotation: str, default: bool = True)

def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
    """
    Check the value of given annotation (argument name) in the prim_func decorator.
    Returns the value of the annotation if present, otherwise giving the default value.
    """
    # look for the named argument in the prim_func / jit decorator
    for dec in node.decorator_list:
        if not isinstance(dec, doc.Call) or dec.func.attr not in ("prim_func", "jit"):
            continue
        for keyword in dec.keywords:
            if keyword.arg == annotation:
                return keyword.value.value
    return default


@dispatch.register(token="tirx", type_name="For")