Как сделать трансляцию numpy matmul между двумя тензорами numpy?

У меня есть матрицы Паули (2x2) и сложные

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)

и функция depolarizing_error, которая принимает нормально распределенное случайное число param, сгенерированное 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]])

Теперь, если я подам одно число для param, скажем, a, моя функция должна вернуть вывод np.array([np.sqrt(1-a)*II, a*X, a*Y, a*Z]), где a — это float, а * обозначает поэлементное умножение между a и элементами (2x2) матриц II, X, Y, Z. Теперь для целей векторизации я хочу передать массив param, т.е.

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

снова со всеми a, b, c, ..., n, сгенерированными независимо np.random.normal(noise_mean, noise_sd) (я думаю, это выполнимо с np.random.normal(noise_mean, noise_sd, 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.random.normal(noise_mean, noise_sd, n) как param, предоставление вывода как np.array([a, b, c,...,n]) само разберется и вернет то, что я хочу выше. но мой XYZ = np.sqrt(param/3)*np.array([X, Y, Z]) закончил тем, что выполнял поэлементное скалярное произведение вместо поэлементного умножения. Я попытался использовать параметр как np.array([a, b]) и в итоге получил

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)])

вместо. До сих пор я пробовал что-то вроде

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]])

думал что матмуль@ будет просто транслировать это удобно для меня но тут я совсем запутался в габаритах.

Теперь моя мотивация для желания сделать все это заключается в том, что у меня есть другая матрица, которая задается:

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)

где здесь размер seq_length эквивалентен количеству записей в уравнении (1) param, обозначая, скажем, N = seq_length = |param|. Здесь моя функция unitary_error должна дать мне вывод

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

так что я смогу использовать np.matmul как попытку реализовать векторизацию, подобную этой

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])

наконец дать

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

где здесь @ обозначает поэлементное скалярное произведение