使用 Tensordot 复制矩阵-矩阵乘法

问题描述

在 PyKeops package 中，没有可用的矩阵-矩阵乘法公式。相反，他们实现了类似于 numpy.tensordot 的东西。我有两个大小为 A,B 和 m x n 的矩阵 n x n。有没有办法使用 A @ B 复制 numpy.tensordot。

import numpy as np
m,n = 10,20
A = np.random.random((m,n))
B = np.random.random((n,n))

result1 = A @ B
result2 = np.tensordot(A,B,(1,1))

print(result1 == result2)

# =======================
# array([[False,False,# ....

解决方法

好的，我找到了一个答案，我们可以使用给定 tensordot 的以下参数来解决这个问题：

import numpy as np
m,n = 10,20
A = np.random.random((m,n))
B = np.random.random((n,n))

result1 = A @ B
result2 = np.tensordot(A,B,[(1,),(0,)])

print(result1 == result2)

# =======================
# array([[True,True,# ....