Function custom_function 

pub fn custom_function(
    f: impl Fn(&[Array]) -> Result<Vec<Array>> + Send + Sync + 'static,
) -> CustomFunction

Wraps f 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, so a custom function without rules is a gradient checkpoint (see checkpoint).

use nuvec_core::{Array, Result, custom_function, grad, multiply, sum};

// A straight-through estimator: forward x², gradient as if the identity.
let square = custom_function(|xs| Ok(vec![multiply(&xs[0], &xs[0])?]))
    .with_vjp(|_primals, _outputs, cotangents| Ok(cotangents.to_vec()));
let g = grad(move |x: &Array| sum(&square.call(std::slice::from_ref(x))?[0], &[], false));
let x = Array::from_slice(&[3.0f32, -1.0], &[2])?;
assert_eq!(g(&x)?.to_vec::<f32>()?, [1.0, 1.0]);