java 并行 矩阵乘法_使用tensordot实现批量矩阵乘法

我试图通过使用tensordot，dot和reshaping等来实现与np.matmul并行矩阵乘法相同的行为 .

我正在将其翻译为使用的库没有支持并行乘法的matmul，只有dot和tensordot .

另外，我想避免迭代第一维，并希望使用一组矩阵乘法和重新整形(希望尽可能多地使用BLAS / GPU运行，因为我有大量小矩阵并行计算) .

这是一个例子：

import numpy as np

angles = np.array([np.pi/4, 2*np.pi/4, 2*np.pi/4])

vectors = np.array([ [1,0],[1,-1],[-1,0]])

s = np.sin(angles)

c = np.cos(angles)

rotations = np.array([[c,s],[-s,c]]).T

print rotations

print vectors

print("Correct: %s" % np.matmul(rotations, vectors.reshape(3,2,1)))

# I want to do this using tensordot/reshaping, i.e just gemm BLAS operations underneath

print("Wrong: %s" % np.tensordot(rotations, vectors, axes=(1,1)))

这个输出是：

Correct: [[[ 7.07106781e-01]

[ 7.07106781e-01]]

[[ 1.00000000e+00]

[ 1.00000000e+00]]

[[ -6.12323400e-17]

[ -1.00000000e+00]]]

Wrong: [[[ 7.07106781e-01 1.11022302e-16 -7.07106781e-01]

[ -7.07106781e-01 -1.41421356e+00 7.07106781e-01]]

[[ 6.12323400e-17 -1.00000000e+00 -6.12323400e-17]

[ -1.00000000e+00 -1.00000000e+00 1.00000000e+00]]

[[ 6.12323400e-17 -1.00000000e+00 -6.12323400e-17]

[ -1.00000000e+00 -1.00000000e+00 1.00000000e+00]]]

有没有一种方法可以修改第二个表达式，以获得与第一个表达式相同的结果，只需使用dot / tensordot .

我相信它是可能的，并且已经看过some comments online，但从来没有任何例子