def fuse_expand_bmm(sorted_graph: List[Tensor], workdir: str = None) -> List[Tensor]:
    """
    Transform expand + bmm into a single bmm op.

    Parameters
    ----------
    sorted_graph : List[Tensor]
        Input graph
    workdir : str, optional
        workdir, by default None

    Returns
    -------
    List[Tensor]
        Optimized graph
    """
    for tensor in sorted_graph:
        src_ops = tensor._attrs["src_ops"]
        if len(src_ops) != 1:
            continue
        op = list(src_ops)[0]
        if op._attrs["op"] != "expand":
            continue
        expand_op = op
        expand_output = expand_op._attrs["outputs"][0]
        dst_ops = expand_output._attrs["dst_ops"]
        if len(dst_ops) != 1:
            continue
        next_op = list(dst_ops)[0]
        if not next_op._attrs["op"].startswith("bmm_"):
            continue
        if not _can_fuse(expand_op, next_op):
            continue

        for int_var_tensor in expand_op._attrs["inputs"][1:]:
            int_var_tensor._attrs["dst_ops"].discard(expand_op)
        expand_op._attrs["inputs"] = [expand_op._attrs["inputs"][0]]
        remove_single_tensor_op_from_sorted_graph(expand_op)

        old_tensor_accessors = next_op._attrs["input_accessors"]
        assert (
            old_tensor_accessors[0].stride_dim is None
            and old_tensor_accessors[1].stride_dim is None
        ), f"next_op {next_op._attrs['name']} tensor accessors are expected to be None"
        bmm_inputs = next_op._attrs["inputs"]
        # refresh tensor accessors, which will be used by codegen
        next_op._attrs["input_accessors"] = [TensorAccessor(t) for t in bmm_inputs]

    sorted_graph = toposort(sorted_graph)
    return sanitize_sorted_graph(sorted_graph)
Transform expand + bmm into a single bmm op.

Parameters
----------
sorted_graph : List[Tensor]
Input graph
workdir : str, optional
workdir, by default None

Returns
-------
List[Tensor]
Optimized graph
python/aitemplate/compiler/transform/fuse_expand_bmm.py