def fuse_bmm_permute(sorted_graph: List[Tensor], _: str) -> List[Tensor]:
    """
    Fuse bmm + permute021 ops. The second argument is unused, it's only
    here to make the type of this function the same as the others called in optimize_graph.
    """
    ops_r = [
        bmm_ccr,
        bmm_crr,
        bmm_rcr,
        bmm_rrr,
    ]

    ops_c = [
        bmm_ccc,
        bmm_crc,
        bmm_rcc,
        bmm_rrc,
    ]
    patterns_cr = [((c_op(), permute021()), r_op) for c_op, r_op in zip(ops_c, ops_r)]
    patterns_rc = [((r_op(), permute021()), c_op) for c_op, r_op in zip(ops_c, ops_r)]

    sorted_graph = transform_simple_fusion_patterns(
        sorted_graph, patterns_cr + patterns_rc
    )

    return sorted_graph
Fuse bmm + permute021 ops. The second argument is unused, it's only
here to make the type of this function the same as the others called in optimize_graph.
python/aitemplate/compiler/transform/fuse_bmm_permute.py