跨多个轴的点积

Question

给定两个 numpy 数组，其中前 d 个维度大小相等

import numpy

d = 3
a = numpy.random.rand(2, 2, 2, 12, 3)
b = numpy.random.rand(2, 2, 2, 5)

我想计算第一个维度的点积。这个

a2 = a.reshape(-1, *a.shape[d:])
b2 = b.reshape(-1, *b.shape[d:])
out = numpy.dot(numpy.moveaxis(a2, 0, -1), numpy.moveaxis(b2, 0, -2))

有效，但前提是 b 不是 (2, 2, 2) 的形状。摆弄 reshape 和 moveaxis 似乎也比必要的更复杂。

有没有更优雅的方案？ （也许 tensordot？）

Answer 1

再次使用np.einsum

np.einsum('ijklm,ijkn->lmn',a,b)

Answer 2

事实证明 tensordot 是 毕竟有帮助。这个

numpy.tensordot(a, b, axes=(range(3), range(3)))

成功了。