pub trait ExtensionLinearizeRule:
Debug
+ Send
+ Sync
+ 'static {
// Required methods
fn family_id(&self) -> &'static str;
fn linearize(
&self,
op: &dyn ExtensionOp,
builder: &mut dyn PrimitiveRuleBuilder,
primal_in: &[ValueKey<StdTensorOp>],
primal_out: &[ValueKey<StdTensorOp>],
tangent_in: &[Option<LocalValueId>],
ctx: &mut ShapeGuardContext,
) -> ADRuleResult<Vec<Option<LocalValueId>>>;
}Expand description
Definitional JVP rule provider for an extension family.
Required only for families that appear in primal graphs and must be differentiable.
Required Methods§
Sourcefn linearize(
&self,
op: &dyn ExtensionOp,
builder: &mut dyn PrimitiveRuleBuilder,
primal_in: &[ValueKey<StdTensorOp>],
primal_out: &[ValueKey<StdTensorOp>],
tangent_in: &[Option<LocalValueId>],
ctx: &mut ShapeGuardContext,
) -> ADRuleResult<Vec<Option<LocalValueId>>>
fn linearize( &self, op: &dyn ExtensionOp, builder: &mut dyn PrimitiveRuleBuilder, primal_in: &[ValueKey<StdTensorOp>], primal_out: &[ValueKey<StdTensorOp>], tangent_in: &[Option<LocalValueId>], ctx: &mut ShapeGuardContext, ) -> ADRuleResult<Vec<Option<LocalValueId>>>
Emit the linear (JVP) rule.