grad、vjp、jvp

g = jax.jit(jax.grad(f))(x)
print(jax.make_jaxpr(jax.grad(f))(x))

let ad = ad_context()?;            // AdContext，带 linalg 和 einsum 的规则
let x = TracedTensor::input_concrete_shape(DType::F64, &[N])?;
let y = f(&x)?;
let g = ad.grad(&y, &x)?;          // 也可以 use tenferro_ad::TracedTensorAdExt 后写 y.grad(&x)
let pg_prog = GraphCompiler::new().compile_with_input_specs(&g, &[(&x, DType::F64, &[N][..])])?;

# JAX
max|grad-exact|=4.441e-16 jit first 14691 us steady 10.6 us
{ lambda ; a:f64[1000]. let
    b:f64[1000] = sin a
    c:f64[1000] = cos a
    d:f64[1000] = integer_pow[y=2] a
    e:f64[1000] = mul 2.0:f64[] a
    f:f64[1000] = mul b d
    _:f64[] = reduce_sum[axes=(0,) out_sharding=None] f
    g:f64[1000] = broadcast_in_dim 1.0:f64[]
    h:f64[1000] = mul b g
    i:f64[1000] = mul g d
    j:f64[1000] = mul h e
    k:f64[1000] = mul i c
    l:f64[1000] = add_any j k
  in (l,) }
max|jax - tenferro|=4.441e-16, bitwise equal 761/1000
# tenferro
grad program inputs=2 ops=10: BroadcastInDim Sin Mul Mul Add Mul Mul Cos Mul Add
grad program: input index of x Some(0)
max|grad-exact|=4.441e-16  ext == ctx bitwise: true  steady 14.4 us

# JAX
TypeError: Gradient only defined for scalar-output functions. Output had shape: (1000,). To get the gradient, reduce the output to a scalar with output.sum(), or use jax.jacobian to differentiate a function with a non-scalar output.
# tenferro
Err: grad requires a scalar output, got shape [1000]
  NonScalarGrad { shape: [1000] }

y, pullback = jax.vjp(h, x)          # h(x) = sin(x) * x**2，向量输出
(ct,) = pullback(v)
jax.jit(lambda x, u: jax.vjp(h, x)[1](u)[0])(x, v)

let ct = TracedTensor::from_vec_col_major(vec![N], vs.clone())?;   // 带数据
let hx = h(&x)?;
let vj = ad.vjp(&hx, &x, &ct)?;
let ct_ph = TracedTensor::input_concrete_shape(DType::F64, &[N])?; // 占位符
ad.vjp(&hx, &x, &ct_ph)

# JAX
max|vjp - g*v|=8.882e-16
jit(vjp) with the cotangent as an argument: first 11171 us steady 10.9 us
# tenferro
vjp program inputs=2 ops=9: Sin Mul Mul Add Mul Mul Cos Mul Add
max|vjp - g*v|=8.882e-16
vjp with placeholder cotangent: Err: vjp (GraphBuild): invalid argument cotangent: vjp cotangent must have concrete tensor data

jit_jvp = jax.jit(lambda x, t: jax.jvp(f, (x,), (t,))[1])
jit_jvp(x, v); jit_jvp(x, 2 * v)

let vc = TracedTensor::from_vec_col_major(vec![N], vs.clone())?;
let j = ad.jvp(&y, &x, &vc)?;
let pj_prog = GraphCompiler::new().compile_with_input_specs(&j, &[(&x, DType::F64, &[N][..])])?;
run1(&rt, &mut eng, &pj, &[&xt]);          // 切向量用默认值 v
run1(&rt, &mut eng, &pj, &[&xt, &v2t]);    // 两个输入都给：切向量换成 2v
ad.jvp(&y, &x, &v_placeholder)             // 占位符

# JAX
jvp=1.879649926237039 |.-exact|=4.885e-15 jit first 17463 us steady 9.8 us
same jit, tangent 2v: 3.759299852474078 (42 us, no retrace: cache size 1)
# tenferro
jvp transform 58.5 us; program inputs=2 ops=10: Cos Mul Mul Mul Mul Add Sin Mul Add ReduceSum
jvp=1.879649926237050 |.-exact|=8.216e-15 steady 14.0 us
jvp program run with [x, 2v]: 3.759299852474099; steady 13.8 us
jvp with placeholder tangent: Err: jvp (GraphBuild): invalid argument tangent: jvp tangent must have concrete tensor data

GraphCompiler::new().compile_many(&[&y, &g])                 // x 是占位符
let xd = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![N], xs.clone())?)?;
let yd = f(&xd)?;                                            // x 是带数据的叶子
let gd = ad.grad(&yd, &xd)?;
let pm_prog = GraphCompiler::new().compile_many(&[&yd, &gd])?;
run_all(&rt, &mut eng, &pm, &[]);              // 默认值
run_all(&rt, &mut eng, &pm, &[&x2]);           // 换 x
run_all(&rt, &mut eng, &pm, &[&x2, &seed]);    // 换 x，种子照写 1

# JAX
value=-0.000000000000028 max|grad-exact|=4.441e-16 jit first 18840 us steady 12.5 us
# tenferro
compile_many with a placeholder: Err: placeholder User { id: 0 } has no runtime input binding
compile_many(value, grad) with a data leaf: inputs=2 ops=12: Sin Mul Mul ReduceSum BroadcastInDim Mul Mul Add Mul Cos Mul Add  (x index Some(0))
defaults: value=0.000000000000048 max|grad-exact|=4.441e-16 steady 16.3 us
run with [x+0.5]: Err: Runtime::run (Execution): invalid argument inputs: expected 2 inputs, received 1
run with [x+0.5, seed=1]: value=855.495821415973751 max|grad-exact|=1.332e-15

f2 = lambda x, w: jnp.sum(jnp.sin(x) * w * x)
gx, gw = jax.grad(f2, argnums=(0, 1))(x, v)

let w = TracedTensor::input_concrete_shape(DType::F64, &[N])?;
let y2 = (&(&x.sin()? * &w)? * &x)?.reduce_sum(None)?;
let one = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64])?;
let gs = ad.vjp_many(&y2, &[&x, &w], &one)?;          // Vec<Option<TracedTensor>>
// 每个梯度各自 compile_with_input_specs，规格先写 [x, w]

# JAX
d/dx max|.-exact|=2.220e-16  d/dw max|.-exact|=0.000e+00
# tenferro
d/dx: ops=9 max|.-exact|=4.441e-16
d/dw: specs [x, w]: Err: GraphCompiler::compile_with_input_specs (Compile): invalid argument bindings: binding 1 declares a placeholder the output does not depend on; remove it from the bindings
d/dw: ops=4 max|.-exact|=2.220e-16

let mut rules = tenferro_linalg::semantic_ad_rules()?;
rules.merge(tenferro_einsum::semantic_ad_rules()?)?;
let ad = AdContext::builder().with_semantic_extension_rules(rules)?.build()?;

let (_, s, _) = a.svd()?;            // a 是 2x2 的 [[4, 1], [1, 3]]
let ls = s.reduce_sum(None)?;
ls.grad(&a)                          // TracedTensorAdExt
ad.grad(&ls, &a)                     // AdContext

TracedTensorAdExt::grad: Err: unsupported vjp AD rule for tenferro-linalg.linalg.v1
AdContext::grad: [1.0000000000000007, 5.551115123125783e-17, 1.1102230246251565e-16, 1.0000000000000002]

vjp program run with [x, ones]: max|.-grad|=4.441e-16