def _is_valid_gemm_op(tensor: Tensor, f_check_src_op: Callable) -> bool:
    """check if the src op of tensor is a valid gemm op for parallel fusion.

    Args:
        tensor (Tensor): the output tensor of the gemm op
        f_check_src_op (Callable): a function to check if the src op of the
        input to gemm op is valid for fusion

    Returns:
        bool: True if src_op of tensor is a valid gemm op
    """
    if len(tensor.dst_ops()) != 1 or len(tensor.src_ops()) != 1:
        return False

    # Don't fuse if tensor is an output tensor
    if tensor._attrs["is_output"]:
        return False

    gemm_op = list(tensor.src_ops())[0]
    if gemm_op._attrs["op"] != "gemm_rcr_bias":
        return False

    gemm_input, weight, bias = gemm_op._attrs["inputs"]

    # check gemm weight/bias is available for constant folding
    if not transform_utils.can_be_constant_folded([weight, bias]):
        return False

    if len(gemm_input.dst_ops()) != 1 or len(gemm_input.src_ops()) != 1:
        return False

    # perm102_bmm only supports 3D input, 3D weight, 2D bias
    if gemm_input._rank() != 2 or weight._rank() != 2 or bias._rank() != 1:
        return False

    if not is_static_dimension(gemm_input.shape(), 1):
        return False

    if not is_static_dimension(weight.shape(), 0) or not is_static_dimension(
        weight.shape(), 1
    ):
        return False
    if not is_static_dimension(bias.shape(), 0):
        return False

    src_op = list(gemm_input.src_ops())[0]

    # the new cat must be eliminated with TensorAccessor
    if not f_check_src_op(src_op):
        return False
    return True
check if the src op of tensor is a valid gemm op for parallel fusion.

    Args:
        tensor (Tensor): the output tensor of the gemm op
        f_check_src_op (Callable): a function to check if the src op of the
        input to gemm op is valid for fusion

    Returns:
        bool: True if src_op of tensor is a valid gemm op
    
python/aitemplate/compiler/transform/fuse_parallel_gemms.py