我正在寻找有关解决以下问题的最有效方法的建议：
我有两个数组，称为 A 和 B。它们的形状均为NxNx3。它们表示两个位置的 2D 矩阵，其中每个位置都是 x、y 和 z 坐标的向量。
我想创建一个形状为 NxN 的新数组，称为 C，其中 C[i， j] 是向量 A[i， j] 和 B[i， j] 的点积。
以下是我到目前为止提出的解决方案。第一个使用numpy的einsum函数(这里描述得很好(。第二个使用 numpy 的广播规则及其 sum 函数。
>>> import numpy as np
>>> A = np.random.randint(0, 10, (100, 100, 3))
>>> B = np.random.randint(0, 10, (100, 100, 3))
>>> C = np.einsum("ijk,ijk->ij", A, B)
>>> D = np.sum(A * B, axis=2)
>>> np.allclose(C, D)
True
有没有更快的方法？我听到有人抱怨说numpy的张量点函数可以非常快，但我一直很难理解它。使用 numpy 的点或内部函数怎么样？
对于某些上下文，A 和 B 数组通常具有 100 到 1000 个元素。
任何指导都非常感谢！
通过一些重塑，我们可以使用matmul. 这个想法是将前 2 个维度视为"批处理"维度，并将最后一个维度的dot：
In [278]: E = A[...,None,:]@B[...,:,None]
In [279]: E.shape
Out[279]: (100, 100, 1, 1)
In [280]: E = np.squeeze(A[...,None,:]@B[...,:,None])
In [281]: np.allclose(C,E)
Out[281]: True
In [282]: timeit E = np.squeeze(A[...,None,:]@B[...,:,None])
130 µs ± 2.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [283]: timeit C = np.einsum("ijk,ijk->ij", A, B)
90.2 µs ± 1.53 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
比较时间可能有点棘手。 在当前版本中，einsum可以根据尺寸采用不同的路线。 在某些情况下，它似乎将任务委托给matmul(或至少相同的底层类似 BLAS 的代码(。 虽然einsum在这个测试中更快很好，但我不会一概而论。
tensordot只是重塑(如果需要，还可以转置(数组，以便它可以应用普通的 2Dnp.dot。 实际上它在这里不起作用，因为您将前 2 个轴视为"批处理"，因为它对它们进行outer product。