Struct CustomFunction 

pub struct CustomFunction { /* private fields */ }

A function with custom transform rules; see custom_function.

Implementations

impl CustomFunction

pub fn with_vjp( self, rule: impl Fn(&[Array], &[Array], &[Array]) -> Result<Vec<Array>> + Send + Sync + 'static, ) -> Self

Sets the reverse-mode rule: (primals, outputs, cotangents) → one cotangent per primal.

pub fn with_jvp( self, rule: impl Fn(&[Array], &[Array]) -> Result<Vec<Array>> + Send + Sync + 'static, ) -> Self

Sets the forward-mode rule: (primals, tangents) → one tangent per output (missing input tangents are passed as zeros).

pub fn with_vmap( self, rule: impl Fn(&[Array], &[Option<usize>]) -> Result<(Vec<Array>, usize)> + Send + Sync + 'static, ) -> Self

Sets the batching rule: (inputs, batch axes) (None for unbatched inputs) → the batched outputs and their common batch axis.

pub fn call(&self, inputs: &[Array]) -> Result<Vec<Array>>

Calls the function on inputs. Values equal f(inputs); transforms of the results use the custom rules.

Trait Implementations

impl Clone for CustomFunction

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.