Trait ExtensionPrimalVjpRule 

pub trait ExtensionPrimalVjpRule:
    Debug
    + Send
    + Sync
    + 'static {
    // Required methods
    fn family_id(&self) -> &'static str;
    fn primal_vjp(
        &self,
        op: &dyn ExtensionOp,
        builder: &mut dyn PrimitiveRuleBuilder,
        cotangent_out: &[Option<LocalValueId>],
        inputs: &[ValueRef<StdTensorOp>],
        ctx: &mut ShapeGuardContext,
    ) -> ADRuleResult<Vec<Option<LocalValueId>>>;
}

Optional direct primal VJP rule provider for an extension family.

This is an opt-in performance or compatibility escape hatch. The default reverse-mode path remains linearize followed by linear_transpose.

Required Methods

fn family_id(&self) -> &'static str

The extension family this rule handles.

fn primal_vjp( &self, op: &dyn ExtensionOp, builder: &mut dyn PrimitiveRuleBuilder, cotangent_out: &[Option<LocalValueId>], inputs: &[ValueRef<StdTensorOp>], ctx: &mut ShapeGuardContext, ) -> ADRuleResult<Vec<Option<LocalValueId>>>

Emit a direct VJP from the primal op.

Implementors