同时要值和梯度:has no runtime input binding
主线位置:变换主线的「grad」到「编译」:一个程序同时输出函数值和梯度。
版本:JAX 0.12.0.1(CPU,x64)、tenferro-rs 0.8.0(CPU,native)
flowchart TB
y["值 y 和梯度 g"] --> cm["compile_many"]
cm --> ph{"x 是占位符"}
ph -->|是| err["未绑定报错"]
ph -->|否| ok["一个程序两个输出"]
ok --> runall["run 传满输入"]
y --> two["两个程序各带规格"]
实验代码:JAX 脚本一个,tenferro 程序 p11_value_and_grad(AdContext 带线代规则,CpuEngine::serial(),release)。函数是变换主线的 f(x) = Σ sin(x)·x²,n = 1000,x_i = −2 + 4i/999。平台 M2 Max、macOS 27、rustc 1.99.0。
现象
JAX 里一个 value_and_grad 加 jit,换输入直接调:
vg = jax.jit(jax.value_and_grad(f))
v, g = vg(x)
v2, g2 = vg(x + 0.5)
value=-0.000000000000028 grad[:3]=[1.9726023611141572, 1.9931199187584416, 2.0134492414958083]
x+0.5: value=855.495821415973865
(f 在这组对称的 x 上数学上是 0,两边给的都是舍入误差。)tenferro 的变换主线用占位符当 jit 的参数;值和梯度两个输出,用 compile_many:
let x = TracedTensor::input_concrete_shape(DType::F64, &[N])?;
let y = f(&x)?;
let g = ad.grad(&y, &x)?;
GraphCompiler::new().compile_many(&[&y, &g])
compile_many([y, g]) with a placeholder: Err: placeholder User { id: 0 } has no runtime input binding
UnboundPlaceholder { input_key: "User { id: 0 }" }
编译期就失败。单个输出时这里会换成 compile_with_input_specs(&g, &[(&x, DType::F64, &[N][..])]),但多个输出没有这个版本。
出处
GraphCompiler::compile_many 只收输出,不收规格;不带数据的占位符找不到 dtype 和形状,descriptor_for_input 返回 UnboundPlaceholder。0.8.0 带规格的入口只有单输出的 compile_with_input_specs。这条报错文字说「runtime」却出现在编译期,原因和规格的其他规则见 tenferro 卷的常见坑 has no runtime input binding。value_and_grad 的对照见 grad、vjp、jvp。
修法
一:两个程序,各带规格。 值和梯度分开编译,run 都只传 x:
let py = GraphCompiler::new().compile_with_input_specs(&y, &[(&x, DType::F64, &[N][..])])?;
let pg = GraphCompiler::new().compile_with_input_specs(&g, &[(&x, DType::F64, &[N][..])])?;
x: value=0.000000000000048 grad[:3]=[1.9726023611141572, 1.9931199187584416, 2.0134492414958083]
x+0.5: value=855.495821415973751 grad[:3]=[3.151643663564495, 3.1508730409643624, 3.1499634249544637]
梯度程序里已经含有前向,所以前向算了两遍。函数便宜时这不要紧。
二:带数据的叶子,一个程序两个输出。 用 from_tensor_concrete_shape 建 x,compile_many 就能编:
let xd = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![N], xs.clone())?)?;
let yd = f(&xd)?;
let gd = ad.grad(&yd, &xd)?;
let pm = GraphCompiler::new().compile_many(&[&yd, &gd])?;
inputs = 2, x index Some(0)
defaults: value=0.000000000000048 grad[:3]=[1.9726023611141572, 1.9931199187584416, 2.0134492414958083]
run [x+0.5]: Err: Runtime::run (Execution): invalid argument inputs: expected 2 inputs, received 1
run [x+0.5, seed=1]: value=855.495821415973751 grad[:3]=[3.151643663564495, 3.1508730409643624, 3.1499634249544637]
程序有两个输入:x 和 grad 的种子,都带默认值。只传新的 x 不行(第三行):程序里没有占位符,run 只接受「一个不传」或「全部传满」。全部传满时按位置覆盖,x 在第 0 位,种子传 1。位置怎么确认、顺序错了会怎样,见切向量、种子、常量都能被 run 的输入替换。
两种写法打印出来的值和梯度前三个元素完全相同。和 JAX 比:x + 0.5 处的值差 1.1e-13(量级 8.6e2,求和顺序和 sin 的实现不同),x 处梯度的前三个元素逐位相同。
怎么确认没中招
- 需要多个输出又想在
run时换输入:要么每个输出一个带规格的程序,要么用数据叶子并对input_count()心里有数。 - 用数据叶子时,换一组输入跑一次,确认输出确实变了(而不是悄悄用了默认值)。
小结
compile_many没有带规格的版本,输入是不带数据的占位符时报placeholder ... has no runtime input binding。- 分成两个带规格的程序最简单,代价是前向算两遍。
- 或者用带数据的叶子编成一个程序;之后
run要么不传输入,要么连grad的种子一起传满。