Function custom_function_flat 

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

Wraps f (a function of positional arrays) so that transforms use custom rules for it, set with CustomFunctionFlat::with_vjp, CustomFunctionFlat::with_jvp and CustomFunctionFlat::with_vmap; the flat form of custom_function.

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_flat).

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

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