traced AD

impl TracedTensorAdExt for TracedTensor {
    fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor> {
        let rules = SemanticExtensionRuleSet::default();
        grad_with_optional_rules(self, wrt, &rules, None)
    }
    // ...
}

    if output.rank != 0 {
        return Err(Error::NonScalarGrad {
            shape: error_shape_hint(output),
        });
    }

    let ones = ones_tensor(output.dtype, vec![])?;
    let seed = TracedTensor::from_tensor_concrete_shape(ones)?;
    let wrt_input_key = leaf_input_key(wrt)?;
    vjp_many_with_transform_and_cache(output, &[wrt], &seed, rules, "grad", ad_transform_cache)?
        .pop()
        .flatten()
        .ok_or_else(|| inactive_wrt_error("grad", &wrt_input_key))

traced grad wrt intermediate: Err(Validation { op: "ad_support::leaf_input_key", phase: GraphBuild, source: InvalidArgument { argument: "tensor", message: "expected traced leaf input, got Derived { operation: OperationKey { operation: Mul, inputs: [Input(User { id: 0 }), Input(User { id: 0 })], role: Primary, fingerprint: 8852925308592293207 }, output_slot: 0 }" } })
traced grad of non-scalar:    Err(NonScalarGrad { shape: [3] })
traced grad, unrelated leaf:  Err(Validation { op: "grad", phase: GraphBuild, source: InvalidArgument { argument: "wrt", message: "grad is inactive for User { id: 2 }; use grad_optional to observe the inactive state" } })
traced grad_optional:         Ok(false)

    let requested_keys = wrts
        .iter()
        .map(|wrt| leaf_input_key(wrt))
        .collect::<Result<Vec<_>>>()?;
    let mut compiler = GraphCompiler::new();
    let source = compile_ad_source(&mut compiler, output)?;
    let source_indices: Vec<_> = requested_keys
        .iter()
        .map(|key| source.input_key_index(key))
        .collect();
    let mut active_inputs = vec![false; source.input_count()];
    for source_index in source_indices.iter().flatten().copied() {
        active_inputs[source_index] = true;
    }
    if !active_inputs.iter().any(|active| *active) {
        return Ok(vec![None; wrts.len()]);
    }

    let active_outputs = vec![true; source.output_count()];
    let derivative = semantic_vjp_with_cache(
        source.frozen_program(),
        &active_inputs,
        &active_outputs,
        rules,
        ad_transform_cache,
    )?;

    for operation in operations.into_iter().rev() {
        let cotangent_outputs: Vec<_> = operation
            .outputs()
            .iter()
            .map(|value| {
                cotangents
                    .get(value)
                    .copied()
                    .map_or(AdValue::Absent, AdValue::Value)
            })
            .collect();
        if cotangent_outputs
            .iter()
            .all(|value| matches!(value, AdValue::Absent))
        {
            continue;
        }
        let active_operation_inputs: Vec<_> = operation
            .inputs()
            .iter()
            .map(|value| forward_active.contains(value))
            .collect();
        if active_operation_inputs.iter().all(|active| !active) {
            continue;
        }
        let cotangent_inputs = match operation.op() {
            SemanticOpRef::Extension(op) => {
                // ... 扩展运算族：查规则集
            }
            SemanticOpRef::Core(op) => vjp_core(
                op,
                &mapped_values(operation.inputs(), &values),
                &mapped_values(operation.outputs(), &values),
                &cotangent_outputs,
                &active_operation_inputs,
                &mut builder,
            )?,
            // ...
        };
        for (source, cotangent) in operation.inputs().iter().copied().zip(cotangent_inputs) {
            if let AdValue::Value(cotangent) = cotangent {
                accumulate_cotangent(&mut builder, &mut cotangents, source, cotangent)?;
            }
        }
    }

    for (input_index, input) in frozen.program.inputs().iter().copied().enumerate() {
        let key = if input_index < source.input_keys().len() {
            source.input_keys()[input_index].clone()
        } else {
            allocate_input_key()
        };
        let local = builder.add_input(key.clone());
        value_map.insert(input, local);
        input_keys.push(key);
    }

== 1. program sizes
L          inputs=2 ops= 3: DotGeneral Mul ReduceSum
dL/dA      inputs=3 ops= 5: BroadcastInDim DotGeneral Mul Add DotGeneral
vjp_many A inputs=3 ops= 5: BroadcastInDim DotGeneral Mul Add DotGeneral
compile_many(dA, dB) ops=6: BroadcastInDim DotGeneral Mul Add DotGeneral DotGeneral
raw vjp_program(A) inputs=3 ops=6: DotGeneral BroadcastInDim Mul Mul Add DotGeneral

        if !seen.insert(key.clone()) {
            return Err(Error::invalid_argument(
                "jvp",
                ErrorPhase::GraphBuild,
                "wrt_tangents",
                "wrt_tangents contains duplicate wrt leaves",
            ));
        }
        let tangent_data = tangent.attached_value().cloned().ok_or_else(|| {
            Error::invalid_argument(
                "jvp",
                ErrorPhase::GraphBuild,
                "tangent",
                "jvp tangent must have concrete tensor data",
            )
        })?;

// f(x) = sum(tanh(M.x)); x is 3x1
let f = f_tanh(&m, &x)?;
let g = f.grad(&x)?;
let hv = g.jvp(&x, &v)?;
// 对照：在 x ± 1e-5·v 处各求一次梯度，中心差分

let cube = (&(&s * &s)? * &s)?; // s = 2.5
let d1 = cube.grad(&s)?;
let d2 = d1.grad(&s)?;
let d3 = d2.grad(&s)?;

f          inputs=2 ops= 3: DotGeneral Tanh ReduceSum
grad f     inputs=3 ops= 9: BroadcastInDim Constant BroadcastInDim DotGeneral Tanh Mul Sub Mul DotGeneral
hvp     = [-0.9930593881003207, 1.1987671592011429, -0.2074128493294373]
hvp_fd  = [-0.993059387921491, 1.1987671584723358, -0.20741284927350276]
max |hvp - fd| = 7.288e-10
== 5. second derivative of x^3 at 2.5 by grad(grad)
f' = [18.75], f'' = [15.0], f''' = [6.0]
x^3        inputs=1 ops= 2: Mul Mul
f'         inputs=2 ops= 6: Mul Mul Mul Mul Add Add
f''        inputs=3 ops=10: Add Mul Mul Mul Mul Add Mul Mul Add Add
f'''       inputs=3 ops=10: Add Mul Mul Mul Add Mul Add Mul Mul Add

pub(crate) struct SemanticAdTransformCacheKey {
    kind: SemanticAdTransformKind,
    input_fingerprint: SemanticFingerprint,
    input_metadata: Box<[ProgramValueMetadata]>,
    active_inputs: Box<[bool]>,
    active_outputs: Box<[bool]>,
}

TracedTensorAdExt::grad  first    370.8us  later min   156.0us median   175.3us
AdContext::grad          first    261.7us  later min    77.5us median    86.0us
compile(grad)            first     95.0us  later min    67.7us median    75.5us
AdContextCacheStats { ad_transforms: CacheStats { entries: 1, retained_bytes: 1155, hits: 49, misses: 1, evictions: 0, clears: 0 } }

grad f     inputs=3 ops= 9: BroadcastInDim Constant BroadcastInDim DotGeneral Tanh Mul Sub Mul DotGeneral
hvp        inputs=4 ops=14: DotGeneral Constant BroadcastInDim DotGeneral Tanh Mul Sub Mul Mul Add Neg BroadcastInDim Mul DotGeneral