def _fuse_bmm_mul_or_div_alpha(sorted_graph: List[Tensor]) -> List[Tensor]:
    """This pass fuses bmm and mul (or div) if mul's other operand is a
       constant scalar tensor (i.e. which has a valid "value" attribute.
       In such a case, we turn this constant value into bmm's alpha.
       Note that for div cases, we assign 1/const_val to alpha.

    Parameters
    ----------
    sorted_graph : List[Tensor]
        input sorted graph

    Return
    ----------
    List[Tensor]
        modified sorted graph upon success. Otherwise, the original sorted
        graph will be returned.
    """
    for tensor in sorted_graph:
        src_ops = tensor._attrs["src_ops"]
        if src_ops is None or len(src_ops) != 1:
            continue
        src_op = list(src_ops)[0]
        if src_op is None:
            continue
        if not src_op._attrs["op"].startswith("bmm"):
            continue
        bmm_op = src_op

        dst_ops = list(tensor._attrs["dst_ops"])
        if not dst_ops or len(dst_ops) != 1:
            continue

        next_op = dst_ops[0]
        if next_op._attrs["op"] != "elementwise":
            continue
        if next_op._attrs["func"] == FuncEnum.MUL:
            is_div = False
        elif next_op._attrs["func"] == FuncEnum.DIV:
            is_div = True
        else:
            continue

        elem_op = next_op
        elem_inputs = elem_op._attrs["inputs"]
        if len(elem_inputs) != 1:
            continue
        elem_args = elem_op._attrs["args"]
        if len(elem_args) != 2:
            continue
        # make sure cst_tensor is the divisor of the DIV op
        if is_div and tensor == elem_args[1]:
            continue
        cst_tensor = elem_args[1] if tensor == elem_args[0] else elem_args[1]
        # skip non-constant scalar tensor
        if not cst_tensor.is_a_const_num():
            continue
        cst_val = cst_tensor._attrs["value"]
        # let's only consider int and float builtin types. Seems that it doesn't
        # make any sense to take other scalar types like str and convert it
        # to a float.
        if not isinstance(cst_val, (float, int)):
            continue
        # OK, we are good so let's add cst_val to bmm's alpha attribute
        bmm_op._attrs["alpha"] = 1.0 / float(cst_val) if is_div else float(cst_val)
        # remove this MUL/DIV
        remove_single_tensor_op_from_sorted_graph(elem_op)

    return sanitize_sorted_graph(sorted_graph)
This pass fuses bmm and mul (or div) if mul's other operand is a
constant scalar tensor (i.e. which has a valid "value" attribute.
In such a case, we turn this constant value into bmm's alpha.
Note that for div cases, we assign 1/const_val to alpha.

Parameters
----------
sorted_graph : List[Tensor]
input sorted graph

Return
----------
List[Tensor]
modified sorted graph upon success. Otherwise, the original sorted
graph will be returned.
python/aitemplate/compiler/transform/fuse_mm_elementwise.py