切向量、种子、常量都能被 run 的输入替换

主线位置：变换主线的「grad / vjp / jvp」到「执行」：变换产生的程序在 run 时接收哪些输入。

版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native）

flowchart TB
  r["run 的输入个数"] --> all{"等于程序输入数"}
  all -->|是| pos["全部按位置覆盖"]
  pos --> swap["顺序错且同形状"]
  swap --> bad["静默算错"]
  all -->|否| ph["只填占位符"]
  ph --> dft["其余用默认值"]

实验代码：JAX 脚本一个，tenferro 程序 p10_default_inputs（AdContext 带线代规则，CpuEngine::serial()，release）。平台 M2 Max、macOS 27、rustc 1.99.0。

现象

JAX 里有两条很清楚的规矩。切向量想在调用时换，就把它写成 jit 函数的参数；闭包里的数组是常量，jit 之后改不了，多传参数直接报错：

jv = jax.jit(lambda x, v: jax.jvp(lambda x: x**3, (x,), (v,))[1])
jv(2.0, 5.0)
w = jnp.full(3, 2.0)
fw = jax.jit(lambda x: jnp.sum(jnp.sin(x) * w))
fw(x3); w = jnp.full(3, 3.0); fw(x3); fw(x3, w)

jvp of x^3 at x=2, v=5: 60.0  (3 x^2 v = 60)
closure w=2: 1.1880459082064578
after rebinding w=3, same jitted fn: 1.1880459082064578
fw(x, w): TypeError: <lambda>() takes 1 positional argument but 2 were given

tenferro 的 jvp 要求切向量建图时就带数据（jvp tangent must have concrete tensor data），于是写一个值为 1 的切向量，编译时只给 x 规格：

let x = TracedTensor::input_concrete_shape(DType::F64, &[])?;
let y = (&(&x * &x)? * &x)?;
let v0 = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64])?;
let j = ad.jvp(&y, &x, &v0)?;
let p = GraphCompiler::new().compile_with_input_specs(&j, &[(&x, DType::F64, &[][..])])?;

program inputs = 2
input index: x Some(0), tangent None
run [x=2]          -> [12.0]  (3 x^2 * 1 = 12)
run [x=2, v=5]     -> [60.0]  (3 x^2 v = 60)
run [v=5, x=2]     -> [150.0]

程序有两个输入：占位符 x，和切向量。只传 x 时切向量用建图时的 1。
传满两个输入时，切向量被覆盖：jit(jvp) 把切向量当参数在 tenferro 里也做得到。
把顺序写反，结果是 3·5²·2 = 150，没有任何报错。
切向量在程序里的位置查不到：v0.input_key() 在程序里找不到对应的下标（None），因为变换给切向量造了一个新的输入。

捕获的常量也一样，它不是冻结的：

let w = TracedTensor::from_vec_col_major(vec![3], vec![2.0_f64; 3])?;
let yw = (&xv.sin()? * &w)?.reduce_sum(None)?;     // xv 是 [3] 的占位符

program inputs = 2
run [x]       -> [1.1880459082064578]
run [x, w=3]  -> [1.7820688623096868]
run [w=3, x]  -> [0.08467200483592033]

第三行把 w 和 x 的位置弄反了，两者都是长度 3 的 F64，算出一个完全不同的数。

出处

resolve_input_reads 决定 run 的每个输入从哪来：

    let resolved = if inputs.len() == expected {
        inputs.iter().cloned().collect()
    } else {
        let mut explicit = inputs.iter();
        let mut resolved = RuntimeInputReads::new();
        for value in program.program().inputs() {
            if let Some(retained) = program.bindings().tensor_ref_for_input(*value) {
                resolved.push(retained.tensor_read()?);
            } else {
                let read = explicit.next().ok_or_else(|| {
// ...

个数等于程序的输入数时，所有输入按位置原样使用，带默认值的那些也被覆盖；否则只给没有默认值的占位符填数，其余用建图时留下的张量。之后只核对形状和 dtype，不核对「这个张量是不是你以为的那个输入」。带数据的叶子（from_vec_col_major 建的常量）、jvp 的切向量、vjp 的余切、grad 的种子在程序里都是「有默认值的输入」，这一点和 JAX 的冻结常量正好相反。变换产生的这些输入是什么、怎么替换，见 grad、vjp、jvp；捕获常量的部分见记图与 jit 的「闭包里的常量」。

📝NOTE

没有蠢问题

问： [v=5, x=2] 明明传反了，为什么不报错？

答：两个输入都是 F64 标量，形状和 dtype 都对得上，run 能核对的就只有这两样。程序里的输入没有名字，只有位置。占位符的位置按 compile_with_input_specs 的顺序，其余的由编译器和变换决定，而切向量的位置连隐藏的 input_key_index 都查不到。只有换位置的两个输入形状或 dtype 不同时，run 才会报错。

修法

想在 run 时换的东西，都做成占位符，写进 compile_with_input_specs，run 时只传占位符。规格的顺序就是输入的顺序，带默认值的输入一个都不传，就不存在覆盖错位置的问题。

常量直接换成占位符：

let wp = TracedTensor::input_concrete_shape(DType::F64, &[3])?;
let yw2 = (&xv.sin()? * &wp)?.reduce_sum(None)?;
let pw2 = GraphCompiler::new().compile_with_input_specs(&yw2, &[(&xv, DType::F64, &[3][..]), (&wp, DType::F64, &[3][..])])?;

切向量不能是占位符，但可以把它挪进函数里：对 x + t·v 在 t = 0 处关于标量 t 求导，切向量是常数 1，v 变成普通的占位符：

let xp = TracedTensor::input_concrete_shape(DType::F64, &[])?;
let vp = TracedTensor::input_concrete_shape(DType::F64, &[])?;
let t = TracedTensor::from_vec_col_major(vec![], vec![0.0_f64])?;
let xt = (&xp + &(&t * &vp)?)?;
let yt = (&(&xt * &xt)? * &xt)?;
let one = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64])?;
let jt = ad.jvp(&yt, &t, &one)?;
let pt = GraphCompiler::new().compile_with_input_specs(&jt, &[(&xp, DType::F64, &[][..]), (&vp, DType::F64, &[][..])])?;

tangent as placeholder: inputs = 4, run [x=2, v=5] -> [60.0]
run [x=2] -> Err: Runtime::run (Execution): invalid argument inputs: expected 4 inputs, or one explicit input for each unbound placeholder
w as placeholder: inputs = 2, run [x, w=3] -> [1.7820688623096868]

程序有 4 个输入（x、v、t 和切向量 1），run 只传两个占位符，按规格顺序填；t 和切向量用默认值。少传一个时报错，不会静默错位（第二行）。

怎么确认没中招

对每个程序打印 input_count()，和你打算传的个数比。传的个数等于它时，每一个带默认值的输入都会被覆盖，要确定顺序。
用一组不对称的值测一次顺序（例如 x = 2、v = 5，正确是 60，反了是 150）。
不要依赖 input_key_index 去找切向量、余切、种子的位置：它找不到变换新造的输入，而且是 #[doc(hidden)] 的 API。

小结

tenferro 程序里，捕获的常量、jvp 的切向量、vjp 的余切、grad 的种子都是带默认值的输入，不是 JAX 那样冻结的常量。
run 传满全部输入时按位置覆盖默认值；只传占位符时其余用默认值。顺序错、形状又相同时静默算错。
要在调用时换的东西做成占位符并列进规格，run 只传占位符；切向量用 x + t·v 的写法变成占位符。

切向量、种子、常量都能被 run 的输入替换

现象

出处

修法

怎么确认没中招

小结

💬 关联讨论

登录