Struct CustomFunctionFlat 

pub struct CustomFunctionFlat { /* private fields */ }

A function of positional arrays with custom transform rules; see custom_function_flat.

Implementations

impl CustomFunctionFlat

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

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>, Error> + Send + Sync + 'static, ) -> CustomFunctionFlat

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<i32>]) -> Result<(Vec<Array>, i32), Error> + Send + Sync + 'static, ) -> CustomFunctionFlat

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>, Error>

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

Trait Implementations

impl Clone for CustomFunctionFlat

fn clone(&self) -> CustomFunctionFlat

Returns a duplicate of the value.

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

Performs copy-assignment from source.