如何使用两个张量和一个向量执行快速张量收缩？答案

【问题标题】：How to perform fast tensor contraction with two tensors and a vector?如何使用两个张量和一个向量执行快速张量收缩？
【发布时间】：2022-08-18 16:22:07
【问题描述】：

我正在使用 numpy（理想情况下是 Numba）来执行涉及三个张量的张量收缩，其中一个是一个向量，应该只乘以其他索引的一个索引。例如，

A = np.random.normal(size=(20,20,20,20))
B = np.random.normal(size=(20,20,20,20))
v = np.sqrt(np.arange(20))

# e.g. v on the 3rd index
>>> %timeit np.vdot(A * v[None, None, :, None], B)
125 µs ± 5.14 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

与之比较

C = np.random.normal(size=(20,20,20,20))

>>> %timeit np.vdot(A * C, B)
76.8 µs ± 4.25 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

有没有更有效的方法将产品包含在v 中？感觉它应该比乘以完整张量C 慢。

np.einsum(\'ijkl,k,ijkl\', A, v, B, optimize=\'optimal\').item()，我不知道，为什么结果是一个优化的单项数组。
我应该说我想避免使用einsum，因为 numba 不支持它。是否可以做与 einsum 相同的操作，但没有 einsum？
也许从这篇文章stackoverflow.com/questions/48253210/… 可以理解为什么第一种情况比第二种慢

标签： python numpy numba

【解决方案1】：

我可以通过使用带有parallel=True 的 numba 来压缩一些性能

import numba as nb
import numpy as np

N = 50

@nb.njit('float64(float64[:,:,:,:], float64[:,:,:,:],float64[:])',parallel=True)
def dotABv(a, b,vv):
    res = 0.0
    for i in nb.prange(a.shape[0]):
        for j in range(a.shape[1]):
            for k in range(a.shape[2]):
                res += vv[k]*np.dot(a[i,j,k,:],b[i,j,k,:])
    return res

v = np.sqrt(np.arange(N))
A = np.random.normal(size=(N,N,N,N))
B = np.random.normal(size=(N,N,N,N))
C = np.random.normal(size=(N,N,N,N))

%timeit dotABv(A,B,v)
%timeit np.dot(A.ravel() , B.ravel()) ## just to compare with dot
%timeit np.vdot(A * v[None, None, :, None], B)
# Output :
# 473 µs ± 950 ns per loop (mean ± std. dev. of 7 runs, 100 loops each)
# 1.57 ms ± 597 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# 15.8 ms ± 1.36 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

print(dotABv(A,B,v), np.vdot(A * v[None, None, :, None], B))
# 5105.504508154087 5105.5045081541075

【讨论】：