def _try_extract_one_mm_op(ops: Set[Union[None, Operator]]) -> Union[None, Operator]:
    """
    Helper function that returns the matmul op from src_ops() or dst_ops() call.
    Return None if there's no bmm ops
    """
    if ops is None:
        return None

    for op in ops:
        if op._attrs["op"].startswith("bmm") or op._attrs["op"].startswith("gemm"):
            return op

    return None

    Helper function that returns the matmul op from src_ops() or dst_ops() call.
    Return None if there's no bmm ops
    
python/aitemplate/compiler/transform/fuse_permute_bmm_and_gemm.py