在 Numpy 中并行化类似 matmul 的矩阵计算(张量积)
Parallelize a matmul-like matrix computation (tensor product) in Numpy
我想实现张量塔克积s.t.,
- 输入:
B
形状(m, n)
,C
形状(n, n, n)
(立方体)。
- 输出:
Y
形状 (m, m, m)
,s.t。 Y_ijk = ∑_{0≤a,b,c<n} B_ia * B_jb * B_kc * C_abc
.
原始for循环代码:
def raw_tensor_product(B, C, i, j, k):
'''
@param B: (m, n)
@param C: (n, n, n)
@param i: index in range(0, m)
@param j: index in range(0, m)
@param k: index in range(0, m)
@return: the (i, j, k) entry of the product result Y
'''
m, n = B.shape
return np.sum([B[i, a] * B[j, b] * B[k, c] * C[a, b, c] for a in range(n) for b in range(n) for c in range(n)])
那如何在numpy中写出与上面等价的代码,直接returns (m, m, m)
结果Y
?我试过 tensordot
和 matmul
.
请注意,C
的维度可以概括。如果它是 (n, n)
,那么结果将是 B.dot(cube).dot(B.T)
。另外例如,如果它是 (n, n, n, n)
,结果将是一个 (m, m, m, m)
张量 - 那么如何实现它?
此外,你能推荐一些我可以学习这些技术的资料吗?谢谢。
这里有一个方法:
np.tensordot(np.tensordot(np.tensordot(C, B, (0, 1)), B, (0, 1)), B, (0, 1))
但我不确定这是否是最efficient/elegent的方法。
我想实现张量塔克积s.t.,
- 输入:
B
形状(m, n)
,C
形状(n, n, n)
(立方体)。 - 输出:
Y
形状(m, m, m)
,s.t。Y_ijk = ∑_{0≤a,b,c<n} B_ia * B_jb * B_kc * C_abc
.
原始for循环代码:
def raw_tensor_product(B, C, i, j, k):
'''
@param B: (m, n)
@param C: (n, n, n)
@param i: index in range(0, m)
@param j: index in range(0, m)
@param k: index in range(0, m)
@return: the (i, j, k) entry of the product result Y
'''
m, n = B.shape
return np.sum([B[i, a] * B[j, b] * B[k, c] * C[a, b, c] for a in range(n) for b in range(n) for c in range(n)])
那如何在numpy中写出与上面等价的代码,直接returns (m, m, m)
结果Y
?我试过 tensordot
和 matmul
.
请注意,C
的维度可以概括。如果它是 (n, n)
,那么结果将是 B.dot(cube).dot(B.T)
。另外例如,如果它是 (n, n, n, n)
,结果将是一个 (m, m, m, m)
张量 - 那么如何实现它?
此外,你能推荐一些我可以学习这些技术的资料吗?谢谢。
这里有一个方法:
np.tensordot(np.tensordot(np.tensordot(C, B, (0, 1)), B, (0, 1)), B, (0, 1))
但我不确定这是否是最efficient/elegent的方法。