def transform_odd_alignment(
    sorted_graph: List[Tensor], workdir: str = None
) -> List[Tensor]:
    """Transform odd alignments to even alignments for bmm operators

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

    Returns
    -------
    List[Tensor]
        Optimized graph
    """
    permutable_pairs = {
        "bmm_ccr": (bmm_rcr, bmm_crr, bmm_rrr),
        "bmm_crr": (bmm_rrr, bmm_ccr, bmm_rcr),
        "bmm_rcr": (bmm_ccr, bmm_rrr, bmm_crr),
        "bmm_rrr": (bmm_crr, bmm_rcr, bmm_ccr),
    }

    sorted_graph = _transform_odd_alignment(sorted_graph, permutable_pairs)

    return sorted_graph
Transform odd alignments to even alignments for bmm operators

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

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