Fastest way

Question

I have a three-dimensional tensor A of size (M, N, N). I also have a weight vector p of length M. I want to compute

Dimension N can be large, so I want to implement it effectively. I am using the following code:

import numpy as np
temp=np.array([p[m]*A[m] for m in range(M)])
B=sum(temp);

I want to know if there are faster and more efficient ways to implement this. The reason this is important is because when N is large, the above variable temp stores the M of these N * N matrices. Therefore, I think we can use a "loop cycle", perhaps, but I read that for loops it is slower than lists.

Answer 1

At least from the tensors on which I tested it, this snipp is up to 3 times faster:

sum(np.einsum('ij...,i->ij...', A, p))

, :

base      = sum(np.array([p[m]*A[m] for m in range(M)]))
einsum1   = sum(np.einsum('ij...,i->ij...', A, p))
einsum2   = np.einsum('ijk,i->jk', A, p)
einsum3   = np.einsum('ij...,i->j...', A, p)
dot       = p.dot(A.reshape(A.shape[0],-1)).reshape(A.shape[1],-1)
tensordot = np.tensordot(p, A, (0, 0))

---

M = 1000, N = 100 ( ):

einsum1   = 1.11x
einsum2   = 5.83x
einsum3   = 6.26x
dot       = 8.60x
tensordot = 9.83x

Answer 2

In [24]: A=np.arange(3*2*2).reshape(3,2,2)    
In [25]: p=np.array([1,2,3])

In [26]: temp=np.array([p[m]*A[m] for m in range(len(p))])
In [27]: sum(temp)  # python sum
Out[27]: 
array([[32, 38],
       [44, 50]])

In [28]: np.einsum('i...,i->...',A,p)
Out[28]: 
array([[32, 38],
       [44, 50]])

ijk,i->jk .

, temp . np.sum(..., axis=0) Python sum().

In [30]: (A*p[:,None,None]).sum(axis=0)
Out[30]: 
array([[32, 38],
       [44, 50]])