Rust in Everything
基于 tenferro-runtime 0.8.0

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 固定顺序。

💬 关联讨论

发起讨论
加载中...