ad-delcont 0.2.0.0 → 0.3.0.0

raw patch · 6 files changed

+220/−48 lines, 6 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• Changelog.md +5/−0
• LICENSE +2/−2
• README.md +1/−4
• ad-delcont.cabal +10/−5
• src/Numeric/AD/DelCont.hs +30/−6
• src/Numeric/AD/DelCont/Internal.hs +172/−31

@@ -1,3 +1,8 @@+0.3++* add `grad`+* refactor AD types to retrieve Functor and Applicative instance+ 0.1  First public version

@@ -1,4 +1,4 @@-Copyright Author name here (c) 2021+Copyright Marco Zocca (c) 2021  All rights reserved. @@ -13,7 +13,7 @@       disclaimer in the documentation and/or other materials provided       with the distribution. -    * Neither the name of Author name here nor the names of other+    * Neither the name of Marco Zocca nor the names of other       contributors may be used to endorse or promote products derived       from this software without specific prior written permission. 

@@ -2,7 +2,4 @@  Reverse-mode Automatic Differentiation using delimited continuations (`shift` and `reset`) , as shown in [1] --## References--1) Wang et al, Demystifying Differentiable Programming : Shift/Reset the Penultimate Backpropagator,  https://www.cs.purdue.edu/homes/rompf/papers/wang-icfp19.pdf +Introductory blog post : http://ocramz.github.io/haskell/automatic-differentiation/2021/07/19/ad-delcont.html

@@ -1,15 +1,20 @@ name:                ad-delcont-version:             0.2.0.0+version:             0.3.0.0 synopsis:            Reverse-mode automatic differentiation with delimited continuations-description:         Reverse-mode automatic differentiation using delimited continuations (@shift@/@reset@), inspired by the papers+description:         Reverse-mode automatic differentiation using delimited continuations (@shift@/@reset@).+                     The package exposes a small and easily extensible user interface to automatic differentiation combinators. It's also lightweight as a dependency, since it only requires @base@ and @transformers@.                      .+                     To use the library, import "Numeric.AD.DelCont", which also contains all documentation.+                     .+                     Blog post : http://ocramz.github.io/haskell/automatic-differentiation/2021/07/19/ad-delcont.html+                     .+                     References :+                     .                      * F. Wang et al, Backpropagation with Continuation Callbacks : Foundations for Efficient and Expressive Differentiable Programming, NeurIPS 2018 - https://papers.nips.cc/paper/2018/file/34e157766f31db3d2099831d348a7933-Paper.pdf                      .                      * F. Wang et al, Demystifying Differentiable Programming : Shift\/Reset the Penultimate Backpropagator, ICFP 2019 - https://www.cs.purdue.edu/homes/rompf/papers/wang-icfp19.pdf                      .-                     The package exposes a small and easily extensible user interface to automatic differentiation combinators. It also manages to be very lightweight, as it only depends on @base@ and @transformers@.-                     .-                     To use the library, import "Numeric.AD.DelCont", which also contains all documentation.+                      homepage:            https://github.com/ocramz/ad-delcont license:             BSD3 license-file:        LICENSE

@@ -1,8 +1,15 @@-{-| Reverse-mode automatic differentiation using delimited continuations.+{-|+Module      : Numeric.AD.DelCont+Description : Reverse-mode automatic differentiation using delimited continuations+Copyright   : (c) Marco Zocca, 2021+License     : BSD+Maintainer  : github.com/ocramz+Stability   : experimental+Portability : POSIX  == Quickstart -Most users will only need to import 'rad1', 'rad2' and leverage the 'Num', 'Fractional', 'Floating' instances of the 'AD' type.+Most users will only need to import 'rad1', 'rad2' or `grad` and leverage the 'Num', 'Fractional', 'Floating' instances of the 'AD' type.  Similarly to @ad@, a user supplies a /polymorphic/ function to be differentiated, e.g. @@ -18,6 +25,23 @@ (2.6399999999999997,3.4000000000000004) @ +'grad' computes the gradient of a scalar function of vector argument. For full generality, the argument can be any 'Traversable' container (a list, an array, a dictionary .. )++@+sqNorm :: Num a => [a] -> a+sqNorm xs = sum $ zipWith (*) xs xs++p :: [Double]+p = [4.1, 2]+@++@+>>> 'grad' sqNorm p+(20.81,[8.2,4.0])+@+++ It's important to emphasize that the library cannot differentiate functions of concrete types, e.g. @Double -> Double@. On the other hand, it's easy to experiment with other numerical interfaces that support one, zero and plus.  == Advanced usage@@ -54,16 +78,16 @@ * M. Innes, Don't unroll adjoint: Differentiating SSA-Form Programs  https://arxiv.org/abs/1810.07951 -} module Numeric.AD.DelCont (-- * Quickstart-                            rad1, rad2+                            rad1, rad2, grad                           , auto                           -- * Advanced usage-                          , rad1g, rad2g+                          , rad1g, rad2g, radNg                             -- ** Lift functions into AD                           , op1, op2                             -- *** Num instances                           , op1Num, op2Num                           -- * Types-                          , AD, AD') where+                          , AD0, AD, AD') where -import Numeric.AD.DelCont.Internal (rad1, rad2, auto, rad1g, rad2g, op1, op2, op1Num, op2Num, AD, AD')+import Numeric.AD.DelCont.Internal (rad1, rad2, grad, auto, rad1g, rad2g, radNg, op1, op2, op1Num, op2Num, AD0, AD, AD') 

@@ -1,15 +1,18 @@ {-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE TypeSynonymInstances #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE DeriveFunctor #-}+{-# language GeneralizedNewtypeDeriving #-} {-# LANGUAGE LambdaCase #-}+-- {-# LANGUAGE MultiParamTypeClasses #-}+-- {-# language TypeFamilies #-}+{-# OPTIONS_GHC -Wno-unused-imports -Wno-unused-top-binds #-} module Numeric.AD.DelCont.Internal-  (rad1, rad2,+  (rad1, rad2, grad,    auto,-   rad1g, rad2g,+   rad1g, rad2g, radNg,    op1Num, op2Num,    op1, op2,-   AD, AD')+   AD0, AD, AD')   where  import Control.Monad.ST (ST, runST)@@ -22,8 +25,9 @@  import Prelude hiding (read) + -- | Dual numbers-data D a da = D a da deriving (Show, Functor)+data D a da = D { primal :: a, dual :: da } deriving (Show, Functor) instance Eq a => Eq (D a da) where   D x _ == D y _ = x == y instance Ord a => Ord (D a db) where@@ -47,19 +51,25 @@ -- -- As one expects from a constant, its value will be used for computing the result, but it will be discarded when computing the sensitivities. auto :: a -> AD s a da-auto x = AD $ lift $ var x undefined+auto x = AD0 $ lift $ var x undefined  --- | Mutable references to dual numbers in the continuation monad+-- | Mutable references in the continuation monad+newtype AD0 s a = AD0 { unAD0 :: forall x . ContT x (ST s) a } deriving (Functor)+instance Applicative (AD0 s) where+  AD0 f <*> AD0 x = AD0 (f <*> x)+  pure x = AD0 $ pure x+-- instance Monad (AD s) where -- TODO++-- | A synonym of 'AD0' for the common case of returning a 'DVar' (which is a 'ST' computation that returns a dual number) -- -- Here the @a@ and @da@ type parameters are respectively the /primal/ and /dual/ quantities tracked by the AD computation.-newtype AD s a da = AD { unAD :: forall x dx . ContT (DVar s x dx) (ST s) (DVar s a da) }+type AD s a da = AD0 s (DVar s a da) -- | Like 'AD' but the types of primal and dual coincide type AD' s a = AD s a a --- runAD :: (forall s . AD s a da) -> D a da--- runAD go = runST (evalContT  (unAD go) >>= readSTRef) + -- | Lift a unary function -- -- This is a polymorphic combinator for tracking how primal and adjoint values are transformed by a function.@@ -80,17 +90,18 @@      -> (a -> (b, db -> da)) -- ^ returns : (function result, pullback)      -> ContT x (ST s) (DVar s a da)      -> ContT x (ST s) (DVar s b db)-op1_ zero plusa f ioa = do+op1_ zeroa plusa f ioa = do   ra <- ioa   (D xa _) <- lift $ readSTRef ra   let (xb, g) = f xa -- 1)   shiftT $ \ k -> lift $ do-    rb <- var xb zero -- 2)+    rb <- var xb zeroa -- 2)     ry <- k rb -- 3)     (D _ yd) <- readSTRef rb -- 4)     modifySTRef' ra (withD (\rda0 -> rda0 `plusa` g yd)) -- 5)     pure ry + -- | Lift a unary function -- -- The first two arguments constrain the types of the adjoint values of the output and input variable respectively, see 'op1Num' for an example.@@ -103,7 +114,7 @@     -> (a -> (b, db -> da)) -- ^ returns : (function result, pullback)     -> AD s a da     -> AD s b db-op1 z plusa f (AD ioa) = AD $ op1_ z plusa f ioa+op1 z plusa f (AD0 ioa) = AD0 $ op1_ z plusa f ioa  -- | Helper for constructing unary functions that operate on Num instances (i.e. 'op1' specialized to Num) op1Num :: (Num da, Num db) =>@@ -120,14 +131,14 @@      -> ContT x (ST s) (DVar s a da)      -> ContT x (ST s) (DVar s b db)      -> ContT x (ST s) (DVar s c dc)-op2_ zero plusa plusb f ioa iob = do+op2_ zeroa plusa plusb f ioa iob = do   ra <- ioa   rb <- iob   (D xa _) <- lift $ readSTRef ra   (D xb _) <- lift $ readSTRef rb   let (xc, g, h) = f xa xb   shiftT $ \ k -> lift $ do-    rc <- var xc zero+    rc <- var xc zeroa     ry <- k rc     (D _ yd) <- readSTRef rc     modifySTRef' ra (withD (\rda0 -> rda0 `plusa` g yd))@@ -142,7 +153,7 @@     -> (db -> db -> db) -- ^ plus     -> (a -> b -> (c, dc -> da, dc -> db)) -- ^ returns : (function result, pullbacks)     -> (AD s a da -> AD s b db -> AD s c dc)-op2 z plusa plusb f (AD ioa) (AD iob) = AD $ op2_ z plusa plusb f ioa iob+op2 z plusa plusb f (AD0 ioa) (AD0 iob) = AD0 $ op2_ z plusa plusb f ioa iob  -- | Helper for constructing binary functions that operate on Num instances (i.e. 'op2' specialized to Num) op2Num :: (Num da, Num db, Num dc) =>@@ -191,23 +202,25 @@   acosh = op1Num $ \x -> (acosh x, (/ sqrt (x*x - 1)))   atanh = op1Num $ \x -> (atanh x, (/ (1 - x*x))) --- instance Eq a => Eq (AD s a da) where -- ??? likely impossible--- instance Ord (AD s a da) where -- ??? see above+-- -- instance Eq a => Eq (AD s a da) where -- ??? likely impossible+-- -- instance Ord (AD s a da) where -- ??? see above ++ -- | Evaluate (forward mode) and differentiate (reverse mode) a unary function, without committing to a specific numeric typeclass rad1g :: da -- ^ zero       -> db -- ^ one       -> (forall s . AD s a da -> AD s b db)       -> a -- ^ function argument       -> (b, da) -- ^ (result, adjoint)-rad1g zero one f x = runST $ do-  xr <- var x zero+rad1g zeroa oneb f x = runST $ do+  xr <- var x zeroa   zr' <- evalContT $     resetT $ do       let-        z = f (AD (pure xr))-      zr <- unAD z-      lift $ modifySTRef' zr (withD (const one))+        z = f (AD0 (pure xr))+      zr <- unAD0 z+      lift $ modifySTRef' zr (withD (const oneb))       pure zr   (D z _) <- readSTRef zr'   (D _ x_bar) <- readSTRef xr@@ -222,22 +235,45 @@       -> (forall s . AD s a da -> AD s b db -> AD s c dc)       -> a -> b       -> (c, (da, db)) -- ^ (result, adjoints)-rad2g zeroa zerob one f x y = runST $ do+rad2g zeroa zerob onec f x y = runST $ do   xr <- var x zeroa   yr <- var y zerob   zr' <- evalContT $     resetT $ do       let-        z = f (AD (pure xr)) (AD (pure yr))-      zr <- unAD z-      lift $ modifySTRef' zr (withD (const one))+        z = f (AD0 (pure xr)) (AD0 (pure yr))+      zr <- unAD0 z+      lift $ modifySTRef' zr (withD (const onec))       pure zr   (D z _) <- readSTRef zr'   (D _ x_bar) <- readSTRef xr   (D _ y_bar) <- readSTRef yr   pure (z, (x_bar, y_bar)) +-- | Evaluate (forward mode) and differentiate (reverse mode) a function of a 'Traversable'+--+-- In linear algebra terms, this computes the gradient of a scalar function of vector argument+radNg :: Traversable t =>+         da -- ^ zero+      -> db -- ^ one+      -> (forall s . t (AD s a da) -> AD s b db)+      -> t a -- ^ argument vector+      -> (b, t da) -- ^ (result, gradient vector)+radNg zeroa onea f xs = runST $ do+  xrs <- traverse (`var` zeroa) xs+  zr' <- evalContT $+    resetT $ do+      let+        (AD0 z) = f (fmap pure xrs)+      zr <- z+      lift $ modifySTRef' zr (withD (const onea))+      pure zr+  (D z _) <- readSTRef zr'+  xs_bar <- traverse readSTRef xrs+  let xs_bar_d = dual <$> xs_bar+  pure (z, xs_bar_d) + -- | Evaluate (forward mode) and differentiate (reverse mode) a unary function -- -- >>> rad1 (\x -> x * x) 1@@ -262,21 +298,126 @@      -> (c, (a, b)) -- ^ (result, adjoints) rad2 = rad2g 0 0 1 +-- | Evaluate (forward mode) and differentiate (reverse mode) a function of a 'Traversable'+--+-- In linear algebra terms, this computes the gradient of a scalar function of vector argument+--+--+-- @+-- sqNorm :: Num a => [a] -> a+-- sqNorm xs = sum $ zipWith (*) xs xs+--+-- p :: [Double]+-- p = [4.1, 2]+-- @+--+-- >>> grad sqNorm p+-- (20.81,[8.2,4.0])+grad :: (Traversable t, Num a, Num b) =>+        (forall s . t (AD' s a) -> AD' s b)+     -> t a -- ^ argument vector+     -> (b, t a) -- ^ (result, gradient vector)+grad = radNg 0 1  +-- ======================== EXPERIMENTAL ========================== +data Backprop a da = Backprop {+  zero :: a -> da+  , one :: da -> da+  , add :: da -> da -> da+                              }++bpNum :: (Num a, Num da) => Backprop a da+bpNum = Backprop zeroNum oneNum addNum++-- | backprop typeclass, adapted from https://hackage.haskell.org/package/backprop-0.2.6.4/docs/src/Numeric.Backprop.Class.html+--+-- we use two type parameters to keep the distinction between primal and dual variables++-- class Backprop a da where+--   zero :: a -> da+--   one :: proxy a -> da -> da+--   add :: proxy a -> da -> da -> da++-- | 'zero' for instances of 'Num'. lazy in its argument.+zeroNum :: Num da => a -> da+zeroNum _ = 0+{-# INLINE zeroNum #-}++-- | 'add' for instances of 'Num'.+addNum :: Num da => da -> da -> da+addNum = (+)+{-# INLINE addNum #-}++-- | 'one' for instances of 'Num'. lazy in its argument.+oneNum :: Num da => a -> da+oneNum _ = 1+{-# INLINE oneNum #-}+++-- rad1BP :: Backprop a da+--        -> Backprop b db+--        -> (forall s . AD s a da -> AD s b db)+--        -> a -- ^ function argument+--        -> (b, da) -- ^ (result, adjoint)+-- rad1BP bpa bpb f x = runST $ do+--   xr <- var x (zero bpa x)+--   zr' <- evalContT $+--     resetT $ do+--       let+--         z = f (AD (pure xr))+--       zr <- unAD z+--       lift $ modifySTRef' zr (withD $ one bpb)+--       pure zr+--   (D z _) <- readSTRef zr'+--   (D _ x_bar) <- readSTRef xr+--   pure (z, x_bar)++-- -- rad1BP :: (Backprop a da, Backprop b db)+-- --        => (forall s . AD s a da -> AD s b db)+-- --        -> a -- ^ function argument+-- --        -> (b, da) -- ^ (result, adjoint)+-- -- rad1BP f x = runST $ do+-- --   xr <- var x (zero x)+-- --   zr' <- evalContT $+-- --     resetT $ do+-- --       let+-- --         z = f (AD (pure xr))+-- --       zr <- unAD z+-- --       let+-- --         oneB :: forall b db . Backprop b db => db -> db -> db+-- --         oneB = one (Proxy :: Proxy db)+-- --       lift $ modifySTRef' zr (withD oneB)+-- --       pure zr+-- --   (D z _) <- readSTRef zr'+-- --   (D _ x_bar) <- readSTRef xr+-- --   pure (z, x_bar)+++++ -- -- playground --- -- | Dual numbers DD (alternative take, using a type family for the first variation)--- data DD a = Dd a (Adj a)--- class Diff a where type Adj a :: *--- instance Diff Double where type Adj Double = Double   -- -- product type (simplified version of vinyl's Rec) -- data Rec :: [*] -> * where --   RNil :: Rec '[] --   (:*) :: !a -> !(Rec as) -> Rec (a ': as)++-- -- dual pairing +-- class Dual a where+--   dual :: Num r => a -> (a -> r)++-- -- | Dual numbers DD (alternative take, using a type family for the first variation)+-- data DD a = Dd a (Adj a)+-- class Diff a where type Adj a :: *+-- instance Diff Double where type Adj Double = Double+++  -- data SDRec s as where --   SDNil :: SDRec s '[]