Backprop memory optimization: Count-trailing-zero / gradient checkpointing
Reference:

- https://rufflewind.com/2016-12-30/reverse-mode-automatic-differentiation
![2018-12-15_23-39-12](https://user-images.githubusercontent.com/22738317/50048047-a5649a80-00c2-11e9-8b70-375a3199906e.png)

- https://rufflewind.com/2016-12-30/reverse-mode-automatic-differentiation and https://github.com/Rufflewind/revad/blob/de509269fe878bc9d564775abc25c4fa663d8a5e/src/chain.rs

```Rust
/// Maintains a partial tape using the count-trailing-zeros (CTZ) eviction
/// strategy.  This results in space usage that is logarithmic in the number
/// of steps, while also keeping a logarithmic amount of recomputation time.
///
/// FIXME: `sweep_mut` is not yet implemented.
pub struct CtzChain<S, J, R> {
    snapshots: Vec<(usize, S)>,
    adjoint: J,
    restore: R,
}

impl<S, J, R> CtzChain<S, J, R> {
    pub fn new<I>(snapshots: I, adjoint: J, restore: R) -> Self
        where I: Iterator<Item=S> {
        let mut chain = CtzChain {
            snapshots: Vec::new(),
            adjoint: adjoint,
            restore: restore,
        };
        ctz_extend(&mut chain.snapshots, 0, snapshots);
        chain
    }

    pub fn sweep<G>(&self, mut x: G) -> G
        where J: Fn(&S, G) -> G, R: Fn(&S) -> S {

        fn extend<S, R>(restored_snapshots: &mut Vec<(usize, S)>,
                        mut num_missing: usize, i: usize,
                        s_old: &mut Option<S>, restore: &R)
            where R: Fn(&S) -> S {
            ctz_extend(restored_snapshots, i + 1, Generator(|| {
                let mut s_new = match s_old {
                    &mut None => {
                        None
                    },
                    &mut Some(ref mut s_old) => {
                        num_missing -= 1;
                        if num_missing == 0 {
                            None
                        } else {
                            Some(restore(s_old))
                        }
                    },
                };
                ::std::mem::swap(s_old, &mut s_new);
                s_new
            }));
        }

        let mut j = match self.snapshots.last() {
            None => return x,
            Some(&(i, _)) => i + 1,
        };
        let mut restored_snapshots = Vec::new();
        for &(i, ref s) in self.snapshots.iter().rev() {
            loop {
                match restored_snapshots.pop() {
                    Some((i, s)) => {
                        let num_missing = j - i - 1;
                        if num_missing != 0 {
                            let mut s_old = Some((self.restore)(&s));
                            restored_snapshots.push((i, s));
                            extend(&mut restored_snapshots, num_missing,
                                   i, &mut s_old, &self.restore);
                        } else {
                            x = (self.adjoint)(&s, x);
                            j -= 1;
                        }
                    },
                    None => {
                        let num_missing = j - i - 1;
                        if num_missing != 0 {
                            let mut s_old = Some((self.restore)(&s));
                            extend(&mut restored_snapshots, num_missing,
                                   i, &mut s_old, &self.restore);
                        } else {
                            x = (self.adjoint)(&s, x);
                            j -= 1;
                            break;
                        }
                    },
                }
            }
        }
        x
    }

    pub fn sweep_once<G>(mut self, mut x: G) -> G
        where J: FnMut(S, G) -> G, R: FnMut(&S) -> S {
        let mut i = match self.snapshots.last() {
            None => return x,
            Some(&(j, _)) => j + 1,
        };
        loop {
            match self.snapshots.pop() {
                Some((j, s)) => {
                    let mut num_missing = i - j - 1;
                    if num_missing == 0 {
                        x = (self.adjoint)(s, x);
                        i -= 1;
                    } else {
                        let mut s_old = Some((self.restore)(&s));
                        self.snapshots.push((j, s));
                        let mut restore = self.restore;
                        ctz_extend(&mut self.snapshots, j + 1, Generator(|| {
                            let mut s_new = match &mut s_old {
                                &mut None => {
                                    None
                                },
                                &mut Some(ref mut s_old) => {
                                    num_missing -= 1;
                                    if num_missing == 0 {
                                        None
                                    } else {
                                        Some(restore(s_old))
                                    }
                                },
                            };
                            ::std::mem::swap(&mut s_old, &mut s_new);
                            s_new
                        }));
                        self.restore = restore;
                    }
                },
                None => {
                    assert_eq!(i, 0);
                    return x;
                },
            }
        }
    }
}

/// Wrap a `next` function into an `Iterator`.
pub struct Generator<F>(pub F);

impl<F, T> Iterator for Generator<F> where F: FnMut() -> Option<T> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        (self.0)()
    }
}
```

Yu. M. Volin and G. M. Ostrovskii. 1985. Automatic computation of derivatives with the use of the multilevel differentiating technique — I: Algorithmic basis. Computers and Mathematics with Applications 11 (1985), 1099–1114.

- https://medium.com/tensorflow/fitting-larger-networks-into-memory-583e3c758ff9
- https://github.com/openai/gradient-checkpointing
[Training Deep Nets with Sublinear Memory Cost, by Chen et al. (2016)](https://arxiv.org/pdf/1604.06174.pdf), [Memory-Efficient Backpropagation Through Time, by Gruslys et al. (2016)](https://arxiv.org/abs/1606.03401v1), [Exact Alpha-Beta Computation in
Logarithmic Space with Application to MAP Word Graph Construction, by Zweig et al. (2000)](https://www.microsoft.com/en-us/research/publication/exact-alpha-beta-computation-in-logarithmic-space-with-application-to-map-word-graph-construction/), [Debugging Backwards in Time, by Bil Lewis (2003)](https://arxiv.org/abs/cs/0310016), [Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, by Griewank and Walther (2008)](https://epubs.siam.org/doi/book/10.1137/1.9780898717761)

