unsupported vjp AD rule for tenferro-linalg.linalg.v1

主线位置：traced 主线的 grad 和 eager 主线的 backward(&mut engine)，经过运算族（linalg、FFT、自定义扩展）的扩展算子时。运算族的求导规则不在核心里，要单独注册，机制见 semantic AD 与扩展；线性代数的规则本身见线性代数。

flowchart TB
  op["a.inv() 等运算族"] --> ext["扩展算子"]
  ext --> t1["AdExt grad"]
  ext --> e1["EagerRuntime new"]
  t1 --> empty["空规则集"]
  e1 --> empty
  empty --> err["报错"]
  ext --> ctx["AdContext 带规则"]
  ctx --> ok["梯度"]

实验 crate exp-tenferro8/pit2-exp，bin p17_linalg_ad，CPU（CpuEngine::serial()），debug，M2 Max、macOS 27、rustc 1.99.0。

现象

损失是 sum(inv(A))，A = [[4, 1], [1, 3]]。traced 一侧用 use tenferro_ad::TracedTensorAdExt 带进来的 grad，eager 一侧用 EagerRuntime::new()：

let a = TracedTensor::from_vec_col_major(vec![2, 2], vec![4.0_f64, 1.0, 1.0, 3.0])?;
let loss = a.inv()?.reduce_sum(None)?;
match loss.grad(&a) { /* 打印 Display 和 Debug */ }

let mut engine = CpuEngine::serial();
let ert = EagerRuntime::new()?;
let x = ert.variable_from(Tensor::from_vec_col_major(vec![2, 2], vec![4.0_f64, 1.0, 1.0, 3.0])?)?;
let l = ert.with_eager_session(&mut engine, |s| {
  let i = s.inv(&x)?;
  s.reduce_sum(&i, None)
})?;
match l.backward(&mut engine) { /* 同上 */ }

traced forward: [0.4545454545454545]
traced grad: Err(unsupported vjp AD rule for tenferro-linalg.linalg.v1)
  UnsupportedAdRule { transform: "vjp", op: "tenferro-linalg.linalg.v1" }
eager forward: [0.4545454545454545]
eager backward: Err(semantic_eager_vjp (GraphBuild): runtime state failure: semantic extension AD failed: semantic extension family "tenferro-linalg.linalg.v1" has no Linearize AD rule)
  RuntimeStateSource { op: "semantic_eager_vjp", phase: GraphBuild, source: Extension(MissingRule { family_id: "tenferro-linalg.linalg.v1", role: Linearize }) }

前向都正常，inv 已经算出来了；只有求导失败。
同一个原因，两种错误类型、两种措辞。traced 说「unsupported」，像是这个运算不可导；eager 的「has no Linearize AD rule」才是实情：规则没注册。
报错里只有运算族的 id tenferro-linalg.linalg.v1，不说是 inv。linalg 的所有运算共用这一个 id。

同样的事会发生在 FFT（tenferro-fft.fft.v1）和自己写的扩展算子上。eager 的 einsum 多半不走扩展算子，所以通常不受影响（semantic AD 与扩展的「einsum 在 eager 下多半不走扩展算子」）。

出处

TracedTensorAdExt 是给核心算子准备的便捷入口，impl TracedTensorAdExt for TracedTensor 的每个方法都现造一个空规则集：

    fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor> {
        let rules = SemanticExtensionRuleSet::default();
        grad_with_optional_rules(self, wrt, &rules, None)
    }

EagerRuntime::new 用的也是 SemanticExtensionRuleSet::default()。变换遇到扩展算子时按 family id 查规则，查不到得到 SemanticAdError::MissingRule。traced 的 semantic_transform_error 把 MissingRule 和 Unsupported 一起改写成不带角色的 Error::UnsupportedAdRule；eager 一侧原样包进 RuntimeStateSource，所以保留了 role: Linearize。linalg 的 semantic_ad_rules 只注册了 linearize 和 linear_transpose，反向时先找 primal_vjp、找不到退回「先线性化」，第一步就缺 Linearize。

扩展算子的执行和求导是两件事：Runtime 装了 tenferro_linalg::extension_module 只解决「能算」，不解决「能求导」。缺执行模块时报的是另一条错，见 no engine is registered。

修法

把要用到的运算族的规则放进 AdContext，traced 用它的 grad，eager 用它构造运行时：

let ad = AdContext::builder()
  .with_semantic_extension_rules(tenferro_linalg::semantic_ad_rules()?)?
  .build()?;
let g = ad.grad(&loss, &a)?;                        // traced
let ert = EagerRuntime::with_ad_context(&ad)?;      // eager
let x = ert.variable_from(Tensor::from_vec_col_major(vec![2, 2], vec![4.0_f64, 1.0, 1.0, 3.0])?)?;
// ... 同上算 l，再 l.backward(&mut engine)?

fixed traced grad: [-0.033057851239669415, -0.04958677685950412, -0.049586776859504134, -0.07438016528925619]
fixed eager x.grad: [-0.033057851239669415, -0.04958677685950412, -0.049586776859504134, -0.07438016528925619]
merge linalg again: Err("semantic extension AD Linearize rule for family \"tenferro-linalg.linalg.v1\" is already registered")

两边结果逐位相同；和中心差分的比较见 semantic AD 与扩展。
用到几个运算族就合并几份：let mut rules = tenferro_linalg::semantic_ad_rules()?; rules.merge(tenferro_fft::semantic_ad_rules()?)?;。同一个 family 合并两次报错（最后一行）。
eager 的规则集在 with_ad_context 时定下，没有事后追加的方法；要换规则就新建一个 EagerRuntime。

怎么确认没中招

搜 use tenferro_ad::TracedTensorAdExt 和 EagerRuntime::new()。图里只有核心算子时它们没问题；一旦出现 inv、svd、solve、det、fft，就换成 AdContext 和 EagerRuntime::with_ad_context。
错误匹配别只认 Error::UnsupportedAdRule：eager 缺规则返回的是 Error::RuntimeStateSource，而 EagerTensor::backward 的文档说「a graph operation lacks a reverse rule」时返回的是 UnsupportedAdRule。
程序启动时建一次 AdContext，把整个项目用到的运算族一次合并进去，所有求导都从它走。

小结

TracedTensorAdExt 的 grad 和 EagerRuntime::new() 都用空规则集，经过运算族的扩展算子求导就失败，前向不受影响。
traced 报 unsupported vjp AD rule for <family>，eager 报 ... has no Linearize AD rule；措辞不同，原因相同。
修法：AdContext::builder().with_semantic_extension_rules(...)，traced 调 ad.grad，eager 用 EagerRuntime::with_ad_context(&ad)。

unsupported vjp AD rule for tenferro-linalg.linalg.v1

现象

出处

修法

怎么确认没中招

小结

💬 关联讨论

登录