Скупистая редкая кончина

Предположим, что у меня есть scipy.sparse.csr_matrix, представляющий значения ниже

[[0 0 1 2 0 3 0 4]
 [1 0 0 2 0 3 4 0]]

Я хочу вычислить кумулятивную сумму ненулевых значений на месте, что изменило бы массив на:

[[0 0 1 3 0 6 0 10]
 [1 0 0 3 0 6 10 0]]

Фактические значения не равны 1, 2, 3,...

Количество ненулевых значений в каждой строке вряд ли будет одинаковым.

Как это сделать быстро?

Текущая программа:

import scipy.sparse
import numpy as np

# sparse data
a = scipy.sparse.csr_matrix(
    [[0,0,1,2,0,3,0,4],
     [1,0,0,2,0,3,4,0]], 
    dtype=int)

# method
indptr = a.indptr
data = a.data
for i in range(a.shape[0]):
    st = indptr[i]
    en = indptr[i + 1]
    np.cumsum(data[st:en], out=data[st:en])

# print result
print(a.todense())

Результат:

[[ 0  0  1  3  0  6  0 10]
 [ 1  0  0  3  0  6 10  0]]