def _transform_gemm_rrr_small_nk(sorted_graph: List[Tensor]) -> List[Tensor]:
    """
    Replace kernel gemm_rrr with N <= 8, K <= 16 with gemm_rrr_small_nk

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

    Returns
    ----------
    List[Tensor]
        AIT graph with suitable gemm_rrr substituted with gemm_rrr_small_nk
    """
    matchFunc, replaceFunc = _single_source_shape_constraint_funcs(
        gemm_rrr, gemm_rrr_small_nk
    )

    return _simple_transform_with_constraint(sorted_graph, matchFunc, replaceFunc)
Replace kernel gemm_rrr with N <= 8, K <= 16 with gemm_rrr_small_nk

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

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