numpy (n, m) and (n, k) to (n, m, k)

Question

x是形状(n, m)的np.array。

让y成为形状(n, k)的np.array。

计算形状(n, m, k)张量z的正确方法是什么，以便

for all i in [0, n - 1]
z[i] = np.dot(x[i][:, np.newaxis], y[i][np.newaxis, :])

？

换句话说，每对行(x_i, y_i)给出一个形状(m, k)矩阵。

我看了看np.tensordot但经过多次试验，我找不到它axes论点的正确值。我不确定它是否适合这项工作。

Answer 1

你可以像这样使用np.einsum()：

z = np.einsum('ij,ik->ijk', x, y)

---

从快速测试来看，这也比基于np.matmul()的方法更快(除了非常小的输入(：

import numpy as np

x = np.random.randint(1, 100, (2, 3))
y = np.random.randint(1, 100, (2, 4))
%timeit np.einsum('ij,ik->ijk', x, y)
# 100000 loops, best of 3: 3.14 µs per loop
%timeit np.matmul(x[:, :, None], y[:, None, :])
# 100000 loops, best of 3: 2,07 µs per loop

x = np.random.randint(1, 100, (20, 30))
y = np.random.randint(1, 100, (20, 40))
%timeit np.einsum('ij,ik->ijk', x, y)
# 10000 loops, best of 3: 32.1 µs per loop
%timeit np.matmul(x[:, :, None], y[:, None, :])
# 10000 loops, best of 3: 76.8 µs per loop

x = np.random.randint(1, 100, (200, 300))
y = np.random.randint(1, 100, (200, 400))
%timeit np.einsum('ij,ik->ijk', x, y)
# 10 loops, best of 3: 48.7 ms per loop
%timeit np.matmul(x[:, :, None], y[:, None, :])
# 10 loops, best of 3: 68.2 ms per loop

---

np.dot()应用于可广播视图，如下所示：

np.dot(x[:, :, None], y[:, None, :])

行不通(它甚至不会达到正确的形状(。

(已编辑(

Answer 2

一个简单的np.matmul(x[:, :, None], y[:, None, :])就可以了。

来自numpy.matmul的文档：

如果任一参数为 N-D、N> 2，则将其视为驻留在最后两个索引中的矩阵堆栈并相应地广播。

这正是我想做的。