Multiplying numpy array stack with hermitian transpose of itself without loop

Question

I want to completely get rid of for loops in my code.

I have a complex numpy array stack1 of dimension OxMxN This is a stack of MxN arrays stacked in the 1st dimension. For each MxN array that we call A I want to compute the matrix multiplication:

for k in range(stack1.shape[0]):
    A=stack1[k,:,:]
    newstack[k,:,:]=A.dot(  numpy.conj(numpy.transpose(A))  )

I tried

newstack = stack1 @ np.conj(stack1.T)

but I run in an issue because the dimensions won't match

Answer 1

We can use einsum -

np.einsum('ijk,ilk->ijl',stack1,np.conj(stack1))

We can also use np.matmul -

np.matmul(stack1,np.conj(stack1).swapaxes(1,2))

On Python 3.x, simplifies with @ operator -

stack1 @ np.conj(stack1).swapaxes(1,2)