semantic AD 与扩展

    // AD rules are registered separately in `tenferro-ad`.

    let Some((prefix, version)) = family_id.rsplit_once('.') else {
        return false;
    };
    let Some(version) = version.strip_prefix('v') else {
        return false;
    };
    let Some((crate_name, op_name)) = prefix.split_once('.') else {
        return false;
    };

/// Definitional JVP rule for one extension family.
pub trait SemanticLinearizeRule: Debug + Send + Sync + 'static {
    /// Return the versioned extension family handled by this rule.
    fn family_id(&self) -> &'static str;

    // ...
    fn linearize(
        &self,
        request: SemanticLinearizeRequest<'_>,
        builder: &mut SemanticProgramBuilder,
    ) -> Result<SemanticLinearizeResult, SemanticAdError>;
}

/// Transpose rule for an extension viewed as a linear map.
pub trait SemanticLinearTransposeRule: Debug + Send + Sync + 'static {
    fn family_id(&self) -> &'static str;

    /// Declare which primal input/output indices this rule reads as tensor
    /// residuals. Indices not declared may only be accessed through metadata.
    fn residual_mask(&self) -> ResidualSpec;

    // ...
    fn linear_transpose(
        &self,
        request: SemanticLinearTransposeRequest<'_>,
        builder: &mut SemanticProgramBuilder,
    ) -> Result<Box<[AdValue]>, SemanticAdError>;
}

/// Optional direct VJP rule expressed against primal semantic values.
pub trait SemanticPrimalVjpRule: Debug + Send + Sync + 'static {
    // family_id、residual_mask 同上
    fn primal_vjp(
        &self,
        request: SemanticPrimalVjpRequest<'_>,
        builder: &mut SemanticProgramBuilder,
    ) -> Result<Box<[AdValue]>, SemanticAdError>;
}

pub fn semantic_ad_rules() -> Result<SemanticExtensionRuleSet, SemanticExtensionRegistryError> {
    SemanticExtensionRuleSet::new()
        .with_linearize(Arc::new(LinalgAdRule))?
        .with_linear_transpose(Arc::new(LinalgAdRule))
}

== 1. family rule sets
SemanticExtensionRuleSet { linearize: ["tenferro-linalg.linalg.v1"], linear_transpose: ["tenferro-linalg.linalg.v1"], primal_vjp: [] }
SemanticExtensionRuleSet { linearize: ["tenferro.einsum.v1"], linear_transpose: ["tenferro.einsum.v1"], primal_vjp: ["tenferro.einsum.v1"] }
SemanticExtensionRuleSet { linearize: ["tenferro-fft.fft.v1"], linear_transpose: ["tenferro-fft.fft.v1"], primal_vjp: ["tenferro-fft.fft.v1"] }
merged: SemanticExtensionRuleSet { linearize: ["tenferro-linalg.linalg.v1", "tenferro.einsum.v1"], linear_transpose: ["tenferro-linalg.linalg.v1", "tenferro.einsum.v1"], primal_vjp: ["tenferro.einsum.v1"] }
merge again -> Err("semantic extension AD Linearize rule for family \"tenferro-linalg.linalg.v1\" is already registered")
default (TracedTensorAdExt) rule set: SemanticExtensionRuleSet { linearize: [], linear_transpose: [], primal_vjp: [] }

            SemanticOpRef::Extension(op) => {
                if rules.lookup_primal_vjp(op.family_id()).is_some() {
                    rules.primal_vjp_operation(
                        operation,
                        &mapped_values(operation.inputs(), &values),
                        &mapped_values(operation.outputs(), &values),
                        &cotangent_outputs,
                        &active_operation_inputs,
                        &mut builder,
                    )?
                } else {
                    let inactive_tangents = vec![AdValue::Absent; operation.inputs().len()];
                    let active_operation_outputs: Vec<_> = cotangent_outputs
                        .iter()
                        .map(|value| matches!(value, AdValue::Value(_)))
                        .collect();
                    let linearized = rules.linearize_operation(
                        operation,
                        &mapped_values(operation.inputs(), &values),
                        &mapped_values(operation.outputs(), &values),
                        &inactive_tangents,
                        &active_operation_outputs,
                        &mut builder,
                    )?;
                    rules.linear_transpose_operation(
                        operation,
                        &mapped_values(operation.inputs(), &values),
                        &mapped_values(operation.outputs(), &values),
                        &cotangent_outputs,
                        &active_operation_inputs,
                        linearized.residuals(),
                        &mut builder,
                    )?
                }
            }

    if index >= values.len() {
        return Err(SemanticAdError::PrimalIndexOutOfBounds {
            family_id,
            kind,
            index,
            len: values.len(),
        });
    }
    let declared = match kind {
        PrimalValueKind::Input => residual_mask.declares_input(index),
        PrimalValueKind::Output => residual_mask.declares_output(index),
    };
    if !declared {
        return Err(SemanticAdError::UndeclaredResidualValue {
            family_id,
            kind,
            index,
        });
    }
    Ok(values[index])

    pub fn with_ad_context(ad: &AdContext) -> Result<Arc<Self>> {
        Ok(Arc::new(Self::with_rules_and_cache(
            ad.semantic_extension_rules().clone(),
            ad.ad_transform_cache(),
        )?))
    }

let (a, l) = inv_loss(&A)?;                      // traced: a.inv()?.reduce_sum(None)
match l.grad(&a) { /* 打印 Ok 或 Err */ }        // TracedTensorAdExt：空规则集
let ad = AdContext::builder().with_semantic_extension_rules(all.clone())?.build()?;
let g = ad.grad(&l, &a)?;                        // 带 linalg + einsum 规则
// eager：同一个 ad 交给 EagerRuntime
let with = EagerRuntime::with_ad_context(&ad)?;
let y = with.variable_from(Tensor::from_vec_col_major(vec![2, 2], A.to_vec())?)?;
let ly = with.with_eager_session(&mut engine, |s| {
  let i = s.inv(&y)?;                            // EagerSessionLinalgExt
  s.reduce_sum(&i, None)
})?;
ly.backward(&mut engine)?;

== 2. traced: inv(A) summed, d/dA
TracedTensorAdExt::grad: Err(unsupported vjp AD rule for tenferro-linalg.linalg.v1)
  UnsupportedAdRule { transform: "vjp", op: "tenferro-linalg.linalg.v1" }
AdContext(linalg) grad = [-0.033057851239669415, -0.04958677685950412, -0.049586776859504134, -0.07438016528925619]
finite difference      = [-0.033057851234197244, -0.049586776851295866, -0.04958677690680702, -0.07438016533245495]
max err = 4.730e-11
== 3. eager: same op, runtime without / with the rule set
...
with AdContext(linalg+einsum): y.grad = Some([-0.033057851239669415, -0.04958677685950412, -0.049586776859504134, -0.07438016528925619])

== 4. eager einsum, 3 operands and 2 operands
EagerRuntime::new(): a.grad = Some([77.0, 77.0, 97.0, 97.0, 117.0, 117.0])
EagerRuntime::new(): 2-operand a.grad = Some([15.0, 15.0, 19.0, 19.0, 23.0, 23.0])
with AdContext: a.grad = Some([77.0, 77.0, 97.0, 97.0, 117.0, 117.0])
with AdContext: 2-operand a.grad = Some([15.0, 15.0, 19.0, 19.0, 23.0, 23.0])

EagerRuntime::new() ij,jk->ik loss=[63.0] backward: Err("semantic_eager_vjp (GraphBuild): runtime state failure: semantic extension AD failed: semantic extension family \"tenferro.einsum.v1\" has no Linearize AD rule")
== ij,jk->ik
  eager tracked: out.shape=[2, 2]
  a.grad shape=[2, 3] value=[5.0, 5.0, 7.0, 7.0, 9.0, 9.0]
  rt.grad shape=[2, 3] value=[5.0, 5.0, 7.0, 7.0, 9.0, 9.0]
  finite difference = [21.000000000270802, 21.000000000270802]
  eager untracked: out.shape=[2, 2] value=[6.0, 12.0, 15.0, 30.0]
  traced: Err(extension family "tenferro.einsum.v1" shape constraint at instruction None failed: Const(1) (1) Equal Const(3) (3))
  concrete: shape=[2, 2] value=[6.0, 12.0, 15.0, 30.0]
...
== ...j,jk->...k
  eager tracked: out.shape=[2, 2]
  a.grad shape=[2, 3] value=[5.0, 5.0, 7.0, 7.0, 9.0, 9.0]
  rt.grad shape=[2, 3] value=[5.0, 5.0, 7.0, 7.0, 9.0, 9.0]
  finite difference = [21.000000000270802, 21.000000000270802]
  eager untracked: out.shape=[2, 2] value=[6.0, 12.0, 15.0, 30.0]
  traced: shape=[2, 2] value=[6.0, 12.0, 15.0, 30.0]
  concrete: shape=[2, 2] value=[6.0, 12.0, 15.0, 30.0]

concrete: Err(CPU einsum: backend failure: label 106 has extents 1 and 3)
eager: Err(CPU einsum: backend failure: label 106 has extents 1 and 3)
  Runtime(TensorRuntime(BackendSource { op: "CPU einsum", source: Shape(Extent { label: 106, first: 1, second: 3 }) }))

EagerRuntime::new(): forward loss = [0.4545454545454545]
EagerRuntime::new(): backward Err(semantic_eager_vjp (GraphBuild): runtime state failure: semantic extension AD failed: semantic extension family "tenferro-linalg.linalg.v1" has no Linearize AD rule)
  RuntimeStateSource { op: "semantic_eager_vjp", phase: GraphBuild, source: Extension(MissingRule { family_id: "tenferro-linalg.linalg.v1", role: Linearize }) }
functional grad: Err(semantic_eager_vjp (GraphBuild): runtime state failure: semantic extension AD failed: semantic extension family "tenferro-linalg.linalg.v1" has no Linearize AD rule)