(
    inputs: List[str],
    result: str,
    idx_removed: ArrayIndexType,
    shapes: Sequence[Tuple[int]] = None,
)
Checks if we can use a BLAS call.

    Parameters
    ----------
    inputs : list of str
        Specifies the subscripts for summation.
    result : str
        Resulting summation.
    idx_removed : set
        Indices that are removed in the summation
    shapes : sequence of tuple[int], optional
        If given, check also that none of the indices are broadcast dimensions.

    Returns
    -------
    type : str or bool
        The type of BLAS call to be used or False if none.

    Notes
    -----
    We assume several operations are not efficient such as a transposed
    DDOT, therefore 'ijk,jki->' should prefer einsum. These return the blas
    type appended with "/EINSUM" to differentiate when they can still be done
    with tensordot if required, e.g. when a backend has no einsum.

    Examples
    --------
    >>> can_blas(['ij', 'jk'], 'ik', set('j'))
    'GEMM'

    >>> can_blas(['ijj', 'jk'], 'ik', set('j'))
    False

    >>> can_blas(['ab', 'cd'], 'abcd', set())
    'OUTER/EINSUM'

    >>> # looks like GEMM but actually 'j' is broadcast:
    >>> can_blas(['ij', 'jk'], 'ik', set('j'), shapes=[(4, 1), (5, 6)])
    False
Checks if we can use a BLAS call.
Checks if we can use a BLAS call .
def can_blas(
    inputs: List[str],
    result: str,
    idx_removed: ArrayIndexType,
    shapes: Sequence[Tuple[int]] = None,
) -> Union[str, bool]:
    """
    Checks if we can use a BLAS call.

    Parameters
    ----------
    inputs : list of str
        Specifies the subscripts for summation.
    result : str
        Resulting summation.
    idx_removed : set
        Indices that are removed in the summation
    shapes : sequence of tuple[int], optional
        If given, check also that none of the indices are broadcast dimensions.

    Returns
    -------
    type : str or bool
        The type of BLAS call to be used or False if none.

    Notes
    -----
    We assume several operations are not efficient such as a transposed
    DDOT, therefore 'ijk,jki->' should prefer einsum. These return the blas
    type appended with "/EINSUM" to differentiate when they can still be done
    with tensordot if required, e.g. when a backend has no einsum.

    Examples
    --------
    >>> can_blas(['ij', 'jk'], 'ik', set('j'))
    'GEMM'

    >>> can_blas(['ijj', 'jk'], 'ik', set('j'))
    False

    >>> can_blas(['ab', 'cd'], 'abcd', set())
    'OUTER/EINSUM'

    >>> # looks like GEMM but actually 'j' is broadcast:
    >>> can_blas(['ij', 'jk'], 'ik', set('j'), shapes=[(4, 1), (5, 6)])
    False
    """
    # Can only do two
    if len(inputs) != 2:
        return False

    input_left, input_right = inputs

    for c in set(input_left + input_right):
        # can't deal with repeated indices on same input or more than 2 total
        nl, nr = input_left.count(c), input_right.count(c)
        if (nl > 1) or (nr > 1) or (nl + nr > 2):
            return False

        # can't do implicit summation or dimension collapse e.g.
        #     "ab,bc->c" (implicitly sum over 'a')
        #     "ab,ca->ca" (take diagonal of 'a')
        if nl + nr - 1 == int(c in result):
            return False

    # check for broadcast indices e.g:
    #     "ij,jk->ik" (but one of the 'j' dimensions is broadcast up)
    if shapes is not None:
        for c in idx_removed:
            if shapes[0][input_left.find(c)] != shapes[1][input_right.find(c)]:
                return False

    # Prefer einsum if not removing indices
    #     (N.B. tensordot outer faster for large arrays?)
    if len(idx_removed) == 0:
        return "OUTER/EINSUM"

    # Build a few temporaries
    sets = [set(x) for x in inputs]
    keep_left = sets[0] - idx_removed
    keep_right = sets[1] - idx_removed
    rs = len(idx_removed)

    # DDOT
    if inputs[0] == inputs[1]:
        return "DOT"

    # DDOT does not make sense if you have to transpose - prefer einsum
    elif sets[0] == sets[1]:
        return "DOT/EINSUM"

    # GEMM no transpose
    if input_left[-rs:] == input_right[:rs]:
        return "GEMM"

    # GEMM transpose both
    elif input_left[:rs] == input_right[-rs:]:
        return "GEMM"

    # GEMM transpose right
    elif input_left[-rs:] == input_right[-rs:]:
        return "GEMM"

    # GEMM transpose left
    elif input_left[:rs] == input_right[:rs]:
        return "GEMM"

    # Einsum is faster than vectordot if we have to copy
    elif (len(keep_left) == 0) or (len(keep_right) == 0):
        return "GEMV/EINSUM"

    # Conventional tensordot
    else:
        return "TDOT"
def can_blas ( inputs : List [ str ] , result : str , idx_removed : ArrayIndexType , shapes : Sequence [ Tuple [ int ] ] = None , ) -> Union [ str , bool ] : # Can only do two if len ( inputs ) != 2 : return False input_left , input_right = inputs for c in set ( input_left + input_right ) : # can't deal with repeated indices on same input or more than 2 total nl , nr = input_left . count ( c ) , input_right . count ( c ) if ( nl > 1 ) or ( nr > 1 ) or ( nl + nr > 2 ) : return False # can't do implicit summation or dimension collapse e.g. #     "ab,bc->c" (implicitly sum over 'a') #     "ab,ca->ca" (take diagonal of 'a') if nl + nr - 1 == int ( c in result ) : return False # check for broadcast indices e.g: #     "ij,jk->ik" (but one of the 'j' dimensions is broadcast up) if shapes is not None : for c in idx_removed : if shapes [ 0 ] [ input_left . find ( c ) ] != shapes [ 1 ] [ input_right . find ( c ) ] : return False # Prefer einsum if not removing indices #     (N.B. tensordot outer faster for large arrays?) if len ( idx_removed ) == 0 : return "OUTER/EINSUM" # Build a few temporaries sets = [ set ( x ) for x in inputs ] keep_left = sets [ 0 ] - idx_removed keep_right = sets [ 1 ] - idx_removed rs = len ( idx_removed ) # DDOT if inputs [ 0 ] == inputs [ 1 ] : return "DOT" # DDOT does not make sense if you have to transpose - prefer einsum elif sets [ 0 ] == sets [ 1 ] : return "DOT/EINSUM" # GEMM no transpose if input_left [ - rs : ] == input_right [ : rs ] : return "GEMM" # GEMM transpose both elif input_left [ : rs ] == input_right [ - rs : ] : return "GEMM" # GEMM transpose right elif input_left [ - rs : ] == input_right [ - rs : ] : return "GEMM" # GEMM transpose left elif input_left [ : rs ] == input_right [ : rs ] : return "GEMM" # Einsum is faster than vectordot if we have to copy elif ( len ( keep_left ) == 0 ) or ( len ( keep_right ) == 0 ) : return "GEMV/EINSUM" # Conventional tensordot else : return "TDOT"