Rust in Everything
基于 tenferro-ad 0.8.0

placeholder semantic input 2 has no runtime input binding

主线位置:变换主线的「grad / vjp」到「编译、准备 → 执行」,走程序级变换 AdContext::vjp_program 的那条路。

版本:JAX 0.12.0.1(CPU,x64)、tenferro-rs 0.8.0(CPU,native,单线程)
flowchart TB
  tc["TraceContext 程序"] --> vp["vjp_program"]
  vp --> ins["原输入 加 种子"]
  ins --> pa["只传原输入"]
  pa --> e1["输入个数不符"]
  ins --> pn["什么都不传"]
  pn --> e2["种子没有绑定"]
  ins --> ps["补上种子 1"]
  ps --> ok["梯度"]

实验代码:tenferro 一侧是 bin p16_vjp_seed(release,CpuEngine::serial()),JAX 一侧是 p16_vjp_seed.py。M2 Max、macOS 27、rustc 1.99.0。

现象

损失是 L(x) = Σ (A x)²,A = [[1, 2, 3], [4, 5, 6]],x = [1, −1, 0.5],精确梯度 2Aᵀ A x = [17, 22, 27]。JAX 里求梯度不用管种子,grad 隐含 1;要种子的是 vjp 返回的那个函数:

f = lambda x: jnp.sum(jnp.einsum("ij,j->i", A, x) ** 2)
jax.grad(f)(x)
y, pull = jax.vjp(f, x)
pull(1.0); pull(2.0); pull()
grad(f)(x)         = [17. 22. 27.]
vjp(f, x) pull(1.0) = [17. 22. 27.]
vjp(f, x) pull(2.0) = [34. 44. 54.]
pull()  TypeError: VJP.__call__() missing 1 required positional argument: 'out_ct'

tenferro 的 TracedTensor 上没有 einsum,用 einsum 就要用 TraceContext 记图,求导走程序级的 AdContext::vjp_program(见 einsum 的求导一节)。照 jax.grad(f)(A, x) 的习惯,编译好导数程序后只传 A 和 x:

let y = tc.einsum(&[va, vx], "ij,j->i")?;
let yy = tc.add_op(CoreSemanticOp::Mul, &[y, y])?[0];
let l = tc.add_op(CoreSemanticOp::ReduceSum { axes: vec![0] }, &[yy])?[0];
let prog = GraphCompiler::new().compile_traced_graph(&tc.finish(&[l])?)?;
let sem = ad.vjp_program(prog.frozen_program(), &[false, true], &[true])?;   // 对 x 求导,给 L 种子
let gprog = GraphCompiler::new().compile_frozen_program(sem.frozen())?;
prepare(&rt, &eng, &gprog, &[&a, &x])   // 只传原输入
prepare(&rt, &eng, &gprog, &[])         // 或者都用默认值
vjp_program: inputs=3 seed index [Some(2)] grad index [None, Some(0)]
prepare(a, x): Err(Runtime::run (Execution): invalid argument inputs: expected 3 inputs, received 2)
prepare(): Err(placeholder semantic input 2 has no runtime input binding)
  • 导数程序有 3 个输入,第 2 个是种子。两个原输入在 TraceContext 里是 input_with_default,带着默认值;种子没有。
  • 传 2 个输入,报个数不符;一个都不传,报种子没有绑定。两条错误都出在 prepare,第一条却标成 Runtime::run。

出处

AdContext::vjp_program 的文档写着 active_outputs 选「receive appended seeds」的输出:种子追加在原输入之后,是一个没有默认值的输入。ad.grad 不一样,它的种子是一个带默认值 1 的输入(grad、vjp、jvp),所以 ad.grad 的程序只传 x 就能跑。

两条报错来自 resolve_input_reads:给的张量少于程序输入时,它只把给的张量依次填进没有默认值的输入。这里只有种子一个空位,给 2 个就多出一个,报「expected 3 inputs, received 2」。一个都不给时走 semantic_default_inputs,逐个取默认值,到第 2 个取不到,报 UnboundPlaceholder,input_key 写成 semantic input 2。同一个错误变体在编译期的用法见 tenferro 卷的 has no runtime input binding,那里缺的是占位符的规格,这里缺的是种子。

修法

把种子当成最后一个输入传进去,标量、dtype 和损失相同:

let seed = Tensor::from_vec_col_major(vec![], vec![1.0_f64])?;
let p = prepare(&rt, &eng, &gprog, &[&a, &x, &seed])?;
run_t(&rt, &mut eng, &p, &[&a, &x, &seed])?;     // 种子 1
run_t(&rt, &mut eng, &p, &[&a, &x, &seed2])?;    // 种子 2,相当于 pull(2.0)
let p1 = prepare(&rt, &eng, &gprog, &[&seed])?;  // 只给种子,A、x 用默认值
run(a, x, seed=1): [17.0, 22.0, 27.0]   exact 2 A^T A x = [17, 22, 27]
run(a, x, seed=2): [34.0, 44.0, 54.0]
run(seed=1) only: [17.0, 22.0, 27.0]

和 JAX 逐位相同(整数结果)。种子为 2 时梯度翻倍,和 pull(2.0) 一样:vjp_program 的导数程序就是 JAX 的 pull,只是余切变成了程序的最后一个输入。只给种子也行,它正好填进唯一的空位;但只要换了 x,就得把三个输入都给齐。

怎么确认没中招

  • 变换之后先打印 sem.derivative_input_indices():每个被选中的输出对应一个种子的输入位置,Some(2) 就是第 2 个输入。多个输出时种子按输出的顺序排在后面。
  • sem.frozen().program.inputs().len() 比原程序多几个,run 时就要多传几个张量(或者只传这几个)。
  • 拿一个已知梯度的小例子跑一次,再把种子换成 2,结果应当正好翻倍。

小结

  • vjp_program 的导数程序 = 原输入 + 每个选中输出一个种子;种子在最后,没有默认值。
  • 只传原输入报 expected N inputs, received M,什么都不传报 placeholder semantic input N has no runtime input binding。
  • 补上标量 1.0 就是 grad;换成别的值就是 jax.vjp 的 pull(c)。

💬 关联讨论

发起讨论
加载中...