Trait ExtensionLinearizeRule 

pub trait ExtensionLinearizeRule:
    Debug
    + Send
    + Sync
    + 'static {
    // Required methods
    fn family_id(&self) -> &'static str;
    fn linearize(
        &self,
        op: &(dyn ExtensionOp + 'static),
        builder: &mut dyn PrimitiveRuleBuilder,
        primal_in: &[ValueKey<StdTensorOp>],
        primal_out: &[ValueKey<StdTensorOp>],
        tangent_in: &[Option<usize>],
        ctx: &mut ShapeGuardContext,
    ) -> Result<Vec<Option<usize>>, ADRuleError>;
}

Definitional JVP rule provider for an extension family.

Required only for families that appear in primal graphs and must be differentiable.

Required Methods

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

The extension family this rule handles.

fn linearize( &self, op: &(dyn ExtensionOp + 'static), builder: &mut dyn PrimitiveRuleBuilder, primal_in: &[ValueKey<StdTensorOp>], primal_out: &[ValueKey<StdTensorOp>], tangent_in: &[Option<usize>], ctx: &mut ShapeGuardContext, ) -> Result<Vec<Option<usize>>, ADRuleError>

Emit the linear (JVP) rule.

Implementors