Function checkpoint_flat 

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

Returns f (a function of positional arrays) with gradient checkpointing; the flat form of checkpoint.

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

let block = checkpoint_flat(|xs| Ok(vec![xs[0].sin().sin()]));
let g = grad(move |x: &Array| Ok(block(std::slice::from_ref(x))?[0].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);