Function custom_function 

pub fn custom_function<I, O, F, R>(f: F) -> CustomFunction<I, O>
where
    I: Tree + Send + Sync,
    O: Tree + Send + Sync,
    F: Fn(&I) -> R + Send + Sync + 'static,
    R: IntoResult<Ok = O>,

Wraps f (a function of a tree, returning a tree) so that transforms use custom rules for it, set with CustomFunction::with_vjp, CustomFunction::with_jvp and CustomFunction::with_vmap.

Without a rule, a transform differentiates or batches f itself, re-running it; f’s intermediates are never kept for a transform. The rules are stored in the graph (and may run later, on another thread), so f and the rules are Send + Sync + 'static.

use nuvec_core::{Array, Result, array, custom_function, grad};

// A straight-through estimator: forward x², gradient as if the identity.
let square = custom_function(|x: &Array| x.multiply(x))
    .with_vjp(|_x: &Array, _y: &Array, cotangent: &Array| Ok(cotangent.clone()));
let g = grad(move |x: &Array| Ok(square.call(x)?.sum(..)));
assert_eq!(g(&array![3.0, -1.0])?.to_vec::<f32>()?, [1.0, 1.0]);