时间戳

首页 > 解决方案 > 关于 np.einsum

python - 关于 np.einsum

问题描述

不明白下面的代码是如何实现维度的转换的？C 的形状是 [2, 3, 3, 4]。没有einsum函数如何实现下面的矩阵运算？

import numpy as np
a = np.random.randint(0, 10, (2,3,4))
b = np.random.randint(0, 10, (3, 6, 4))
c = np.einsum('bld,hid-> blhd', a,b)

标签： pythonnumpypytorch

解决方案

您可以在有关爱因斯坦符号的维基百科中找到更多详细信息

这意味着您有索引，b,l,h,i,d 这将迭代索引以覆盖所有输入并构建输入

我将在这里的数组使用大写字母以与索引区分开来。

C[b,l,h,d] += A[b,l,d] * B[h,i,d]

输出的形状可以如下确定。您获取每个输出轴的索引并在输入中查找相同的索引。例如，第一个轴C被索引b，也用于索引 的第一个轴A，因此assert C.shape[0] == A.shape[0]。对我们拥有的其他轴重复assert C.shape[1] == A.shape[1], assert C.shape[2] == B.shape[0], 和assert C.shape[3] == A.shape[2], 还有assert C.shape[3] == B.shape[2]。

请注意，索引i不影响将添加术语的位置，输出的每个元素都可以写为

C[b,l,h,d] = sum(A[b,l,d] * B[h,i,d] for i in range(B.shape[1]))

另请注意，i它不用于索引A。所以这也可以写成

C[b,l,h,d] = A[b,l,d] * B[h,:,d].sum();

或者如果你想使用矢量化操作

先扩大后缩小

C = A[:,:,None,:] * B[None,None,:,:,:].sum(-2)

扩展减少然后扩展，可能因为A不使用i

C = A[:,:,None,:] * B.sum(-2)[None,None,:,:]

推荐阅读