Function value_and_grad 

pub fn value_and_grad<P, F, R>(f: F) -> impl Fn(&P) -> Result<(R::Ok, P)>
where
    P: Tree,
    F: Fn(&P) -> R,
    R: IntoResult,
    R::Ok: GradOutput,

Returns a function computing f’s value and its gradient with respect to the trainable leaves of its argument, as a tree of the argument’s type (grads.layers[0].weight).

f takes the differentiated tree (a tuple, a struct, a model) and returns the 0-D loss, or (loss, aux) to also return auxiliary outputs (the return type decides); it may return the value directly or a Result. Constants, such as a batch, are captured by the closure. Write the closure’s parameter type: |p: &(Array, Array)|, |m: &Mlp| (Rust cannot infer it through the generic transform: E0282 otherwise).

Frozen parameters, buffers and integer parameters get zero gradients that allocate no full-size memory (stride-0 broadcasts of one element); a plain integer or boolean Array leaf is an error (capture such data in the closure instead).

use nuvec_core::{Array, Result, array, value_and_grad};

let x = array![1.0, 2.0];
// loss(w, b) = sum((w x + b)^2), with `x` captured.
let loss = |(w, b): &(Array, Array)| {
    let r = w.multiply(&x).add(b);
    r.multiply(&r).sum(..)
};
let (value, (gw, gb)) = value_and_grad(loss)(&(array![1.0, 1.0], array![0.0, 0.0]))?;
assert_eq!(value.item::<f32>()?, 5.0);
assert_eq!(gw.to_vec::<f32>()?, [2.0, 8.0]); // 2 (w x + b) x
assert_eq!(gb.to_vec::<f32>()?, [2.0, 4.0]);