Использование einsum без явной репликации

Question

У меня есть матрица A формы (n, n) и другая матрица b формы (p, n). Мне нужно, чтобы получить матрицу С, что

C[i] = (A * b[i, np.newaxis, :]) * b[i, :, np.newaxis] 

В настоящее время я делаю что делает np.matlib.repmat (A), а затем с помощью np.einsum следующего

A1 = np.matlib.repmat(A, p, 1).reshape(p, n, n) 
C = np.einsum('ijk, ij..., ik... -> ijk', A1, b, b) 

Но я могу что создание А1 является пустой тратой ресурсов, поскольку она использует одни и те же значения снова и снова. В любом случае я могу это сделать без создания промежуточной матрицы A1?

Answer 1

Использование np.einsum как так -

np.einsum('ij,kij,kij->kij',A,b[:,None,:],b[...,None])