Function checkpoint 

pub fn checkpoint<I, O, F, R>(f: F) -> impl Fn(&I) -> Result<O> + Send + Sync
where
    I: Tree + Send + Sync,
    O: Tree + Send + Sync,
    F: Fn(&I) -> R + Send + Sync + 'static,
    R: IntoResult<Ok = O>,

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. Random numbers drawn inside f (dropout) are drawn again with the same key, so the recomputation sees the same masks.

f is stored in the graph for the recomputation, so it is Send + Sync + 'static (it cannot borrow locals: move them in).

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

let block = checkpoint(|x: &Array| x.sin().sin());
let g = grad(move |x: &Array| Ok(block(x)?.sum(..)));
let x = Array::from(0.5f32);
let expect = 0.5f32.sin().cos() * 0.5f32.cos();
assert!((g(&x)?.item::<f32>()? - expect).abs() < 1e-6);