Function get_op

python/tvm/script/parser/core/dispatch.py:127–157

Register a method for a operand type, AST operator node and operand index. Parameters ---------- operand_type : Type The type of operands, e.g., tirx.PrimExpr, tirx.IterVar. op_node_type : AST The doc AST operator node type, e.g., doc.Add, doc.Eq. operand_index

(
    operand_type: type,
    op_node_type: type,
    operand_index: int,
    default: OpMethod | None = None,
)

def get_op(
    operand_type: type,
    op_node_type: type,
    operand_index: int,
    default: OpMethod | None = None,
) -> OpMethod | None:
    """Register a method for a operand type, AST operator node and operand index.

    Parameters
    ----------
    operand_type : Type
        The type of operands, e.g., tirx.PrimExpr, tirx.IterVar.

    op_node_type : AST
        The doc AST operator node type, e.g., doc.Add, doc.Eq.

    operand_index : int
        The operand index, i.e., 0 for left operand and 1 for right operand.


    default : Optional[OpMethod]
        The default method when no registered methods with this operand type,
        AST operator node and operand index.

    Returns
    -------
    func : Optional[OpMethod]
        The function to register dispatched method of parsing
        corresponding a operand type, AST operator node and operand index.
    """
    return OpVTable.get((operand_type, op_node_type, operand_index), default)