How to do numpy matmul broadcasting between two nu

Question

I have the Pauli matrices which are (2x2) and complex

II = np.identity(2, dtype=complex)
X = np.array([[0, 1], [1, 0]], dtype=complex)
Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
Z = np.array([[1, 0], [0, -1]], dtype=complex)

and a depolarizing_error function which takes in a normally distributed random number param, generated by np.random.normal(noise_mean, noise_sd)

def depolarizing_error(param):
    XYZ = np.sqrt(param/3)*np.array([X, Y, Z])
    return np.array([np.sqrt(1-param)*II, XYZ[0], XYZ[1], XYZ[2]])

Now if I feed in a single number for param of let's say a, my function should return an output of np.array([np.sqrt(1-a)*II, a*X, a*Y, a*Z]) where a is a float and * denotes the element-wise multiplication between a and the entries of the (2x2) matrices II, X, Y, Z. Now for vectorization purposes, I wish to feed in an array of param i.e.

param = np.array([a, b, c, ..., n])   Eqn(1)

again with all a, b, c, ..., n generated independently by np.random.normal(noise_mean, noise_sd) (I think it's doable with np.random.normal(noise_mean, noise_sd, n) or something) such that my function now returns:

np.array([[np.sqrt(1-a)*II, a*X, a*Y, a*Z],
          [np.sqrt(1-b)*II, b*X, b*Y, b*Z],
          ................................,
          [np.sqrt(1-n)*II, n*X, n*Y, n*Z]])

I thought feeding in something like np.random.normal(noise_mean, noise_sd, n) as param, giving output as np.array([a, b, c,...,n]) would sort itself out and return what I want above. but my XYZ = np.sqrt(param/3)*np.array([X, Y, Z]) ended up doing element-wise dot product instead of element-wise multiplication. I tried using param as np.array([a, b]) and ended up with

np.array([np.dot(np.sqrt(1-[a, b]), II),
          np.dot(np.sqrt([a, b]/3),  X),
          np.dot(np.sqrt([a, b]/3),  Y),
          np.dot(np.sqrt([a, b]/3),  Z)])

instead. So far I've tried something like

def depolarizing_error(param):
    XYZ = np.sqrt(param/3)@np.array([X, Y, Z])
    return np.array([np.sqrt(1-param)*II, XYZ[0], XYZ[1], XYZ[2]])

thinking that the matmul @ will just broadcast it conveniently for me but then I got really bogged down by the dimensions.

Now my motivation for wanting to do all this is because I have another matrix that's given by:

def random_angles(sd, seq_length):
    return np.random.normal(0, sd, (seq_length,3))

def unitary_error(params):
    e_1 = np.exp(-1j*(params[:,0]+params[:,2])/2)*np.cos(params[:,1]/2)
    e_2 = np.exp(-1j*(params[:,0]-params[:,2])/2)*np.sin(params[:,1]/2)
    return np.array([[e_1, e_2], [-e_2.conj(), e_1.conj()]],
                     dtype=complex).transpose(2,0,1)

where here the size of seq_length is equivalent to the number of entries in Eqn(1) param, denoting N = seq_length = |param| say. Here my unitary_error function should give me an output of

np.array([V_1, V_2, ..., V_N])

such that I'll be able to use np.matmul as an attempt to implement vectorization like this

np.array([V_1, V_2, ..., V_N])@np.array([[np.sqrt(1-a)*II, a*X, a*Y, a*Z],
                                         [np.sqrt(1-b)*II, b*X, b*Y, b*Z],
                                         ................................,
                                         [np.sqrt(1-n)*II, n*X, n*Y, n*Z]])@np.array([V_1, V_2, ..., V_N])

to finally give

np.array([[[e-mail](1-a)*II@V_1, V_1@a*X@V_1, V_1@a*Y@V_1, V_1@a*Z@V_1],
          [[e-mail](1-b)*II@V_2, V_2@b*X@V_2, V_2@b*Y@V_2, V_2@b*Z@V_2],
          ................................,
          [[e-mail](1-n)*II@V_N, V_N@n*X@V_N, V_N@n*Y@V_N, V_N@n*Z@V_N]])

where here @ denotes the element-wise dot-product