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)。