Умножение разреженной матрицы на матрицу с использованием SciPy и Numba

Я пытаюсь ускорить некоторые разреженные матрично-матричные умножения в Python, используя Numba и его JIT-компилятор. К сожалению, он не поддерживает библиотеку SciPy, которая мне нужна.

Мое решение состоит в том, чтобы перевести функции csr_matmat_pass1() и csr_matmat_pass2() из здесь в код Python. Мой код, кажется, работает для матриц меньше ~ 80x80 и дает правильные результаты. Вот мое решение:

import scipy.sparse as sparse
import numpy as np

def csr_matmat_pass1(n_row, n_col, Ap, Aj, Bp, Bj):   
   mask = np.ones(n_col, dtype="int") * -1
   Cp = np.zeros(n_row+1, dtype="int")

   nnz = 0
   for i in range(n_row):
       row_nnz = 0

       for jj in range(Ap[i],Ap[i+1]):
           j = Aj[jj]
           for kk in range(Bp[j],Bp[j+1]):
               k = Bj[kk]

               if(mask[k] != i):
                   mask[k] = i
                   row_nnz += 1


       next_nnz = nnz + row_nnz;

       nnz = next_nnz;
       Cp[i+1] = nnz;
   return Cp            

def csr_matmat_pass2(n_row, n_col, Ap, Aj, Ax, Bp, Bj, Bx, Cp):
   nextV = np.ones(n_col, dtype="int") * -1
   sums = np.zeros(n_col)

   nnz = 0

   Cp[0] = 0
   #preallocate space
   sizeC = max(len(Ax),len(Bx)) 
   Cj = np.zeros(sizeC, dtype="int")
   Cx = np.zeros(sizeC)

   for i in range(n_row):
       head   = -2
       length =  0

       jj_start = Ap[i]
       jj_end   = Ap[i+1]
       for jj in range(jj_start,jj_end):   
           j = Aj[jj]
           v = Ax[jj]

           kk_start = Bp[j]
           kk_end   = Bp[j+1]
           for kk in range(kk_start,kk_end):
               k = Bj[kk]

               sums[k] += v*Bx[kk]

               if(nextV[k] == -1):
                   nextV[k] = head
                   head  = k
                   length += 1

       for jj in range(length):

           if(sums[head] != 0.0):
               Cj[nnz] = head
               Cx[nnz] = sums[head]
               nnz += 1


           temp = head     
           head = nextV[head]

           nextV[temp] = -1
           sums[temp] =  0


       Cp[i+1] = nnz
   return Cp, Cj, Cx

#calculate random sparse matrices A,B
mSize = 50
A = sparse.random(mSize, mSize, 0.01).tocsr()
B = sparse.random(mSize, mSize, 0.01).tocsr() 

#calculate sparse C  
Cp = csr_matmat_pass1(np.shape(A)[0], np.shape(B)[1], A.indptr, A.indices, B.indptr, B.indices)
Cp, Cj, Cx = csr_matmat_pass2(np.shape(A)[0], np.shape(B)[1], A.indptr, A.indices, A.data, B.indptr, B.indices, B.data, Cp)
#generate numpy sparse matrix from Cx, Cj, Cp    
C = sparse.csr_matrix((Cx,Cj,Cp),shape=(nRow,nCol))


diffC = A.dot(B) - C

#validate function -> check if any nonzero element is present. If so -> calc is wrong
if np.any(diffC.todense()): UserWarning('Calculations are wrong')

При увеличении размера матриц (скажем, mSize=100) я получаю следующую ошибку:

line 168, in csr_matmat_pass2 Cj[nnz] = head

IndexError: index 72 is out of bounds for axis 0 with size 72

Я предполагаю, что ошибка в моем переводе на Python, а не в коде C++ (поскольку он из библиотеки scipy). Также Cp имеет больше элементов, чем размер матриц A, B. По этой причине в переводе csr_matmat_pass1() должна быть ошибка. К сожалению, я не могу найти никаких синтаксических ошибок и не знаю, почему nnz становится больше, чем должно.