Broadcasting (N,2,2) array and (N,2) array through a for loop

  arrays, broadcasting, matlab, python

I am trying to migrate some code from MATLAB to python and I am having trouble understanding why the following code does not work.

import numpy as np

ngrid = 56
A = np.random.randint(10, size =(ngrid*ngrid,2))
A_tmp = A
B = np.random.randint(10,size =(ngrid*ngrid,2,2) )

for jj in range(ngrid*ngrid):
    A[jj,:] = A_tmp[jj,:]*B[jj,:,:].conj()

When I execute this code I get the error.

ValueError: could not broadcast input array from shape (2,2) into shape (2,)

I am not understanding why this is giving me a one dimensional array vs a (ngrid*ngrid,2) array.

The MATLAB code I am trying to recreate for is

for jj = 1:ngrid^2
    Psi0(jj, :) = Psi0_tmp(jj, :)*dia2adi(:,:, jj)';

Any guidance on theory and how to correct my code would be very helpful.


Source: Python Questions