Skip to main content

checkpoint

Function checkpoint 

Source
pub fn checkpoint(
    f: impl Fn(&[Array]) -> Result<Vec<Array>> + Send + Sync + 'static,
) -> impl Fn(&[Array]) -> Result<Vec<Array>> + Send + Sync
Expand description

Returns f with gradient checkpointing: values and derivatives equal f’s, but f’s intermediates are not kept for the backward pass; f runs again to differentiate.

use nuvec_core::{Array, Result, checkpoint, grad, sin, sum};

let block = checkpoint(|xs| Ok(vec![sin(&sin(&xs[0])?)?]));
let g = grad(move |x: &Array| sum(&block(std::slice::from_ref(x))?[0], &[], false));
let x = Array::scalar(0.5f32);
let expect = 0.5f32.sin().cos() * 0.5f32.cos();
assert!((g(&x)?.item::<f32>()? - expect).abs() < 1e-6);