pub fn value_and_grad_flat<F>(
f: F,
argnums: &[usize],
) -> impl Fn(&[Array]) -> Result<(Array, Vec<Array>)>Expand description
Returns a function computing f’s value and its gradients with respect to the
positional arguments at argnums (the flat form of value_and_grad).
f must return a 0-D array. Only the selected arguments are traced; the others are
treated as constants.
use nuvec_core::{Array, Result, array, value_and_grad_flat};
// loss(w, x) = sum((w * x - 1)^2)
let loss = |xs: &[Array]| {
let r = xs[0].multiply(&xs[1]).subtract(1.0);
Ok(r.multiply(&r).sum(..))
};
let (w, x) = (array![2.0], array![3.0]);
let (value, grads) = value_and_grad_flat(loss, &[0])(&[w, x])?;
assert_eq!(value.item::<f32>()?, 25.0);
assert_eq!(grads[0].to_vec::<f32>()?, [30.0]); // 2 (w x - 1) x