checkpointing with 'ad'

checkpoint.hs

{-# language RankNTypes #-}
{-# language FlexibleContexts #-}
{-# language ScopedTypeVariables #-}
import Numeric.AD.Internal.Identity
import Numeric.AD.Internal.Reverse
import Numeric.AD.Internal.Or
import Numeric.AD.Mode
import Numeric.AD.Mode.Reverse (grad)
import Data.Reflection (Reifies)
import Data.Foldable (toList)
import Data.Proxy

-- based on an implementation by Sofus Mortensen
grad_cp
  :: forall f g a. (Traversable f, Traversable g, Num a)
  => (forall s. Reifies s Tape => f (Reverse s a) -> g (Reverse s a))
  -> (forall s. Reifies s Tape => g (Reverse s a) -> Reverse s a)
  -> f a -> f a
grad_cp f g xs = grad (sum . zipWith ((*) . auto) z' . toList . f) xs where
  z' = toList $ grad g fxs
  fxs = reifyTape 0 $ \(_ :: Proxy s) -> primal <$> f (auto <$> xs :: f (Reverse s a)) -- should have an empty tape when done

Nice!