einsum 完整详解爱因斯坦求和约定一、基本定义einsum Einstein Summation即爱因斯坦求和约定是一套用极简字符串描述任意张量收缩、相乘、求和、转置的语法规则。主流库实现numpy.einsum数值计算torch.einsum/tf.einsum深度学习注意力机制高频使用核心两大铁律记这两条就够用相同字母下标对应维度相乘并自动求和收缩掉该维度箭头-右侧只保留的字母作为输出维度右侧消失的字母全部求和消去逗号,分隔不同输入张量字母只是维度别名顺序代表轴顺序公式模板np.einsum(输入1下标,输入2下标-输出下标,arr1,arr2)二、语法规则拆解1. 隐式写法不写-输入端重复下标自动求和剩余下标按字母顺序输出np.einsum(ij,jk,A,B)# 等价 ij,jk-ik 矩阵乘法2. 显式写法推荐可读性极强用-强制指定输出维度可自由调换顺序实现转置np.einsum(ij,jk-ki,A,B)# 矩阵乘法后再转置3. 省略号...万能通配符处理高维批量张量比如[B, H, N, d]...代表前面/后面所有未标注维度# 批量矩阵乘B个矩阵 [B,m,n] × [B,n,k] → [B,m,k]einsum(...mn,...nk-...mk,X,Y)三、最常用运算对照表直接抄1. 一维向量运算importnumpyasnp anp.array([1,2,3])bnp.array([4,5,6])# 向量点积内积 i,i- 输出标量np.einsum(i,i-,a,b)# 32# 向量外积笛卡尔积 i,j-ijnp.einsum(i,j-ij,a,b)# shape (3,3)2. 二维矩阵经典操作设 A:(i,j)B:(j,k)功能einsum表达式等价numpy矩阵乘法ij,jk-ikA B / np.dot矩阵转置ij-jiA.T矩阵逐元素相乘ij,ij-ijA * B矩阵所有元素求和ij-np.sum(A)按行求和axis1ij-iA.sum(1)按列求和axis0ij-jA.sum(0)矩阵迹对角线和ii-np.trace(A)提取对角线ii-inp.diag(A)3. 高阶/深度学习高频写法批量矩阵乘法batch matmul# [B,m,n] [B,n,k] - [B,m,k]einsum(bmn,bnk-bmk,X,Y)Transformer多头注意力核心# Q:bhqd, K:bhkd → bhqk 注意力分数einsum(bhqd,bhkd-bhqk,Q,K)外积/格拉姆矩阵einsum(bi,bj-bij,x,x)四、极简实例跑一遍例1矩阵乘法importnumpyasnp Anp.array([[1,2],[3,4]])# (2,2) ijBnp.array([[5,6],[7,8]])# (2,2) jkresnp.einsum(ij,jk-ik,A,B)# 计算逻辑∑j A[i,j] * B[j,k]print(res)# [[19 22]# [43 50]]例2矩阵转置一行搞定np.einsum(ij-ji,A)五、优缺点 使用建议优点表达统一点积、矩阵乘、张量缩并、求和、转置全部一套语法避免冗余代码不用频繁.transpose()、.reshape()、keepdims维度含义直观字母直接对应业务维度b批次、h头、q查询、k键读代码秒懂优化开关加optimizeTrue自动做计算图优化速度接近底层BLAS缺点新手需要短暂适应下标规则过于复杂的超高阶张量可读性反而下降最佳实践矩阵、批量矩阵乘、注意力张量收缩优先用einsum简单加减、广播加法直接用原生运算符公式一定加上-显式指定输出方便排查维度错误六、一句话口诀同字母相乘求和消掉箭头留下就是输出逗号分开两个张量点点省略任意批量。