Function _check_base_py_module_inheritance

python/tvm/script/parser/ir/parser.py:180–212

Check if a class inherits from BasePyModule. Parameters ---------- node : doc.ClassDef The class definition node to check. Returns ------- bool True if the class inherits from BasePyModule, False otherwise.

(node: doc.ClassDef)

def _check_base_py_module_inheritance(node: doc.ClassDef) -> bool:
    """Check if a class inherits from BasePyModule.

    Parameters
    ----------
    node : doc.ClassDef
        The class definition node to check.

    Returns
    -------
    bool
        True if the class inherits from BasePyModule, False otherwise.
    """
    if not node.bases:
        return False

    # Check each base class
    for base in node.bases:
        if hasattr(base, "id"):
            if base.id == "BasePyModule":
                return True
        elif hasattr(base, "attr"):
            if base.attr == "BasePyModule":
                return True
        elif hasattr(base, "value") and hasattr(base.value, "id"):
            if (
                base.value.id in ["BasePyModule", "tvm", "relax"]
                and hasattr(base, "attr")
                and base.attr == "BasePyModule"
            ):
                return True

    return False