沿指定轴的numpy kron

Question

是否有一种函数可以沿着指定的轴应用Kronecker乘积？例如，给定形状为a.shape == (n, k)和b.shape == (n, l)的2D数组a和b，计算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])

- Lukasz Wiklendt

---

你能看一下 np.kron 的代码吗？我不记得它是编译的还是Python的？ - hpaulj

Answer 1

虽然没有内置的方法，但我们可以使用 外部逐元素乘法 使它们的第一个轴对齐，然后进行重塑 -

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

或者，我们可以使用einsum -

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

- Divakar

---

太棒了，谢谢。我喜欢 einsum 的解决方案，它是一个非常多才多艺的工具。 - Lukasz Wiklendt

Answer 2

如果需要稀疏性：

from scipy.sparse import kron


def kron_prod_in_axis1(a, b):
    return kron(a, b, format="csr")[[i * a.shape[0] + i for i in range(a.shape[0])], :]

- alan