Replacing for loops with numpy operations in case of 3D matrices

Question

I have this mathematical formula to implement ![https://i.sstatic.net/S28BA.png] where for example w_fk denotes matrix of shape (F, K). I have implemented this as

gamma_dashed_lft = np.zeros((L, F, T))
for l in range(L):
    for f in range(F):
        for t in range(T):
            temp = 0
            for k in range(K):
                temp = temp + (q_lk[l, k] * w_fk[f, k] * h_kt[k, t])
            gamma_dashed_lft[l, f, t] = temp
return gamma_dashed_lft

What would be the way to replace for loops with matrix multiplication in the case of given formula?

Answer 1

You should have provided a concrete example. Fortunately it's not too hard to read the dimensions and create:

In [302]: L,F,T,K=2,3,4,5
In [303]: q_lk=np.arange(L*K).reshape(L,K)
In [304]: w_fk=np.arange(F*K).reshape(F,K)
In [305]: h_kt=np.arange(K*T).reshape(K,T)

which when applied to your code produces:

In [306]: gamma_dashed_lft = np.zeros((L, F, T))
     ...: for l in range(L):
     ...:     for f in range(F):
     ...:         for t in range(T):
     ...:             temp = 0
     ...:             for k in range(K):
     ...:                 temp = temp + (q_lk[l, k] * w_fk[f, k] * h_kt[k, t])
     ...:             gamma_dashed_lft[l, f, t] = temp
     ...: 

In [308]: gamma_dashed_lft
Out[308]: 
array([[[ 400.,  430.,  460.,  490.],
        [1000., 1080., 1160., 1240.],
        [1600., 1730., 1860., 1990.]],

       [[1000., 1080., 1160., 1240.],
        [2600., 2855., 3110., 3365.],
        [4200., 4630., 5060., 5490.]]])

An equivalent expression making full use of broadcasting is:

In [309]: arr =(q_lk[:,None,:,None]*w_fk[None,:,:,None]*h_kt[None,None,:,:]).sum(axis=2)
In [310]: arr.shape
Out[310]: (2, 3, 4)
In [311]: np.allclose(arr,gamma_dashed_lft)
Out[311]: True

In setting up the broadcasting I was aiming for an array with shape (L,F,K,T) with the sum reduction on the K.

Since you made me create the test case, I let you work out the broadcasting details. It'll be a good exercise for you.

einsum

In [446]: D=np.einsum('lk,fk,kt->lft', q_lk, w_fk, h_kt)
In [447]: D.shape
Out[447]: (2, 3, 4)
In [448]: arr =(q_lk[:,None,:,None]*w_fk[None,:,:,None]*h_kt[None,None,:,:]).sum
     ...: (axis=2)
In [449]: np.allclose(arr,D)
Out[449]: True
In [450]: timeit arr =(q_lk[:,None,:,None]*w_fk[None,:,:,None]*h_kt[None,None,:,
     ...: :]).sum(axis=2)
22.4 µs ± 2.02 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [451]: timeit D=np.einsum('lk,fk,kt->lft', q_lk, w_fk, h_kt)
12.2 µs ± 40.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

matmul

In [458]: timeit E=((q_lk[:,None,None,:]*w_fk[None,:,None,:])@h_kt[None,None,:,:,: ]).squeeze()
10.6 µs ± 44.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

For this I'm making a (L,F,1,K) array to @ with a (1,1,K,T), resulting in a (L,F,1,T). LF are the matmul 'batch' dimensions, while K is the sum-of-products dimension.