einsum notation must contain exactly one '->'
主线位置:einsum 线的第 1 步。下标字符串先被解析成整数标签,解析失败就到不了规划和执行。三种下标写法和它们的内部表示见 einsum。
flowchart TB
s["下标字符串"] --> arrow{"恰好一个箭头"}
arrow -->|否| e1["must contain exactly one"]
arrow -->|是| paren{"有括号"}
paren -->|是| e2["does not accept parentheses"]
paren -->|否| n{"输入个数对"}
n -->|否| e3["expects 2 inputs"]
n -->|是| ok["规划 执行"]
nest["NestedEinsum 加 einsum_with"] --> ok
实验 crate exp-tenferro8/pit1-pit,bin p10_arrow,release;M2 Max、macOS 27、rustc 1.99.0。
现象
从 NumPy 或 opt_einsum 搬过来的下标,常见三种写法:省略输出("ij,jk",NumPy 按字母顺序推出 ik)、用括号指定收缩顺序(Julia 的 OMEinsum 写法)、省略号批维。在具体张量上逐个试(x 是 2×3,y 是 3×4,z 是 4×5,bx、by 带一个长 7 的批维,全 1):
let mut engine = CpuEngine::serial();
let r = [&x, &y].as_slice().einsum("ij,jk", &mut engine);
no arrow ij,jk Err(invalid einsum subscripts: einsum notation must contain exactly one '->', got: ij,jk)
arrow ij,jk->ik Ok([2, 4])
trace-out ij,jk-> Ok([])
two arrows ij,jk->ik->ki Err(invalid einsum subscripts: einsum notation must contain exactly one '->', got: ij,jk->ik->ki)
parens (ij,jk),kl->il Err(invalid einsum subscripts: EinsumNotation::parse does not accept parentheses; use NestedEinsum::parse for parenthesized contraction order)
ellipsis ...ij,...jk->...ik Ok([7, 2, 4])
arity ij,jk->ik Err(einsum: invalid argument inputs: einsum notation expects 2 inputs, got 1)
(einsum 是 TensorEinsumExt 给 [&Tensor] 实现的方法,第二个参数是 &mut dyn Engine。)照报错的提示去解析带括号的写法:
NestedEinsum::parse (ij,jk),kl->il Ok(Node { subscripts: Subscripts { inputs: [[105, 107], [107, 108]], output: [105, 108] }, children: [Node { subscripts: Subscripts { inputs: [[105, 106], [106, 107]], output: [105, 107] }, children: [Leaf(0), Leaf(1)] }, Leaf(2)] })
NestedEinsum::parse (...ij,jk),kl->...il Err(invalid einsum subscripts: ellipsis with parenthesized contraction order is not supported)
NestedEinsum::parse (ij,jk),kl Err(invalid einsum subscripts: einsum notation must contain exactly one '->', got: (ij,jk),kl)
解析是成功了,可具体张量的 einsum 只收字符串,拿到一个 NestedEinsum 也不知道交给谁。
出处
箭头的检查在 split_and_validate_notation:
let parts: Vec<&str> = notation.split("->").collect();
if parts.len() != 2 {
return Err(Error::invalid_subscripts(format!(
"einsum notation must contain exactly one '->', got: {notation}"
)));
}
没有箭头和有两个箭头是同一条错误。输出为空("ij,jk->")是合法的,表示全部求和成标量。README 的约定只写了「Equations need the explicit arrow ("ij,jk->ik")」,没有解释原因;结果就是 NumPy 的隐式输出规则(只出现一次的标签按字母顺序组成输出)在 tenferro 里不存在,移植时每个下标都要补上输出。
平铺的写法由 EinsumNotation::parse 解析,见到括号就拒绝,并在错误里指向 NestedEinsum::parse。NestedEinsum::parse 得到一棵树(105、106…… 是 i、j 的字符码,标签在内部是整数),但消费这棵树的只有 traced 入口的 EinsumOptimize::Nested:具体张量的 einsum、ConcreteEinsumPlan 和 eager 的 einsum 都不收收缩顺序。括号和省略号不能同时用,这是 README 说的「parenthesized ellipsis remains unsupported」。
输入个数和下标里逗号分出的项数不一致,在省略号展开那一步报 invalid argument,和下标语法错误是两种错误类型(InvalidArgument 对 InvalidSubscripts)。
修法
平铺写法补上箭头:
let out = [&x, &y].as_slice().einsum("ij,jk->ik", &mut engine)?;
要指定收缩顺序,走 traced:输入 trace 进图,下标写平铺形式,顺序用 EinsumOptimize::Nested 传:
let nested = NestedEinsum::parse("(ij,jk),kl->il")?;
let out = trace.einsum_with(&vals, "ij,jk,kl->il", EinsumOptimize::Nested(nested))?;
let program = GraphCompiler::new().compile_traced_graph(&trace.finish(&[out])?)?;
// Runtime 要注册 CPU engine 并装上 tenferro_einsum::extension_module
traced einsum_with Nested((ij,jk),kl->il): [2, 5] first=12
(三个全 1 矩阵,il 的每个元素是 3 × 4 = 12。)同样能用 EinsumOptimize::Path(JAX 式的位置路径)或 EinsumOptimize::Tree。按 Path 的文档注释,它与具体维度无关,还能用在符号形状上。带省略号又想固定顺序时用 Path,同一个程序里:
traced einsum_with Path[(1,2),(0,1)] ...ij,...jk,...kl->...il: [7, 2, 5] first=12
为什么不用 EinsumOptimize::False 来「按写的顺序算」,见 EinsumOptimize::False 反而慢几百倍。
怎么确认没中招
- 错误类型是
InvalidSubscripts的,都在解析这一步,和张量的形状无关;InvalidArgument的多半是输入个数或某一维长度对不上。 - 从 NumPy 移植代码时,搜一遍
einsum(的下标里有没有不带->的。 - 自己写一个 einsum 时方言怎么取舍,见 11 动手。
小结
- 下标必须恰好有一个
->;没有隐式输出,"ij,jk->"表示求和成标量。 - 平铺写法不收括号;括号写法用
NestedEinsum::parse解析后只能交给 traced 的einsum_with(.., EinsumOptimize::Nested(..))。 - 括号和省略号不能同用;带省略号时用
Path固定顺序。