numpy elementwise外部产品

Question

我想在numpy中做两个2d数组的元素外部产品.

A.shape = (100, 3) # A numpy ndarray
B.shape = (100, 5) # A numpy ndarray

C = element_wise_outer_product(A, B) # A function that does the trick
C.shape = (100, 3, 5) # This should be the result
C[i] = np.outer(A[i], B[i]) # This should be the result

一个天真的实现可以如下.

tmp = []
for i in range(len(A):
    outer_product = np.outer(A[i], B[i])
    tmp.append(outer_product)
C = np.array(tmp)

从堆栈溢出中获得更好的解决方案.

big_outer = np.multiply.outer(A, B)
tmp = np.swapaxes(tmp, 1, 2)
C_tmp = [tmp[i][i] for i in range(len(A)]
C = np.array(C_tmp)

我正在寻找一个摆脱for循环的矢量化实现.有没有人有想法？谢谢!

Answer 1

延伸A并B以3D保持他们的第一轴线对齐,并且沿所述第三和第二个分别与引入新的轴None/np.newaxis,然后彼此相乘.这将允许broadcasting为矢量化解决方案发挥作用.

因此,实施将是 -

A[:,:,None]*B[:,None,:]

我们可以通过使用ellipsisA 来缩短它::,:并跳过列出剩余的最后一个轴B,就像这样 -

A[...,None]*B[:,None]

作为另一种矢量化方法,我们也可以使用np.einsum,一旦我们通过字符串表示法语法并且认为那些符号是参与天真循环实现的迭代器的代表,这可能更直观,如此 -

np.einsum('ij,ik->ijk',A,B)