将表达式转换为 einsum

Question

我需要计算一些矩阵的以下表达式：

我当然可以使用 for 循环来完成此操作，但我尝试使用 torch.einsum 函数以矢量化方式计算此值。我无法理解如何将两个求和合并为一个 einsum 表达式。

供参考：

• c 是一个 nxm 矩阵
• w 是一个 mxn 矩阵
• X 是一个 n 维向量
• b 是 m 维向量
• X_tilde 是一个 nxmxn 矩阵

其整体输出应该是一个 mxn 矩阵。

单独计算每个总和非常简单。例如，内部总和（l 上）可以轻松写为：

torch.einsum("l..., blmn -> mn", w[k], X_tilde)

但是，我看不到将双重求和合并为单个 einsum 表达式的方法。

Answer 1

w*x+b

无法用

einsum

完成。

c (n,m)
w (mxn)
X (n,)
b (m,)
X_tilde (n,m,n)

我能做的最好的事情就是使用两个临时表达式。

l

总和有点令人困惑。我假设

l

对应于

n

之一，所以我猜测

sum(l) w, X_tilde

是：

temp1 = np.einsum('ij,kij->ki', w,X_tilde) # (n,m)

表达式是

b

看起来更简单：

w*x+b

temp2 = np.einsum('mn,n->m', w, x) + b   # (m,)

将所有 3 个结合起来，我假设

k

的总和超过

m

维度。

sum(k) c * sigma * temp1 * temp2

res = np.einsum('im,m,jm->ji', c*sigma, temp2, temp1)

我认为仅使用一个

einsum

不会带来性能优势。最初

einsum

构造了一个由所有变量组成的迭代空间。如果有很多变量和两个以上的指数，这个数字可能会很大。较新的版本将计算分成几部分（如果可能的话，

matmul

），您可以使用

einsum_path

进行探索。