How to vectorize a numpy for loop that has a multiple indexed access-CodePudding

unigram is an array shape (N, M, 100)

I would like to remove the for loop and perform all the calculations.

seq is a 1D array of size M, and the size of M maybe up to 10000.

I would like to remove the for loop and vectorize it for easier computation.

batch_size, seq_len, num_labels = unigram_scores.shape
broadcast = np.broadcast_to(seq, (batch_size, seq_len))


for i in range(0, broadcast.shape[1]):
    n_seq[i] = unigram_scores[np.arange(batch_size), i , broadcast[:,i]]

edit: answer by @hpaulj worked perfectly and also has the advantage of not having to install any extra dependency
the speed up was much lower than I expected

I ended up finally installing numba

import numpy as np
from numba import njit, prange

@njit(parallel=True)
def calculate_unigram_probability(unigram_scores,seq):
    
    batch_size, seq_len, num_labels = unigram_scores.shape
    broadcast = np.broadcast_to(seq, (batch_size, seq_len))

    for i in prange( broadcast.shape[1]):
       n_seq[i] = unigram_scores[np.arange(batch_size), i , broadcast[:,i]]


    return n_seq

which is also taking a a bit too long, Currently I am trying to move it from the cpu to cuda which should bring about the speedup I am hoping for

CodePudding user response：

In [129]: N,M = 5,3
In [130]: unigram=np.arange(N*M*4).reshape(N,M,4)
In [131]: seq = np.arange(M)
In [132]: b_seq = np.broadcast_to(seq, (N,M))

For a single i:

In [133]: i=0; unigram[np.arange(N),i,b_seq[:,i]]
Out[133]: array([ 0, 12, 24, 36, 48])

For all i in the range:

In [136]: i=np.arange(M)[:,None]
In [137]: unigram[np.arange(N),i,b_seq[:,i]]
Out[137]: 
array([[[ 0, 12, 24, 36, 48],
        [ 5, 17, 29, 41, 53],
        [10, 22, 34, 46, 58]],

        ...
       [[ 0, 12, 24, 36, 48],
        [ 5, 17, 29, 41, 53],
        [10, 22, 34, 46, 58]]])

A (5,3,5) array. This (5,3) might be better)

In [141]: i=np.arange(M); unigram[np.arange(N)[:,None],i,b_seq[:,i]]
Out[141]: 
array([[ 0,  5, 10],
       [12, 17, 22],
       [24, 29, 34],
       [36, 41, 46],
       [48, 53, 58]])

We don't need to index b_seq: unigram[np.arange(N)[:,None],i,b_seq]

Or even use; let the indexing broadcast seq:

unigram[np.arange(N)[:,None],i,seq]

and with the help of ix_:

In [145]: I,J=np.ix_(np.arange(N), np.arange(M))
In [146]: unigram[I,J,seq]

To get a visual idea of what this indexing does, look at unigram. It's pull 'diagonals' from successive blocks/batches:

In [147]: unigram
Out[147]: 
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]],
        ...

CodePudding user response：

you can use x.flatten() to reshape a 3d array to 1d array (x must be a numpy array )

in your case :

broadcast = broadcast.flatten()

this will transform an array of shape (NM1000) to an array of one dimension