Estou usando a distribuição Anaconda do Python, junto com o Numba, e escrevi a seguinte função Python que multiplica uma matriz esparsaA (armazenado em formato CSR) por um vetor densox:

@jit
def csrMult( x, Adata, Aindices, Aindptr, Ashape ):

    numRowsA = Ashape[0]
    Ax       = numpy.zeros( numRowsA )

    for i in range( numRowsA ):
        Ax_i = 0.0
        for dataIdx in range( Aindptr[i], Aindptr[i+1] ):

            j     = Aindices[dataIdx]
            Ax_i +=    Adata[dataIdx] * x[j]

        Ax[i] = Ax_i

    return Ax 

AquiA é um grandescipy matriz esparsa,

>>> A.shape
( 56469, 39279 )
#                  having ~ 142,258,302 nonzero entries (so about 6.4% )
>>> type( A[0,0] )
dtype( 'float32' )

ex é umnumpy array. Aqui está um trecho de código que chama a função acima:

x       = numpy.random.randn( A.shape[1] )
Ax      = A.dot( x )   
AxCheck = csrMult( x, A.data, A.indices, A.indptr, A.shape )

Observe o@jit-decorator que instrui o Numba a fazer uma compilação just-in-time para ocsrMult() função.

Nas minhas experiências, minha funçãocsrMult() é sobreduas vezes mais rápido Enquanto oscipy .dot() método. Esse é um resultado bastante impressionante para o Numba.

No entanto, o MATLAB ainda realiza essa multiplicação vetor-matriz sobre6 vezes mais rápido do quecsrMult(). Acredito que isso ocorre porque o MATLAB usa multithreading ao realizar a multiplicação esparsa de vetores matriz.

Pergunta, questão:

Como posso paralelizar o exteriorfor-loop ao usar o Numba?

Numba costumava ter umprange() função, que simplificou o paralelismo embaraçosamente paralelofor-rotações. Infelizmente, o Numba não tem maisprange() [na verdade, isso é falso, veja a edição abaixo]Então, qual é a maneira correta de paralelizar issofor-loop agora, que Numbaprange() função se foi?

Quandoprange() foi removido do Numba, que alternativa os desenvolvedores do Numba tinham em mente?

Editar 1:
Atualizei para a versão mais recente do Numba, que é .35, eprange() está de volta! Não estava incluído na versão .33, a versão que eu estava usando.
É uma boa notícia, mas infelizmente estou recebendo uma mensagem de erro quando tento paralelizar meu loop for usandoprange(). Aqui está um paralelo para loopexemplo da documentação do Numba (consulte a seção 1.9.2 "Loops paralelos explícitos") e abaixo está o meu novo código:

from numba import njit, prange
@njit( parallel=True )
def csrMult_numba( x, Adata, Aindices, Aindptr, Ashape):

    numRowsA = Ashape[0]    
    Ax       = np.zeros( numRowsA )

    for i in prange( numRowsA ):
        Ax_i = 0.0        
        for dataIdx in range( Aindptr[i],Aindptr[i+1] ):

            j     = Aindices[dataIdx]
            Ax_i +=    Adata[dataIdx] * x[j]

        Ax[i] = Ax_i            

    return Ax 

Quando eu chamo essa função, usando o trecho de código fornecido acima, recebo o seguinte erro:

AttributeError: falha no nopython (converter para parfors) O objeto 'SetItem' não tem atributo 'get_targets'

Dado
a tentativa acima de usarprange falhas, minha pergunta é:

Qual é a maneira correta (usandoprange ou um método alternativo)paralelizar esse Pythonfor-ciclo?

Como observado abaixo, foi trivial paralelizar um loop for similar em C ++ e obter um8x aceleração, tendo sido executado em20-omp-threads. Deve haver uma maneira de fazê-lo usando o Numba, já que o loop for é embaraçosamente paralelo (e uma vez que a multiplicação esparsa de vetores de matriz é uma operação fundamental na computação científica).

Edição 2:
Aqui está minha versão C ++ docsrMult(). Paralelamente aofor() loop na versão C ++ torna o código 8x mais rápido nos meus testes. Isso sugere que uma aceleração semelhante deve ser possível para a versão Python ao usar o Numba.

void csrMult(VectorXd& Ax, VectorXd& x, vector<double>& Adata, vector<int>& Aindices, vector<int>& Aindptr)
{
    // This code assumes that the size of Ax is numRowsA.
    #pragma omp parallel num_threads(20)
    {       
        #pragma omp for schedule(dynamic,590) 
        for (int i = 0; i < Ax.size(); i++)
        {
            double Ax_i = 0.0;
            for (int dataIdx = Aindptr[i]; dataIdx < Aindptr[i + 1]; dataIdx++)
            {
                Ax_i += Adata[dataIdx] * x[Aindices[dataIdx]];
            }

            Ax[i] = Ax_i;
        }
    }
}