时间戳

首页 > 解决方案 > 沿给定轴的 numpy kron

python - 沿给定轴的 numpy kron

问题描述

是否有沿给定轴应用克罗内克积的功能？例如，给定 2D 数组a和b的形状a.shape == (n, k)和b.shape == (n, l)，计算c形状c.shape == (n, k*l)的等效结果为：

c = np.empty((a.shape[0], a.shape[1] * b.shape[1]))
for i in range(c.shape[0]):
    c[i,:] = np.kron(a[i], b[i])

标签： pythonnumpy

解决方案

没有内置的，但我们可以使用outer elementwise-multiplication保持它们的第一个轴对齐然后重塑 -

c = (a[:,:,None]*b[:,None,:]).reshape(a.shape[0],-1)

或者，我们可以使用einsum-

c = np.einsum('nk,nl->nkl',a,b).reshape(a.shape[0],-1)

推荐阅读