Skip to main content

value_and_grad_flat

Function value_and_grad_flat 

Source
pub fn value_and_grad_flat<F>(
    f: F,
    argnums: &[usize],
) -> impl Fn(&[Array]) -> Result<(Array, Vec<Array>)>
where F: Fn(&[Array]) -> Result<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