Skip to main content

linearize

Function linearize 

Source
pub fn linearize(
    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>>>
Expand description

Forward-mode AD (JVP) for StdTensorOp: given the primal op and its tangent inputs, emit the linearized graph into builder and return the output tangents.

Rules per op live in the category submodules (semiring, analytic, elementwise, structural, contraction, indexing, diagonal, dynamic). StdTensorOp::Extension(_) delegates to the trait.