def _transform_bmm_rcr_n1(sorted_graph: List[Tensor]) -> List[Tensor]:
    """
    Replace kernel bmm_rcr with N == 1 and K % 8 == 0 with bmm_rcr_n1

    Parameters
    ----------
    sorted_graph : List[Tensor]
        original AIT graph

    Returns
    ----------
    List[Tensor]
        AIT graph with suitable bmm_rcr substituted with bmm_rcr_n1
    """
    matchFunc, replaceFunc = _single_source_shape_constraint_funcs(
        bmm_rcr, ops.gemm_special.bmm_rcr_n1
    )

    return _simple_transform_with_constraint(sorted_graph, matchFunc, replaceFunc)
Replace kernel bmm_rcr with N == 1 and K % 8 == 0 with bmm_rcr_n1

Parameters
----------
sorted_graph : List[Tensor]
original AIT graph

Returns
----------
List[Tensor]
AIT graph with suitable bmm_rcr substituted with bmm_rcr_n1
python/aitemplate/compiler/transform/transform_special_ops.py