primitive AD 规则

$ git grep -nE 'ad::linearize|ad::transpose_rule|primitive_residual_spec' v0.8.0 -- crates ':!crates/tenferro-internal-ops/src/ad'
v0.8.0:crates/tenferro-ad/src/semantic_transform.rs:647:        _ => tenferro_ops::ad::primitive_residual_spec(op).unwrap_or_default(),
v0.8.0:crates/tenferro-internal-ops/src/std_tensor_op.rs:487:        crate::ad::linearize(self, builder, primal_in, primal_out, tangent_in, ctx)
v0.8.0:crates/tenferro-internal-ops/src/std_tensor_op.rs:510:        crate::ad::transpose_rule(self, builder, cotangent_out, inputs.as_slice(), mode, ctx)

pub(crate) trait PrimitiveAdRule: Send + Sync {
    fn kind(&self) -> PrimitiveOpKind;

    /// Declare which primal inputs this rule's transpose reads as tensor
    /// residuals. Inputs not declared may only be accessed through metadata.
    fn residual_mask(&self) -> ResidualSpec;

    fn linearize(
        &self,
        op: &StdTensorOp,
        builder: &mut dyn PrimitiveRuleBuilder,
        primal_in: &[ValueKey<StdTensorOp>],
        primal_out: &[ValueKey<StdTensorOp>],
        tangent_in: &[Option<LocalValueId>],
        ctx: &mut ShapeGuardContext,
    ) -> ADRuleResult<Vec<Option<LocalValueId>>>;

    fn transpose_rule(
        &self,
        op: &StdTensorOp,
        builder: &mut dyn PrimitiveRuleBuilder,
        cotangent_out: &[Option<LocalValueId>],
        inputs: &[TransposeInputRef<'_>],
        mode: &OperationRole,
        ctx: &mut ShapeGuardContext,
    ) -> ADRuleResult<Vec<Option<LocalValueId>>>;
}

static PRIMITIVE_AD_RULES: [&'static dyn PrimitiveAdRule; PrimitiveOpKind::COUNT] = [
    &FunctionPrimitiveAdRule {
        kind: PrimitiveOpKind::Add,
        linearize: linearize_add,
        transpose_rule: transpose_add,
        residual_mask: ResidualSpec::none(),
    },
    // ...
    &FunctionPrimitiveAdRule {
        kind: PrimitiveOpKind::Erf,
        linearize: linearize_erf,
        transpose_rule: transpose_erf,
        residual_mask: ResidualSpec::input(0),
    },
    // ...
];

residual_spec Add        = Some(ResidualSpec { input_mask: 0, output_mask: 0 })
residual_spec Mul        = Some(ResidualSpec { input_mask: 18446744073709551615, output_mask: 0 })
residual_spec Exp        = Some(ResidualSpec { input_mask: 1, output_mask: 0 })
residual_spec Tanh       = Some(ResidualSpec { input_mask: 1, output_mask: 0 })
residual_spec Erf        = Some(ResidualSpec { input_mask: 1, output_mask: 0 })
residual_spec Maximum    = Some(ResidualSpec { input_mask: 18446744073709551615, output_mask: 0 })
residual_spec Clamp      = Some(ResidualSpec { input_mask: 18446744073709551615, output_mask: 0 })
residual_spec ReduceSum  = Some(ResidualSpec { input_mask: 1, output_mask: 0 })
residual_spec DotGeneral = Some(ResidualSpec { input_mask: 18446744073709551615, output_mask: 0 })

    fn check_declared(&self, op: &str, use_case: &str) {
        debug_assert!(
            self.mask.declares_input(self.index),
            "transpose rule {op} accessed input {} as {use_case} but its residual mask does \
             not declare input {}; declare it in the rule's residual mask",
            self.index,
            self.index
        );
    }

// These direct semantic VJPs consume outputs, unlike the lower-level linear
// transpose adapters, whose residual contract exposes only their inputs.
pub(crate) fn eager_core_residual_spec(
    op: &tenferro_ops::std_tensor_op::StdTensorOp,
) -> crate::semantic_extension::ResidualSpec {
    use tenferro_ops::std_tensor_op::StdTensorOp;
    match op {
        StdTensorOp::Exp | StdTensorOp::Tanh => crate::semantic_extension::ResidualSpec::output(0),
        _ => tenferro_ops::ad::primitive_residual_spec(op).unwrap_or_default(),
    }
}

    if let Some(dx) = tangent_in[0] {
        let dx = convert_linear_to_dtype(builder, dx, lhs_dtype, output_dtype);
        let rhs = convert_fixed_ref_to_dtype(
            builder,
            ValueRef::External(primal_in[1].clone()),
            rhs_dtype,
            output_dtype,
        );
        let term = builder.add_operation(
            StdTensorOp::Mul,
            vec![ValueRef::Local(dx), rhs],
            OperationRole::Linearized {
                active_mask: vec![true, false],
            },
        );
        terms.push(term[0]);
    }
    // ... dy 一项对称
    match terms.as_slice() {
        [] => vec![None],
        [only] => vec![Some(*only)],
        [lhs, rhs] => {
            let sum = builder.add_operation(
                StdTensorOp::Add,
                vec![ValueRef::Local(*lhs), ValueRef::Local(*rhs)],
                OperationRole::Linearized {
                    active_mask: vec![true, true],
                },
            );
            vec![Some(sum[0])]
        }
        _ => unreachable!("mul linearization creates at most two terms"),
    }

    let active_mask = match mode {
        OperationRole::Linearized { active_mask } => active_mask,
        OperationRole::Primary => return Ok(vec![None, None]),
    };
    // ...
    if active_mask[0] {
        let rhs = inputs[1].fixed_value("mul", 1)?;
        let rhs_conj = conjugate_primal_if_complex(builder, rhs, ctx)?;
        let rhs_conj = convert_fixed_ref_to_dtype(builder, rhs_conj, rhs_dtype, output_dtype);
        let out = builder.add_operation(
            StdTensorOp::Mul,
            vec![rhs_conj, ValueRef::Local(ct)],
            OperationRole::Linearized {
                active_mask: vec![false, true],
            },
        );
        // ...
    }

== linearize mul, both tangents: returned [Some(4)]
  op0: Mul            inputs=["L0", "x2"] role=Linearized { active_mask: [true, false] }
  op1: Mul            inputs=["x1", "L1"] role=Linearized { active_mask: [false, true] }
  op2: Add            inputs=["L2", "L3"] role=Linearized { active_mask: [true, true] }
== transpose mul, active_mask [true,false]: returned [Some(1), None]
  op0: Mul            inputs=["x2", "L0"] role=Linearized { active_mask: [false, true] }
== transpose mul, role Primary: returned [None, None]

        CoreSemanticOp::Mul => {
            let rhs_coefficient = conjugate_if_complex(builder, primal_inputs[1])?;
            let lhs_coefficient = conjugate_if_complex(builder, primal_inputs[0])?;
            let lhs = multiply_ad_value(builder, cotangent, rhs_coefficient)?;
            let rhs = multiply_ad_value(builder, cotangent, lhs_coefficient)?;
            vec![
                normalize_ad_value(builder, lhs, active_inputs[0], primal_inputs[0])?,
                normalize_ad_value(builder, rhs, active_inputs[1], primal_inputs[1])?,
            ]
        }

    match tangent_in[0] {
        Some(dx) => {
            let out = builder.add_operation(
                op.clone(),
                vec![ValueRef::Local(dx)],
                OperationRole::Linearized {
                    active_mask: vec![true],
                },
            );
            vec![Some(out[0])]
        }
        None => vec![None],
    }

            let input_key = ValueRef::External(input.key().clone());
            let input_rank = ctx.rank_of(&input_key)?;
            let kept_dims = kept_dims(input_rank, axes);
            // ... 全部轴都被求和时，先 reshape 成标量
            let (shape, shape_sources) = input.shape_operand(input_rank, 1, ctx)?;
            let (op_inputs, active_mask) =
                linearized_inputs_with_inactive_shape_sources(cotangent, shape_sources);
            let out = builder.add_operation(
                StdTensorOp::BroadcastInDim {
                    shape,
                    dims: kept_dims,
                },
                op_inputs,
                OperationRole::Linearized { active_mask },
            );

== transpose reduce_sum axes [0] of [2,3]: returned [Some(1)]
  op0: BroadcastInDim { shape: [Const(2), Const(3)], dims: [1] } role=Linearized { active_mask: [true] }

    let new_config = DotGeneralConfig {
        lhs_contracting_dims: ct_rhs_free_positions,
        rhs_contracting_dims: rhs_free.into(),
        lhs_batch_dims: (lhs_free.len() + rhs_free.len()..output_rank).collect(),
        rhs_batch_dims: config.rhs_batch_dims.clone(),
    };
    Ok((
        new_config,
        output_rank,
        rhs_rank,
        permutation_to_original_order(lhs_rank, &result_order),
    ))

== transpose dot_general [2,3].[3,4]: returned [Some(1), Some(2)]
  op0: DotGeneral { config: DotGeneralConfig { lhs_contracting_dims: [1], rhs_contracting_dims: [1], lhs_batch_dims: [], rhs_batch_dims: [] } }
  op1: DotGeneral { config: DotGeneralConfig { lhs_contracting_dims: [0], rhs_contracting_dims: [0], lhs_batch_dims: [], rhs_batch_dims: [] } }

        CoreSemanticOp::Erf => {
            // d erf(x) = 2/sqrt(pi) * exp(-x^2)
            let square = builder.add_op(CoreSemanticOp::Mul, &[primal_input, primal_input])?[0];
            let negated = builder.add_op(CoreSemanticOp::Neg, &[square])?[0];
            let gaussian = builder.add_op(CoreSemanticOp::Exp, &[negated])?[0];
            let scale = real_float_constant_like(
                builder,
                primal_input,
                std::f64::consts::FRAC_2_SQRT_PI,
                role,
            )?;
            builder.add_op(CoreSemanticOp::Mul, &[scale, gaussian])?[0]
        }

== linearize erf: returned [Some(7)]
  op0: Constant       inputs=[] role=Primary
  op1: BroadcastInDim inputs=["L1", "x7"] role=Primary
  op2: Mul            inputs=["x7", "x7"] role=Linearized { active_mask: [false, false] }
  op3: Neg            inputs=["L3"] role=Linearized { active_mask: [false] }
  op4: Exp            inputs=["L4"] role=Linearized { active_mask: [false] }
  op5: Mul            inputs=["L2", "L5"] role=Linearized { active_mask: [false, false] }
  op6: Mul            inputs=["L6", "L0"] role=Linearized { active_mask: [false, true] }

pub fn gelu<O: CompositeOps>(ops: &mut O, x: &O::Value) -> Result<O::Value, O::Error> {
    let dtype = ops.dtype(x);
    require_real_float("gelu", dtype)?;
    let shape = ops.shape(x)?;
    let scaled = scale_by(ops, x, std::f64::consts::FRAC_1_SQRT_2, &shape)?;
    let erf = ops.unary(CompositeUnary::Erf, &scaled)?;
    drop(scaled);
    let one = splat(ops, dtype, 1.0, &shape)?;
    let gate = ops.binary(CompositeBinary::Add, &one, &erf)?;
    drop((one, erf));
    let gated = ops.binary(CompositeBinary::Mul, x, &gate)?;
    drop(gate);
    scale_by(ops, &gated, 0.5, &shape)
}

erf        inputs=1 ops= 1: Erf
d erf      inputs=3 ops= 9: BroadcastInDim Mul Constant BroadcastInDim Mul Neg Exp Mul Mul
sigmoid    inputs=1 ops= 9: Neg Compare Constant BroadcastInDim Select Exp Select Add Div
d sigmoid  inputs=3 ops=24: Neg Compare Constant BroadcastInDim BroadcastInDim Mul Constant BroadcastInDim Select Exp Select Mul Add Mul Div Neg Div Select Add Mul Select Select Neg Add
gelu       inputs=1 ops=11: Constant BroadcastInDim Constant BroadcastInDim Mul Erf Add Mul Constant BroadcastInDim Mul
d gelu     inputs=3 ops=23: BroadcastInDim Mul Constant BroadcastInDim Mul Constant BroadcastInDim Constant BroadcastInDim Mul Erf Add Mul Mul Constant BroadcastInDim Mul Neg Exp Mul Mul Mul Add
softmax    inputs=1 ops=11: ReduceMax Constant Maximum BroadcastInDim Sub Exp ReduceSum Constant Maximum BroadcastInDim Div
d softmax  inputs=3 ops=45: BroadcastInDim Mul ReduceMax Constant Maximum BroadcastInDim Sub Exp ReduceSum Constant Maximum BroadcastInDim Div Compare Compare Mul Mul Div Neg ReduceSum Constant Select Constant Add Div Select BroadcastInDim Add Mul BroadcastInDim Compare Convert ReduceSum BroadcastInDim Div Compare Compare Neg ReduceSum Select Div Select BroadcastInDim Mul Add
layer_norm inputs=1 ops=12: ReduceSum Constant Mul BroadcastInDim Sub ReduceSumSquares Mul Constant Add Rsqrt BroadcastInDim Mul
d layer_norm inputs=3 ops=30: BroadcastInDim Mul ReduceSum Constant Mul BroadcastInDim Sub ReduceSumSquares Mul Constant Add Rsqrt BroadcastInDim Mul Mul ReduceSum Neg Add Div Mul Mul BroadcastInDim Mul Add Add Neg ReduceSum Mul BroadcastInDim Add

erf         max|traced-fd|=3.1e-11 max|eager-traced|=0.0e0 grad=[0.03567908676708881, -1.084134787104863, 2.256758334191025, 0.2974928931287344]
sigmoid     max|traced-fd|=4.7e-11 max|eager-traced|=0.0e0 grad=[0.04474393562109985, -0.24751657271185992, 0.5, 0.10695484826014721]
silu        max|traced-fd|=2.8e-11 max|eager-traced|=0.0e0 grad=[-0.012388246289742881, -0.4006626881451501, 1.0, 0.430551119171924]
softplus    max|traced-fd|=1.0e-10 max|eager-traced|=0.0e0 grad=[0.054727657141906895, -0.4501660026875221, 1.0, 0.34498724056380625]
gelu        max|traced-fd|=3.0e-11 max|eager-traced|=0.0e0 grad=[-0.03824075766899386, -0.34253175176580575, 1.0, 0.5099489218128948]
gelu_tanh   max|traced-fd|=2.6e-11 max|eager-traced|=0.0e0 grad=[-0.03831323794542993, -0.34254185080430055, 1.0, 0.5097918655286516]
softmax     max|traced-fd|=4.2e-11 max|eager-traced|=0.0e0 grad=[-0.013242187672278883, -0.29800356572608855, 0.3390214906371875, -0.027775737238820097]
log_softmax max|traced-fd|=3.6e-10 max|eager-traced|=0.0e0 grad=[0.20588318064150288, -1.345342531639056, 1.5781976793460415, -0.43873832834848847]
layer_norm  max|traced-fd|=7.9e-11 max|eager-traced|=0.0e0 grad=[0.13252838602003836, -1.762111528752393, 1.8216079804071303, -0.19202483767477602]

    let AdValue::Value(active) = active else {
        return Ok(AdValue::Absent);
    };
    let zero = zero_constant_like(builder, active, role)?;
    let selected = builder.add_op(CoreSemanticOp::Select, &[self_eq_output, active, zero])?[0];
    let one = one_like(builder, active, role)?;
    let two = builder.add_op(CoreSemanticOp::Add, &[one, one])?[0];
    let half = builder.add_op(CoreSemanticOp::Div, &[selected, two])?[0];
    Ok(AdValue::Value(
        builder.add_op(CoreSemanticOp::Select, &[other_eq_output, half, selected])?[0],
    ))

traced d/dx  = [-inf, 0.0]
traced d/dlo = [0.0, -1.0]
clamp d/dlo  = [0.0, -1.0], d/dhi = [0.0, 0.0]
eager  x.grad = [-inf, 0.0], lo.grad = [0.0, -1.0]
tie: d/dx and d/dlo of sum(maximum(x, lo)) at x = lo = [1, 1]
  [0.5, 0.5] [0.5, 0.5]

// L(A, B) = sum((A.B) * (A.B)); A is 2x3, B is 3x2 (column-major data)
fn loss(a: &TracedTensor, b: &TracedTensor) -> R<TracedTensor> {
  let c = a.matmul(b)?;
  Ok((&c * &c)?.reduce_sum(None)?)
}
// 对 A、B 的每个元素：用 ±1e-6 扰动后的数据重建图、编译、prepare、run，中心差分，
// 与 l.grad(&a) 和 AdContext::vjp_many(&l, &[&a, &b], ..) 的结果比

grad A      = [5.07, -6.91, 4.24, -4.995, -15.719999999999999, 18.86]
vjp_many A  = [5.07, -6.91, 4.24, -4.995, -15.719999999999999, 18.86]
vjp_many B  = [10.75, 10.0625, -16.125, -4.35, -5.499999999999999, 6.524999999999999]
max |grad - fd| A = 2.162e-9, B = 2.313e-9
== 3. jvp vs directional difference
jvp_A      inputs=3 ops= 5: DotGeneral DotGeneral Mul Add ReduceSum
jvp = -13.575000000000, fd = -13.575000000543, <grad, dA> = -13.575000000000