diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for downhill + +## 0.1.0.0 -- 2021-12-12 + +* First version diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,14 @@ +Copyright 2021 Andrius Stankevičius + +Permission is hereby granted, free of charge, to any person obtaining a copy of this +software and associated documentation files (the "Software"), to deal in the Software +without restriction, including without limitation the rights to use, copy, modify, +merge, publish, distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, +INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT +HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/downhill.cabal b/downhill.cabal new file mode 100644 --- /dev/null +++ b/downhill.cabal @@ -0,0 +1,58 @@ +cabal-version: 2.4 + +name: downhill +version: 0.1.0.0 +synopsis: Reverse mode automatic differentiation +homepage: https://andriusstank.github.io/downhill/ +description: + Simple and well typed implementation of reverse mode automatic differentiation. + See home page for more detailed + description. +bug-reports: https://github.com/andriusstank/downhill/issues +license: MIT +license-file: LICENSE +author: Andrius Stankevičius +maintainer: floppycat@gmail.com +-- copyright: +category: Math +extra-source-files: CHANGELOG.md + +library + exposed-modules: Downhill.Linear.Expr, + Downhill.Linear.BackGrad, + Downhill.Linear.Backprop, + Downhill.Linear.Lift, + Downhill.Linear.Prelude, + Downhill.Internal.Graph.Types, + Downhill.Internal.Graph.OpenMap, + Downhill.Internal.Graph.NodeMap, + Downhill.Internal.Graph.OpenGraph, + Downhill.Internal.Graph.Graph, + Downhill.Grad, + Downhill.BVar, + Downhill.BVar.Num + Downhill.BVar.Prelude, + Downhill.BVar.Traversable, + Downhill.TH + -- other-modules: + -- other-extensions: + build-depends: base >= 4.12.0.0 && <4.17, + containers >= 0.6.5 && < 0.7, + reflection >= 2.1.6 && < 2.2, + template-haskell >= 2.16.0 && < 2.19, + transformers >= 0.5.6 && < 0.6, + th-abstraction >= 0.4.3 && < 0.5, + unordered-containers >= 0.2.14 && < 0.3, + vector-space >= 0.16 && < 0.17, + hs-source-dirs: src + other-modules: + default-language: Haskell2010 + ghc-options: -Wall + +test-suite downhill-test + type: exitcode-stdio-1.0 + main-is: Main.hs + other-modules: DownhillTest.Point, DownhillTest.Traversable, DownhillTest.TH, DownhillTest.TestTHOptions + build-depends: base, downhill, tasty, tasty-hunit, vector-space + hs-source-dirs: test + default-language: Haskell2010 diff --git a/src/Downhill/BVar.hs b/src/Downhill/BVar.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/BVar.hs @@ -0,0 +1,133 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +module Downhill.BVar + ( BVar (..), + var, + constant, + backprop, + ) +where + +import Data.AdditiveGroup (AdditiveGroup) +import Data.AffineSpace (AffineSpace ((.+^), (.-.))) +import qualified Data.AffineSpace as AffineSpace +import Data.VectorSpace + ( AdditiveGroup (..), + VectorSpace ((*^)), + ) +import qualified Data.VectorSpace as VectorSpace +import Downhill.Grad + ( Dual (evalGrad), + HasFullGrad, + HasGrad (Grad, MScalar, Tang), + HasGradAffine, + ) +import Downhill.Linear.BackGrad + ( BackGrad (..), + realNode, + ) +import qualified Downhill.Linear.Backprop as BP +import Downhill.Linear.Expr (BasicVector, Expr (ExprVar), FullVector) +import Downhill.Linear.Lift (lift2_dense) +import Prelude hiding (id, (.)) + +-- | Variable is a value paired with derivative. +data BVar r a = BVar + { bvarValue :: a, + bvarGrad :: BackGrad r (Grad a) + } + +instance (AdditiveGroup b, HasFullGrad b) => AdditiveGroup (BVar r b) where + zeroV = BVar zeroV zeroV + negateV (BVar y0 dy) = BVar (negateV y0) (negateV dy) + BVar y0 dy ^-^ BVar z0 dz = BVar (y0 ^-^ z0) (dy ^-^ dz) + BVar y0 dy ^+^ BVar z0 dz = BVar (y0 ^+^ z0) (dy ^+^ dz) + +instance (Num b, HasFullGrad b, MScalar b ~ b) => Num (BVar r b) where + (BVar f0 df) + (BVar g0 dg) = BVar (f0 + g0) (df ^+^ dg) + (BVar f0 df) - (BVar g0 dg) = BVar (f0 - g0) (df ^-^ dg) + (BVar f0 df) * (BVar g0 dg) = BVar (f0 * g0) (f0 *^ dg ^+^ g0 *^ df) + negate (BVar f0 df) = BVar (negate f0) (negateV df) + abs (BVar f0 df) = BVar (abs f0) (signum f0 *^ df) -- TODO: ineffiency: multiplication by 1 + signum (BVar f0 _) = BVar (signum f0) zeroV + fromInteger x = BVar (fromInteger x) zeroV + +sqr :: Num a => a -> a +sqr x = x * x + +rsqrt :: Floating a => a -> a +rsqrt x = recip (sqrt x) + +instance (Fractional b, HasFullGrad b, MScalar b ~ b) => Fractional (BVar r b) where + fromRational x = BVar (fromRational x) zeroV + recip (BVar x dx) = BVar (recip x) (df *^ dx) + where + df = negate (recip (sqr x)) + BVar x dx / BVar y dy = BVar (x / y) ((recip y *^ dx) ^-^ ((x / sqr y) *^ dy)) + +instance (Floating b, HasFullGrad b, MScalar b ~ b) => Floating (BVar r b) where + pi = BVar pi zeroV + exp (BVar x dx) = BVar (exp x) (exp x *^ dx) + log (BVar x dx) = BVar (log x) (recip x *^ dx) + sin (BVar x dx) = BVar (sin x) (cos x *^ dx) + cos (BVar x dx) = BVar (cos x) (negate (sin x) *^ dx) + asin (BVar x dx) = BVar (asin x) (rsqrt (1 - sqr x) *^ dx) + acos (BVar x dx) = BVar (acos x) (negate (rsqrt (1 - sqr x)) *^ dx) + atan (BVar x dx) = BVar (atan x) (recip (1 + sqr x) *^ dx) + sinh (BVar x dx) = BVar (sinh x) (cosh x *^ dx) + cosh (BVar x dx) = BVar (cosh x) (sinh x *^ dx) + asinh (BVar x dx) = BVar (asinh x) (rsqrt (1 + sqr x) *^ dx) + acosh (BVar x dx) = BVar (acosh x) (rsqrt (sqr x - 1) *^ dx) + atanh (BVar x dx) = BVar (atanh x) (recip (1 - sqr x) *^ dx) + +instance + ( VectorSpace v, + HasFullGrad v, + Tang v ~ v, + FullVector (MScalar v), + Grad (MScalar v) ~ MScalar v + ) => + VectorSpace (BVar r v) + where + type Scalar (BVar r v) = BVar r (MScalar v) + BVar a da *^ BVar v dv = BVar (a *^ v) (lift2_dense bpA bpV da dv) + where + bpA :: Grad v -> MScalar v + bpA dz = evalGrad dz v + bpV :: Grad v -> Grad v + bpV dz = a *^ dz + +instance (HasFullGrad p, HasGradAffine p) => AffineSpace (BVar r p) where + type Diff (BVar r p) = BVar r (Tang p) + BVar y0 dy .+^ BVar z0 dz = BVar (y0 .+^ z0) (dy ^+^ dz) + BVar y0 dy .-. BVar z0 dz = BVar (y0 .-. z0) (dy ^-^ dz) + +-- | A variable with derivative of zero. +constant :: forall r a. FullVector (Grad a) => a -> BVar r a +constant x = BVar x zeroV + +-- | A variable with identity derivative. +var :: a -> BVar (Grad a) a +var x = BVar x (realNode ExprVar) + +--backprop :: forall a p. (HasGrad p, BasicVector a) => BVar a p -> GradBuilder p -> a +--backprop (BVar _y0 x) = BP.backprop x + +-- | Reverse mode differentiation. +-- +-- +backprop :: forall r a. (HasGrad a, FullVector (Grad a), BasicVector r) => BVar r a -> Grad a -> r +backprop (BVar _y0 x) = BP.backprop x diff --git a/src/Downhill/BVar/Num.hs b/src/Downhill/BVar/Num.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/BVar/Num.hs @@ -0,0 +1,99 @@ +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +module Downhill.BVar.Num + ( -- | Automatic differentiation for @Num@ hierarchy. + -- + -- Polymorphic functions of type such as @Num a => a -> a@ + -- can't be differentiated directly, because 'backprop' needs some additional instances. + -- 'AsNum' wrapper provides those instances. + -- + -- @ + -- derivative :: (forall b. Floating b => b -> b) -> (forall a. Floating a => a -> a) + -- derivative fun x0 = backpropNum (fun (var (AsNum x0))) + -- @ + + AsNum (..), + NumBVar, + numbvarValue, + var, + constant, + backpropNum + ) +where + +import Data.AffineSpace (AffineSpace (..)) +import Data.Semigroup (Sum (Sum, getSum)) +import Data.VectorSpace (AdditiveGroup (..), VectorSpace (..), zeroV) +import Downhill.BVar (BVar (bvarValue), backprop) +import qualified Downhill.BVar as BVar +import Downhill.Grad + ( Dual (evalGrad), + HasGrad (Grad, Metric, MScalar, Tang), + MetricTensor (MtCovector, MtVector, evalMetric), + ) +import Downhill.Linear.Expr (BasicVector (..), FullVector (identityBuilder, negateBuilder, scaleBuilder)) + +-- | @AsNum a@ implements many instances in terms of @Num a@ instance. +newtype AsNum a = AsNum {unAsNum :: a} + deriving (Show) + deriving (Num) via a + deriving (Fractional) via a + deriving (Floating) via a + +instance Num a => Dual (AsNum a) (AsNum a) (AsNum a) where + evalGrad = (*) + +instance Num a => HasGrad (AsNum a) where + type MScalar (AsNum a) = AsNum a + type Grad (AsNum a) = AsNum a + type Tang (AsNum a) = AsNum a + type Metric (AsNum a) = AsNum a + +instance Num a => MetricTensor (AsNum a) where + type MtVector (AsNum a) = AsNum a + type MtCovector (AsNum a) = AsNum a + evalMetric (AsNum m) (AsNum x) = AsNum (m * x) + +instance Num a => AdditiveGroup (AsNum a) where + zeroV = 0 + (^+^) = (+) + (^-^) = (-) + negateV = negate + +instance Num a => VectorSpace (AsNum a) where + type Scalar (AsNum a) = AsNum a + (*^) = (*) + +instance Num a => BasicVector (AsNum a) where + type VecBuilder (AsNum a) = Sum a + sumBuilder = AsNum . getSum + +instance Num a => FullVector (AsNum a) where + identityBuilder = Sum . unAsNum + negateBuilder = Sum . negate . unAsNum + scaleBuilder (AsNum x) (AsNum y) = Sum $ x * y + +instance Num a => AffineSpace (AsNum a) where + type Diff (AsNum a) = AsNum a + AsNum x .-. AsNum y = AsNum (x - y) + AsNum x .+^ AsNum y = AsNum (x + y) + +type NumBVar a = BVar (AsNum a) (AsNum a) + +constant :: forall a. Num a => a -> NumBVar a +constant = BVar.constant @(AsNum a) @(AsNum a) . AsNum + +var :: Num a => a -> NumBVar a +var = BVar.var . AsNum + +backpropNum :: forall a. Num a => NumBVar a -> a +backpropNum x = unAsNum $ backprop @(AsNum a) @(AsNum a) x (AsNum 1) + +numbvarValue :: NumBVar a -> a +numbvarValue = unAsNum . bvarValue diff --git a/src/Downhill/BVar/Prelude.hs b/src/Downhill/BVar/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/BVar/Prelude.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ViewPatterns #-} + +module Downhill.BVar.Prelude + ( -- * Tuples + + -- | Pattern synonyms @T2@, @T3@ pack and unpack tuples: + -- + -- @ + -- fstBVar :: (HasGrad a, HasGrad b) => BVar r (a, b) -> BVar r a + -- fstBVar (T2 a _b) = a + -- + -- tieBVar :: (HasGrad a, HasGrad b) => BVar r a -> BVar r b -> BVar r (a, b) + -- tieBVar a b = T2 a b + -- @ + pattern T2, + pattern T3, + ) +where + +import Downhill.BVar (BVar (BVar)) +import Downhill.Grad (HasGrad) +import qualified Downhill.Linear.Prelude as Linear +import Prelude () + +toPair :: (HasGrad a, HasGrad b) => BVar r (a, b) -> (BVar r a, BVar r b) +toPair (BVar (x, y) (Linear.T2 dx dy)) = (BVar x dx, BVar y dy) + +{-# COMPLETE T2 #-} + +pattern T2 :: (HasGrad a, HasGrad b) => BVar r a -> BVar r b -> BVar r (a, b) +pattern T2 a b <- + (toPair -> (a, b)) + where + T2 (BVar a da) (BVar b db) = BVar (a, b) (Linear.T2 da db) + +toTriple :: (HasGrad a, HasGrad b, HasGrad c) => BVar r (a, b, c) -> (BVar r a, BVar r b, BVar r c) +toTriple (BVar (x, y, z) (Linear.T3 dx dy dz)) = (BVar x dx, BVar y dy, BVar z dz) + +{-# COMPLETE T3 #-} + +pattern T3 :: (HasGrad a, HasGrad b, HasGrad c) => BVar r a -> BVar r b -> BVar r c -> BVar r (a, b, c) +pattern T3 a b c <- + (toTriple -> (a, b, c)) + where + T3 (BVar a da) (BVar b db) (BVar c dc) = BVar (a, b, c) (Linear.T3 da db dc) diff --git a/src/Downhill/BVar/Traversable.hs b/src/Downhill/BVar/Traversable.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/BVar/Traversable.hs @@ -0,0 +1,284 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Easy backpropagation when all variables have the same type. +-- +-- @ +-- data MyRecord a = ... +-- deriving (Functor, Foldable, Traversable) +-- +-- deriving via (TraversableVar MyRecord a) instance HasGrad a => HasGrad (MyRecord a) +-- @ +-- +-- = Gradient type +-- One might excect gradient type to be @type Grad (MyRecord a) = MyRecord (Grad a)@, but it's not +-- the case, because record could contain additional members apart from @a@s, for example: +-- +-- @ +-- data MyPoint a = MyPoint +-- { +-- , pointLabel :: String +-- , pointX :: a +-- , pointY :: a +-- } +-- @ +-- +-- and @MyPoint (Grad a)@ can't be made @VectorSpace@. Gradient type @Grad (MyRecord a)@ +-- is a newtype wrapper over @IntMap@ +-- that is not exported. +-- + +module Downhill.BVar.Traversable + ( -- * Backpropagate + backpropTraversable, + backpropTraversable_GradOnly, + backpropTraversable_ValueAndGrad, + + -- * Split + splitTraversable, + + -- * TraversableVar + TraversableVar (..), + ) +where + +import Control.Monad.Trans.State.Strict (State, evalState, get, put) +import Data.AdditiveGroup (AdditiveGroup, sumV) +import Data.Foldable (toList) +import Data.IntMap (IntMap) +import qualified Data.IntMap as IntMap +import Data.Maybe (fromMaybe) +import Data.VectorSpace (AdditiveGroup (negateV, zeroV, (^+^), (^-^)), VectorSpace (Scalar, (*^))) +import qualified Data.VectorSpace as VectorSpace +import Downhill.BVar (BVar (BVar, bvarGrad, bvarValue), backprop, var) +import Downhill.Grad + ( Dual (evalGrad), + HasGrad (Grad, MScalar, Metric, Tang), + MetricTensor + ( MtCovector, + MtVector, + evalMetric + ), + ) +import Downhill.Linear.BackGrad (BackGrad (BackGrad), castBackGrad, realNode) +import Downhill.Linear.Expr + ( BasicVector (VecBuilder, sumBuilder), + Expr (ExprSum), + FullVector, + SparseVector (unSparseVector), + Term, + ) +import Downhill.Linear.Lift (lift1_sparse) +import GHC.Generics (Generic) + +-- | Provides HasGrad instance for use in deriving via +newtype TraversableVar f a = TraversableVar {unTraversableVar :: f a} + deriving stock (Functor, Foldable, Traversable) + +newtype TraversableMetric f a = TraversableMetric (Metric a) + deriving (Generic) + +instance AdditiveGroup (Metric a) => AdditiveGroup (TraversableMetric f a) + +instance VectorSpace (Metric a) => VectorSpace (TraversableMetric f a) where + type Scalar (TraversableMetric f a) = Scalar (Metric a) + +instance + ( MetricTensor (Metric a), + MtVector (Metric a) ~ Tang a, + MtCovector (Metric a) ~ Grad a, + Dual s (Tang a) (Grad a) + ) => + MetricTensor (TraversableMetric f a) + where + type MtVector (TraversableMetric f a) = IntmapVector f (Tang a) + type MtCovector (TraversableMetric f a) = IntmapVector f (Grad a) + evalMetric (TraversableMetric m) (IntmapVector da) = IntmapVector (IntMap.map (evalMetric m) da) + +instance HasGrad a => HasGrad (TraversableVar f a) where + type MScalar (TraversableVar f a) = MScalar a + type Tang (TraversableVar f a) = IntmapVector f (Tang a) + type Grad (TraversableVar f a) = IntmapVector f (Grad a) + type Metric (TraversableVar f a) = TraversableMetric f a + +-- | @IntmapVector@ serves as a gradient of 'TraversableVar'. +newtype IntmapVector f v = IntmapVector {unIntmapVector :: IntMap v} + deriving (Show) + +instance AdditiveGroup a => AdditiveGroup (IntmapVector f a) where + zeroV = IntmapVector IntMap.empty + negateV (IntmapVector v) = IntmapVector (negateV <$> v) + IntmapVector u ^+^ IntmapVector v = IntmapVector (IntMap.unionWith (^+^) u v) + IntmapVector u ^-^ IntmapVector v = IntmapVector (IntMap.mergeWithKey combine only1 only2 u v) + where + combine _key x y = Just (x ^-^ y) + only1 = id + only2 = fmap negateV + +instance VectorSpace v => VectorSpace (IntmapVector f v) where + type Scalar (IntmapVector f v) = VectorSpace.Scalar v + a *^ (IntmapVector v) = IntmapVector (fmap (a *^) v) + +instance Dual s dv v => Dual s (IntmapVector f dv) (IntmapVector f v) where + evalGrad (IntmapVector dv) (IntmapVector v) = sumV $ IntMap.intersectionWith evalGrad dv v + +deriving via (IntMap v) instance Semigroup v => Semigroup (IntmapVector f v) + +deriving via (IntMap v) instance Monoid v => Monoid (IntmapVector f v) + +instance BasicVector v => BasicVector (IntmapVector f v) where + type VecBuilder (IntmapVector f v) = IntmapVector f (VecBuilder v) + sumBuilder (IntmapVector v) = IntmapVector (fmap sumBuilder v) + +imap :: + forall t a b. + Traversable t => + (Int -> a -> b) -> + t a -> + t b +imap mkBVar' xs' = evalState (traverse getmkvar xs') 0 + where + getmkvar :: a -> State Int b + getmkvar x = do + index <- get + put (index + 1) + return (mkBVar' index x) + +-- | Note that @splitTraversable@ won't be useful +-- for top level @BVar@, because the type @Grad (f a)@ is not exposed. +splitTraversable :: + forall f r a. + ( Traversable f, + Grad (f a) ~ Grad (TraversableVar f a), + HasGrad a + ) => + BVar r (f a) -> + f (BVar r a) +splitTraversable (BVar xs dxs) = vars + where + vars :: f (BVar r a) + vars = imap mkBVar xs + mkBVar :: Int -> a -> BVar r a + mkBVar index x = + let mkBuilder :: VecBuilder (Grad a) -> IntmapVector f (VecBuilder (Grad a)) + mkBuilder dx = IntmapVector (IntMap.singleton index dx) + in BVar x (lift1_sparse mkBuilder dxs) + +lift1_sparseT :: + forall r a z. + BasicVector z => + (VecBuilder z -> VecBuilder a) -> + BackGrad r a -> + Term r (SparseVector z) +lift1_sparseT fa (BackGrad f) = f (fa . unSparseVector) + +-- Not exported, because it is untested and hardly useful. +_joinTraversable :: + forall f r a. + ( Traversable f, + Grad (f a) ~ Grad (TraversableVar f a), + HasGrad a, + FullVector (Grad a) + ) => + f (BVar r a) -> + BVar r (f a) +_joinTraversable x = BVar values (castBackGrad node) + where + values :: f a + values = bvarValue <$> x + grads :: f (BackGrad r (Grad a)) + grads = bvarGrad <$> x + terms :: [Term r (SparseVector (IntmapVector f (Grad a)))] + terms = toList (imap mkTerm grads) + mkTerm :: Int -> BackGrad r (Grad a) -> Term r (SparseVector (IntmapVector f (Grad a))) + mkTerm index = lift1_sparseT (lookupIntMap index) + lookupIntMap :: Int -> IntmapVector f x -> x + lookupIntMap key (IntmapVector intmap) = case IntMap.lookup key intmap of + Nothing -> error "Downhill BUG: Bad index in joinTraversable" + Just value -> value + node :: BackGrad r (SparseVector (IntmapVector f (Grad a))) + node = realNode (ExprSum terms) + +-- | @backpropTraversable one combine fun@ +-- +-- @one@ is a value to be backpropagated. In case of @p@ being scalar, set @one@ +-- to 1 to compute unscaled gradient. +-- +-- @combine@ is given value of a parameter and its gradient to construct result, +-- just like @zipWith@. +-- +-- @fun@ is the function to be differentiated. +backpropTraversable :: + forall f a b p. + ( Traversable f, + Grad (f a) ~ Grad (TraversableVar f a), + HasGrad a, + HasGrad p, + FullVector (Grad p) + ) => + Grad p -> + (a -> Grad a -> b) -> + (forall r. f (BVar r a) -> BVar r p) -> + f a -> + f b +backpropTraversable one combine fun x = imap makeResult x + where + splitX :: f (BVar (Grad (f a)) a) + splitX = splitTraversable (var x) + + y :: BVar (Grad (f a)) p + y = fun splitX + + grad :: IntMap (Grad a) + IntmapVector grad = backprop y one + + lookupGrad i = fromMaybe zeroV (IntMap.lookup i grad) + + makeResult :: Int -> a -> b + makeResult i x' = combine x' (lookupGrad i) + +{-# ANN backpropTraversable_GradOnly "HLint: ignore Use camelCase" #-} + +-- | Like 'backpropTraversable', but returns gradient only. +backpropTraversable_GradOnly :: + forall f a p. + ( Traversable f, + Grad (f a) ~ Grad (TraversableVar f a), + HasGrad a, + HasGrad p, + FullVector (Grad p) + ) => + Grad p -> + (forall r. f (BVar r a) -> BVar r p) -> + f a -> + f (Grad a) +backpropTraversable_GradOnly one = backpropTraversable one gradOnly + where + gradOnly _value grad = grad + +-- | 'backpropTraversable' specialized to return a pair of value and gradient. +{-# ANN backpropTraversable_ValueAndGrad "HLint: ignore Use camelCase" #-} +backpropTraversable_ValueAndGrad :: + forall f a p. + ( Traversable f, + Grad (f a) ~ Grad (TraversableVar f a), + HasGrad a, + HasGrad p, + FullVector (Grad p) + ) => + Grad p -> + (forall r. f (BVar r a) -> BVar r p) -> + f a -> + f (a, Grad a) +backpropTraversable_ValueAndGrad one = backpropTraversable one (,) diff --git a/src/Downhill/Grad.hs b/src/Downhill/Grad.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Grad.hs @@ -0,0 +1,218 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +module Downhill.Grad + ( Dual (..), + MetricTensor (..), + HasGrad (..), + GradBuilder, + HasFullGrad, + HasGradAffine, + ) +where + +import Data.AffineSpace (AffineSpace (Diff)) +import Data.Kind (Type) +import Data.VectorSpace (AdditiveGroup ((^+^)), VectorSpace (Scalar, (*^))) +import qualified Data.VectorSpace as VectorSpace +import Downhill.Linear.Expr (BasicVector (VecBuilder), FullVector) +import GHC.Generics (Generic) + +-- | Dual of a vector @v@ is a linear map @v -> Scalar v@. +class + ( AdditiveGroup s, + VectorSpace v, + VectorSpace dv, + VectorSpace.Scalar v ~ s, + VectorSpace.Scalar dv ~ s + ) => + Dual s v dv + where + -- if evalGrad goes to HasGrad class, parameter p is ambiguous + evalGrad :: dv -> v -> s + +-- | @MetricTensor@ converts gradients to vectors. +-- +-- It is really inverse of a metric tensor, because it maps cotangent +-- space into tangent space. Gradient descent doesn't need metric tensor, +-- it needs inverse. + +class + ( Dual (Scalar g) (MtVector g) (MtCovector g), + VectorSpace g + ) => + MetricTensor g + where + type MtVector g :: Type + type MtCovector g :: Type + + -- | @m@ must be symmetric: + -- + -- @evalGrad x (evalMetric m y) = evalGrad y (evalMetric m x)@ + evalMetric :: g -> MtCovector g -> MtVector g + + -- | @innerProduct m x y = evalGrad x (evalMetric m y)@ + innerProduct :: g -> MtCovector g -> MtCovector g -> Scalar g + innerProduct g x y = evalGrad x (evalMetric g y) + + -- | @sqrNorm m x = innerProduct m x x@ + sqrNorm :: g -> MtCovector g -> Scalar g + sqrNorm g x = innerProduct g x x + +-- | @HasGrad@ is a collection of types and constraints that are useful +-- in many places. It helps to keep type signatures short. + +-- TODO: FullVector or not? +-- TODO: Metric or not? +class + ( Dual (MScalar p) (Tang p) (Grad p), + MetricTensor (Metric p), + MtVector (Metric p) ~ Tang p, + MtCovector (Metric p) ~ Grad p, + BasicVector (Tang p), + BasicVector (Grad p) + ) => + HasGrad p + where + -- | Scalar of @Tang p@ and @Grad p@. + type MScalar p :: Type + + -- | Tangent vector of manifold @p@. If p is 'AffineSpace', @Tang p@ should + -- be @'Diff' p@. If @p@ is 'VectorSpace', @Tang p@ might be the same as @p@ itself. + type Tang p :: Type + + -- | Dual of tangent space of @p@. + type Grad p :: Type + + -- | A 'MetricTensor'. + type Metric p :: Type + +type GradBuilder v = VecBuilder (Grad v) + +type HasFullGrad p = (HasGrad p, FullVector (Grad p)) + +type HasGradAffine p = + ( AffineSpace p, + HasGrad p, + HasGrad (Tang p), + Tang p ~ Diff p, + Tang (Tang p) ~ Tang p, + Grad (Tang p) ~ Grad p + ) + +instance Dual Integer Integer Integer where + evalGrad = (*) + +instance MetricTensor Integer where + type MtVector Integer = Integer + type MtCovector Integer = Integer + evalMetric m x = m * x + +instance HasGrad Integer where + type MScalar Integer = Integer + type Tang Integer = Integer + type Grad Integer = Integer + type Metric Integer = Integer + +instance (Dual s a da, Dual s b db) => Dual s (a, b) (da, db) where + evalGrad (a, b) (x, y) = evalGrad a x ^+^ evalGrad b y + +instance (Dual s a da, Dual s b db, Dual s c dc) => Dual s (a, b, c) (da, db, dc) where + evalGrad (a, b, c) (x, y, z) = evalGrad a x ^+^ evalGrad b y ^+^ evalGrad c z + +instance (MetricTensor ma, MetricTensor mb, Scalar ma ~ Scalar mb) => MetricTensor (ma, mb) where + type MtVector (ma, mb) = (MtVector ma, MtVector mb) + type MtCovector (ma, mb) = (MtCovector ma, MtCovector mb) + evalMetric (ma, mb) (a, b) = (evalMetric ma a, evalMetric mb b) + sqrNorm (ma, mb) (a, b) = sqrNorm ma a ^+^ sqrNorm mb b + +instance + ( HasGrad a, + HasGrad b, + MScalar b ~ MScalar a + ) => + HasGrad (a, b) + where + type MScalar (a, b) = MScalar a + type Grad (a, b) = (Grad a, Grad b) + type Tang (a, b) = (Tang a, Tang b) + type Metric (a, b) = (Metric a, Metric b) + +instance + ( MetricTensor ma, + MetricTensor mb, + MetricTensor mc, + Scalar ma ~ Scalar mb, + Scalar ma ~ Scalar mc + ) => + MetricTensor (ma, mb, mc) + where + type MtVector (ma, mb, mc) = (MtVector ma, MtVector mb, MtVector mc) + type MtCovector (ma, mb, mc) = (MtCovector ma, MtCovector mb, MtCovector mc) + evalMetric (ma, mb, mc) (a, b, c) = (evalMetric ma a, evalMetric mb b, evalMetric mc c) + sqrNorm (ma, mb, mc) (a, b, c) = sqrNorm ma a ^+^ sqrNorm mb b ^+^ sqrNorm mc c + +instance + ( HasGrad a, + HasGrad b, + HasGrad c, + MScalar b ~ MScalar a, + MScalar c ~ MScalar a + ) => + HasGrad (a, b, c) + where + type MScalar (a, b, c) = MScalar a + type Grad (a, b, c) = (Grad a, Grad b, Grad c) + type Tang (a, b, c) = (Tang a, Tang b, Tang c) + type Metric (a, b, c) = (Metric a, Metric b, Metric c) + +instance Dual Float Float Float where + evalGrad = (*) + +instance MetricTensor Float where + type MtVector Float = Float + type MtCovector Float = Float + evalMetric m dv = m * dv + +instance HasGrad Float where + type MScalar Float = Float + type Grad Float = Float + type Tang Float = Float + type Metric Float = Float + +instance Dual Double Double Double where + evalGrad = (*) + +instance MetricTensor Double where + type MtVector Double = Double + type MtCovector Double = Double + evalMetric m dv = m * dv + +instance HasGrad Double where + type MScalar Double = Double + type Grad Double = Double + type Tang Double = Double + type Metric Double = Double + +newtype L2 v = L2 (Scalar v) + deriving (Generic) + +instance AdditiveGroup (Scalar v) => AdditiveGroup (L2 v) + +instance (AdditiveGroup (Scalar v), Num (Scalar v)) => VectorSpace (L2 v) where + type Scalar (L2 v) = Scalar v + x *^ L2 y = L2 (x * y) + +instance (AdditiveGroup a, Num a, a ~ Scalar v, Dual a v v) => MetricTensor (L2 v) where + type MtVector (L2 v) = v + type MtCovector (L2 v) = v + evalMetric (L2 a) u = a *^ u + innerProduct (L2 a) x y = a * evalGrad x y + sqrNorm g x = innerProduct g x x diff --git a/src/Downhill/Internal/Graph/Graph.hs b/src/Downhill/Internal/Graph/Graph.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Internal/Graph/Graph.hs @@ -0,0 +1,188 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +module Downhill.Internal.Graph.Graph + ( -- * Graph type + Graph (..), Node(..), + SomeGraph (..), + -- * Evaluate + evalGraph, + -- * Transpose + transposeGraph, + --transposeFwdGraph, + --transposeBackGraph, + -- * Construct + unsafeFromOpenGraph, + ) +where + +import Data.Either (partitionEithers) +import Data.Functor.Identity (Identity (Identity, runIdentity)) +import Downhill.Internal.Graph.NodeMap + ( IsNodeSet, + NodeKey, + NodeMap, + KeyAndValue (KeyAndValue), + SomeNodeMap (SomeNodeMap), + ) +import qualified Downhill.Internal.Graph.NodeMap as NodeMap +import Downhill.Internal.Graph.OpenGraph (OpenGraph (OpenGraph), OpenNode (OpenNode), OpenEdge (OpenEdge), OpenEndpoint (OpenSourceNode, OpenInnerNode)) +import Downhill.Internal.Graph.Types (FwdFun (FwdFun), BackFun) +import Downhill.Linear.Expr (BasicVector (VecBuilder, sumBuilder)) +import Prelude hiding (head, tail) +import GHC.Stack (callStack, prettyCallStack, HasCallStack) + +data Endpoint s a v where + SourceNode :: Endpoint s a a + InnerNode :: NodeKey s v -> Endpoint s a v + +data Edge s e a v where + Edge :: e u v -> Endpoint s a u -> Edge s e a v + +{-| Inner node. This does not include initial node. Contains a list +of ingoing edges. -} +data Node s e a v = BasicVector v => Node [Edge s e a v] + +data Graph s e a z = BasicVector a => + Graph + { graphInnerNodes :: NodeMap s (Node s e a), + graphFinalNode :: Node s e a z + } + +data SomeGraph e a z where + SomeGraph :: IsNodeSet s => Graph s e a z -> SomeGraph e a z + +{- `Edge` stores head endpoint only. `AnyEdge` stores both endpoints. -} +data AnyEdge s e a z = forall u v. + AnyEdge + { _edgeTail :: Endpoint s z v, + _edgeLabel :: e u v, + _edgeHead :: Endpoint s a u + } + +-- | Forward mode evaluation +evalGraph :: forall s x z. Graph s FwdFun z x -> z -> x +evalGraph (Graph nodes finalNode) dz = evalNode finalNode + where + evalParent :: forall v. Endpoint s z v -> v + evalParent = \case + SourceNode -> dz + InnerNode nodeName -> runIdentity (NodeMap.lookup innerValues nodeName) + evalEdge :: Edge s FwdFun z v -> VecBuilder v + evalEdge (Edge (FwdFun f) tail) = f $ evalParent tail + evalNode :: Node s FwdFun z v -> v + evalNode (Node xs) = sumBuilder (mconcat [evalEdge x | x <- xs]) + innerValues :: NodeMap s Identity + innerValues = NodeMap.map (Identity . evalNode) nodes + +nodeEdges :: forall s f a z x. NodeKey s x -> Node s f a x -> [AnyEdge s f a z] +nodeEdges name (Node xs) = go <$> xs + where + go :: Edge s f a x -> AnyEdge s f a z + go (Edge f head) = AnyEdge (InnerNode name) f head + +allGraphEdges :: forall s f a z. Graph s f a z -> [AnyEdge s f a z] +allGraphEdges (Graph innerNodes (Node es)) = finalEdges ++ innerEdges + where + innerEdges :: [AnyEdge s f a z] + innerEdges = concat (NodeMap.toListWith nodeEdges innerNodes) + finalEdges :: [AnyEdge s f a z] + finalEdges = wrapFinalEdge <$> es + where + wrapFinalEdge :: Edge s f a z -> AnyEdge s f a z + wrapFinalEdge (Edge f head) = AnyEdge SourceNode f head + +sortByTail :: + forall s f da dz. + AnyEdge s f da dz -> + Either (Edge s f da dz) (KeyAndValue s (Edge s f da)) +sortByTail (AnyEdge tail f head) = case tail of + SourceNode -> Left (Edge f head) + InnerNode x -> Right (KeyAndValue x (Edge f head)) + +flipAnyEdge :: (forall u v. f u v -> g v u) -> AnyEdge s f a z -> AnyEdge s g z a +flipAnyEdge flipF (AnyEdge tail f head) = AnyEdge head (flipF f) tail + +{- BasicVector constraint is needed to construct a node. + `NodeMap s NodeDict` is a list of all nodes. +-} +data NodeDict x = BasicVector x => NodeDict + +emptyNodeMap :: forall s e z. NodeMap s NodeDict -> NodeMap s (Node s e z) +emptyNodeMap = NodeMap.map emptyNode + where + emptyNode :: forall x. NodeDict x -> Node s e z x + emptyNode = \case + NodeDict -> Node [] + +edgeListToGraph :: + forall s e a z. + (IsNodeSet s, BasicVector a, BasicVector z) => + NodeMap s NodeDict -> + [AnyEdge s e z a] -> + Graph s e z a +edgeListToGraph nodes flippedEdges = Graph innerNodes (Node initialEdges) + where + initialEdges :: [Edge s e z a] + innerEdges :: [KeyAndValue s (Edge s e z)] + (initialEdges, innerEdges) = partitionEithers (sortByTail <$> flippedEdges) + prependToMap :: KeyAndValue s (Edge s e z) -> NodeMap s (Node s e z) -> NodeMap s (Node s e z) + prependToMap (KeyAndValue key edge) = NodeMap.adjust prependToNode key + where + prependToNode (Node edges) = Node (edge : edges) + innerNodes = foldr prependToMap (emptyNodeMap nodes) innerEdges + +graphNodes :: Graph s f da dz -> NodeMap s NodeDict +graphNodes (Graph env _) = NodeMap.map go env + where + go :: Node s f da dv -> NodeDict dv + go = \case + Node _ -> NodeDict + +-- | Reverse edges. Turns reverse mode evaluation into forward mode. +transposeGraph :: forall s f g a z. IsNodeSet s => (forall u v. f u v -> g v u) -> Graph s f a z -> Graph s g z a +transposeGraph flipEdge g@(Graph _ (Node _)) = edgeListToGraph (graphNodes g) flippedEdges + where edges :: [AnyEdge s f a z] + edges = allGraphEdges g + flippedEdges :: [AnyEdge s g z a] + flippedEdges = flipAnyEdge flipEdge <$> edges + +_mapEdges :: forall s f g a z. (forall u v. f u v -> g u v) -> Graph s f a z -> Graph s g a z +_mapEdges f (Graph inner final) = Graph (NodeMap.map go inner) (go final) + where + go :: Node s f a v -> Node s g a v + go (Node xs) = Node [goEdge x | x <- xs] + goEdge :: Edge p f a x -> Edge p g a x + goEdge (Edge e x) = Edge (f e) x + +unsafeConstructGraph :: forall s a v. (IsNodeSet s, BasicVector a, HasCallStack) => NodeMap s (OpenNode a) -> OpenNode a v -> Graph s BackFun a v +unsafeConstructGraph m x = Graph (NodeMap.map mkExpr m) (mkExpr x) + where + mkExpr :: forall x. OpenNode a x -> Node s BackFun a x + mkExpr = \case + OpenNode terms -> Node (mkTerm <$> terms) + mkTerm :: forall x. OpenEdge a x -> Edge s BackFun a x + mkTerm = \case + OpenEdge f x' -> Edge f (mkArg x') + mkArg :: forall u. OpenEndpoint a u -> Endpoint s a u + mkArg = \case + OpenSourceNode -> SourceNode + OpenInnerNode key -> case NodeMap.tryLookup m key of + Just (key', _value) -> InnerNode key' + Nothing -> error ("Downhill: invalid key in constructGraph\n" ++ prettyCallStack callStack) + +-- | Will crash if graph has invalid keys +unsafeFromOpenGraph :: (BasicVector a, HasCallStack) => OpenGraph a v -> SomeGraph BackFun a v +unsafeFromOpenGraph (OpenGraph x m) = + case NodeMap.fromOpenMap m of + SomeNodeMap m' -> SomeGraph (unsafeConstructGraph m' x) diff --git a/src/Downhill/Internal/Graph/NodeMap.hs b/src/Downhill/Internal/Graph/NodeMap.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Internal/Graph/NodeMap.hs @@ -0,0 +1,120 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} + +module Downhill.Internal.Graph.NodeMap + ( -- * NodeMap + NodeMap, + NodeKey, + -- * Construction + fromOpenMap, + generate, + -- * Query + lookup, + tryLookup, + toList, + toListWith, + elems, + -- * Modify + map, + mapWithKey, + adjust, + zipWith, + -- * Node Set + IsNodeSet, + SomeNodeMap (..), + KeyAndValue (..), + ) +where + +import Control.Applicative (Const) +import Data.Data (Proxy (Proxy)) +import Data.Reflection (Reifies (reflect), reify) +import Downhill.Internal.Graph.OpenMap (OpenKey, OpenMap, SomeOpenItem (SomeOpenItem)) +import qualified Downhill.Internal.Graph.OpenMap as OpenMap +import Prelude (Maybe (Just, Nothing), const, error, (.), (<$>)) + +type role NodeKey nominal nominal + +-- | Valid key, guaranteed to be a member of @s@ +newtype NodeKey s x = NodeKey (OpenKey x) + +-- | @NodeMap s f@ is a map where value of type @f x@ is associated with key @NodeKey s x@. +-- Type variable `s` tracks the set of nodes. Lookups never fail. Maps can +-- be zipped without losing any nodes. +newtype NodeMap s f = NodeMap {unNodeMap :: OpenMap f} + +data KeyAndValue s f = forall x. KeyAndValue (NodeKey s x) (f x) + +class IsNodeSet s where + allNodes :: OpenMap Proxy + +map :: forall s f g. (forall v. f v -> g v) -> NodeMap s f -> NodeMap s g +map f = NodeMap . OpenMap.map f . unNodeMap + +mapWithKey :: forall s f g. (forall x. NodeKey s x -> f x -> g x) -> NodeMap s f -> NodeMap s g +mapWithKey f (NodeMap x) = NodeMap (OpenMap.mapWithKey f' x) + where + f' :: OpenKey dx -> f dx -> g dx + f' key' = f (NodeKey key') + +toList :: NodeMap s f -> [KeyAndValue s f] +toList = toListWith KeyAndValue + +toListWith :: forall s f r. (forall x. NodeKey s x -> f x -> r) -> NodeMap s f -> [r] +toListWith f (NodeMap m) = wrap <$> OpenMap.toList m + where + wrap :: SomeOpenItem f -> r + wrap (SomeOpenItem key value) = f (NodeKey key) value + +elems :: NodeMap s (Const b) -> [b] +elems (NodeMap m) = OpenMap.elems m + +lookup :: NodeMap s f -> NodeKey s v -> f v +lookup (NodeMap m) (NodeKey key) = + case OpenMap.lookup m key of + Just x -> x + Nothing -> error "oh fuck" + +-- | If key belongs to @s@, @tryLookup@ will return a proof of this fact +-- and a corresponding value from the map. Otherwise returns @Nothing@. +tryLookup :: NodeMap s f -> OpenKey x -> Maybe (NodeKey s x, f x) +tryLookup (NodeMap m) key = + case OpenMap.lookup m key of + Just x -> Just (NodeKey key, x) + Nothing -> Nothing + +generate :: forall s f. IsNodeSet s => (forall x. NodeKey s x -> f x) -> NodeMap s f +generate f = case allNodes @s of + m -> mapWithKey (\key _ -> f key) (NodeMap m) + +zipWith :: forall s f g h. (forall x. f x -> g x -> h x) -> NodeMap s f -> NodeMap s g -> NodeMap s h +zipWith f (NodeMap x) (NodeMap y) = NodeMap (OpenMap.intersectionWith f x y) + +adjust :: forall s f x. (f x -> f x) -> NodeKey s x -> NodeMap s f -> NodeMap s f +adjust f (NodeKey key) (NodeMap m) = NodeMap (OpenMap.adjust f key m) + +data NodeSetWrapper s + +instance Reifies s (OpenMap Proxy) => IsNodeSet (NodeSetWrapper s) where + allNodes = reflect @s Proxy + +-- | 'NodeMap' with existential set of nodes. +data SomeNodeMap f where + SomeNodeMap :: IsNodeSet s => NodeMap s f -> SomeNodeMap f + +fromOpenMap :: forall f. OpenMap f -> SomeNodeMap f +fromOpenMap x = reify nodes go + where + nodes :: OpenMap Proxy + nodes = OpenMap.map (const Proxy) x + go :: forall s. Reifies s (OpenMap Proxy) => Proxy s -> SomeNodeMap f + go _proxy = SomeNodeMap @(NodeSetWrapper s) (NodeMap x) diff --git a/src/Downhill/Internal/Graph/OpenGraph.hs b/src/Downhill/Internal/Graph/OpenGraph.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Internal/Graph/OpenGraph.hs @@ -0,0 +1,86 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralisedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Downhill.Internal.Graph.OpenGraph + ( OpenEdge (..), + OpenEndpoint (..), + OpenNode (..), + OpenGraph (..), + recoverSharing, + ) +where + +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.State.Strict (StateT (..), get, modify) +import Downhill.Internal.Graph.OpenMap (OpenKey, OpenMap) +import qualified Downhill.Internal.Graph.OpenMap as OpenMap +import Downhill.Internal.Graph.Types (BackFun (BackFun)) +import Downhill.Linear.Expr (BasicVector, Expr (ExprSum, ExprVar), Term (..)) +import Prelude hiding (lookup) + +data OpenEndpoint a v where + OpenSourceNode :: OpenEndpoint a a + OpenInnerNode :: OpenKey v -> OpenEndpoint a v + +data OpenEdge a v where + OpenEdge :: BackFun u v -> OpenEndpoint a u -> OpenEdge a v + +data OpenNode a v = BasicVector v => OpenNode [OpenEdge a v] + +-- | Maintains a cache of visited 'Expr's. +newtype TreeBuilder a r = TreeCache {unTreeCache :: StateT (OpenMap (OpenNode a)) IO r} + deriving (Functor, Applicative, Monad) + +insertIntoCache :: OpenKey dv -> OpenNode a dv -> TreeBuilder a () +insertIntoCache name value = TreeCache $ modify (OpenMap.insert name value) + +-- | @buildExpr action key@ will run @action@, associate result with @key@ and +-- store it in cache. If @key@ is already in cache, @action@ will not be run. +buildExpr :: + TreeBuilder a (OpenNode a v) -> + Expr a v -> + TreeBuilder a (OpenKey v, OpenNode a v) +buildExpr action key = do + name <- TreeCache (lift (OpenMap.makeOpenKey key)) + cache <- TreeCache get + case OpenMap.lookup cache name of + Just x -> return (name, x) + Nothing -> do + value <- action + insertIntoCache name value + return (name, value) + +runTreeBuilder :: forall a g dv. TreeBuilder a (g dv) -> IO (g dv, OpenMap (OpenNode a)) +runTreeBuilder rs_x = runStateT (unTreeCache rs_x) OpenMap.empty + +-- | Computational graph under construction. "Open" refers to the set of the nodes – new nodes can be +-- added to this graph. Once the graph is complete the set of nodes will be frozen +-- and the type of the graph will become 'Graph' ("Downhill.Internal.Graph" module). +data OpenGraph a z = OpenGraph (OpenNode a z) (OpenMap (OpenNode a)) + +goEdges :: BasicVector v => [Term a v] -> TreeBuilder a (OpenNode a v) +goEdges xs = do + xs' <- traverse goSharing4term xs + return $ OpenNode xs' + +goSharing4arg :: forall a v. Expr a v -> TreeBuilder a (OpenEndpoint a v) +goSharing4arg key = case key of + ExprVar -> return OpenSourceNode + ExprSum xs -> do + (gRef, _) <- buildExpr (goEdges xs) key + return (OpenInnerNode gRef) + +goSharing4term :: forall a v. Term a v -> TreeBuilder a (OpenEdge a v) +goSharing4term = \case + Term f arg -> do + arg' <- goSharing4arg arg + return (OpenEdge (BackFun f) arg') + +-- | Collects duplicate nodes in 'Expr' tree and converts it to a graph. +recoverSharing :: forall a z. BasicVector z => [Term a z] -> IO (OpenGraph a z) +recoverSharing xs = do + (final_node, graph) <- runTreeBuilder (goEdges xs) + return (OpenGraph final_node graph) diff --git a/src/Downhill/Internal/Graph/OpenMap.hs b/src/Downhill/Internal/Graph/OpenMap.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Internal/Graph/OpenMap.hs @@ -0,0 +1,110 @@ +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RankNTypes #-} + +module Downhill.Internal.Graph.OpenMap + ( -- * OpenMap + OpenMap, + OpenKey, + SomeOpenItem (SomeOpenItem), + -- * Construction + makeOpenKey, + empty, + insert, + -- * Query + lookup, + toList, + elems, + -- * Modify + map, + mapWithKey, + mapMaybe, + adjust, + intersectionWith, + ) +where + +import Control.Applicative (Const (Const)) +import Control.Exception (evaluate) +import Data.HashMap.Lazy (HashMap) +import qualified Data.HashMap.Lazy as HashMap +import Data.Kind (Type) +import GHC.Base (Any, Maybe (Just, Nothing), coerce) +import GHC.StableName (StableName) +import System.Mem.StableName (makeStableName) +import Unsafe.Coerce (unsafeCoerce) +import Prelude (Functor (fmap), IO, Monad (return), (.), (<$>)) + +data SomeExpr f = forall v. SomeExpr (f v) + +-- | A key of @OpenMap@. +newtype OpenKey x = OpenKey (StableName Any) + +-- | Heterogeneous map with 'StableName' as a key. +newtype OpenMap (f :: Type -> Type) = OpenMap {unOpenMap :: HashMap (StableName Any) (SomeExpr f)} + +-- | Key and value. +data SomeOpenItem f = forall x. SomeOpenItem (OpenKey x) (f x) + +empty :: OpenMap f +empty = OpenMap HashMap.empty + +map :: forall f g. (forall x. f x -> g x) -> OpenMap f -> OpenMap g +map f = OpenMap . fmap go . unOpenMap + where + go (SomeExpr y) = SomeExpr (f y) + +mapMaybe :: forall f g. (forall x. f x -> Maybe (g x)) -> OpenMap f -> OpenMap g +mapMaybe f = OpenMap . HashMap.mapMaybe go . unOpenMap + where + go (SomeExpr y) = case f y of + Just fy -> Just (SomeExpr fy) + Nothing -> Nothing + +mapWithKey :: forall f g. (forall d. OpenKey d -> f d -> g d) -> OpenMap f -> OpenMap g +mapWithKey f = OpenMap . HashMap.mapWithKey go . unOpenMap + where + go key (SomeExpr y) = SomeExpr (f (OpenKey key) y) + +lookup :: OpenMap f -> OpenKey x -> Maybe (f x) +lookup (OpenMap m) (OpenKey k) = unsafeCastTypeSomeExpr <$> HashMap.lookup k m + +toList :: OpenMap f -> [SomeOpenItem f] +toList = fmap wrap . HashMap.toList . unOpenMap + where + wrap :: (StableName Any, SomeExpr f) -> SomeOpenItem f + wrap (key, x) = case x of + SomeExpr x' -> SomeOpenItem (OpenKey key) x' + +elems :: OpenMap (Const b) -> [b] +elems = fmap unSomeExpr . HashMap.elems . unOpenMap + where + unSomeExpr :: SomeExpr (Const r) -> r + unSomeExpr (SomeExpr (Const x)) = x + +unsafeCastTypeSomeExpr :: SomeExpr f -> f v +unsafeCastTypeSomeExpr = \case + SomeExpr x -> unsafeCoerce x + +intersectionWith :: forall f g h. (forall x. f x -> g x -> h x) -> OpenMap f -> OpenMap g -> OpenMap h +intersectionWith f (OpenMap x) (OpenMap y) = OpenMap (HashMap.intersectionWith f' x y) + where + f' (SomeExpr x') sy = SomeExpr (f x' y') + where + y' = unsafeCastTypeSomeExpr sy + +insert :: forall f dx. OpenKey dx -> f dx -> OpenMap f -> OpenMap f +insert (OpenKey k) x (OpenMap m) = OpenMap (HashMap.insert k (SomeExpr x) m) + +adjust :: forall f x. (f x -> f x) -> OpenKey x -> OpenMap f -> OpenMap f +adjust f (OpenKey key) (OpenMap m) = OpenMap m' + where + m' = HashMap.adjust f' key m + f' x = SomeExpr (f (unsafeCastTypeSomeExpr x)) + +makeOpenKey :: f v -> IO (OpenKey v) +makeOpenKey x = do + x' <- evaluate x + z <- makeStableName x' + return (OpenKey (coerce z)) diff --git a/src/Downhill/Internal/Graph/Types.hs b/src/Downhill/Internal/Graph/Types.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Internal/Graph/Types.hs @@ -0,0 +1,37 @@ +{-| Types of nodes and edges of the computational graph. + +Parameters: + + * @p@ - is parent node; might be 'OpenKey' or 'NodeKey' + + * @e@ - edge type + + * @a@ - type of the initial node of expression + + * @v@ - type of the node. +-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE GADTs #-} +module Downhill.Internal.Graph.Types +( + -- * Linear functions + BackFun(..), FwdFun(..), + flipBackFun, flipFwdFun +) + where + +import Downhill.Linear.Expr (BasicVector (VecBuilder)) + + +-- | Edge type for backward mode evaluation +newtype BackFun u v = BackFun {unBackFun :: v -> VecBuilder u} + +-- | Edge type for forward mode evaluation +newtype FwdFun u v = FwdFun {unFwdFun :: u -> VecBuilder v} + +flipBackFun :: BackFun u v -> FwdFun v u +flipBackFun (BackFun f) = FwdFun f + +flipFwdFun :: FwdFun u v -> BackFun v u +flipFwdFun (FwdFun f) = BackFun f + diff --git a/src/Downhill/Linear/BackGrad.hs b/src/Downhill/Linear/BackGrad.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Linear/BackGrad.hs @@ -0,0 +1,90 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE TypeApplications #-} + +module Downhill.Linear.BackGrad + ( BackGrad (..), + realNode, + inlineNode, + sparseNode, + castBackGrad, + ) +where + +import Data.VectorSpace + ( AdditiveGroup (..), + Scalar, + VectorSpace (..), + ) +import Downhill.Linear.Expr + ( BasicVector (VecBuilder), + Expr (ExprSum), + FullVector (identityBuilder, negateBuilder, scaleBuilder), + Term (Term), SparseVector (unSparseVector), + ) + +-- | Linear expression, made for backpropagation. +-- It is similar to @'Expr' 'BackFun'@, but has a more flexible form. +newtype BackGrad a v + = BackGrad + ( forall x. + (x -> VecBuilder v) -> + Term a x + ) + +-- | Creates a @BackGrad@ that is backed by a real node. Gradient of type @v@ will be computed and stored +-- in a graph for this node. +{-# ANN module "HLint: ignore Avoid lambda using `infix`" #-} + +realNode :: Expr a v -> BackGrad a v +realNode x = BackGrad (\f -> Term f x) + +-- | @inlineNode f x@ will apply function @f@ to variable @x@ without creating a node. All of the gradients +-- coming to this expression will be forwarded to the parents of @x@. However, if this expression is used +-- more than once, @f@ will be evaluated multiple times, too. It is intended to be used for @newtype@ wrappers. +-- @inlineNode f x@ also doesn't prevent +-- compiler to inline and optimize @x@ +inlineNode :: + forall r u v. + (VecBuilder v -> VecBuilder u) -> + BackGrad r u -> + BackGrad r v +inlineNode f (BackGrad g) = BackGrad go + where + go :: forall x. (x -> VecBuilder v) -> Term r x + go h = g (f . h) + +sparseNode :: + forall r a z. + BasicVector z => + (VecBuilder z -> VecBuilder a) -> + BackGrad r a -> + BackGrad r z +sparseNode fa (BackGrad x) = castBackGrad (realNode node) + where + fa' = fa . unSparseVector + node :: Expr r (SparseVector z) + node = ExprSum [x fa'] + +-- | @BackGrad@ doesn't track the type of the node. Type of @BackGrad@ can be changed freely +-- as long as @VecBuilder@ stays the same. +castBackGrad :: + forall r v z. + VecBuilder z ~ VecBuilder v => + BackGrad r v -> + BackGrad r z +castBackGrad (BackGrad g) = BackGrad g + +instance (FullVector v) => AdditiveGroup (BackGrad r v) where + zeroV = realNode (ExprSum []) + negateV (BackGrad x) = realNode (ExprSum [x negateBuilder]) + BackGrad x ^+^ BackGrad y = realNode (ExprSum [x identityBuilder, y identityBuilder]) + BackGrad x ^-^ BackGrad y = realNode (ExprSum [x identityBuilder, y negateBuilder]) + +instance FullVector v => VectorSpace (BackGrad r v) where + type Scalar (BackGrad r v) = Scalar v + a *^ BackGrad v = realNode (ExprSum [v (scaleBuilder a)]) diff --git a/src/Downhill/Linear/Backprop.hs b/src/Downhill/Linear/Backprop.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Linear/Backprop.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} + +module Downhill.Linear.Backprop + ( -- * Backpropagation + backprop, + + -- * Graph + buildGraph, + --abstractBackprop, + ) +where + +import Downhill.Internal.Graph.Graph + ( SomeGraph (..), + evalGraph, + transposeGraph, + ) +import qualified Downhill.Internal.Graph.Graph as Graph +import Downhill.Internal.Graph.OpenGraph (recoverSharing) +import Downhill.Internal.Graph.Types (BackFun, flipBackFun) +import Downhill.Linear.BackGrad (BackGrad (..), castBackGrad) +import Downhill.Linear.Expr + ( BasicVector (VecBuilder), + FullVector (identityBuilder), + SparseVector (SparseVector, unSparseVector), + Term, + ) +import GHC.IO.Unsafe (unsafePerformIO) + +buildGraph :: + forall a v. + (BasicVector a, BasicVector v) => + [Term a v] -> + IO (SomeGraph BackFun a v) +buildGraph fidentityBuilder = do + og <- recoverSharing fidentityBuilder + return (Graph.unsafeFromOpenGraph og) + +abstractBackprop :: + forall a u v. + (BasicVector a, BasicVector u, BasicVector v) => + BackGrad a u -> + (v -> VecBuilder u) -> + v -> + a +abstractBackprop (BackGrad f) builder x = + case unsafePerformIO (buildGraph [f builder]) of + SomeGraph g -> evalGraph (transposeGraph flipBackFun g) x + +_backprop :: forall a v. (BasicVector a, BasicVector v) => BackGrad a v -> VecBuilder v -> a +_backprop dvar x = + abstractBackprop @a @(SparseVector v) @(SparseVector v) + sparseDVar + unSparseVector + (SparseVector x) + where + sparseDVar :: BackGrad a (SparseVector v) + sparseDVar = castBackGrad dvar + +-- | Purity of this function depends on laws of arithmetic +-- and linearity law of 'Term'. If your addition is approximately +-- associative, then this function is approximately pure. Fair? +backprop :: forall a v. (BasicVector a, FullVector v) => BackGrad a v -> v -> a +backprop dvar = abstractBackprop dvar identityBuilder diff --git a/src/Downhill/Linear/Expr.hs b/src/Downhill/Linear/Expr.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Linear/Expr.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +module Downhill.Linear.Expr + ( -- * Expression + Expr (..), + Term (..), + + -- * Vectors + BasicVector (..), + FullVector (..), + SparseVector (..), + DenseVector (..), + DenseBuilder (..), + toDenseBuilder, + + -- * Misc + maybeToMonoid, + ) +where + +import Data.Kind (Type) +import Data.Maybe (fromMaybe) +import Data.Semigroup (Sum (Sum, getSum)) +import Data.VectorSpace (AdditiveGroup (..), VectorSpace (..)) + +-- | Argument @f@ in @Term f x@ must be /linear/ function. That's a law. +data Term a v where + Term :: (v -> VecBuilder u) -> Expr a u -> Term a v + +-- | @Expr a v@ represents a linear expression of type @v@, containing some free variables of type @a@. +data Expr a v where + ExprVar :: Expr a a + ExprSum :: BasicVector v => [Term a v] -> Expr a v + +class Monoid (VecBuilder v) => BasicVector v where + -- | @VecBuilder v@ is a sparse representation of vector @v@. Edges of a computational graph + -- produce builders, which are then summed into vectors in nodes. Monoid operation '<>' + -- means addition of vectors, but it doesn't need to compute the sum immediately - it + -- might defer computation until 'sumBuilder' is evaluated. + -- + -- @ + -- sumBuilder mempty = zeroV + -- sumBuilder (x <> y) = sumBuilder x ^+^ sumBuilder y + -- @ + -- + -- 'mempty' must be cheap. '<>' must be O(1). + type VecBuilder v :: Type + + sumBuilder :: VecBuilder v -> v + +maybeToMonoid :: Monoid m => Maybe m -> m +maybeToMonoid = fromMaybe mempty + +instance BasicVector Integer where + type VecBuilder Integer = Sum Integer + sumBuilder = getSum + +instance (BasicVector a, BasicVector b) => BasicVector (a, b) where + type VecBuilder (a, b) = Maybe (VecBuilder a, VecBuilder b) + sumBuilder = sumPair . maybeToMonoid + where + sumPair (a, b) = (sumBuilder a, sumBuilder b) + +instance (BasicVector a, BasicVector b, BasicVector c) => BasicVector (a, b, c) where + type VecBuilder (a, b, c) = Maybe (VecBuilder a, VecBuilder b, VecBuilder c) + sumBuilder = sumTriple . maybeToMonoid + where + sumTriple (a, b, c) = (sumBuilder a, sumBuilder b, sumBuilder c) + +instance BasicVector Float where + type VecBuilder Float = Sum Float + sumBuilder = getSum + +instance BasicVector Double where + type VecBuilder Double = Sum Double + sumBuilder = getSum + +-- | Full-featured vector. +-- +-- Gradients are linear functions and form a vector space. +-- @FullVector@ class provides functionality that is needed to +-- make 'VectorSpace' instances. +class (BasicVector v, VectorSpace v) => FullVector v where + identityBuilder :: v -> VecBuilder v + negateBuilder :: v -> VecBuilder v + scaleBuilder :: Scalar v -> v -> VecBuilder v + +instance FullVector Float where + identityBuilder = Sum + negateBuilder = Sum . negate + scaleBuilder x = Sum . (x *) + +instance FullVector Double where + identityBuilder = Sum + negateBuilder = Sum . negate + scaleBuilder x = Sum . (x *) + +instance FullVector Integer where + identityBuilder = Sum + negateBuilder = Sum . negate + scaleBuilder x = Sum . (x *) + +instance (Scalar a ~ Scalar b, FullVector a, FullVector b) => FullVector (a, b) where + identityBuilder (x, y) = Just (identityBuilder x, identityBuilder y) + negateBuilder (x, y) = Just (negateBuilder x, negateBuilder y) + scaleBuilder a (x, y) = Just (scaleBuilder a x, scaleBuilder a y) + +instance (s ~ Scalar a, s ~ Scalar b, s ~ Scalar c, FullVector a, FullVector b, FullVector c) => FullVector (a, b, c) where + identityBuilder (x, y, z) = Just (identityBuilder x, identityBuilder y, identityBuilder z) + negateBuilder (x, y, z) = Just (negateBuilder x, negateBuilder y, negateBuilder z) + scaleBuilder a (x, y, z) = Just (scaleBuilder a x, scaleBuilder a y, scaleBuilder a z) + +-- | Normally graph node would compute the sum of gradients and then +-- propagate it to ancestor nodes. That's the best strategy when +-- some computation needs to be performed for backpropagation. +-- Some operations, like constructing/deconstructing tuples or +-- wrapping/unwrapping, don't need to compute the sum. Doing so only +-- destroys sparsity. A node of type @SparseVector v@ won't sum +-- the gradients, it will simply forward builders to its parents. +newtype SparseVector v = SparseVector + {unSparseVector :: VecBuilder v} + +deriving via (VecBuilder v) instance Semigroup (VecBuilder v) => Semigroup (SparseVector v) + +instance Monoid (VecBuilder v) => BasicVector (SparseVector v) where + type VecBuilder (SparseVector v) = VecBuilder v + sumBuilder = SparseVector + +newtype DenseSemibuilder v = DenseSemibuilder {_unDenseSemibuilder :: v} + +instance AdditiveGroup v => Semigroup (DenseSemibuilder v) where + DenseSemibuilder x <> DenseSemibuilder y = DenseSemibuilder (x ^+^ y) + +newtype DenseBuilder v = DenseBuilder (Maybe v) + deriving (Semigroup, Monoid) via (Maybe (DenseSemibuilder v)) + +toDenseBuilder :: v -> DenseBuilder v +toDenseBuilder = DenseBuilder . Just + +-- | When sparsity is not needed, we can use vector @v@ as a builder of itself. +-- @DenseVector@ takes care of that. +newtype DenseVector v = DenseVector v + deriving (AdditiveGroup, VectorSpace) via v + +instance AdditiveGroup v => BasicVector (DenseVector v) where + type VecBuilder (DenseVector v) = DenseBuilder v + sumBuilder (DenseBuilder Nothing) = DenseVector zeroV + sumBuilder (DenseBuilder (Just x)) = DenseVector x + +instance VectorSpace v => FullVector (DenseVector v) where + identityBuilder (DenseVector v) = DenseBuilder (Just v) + negateBuilder (DenseVector v) = DenseBuilder (Just (negateV v)) + scaleBuilder a (DenseVector v) = DenseBuilder (Just (a *^ v)) + +instance FullVector v => AdditiveGroup (Expr a v) where + zeroV = ExprSum [] + negateV x = ExprSum [Term negateBuilder x] + x ^+^ y = ExprSum [Term identityBuilder x, Term identityBuilder y] + x ^-^ y = ExprSum [Term identityBuilder x, Term negateBuilder y] + +instance FullVector dv => VectorSpace (Expr da dv) where + type Scalar (Expr da dv) = Scalar dv + a *^ v = ExprSum [Term (scaleBuilder a) v] diff --git a/src/Downhill/Linear/Lift.hs b/src/Downhill/Linear/Lift.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Linear/Lift.hs @@ -0,0 +1,136 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PartialTypeSignatures #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | While 'BackGrad' is intended to be simple to construct manually, this module provides a way to do +-- that with a bit less of boilerplate. +module Downhill.Linear.Lift + ( -- * Lifts + lift1, + lift2, + lift3, + + -- * Dense lifts + lift1_dense, + lift2_dense, + lift3_dense, + + -- * Lifts for 'SparseVector' + lift1_sparse, + lift2_sparse, + lift3_sparse, + ) +where + +import Downhill.Linear.BackGrad (BackGrad (..), castBackGrad, realNode) +import Downhill.Linear.Expr (BasicVector (..), Expr (ExprSum), FullVector (identityBuilder), SparseVector (unSparseVector)) +import Prelude hiding (fst, snd, zip) + +lift1 :: + forall z r a. + BasicVector z => + (z -> VecBuilder a) -> + BackGrad r a -> + BackGrad r z +lift1 fa (BackGrad da) = realNode node + where + node = ExprSum [da fa] + +lift2 :: + forall z r a b. + BasicVector z => + (z -> VecBuilder a) -> + (z -> VecBuilder b) -> + BackGrad r a -> + BackGrad r b -> + BackGrad r z +lift2 fa fb (BackGrad da) (BackGrad db) = realNode node + where + node = ExprSum [da fa, db fb] + +lift3 :: + forall z r a b c. + BasicVector z => + (z -> VecBuilder a) -> + (z -> VecBuilder b) -> + (z -> VecBuilder c) -> + BackGrad r a -> + BackGrad r b -> + BackGrad r c -> + BackGrad r z +lift3 fa fb fc (BackGrad da) (BackGrad db) (BackGrad dc) = realNode node + where + node = ExprSum [da fa, db fb, dc fc] + +-- | Same as 'sparseNode', included here for completeness. +lift1_sparse :: + forall r a z. + BasicVector z => + (VecBuilder z -> VecBuilder a) -> + BackGrad r a -> + BackGrad r z +lift1_sparse fa = castBackGrad . lift1 @(SparseVector z) fa' + where + fa' = fa . unSparseVector + +lift2_sparse :: + forall r a b z. + BasicVector z => + (VecBuilder z -> VecBuilder a) -> + (VecBuilder z -> VecBuilder b) -> + BackGrad r a -> + BackGrad r b -> + BackGrad r z +lift2_sparse fa fb a b = castBackGrad $ lift2 @(SparseVector z) fa' fb' a b + where + fa' = fa . unSparseVector + fb' = fb . unSparseVector + +lift3_sparse :: + forall r a b c z. + BasicVector z => + (VecBuilder z -> VecBuilder a) -> + (VecBuilder z -> VecBuilder b) -> + (VecBuilder z -> VecBuilder c) -> + BackGrad r a -> + BackGrad r b -> + BackGrad r c -> + BackGrad r z +lift3_sparse fa fb fc a b c = + castBackGrad $ + lift3 @(SparseVector z) fa' fb' fc' a b c + where + fa' = fa . unSparseVector + fb' = fb . unSparseVector + fc' = fc . unSparseVector + +lift1_dense :: + (BasicVector v, FullVector a) => + ((v -> a) -> BackGrad r a -> BackGrad r v) +lift1_dense fa = lift1 (identityBuilder . fa) + +lift2_dense :: + (BasicVector v, FullVector a, FullVector b) => + (v -> a) -> + (v -> b) -> + BackGrad r a -> + BackGrad r b -> + BackGrad r v +lift2_dense fa fb = lift2 (identityBuilder . fa) (identityBuilder . fb) + +lift3_dense :: + (BasicVector v, FullVector a, FullVector b, FullVector c) => + (v -> a) -> + (v -> b) -> + (v -> c) -> + BackGrad r a -> + BackGrad r b -> + BackGrad r c -> + BackGrad r v +lift3_dense fa fb fc = lift3 (identityBuilder . fa) (identityBuilder . fb) (identityBuilder . fc) diff --git a/src/Downhill/Linear/Prelude.hs b/src/Downhill/Linear/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/Linear/Prelude.hs @@ -0,0 +1,98 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE NoImplicitPrelude #-} + +module Downhill.Linear.Prelude + ( pattern T2, + pattern T3, + ) +where + +import Downhill.Linear.BackGrad (BackGrad) +import Downhill.Linear.Expr (BasicVector (VecBuilder), maybeToMonoid) +import qualified Downhill.Linear.Lift as Lift +import Prelude (Maybe (Just), Monoid (mempty), fmap, (.)) +import qualified Prelude + +splitPair :: forall r a b. (BasicVector a, BasicVector b) => BackGrad r (a, b) -> (BackGrad r a, BackGrad r b) +splitPair x = (bg1, bg2) + where + go1 :: VecBuilder a -> VecBuilder (a, b) + go2 :: VecBuilder b -> VecBuilder (a, b) + go1 da = Just (da, mempty) + go2 db = Just (mempty, db) + bg1 :: BackGrad r a + bg2 :: BackGrad r b + bg1 = Lift.lift1_sparse go1 x + bg2 = Lift.lift1_sparse go2 x + +toTriple :: + forall r a b c. + (BasicVector a, BasicVector b, BasicVector c) => + BackGrad r (a, b, c) -> + (BackGrad r a, BackGrad r b, BackGrad r c) +toTriple x = (bg1, bg2, bg3) + where + go1 :: VecBuilder a -> VecBuilder (a, b, c) + go2 :: VecBuilder b -> VecBuilder (a, b, c) + go3 :: VecBuilder c -> VecBuilder (a, b, c) + go1 da = Just (da, mempty, mempty) + go2 db = Just (mempty, db, mempty) + go3 dc = Just (mempty, mempty, dc) + bg1 :: BackGrad r a + bg2 :: BackGrad r b + bg3 :: BackGrad r c + bg1 = Lift.lift1_sparse go1 x + bg2 = Lift.lift1_sparse go2 x + bg3 = Lift.lift1_sparse go3 x + +-- | +-- +-- @ +-- getFst :: (BasicVector (DualOf a), BasicVector (DualOf b)) => BackGrad r (a, b) -> BackGrad r a +-- getFst (T2 x _) = x +-- @ +-- +-- @ +-- mkPair :: (BasicVector (DualOf a), BasicVector (DualOf b)) => BackGrad r a -> BackGrad r b -> BackGrad r (a, b) +-- mkPair x y = (T2 x y) +-- @ +{-# COMPLETE T2 #-} + +pattern T2 :: forall r a b. (BasicVector a, BasicVector b) => BackGrad r a -> BackGrad r b -> BackGrad r (a, b) +pattern T2 a b <- + (splitPair -> (a, b)) + where + T2 a b = Lift.lift2_sparse go1 go2 a b + where + go1 :: VecBuilder (a, b) -> VecBuilder a + go2 :: VecBuilder (a, b) -> VecBuilder b + go1 = maybeToMonoid . fmap Prelude.fst + go2 = maybeToMonoid . fmap Prelude.snd + +{-# COMPLETE T3 #-} + +pattern T3 :: + forall r a b c. + (BasicVector a, BasicVector b, BasicVector c) => + BackGrad r a -> + BackGrad r b -> + BackGrad r c -> + BackGrad r (a, b, c) +pattern T3 a b c <- + (toTriple -> (a, b, c)) + where + T3 a b c = Lift.lift3_sparse go1 go2 go3 a b c + where + go1 :: VecBuilder (a, b, c) -> VecBuilder a + go2 :: VecBuilder (a, b, c) -> VecBuilder b + go3 :: VecBuilder (a, b, c) -> VecBuilder c + go1 = maybeToMonoid . fmap (\(x, _, _) -> x) + go2 = maybeToMonoid . fmap (\(_, x, _) -> x) + go3 = maybeToMonoid . fmap (\(_, _, x) -> x) diff --git a/src/Downhill/TH.hs b/src/Downhill/TH.hs new file mode 100644 --- /dev/null +++ b/src/Downhill/TH.hs @@ -0,0 +1,917 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Use like this: +-- +-- @ +-- mkHasGradInstances +-- defaultBVarOptions +-- [d| +-- instance HasGrad MyRecord where +-- type MScalar MyRecord = Float +-- |] +-- @ +-- +-- Instance declaration passed to @mkHasGradInstances@ gives two important bits of information: +-- +-- * Type variables for @MyRecord@, which can be concrete types (such as @instance HasGrad (MyRecord Float)@) +-- or regular type variables (@instance HasGrad (MyRecord a)@) +-- +-- * Scalar type. +-- +module Downhill.TH + ( + mkHasGradInstances, + AffineSpaceOptions (..), + RecordNamer (..), + BVarOptions (..), + defaultBVarOptions, + ) +where + +import Control.Monad +import Data.AdditiveGroup ((^+^), (^-^)) +import Data.AffineSpace (AffineSpace (Diff, (.+^), (.-.))) +import Data.Foldable (traverse_) +import qualified Data.Map as Map +import Data.Maybe (catMaybes) +import Data.VectorSpace (AdditiveGroup (negateV, zeroV), VectorSpace (Scalar, (*^))) +import Downhill.BVar (BVar (BVar)) +import Downhill.Grad + ( Dual (evalGrad), + HasGrad (Grad, MScalar, Metric, Tang), + MetricTensor (MtCovector, MtVector, evalMetric, sqrNorm), + ) +import Downhill.Linear.Expr (BasicVector (VecBuilder, sumBuilder)) +import Downhill.Linear.Lift (lift1_sparse) +import GHC.Records (HasField (getField)) +import Language.Haskell.TH + ( Bang (Bang), + Con (NormalC, RecC), + Cxt, + Dec (DataD, InstanceD, NewtypeD, SigD), + Exp (AppE, ConE, InfixE, VarE), + Name, + Pat (VarP), + Q, + SourceStrictness (NoSourceStrictness), + SourceUnpackedness (NoSourceUnpackedness), + Type (AppT, ConT, VarT), + nameBase, + newName, + ) +import Language.Haskell.TH.Datatype (ConstructorInfo (constructorFields, constructorName, constructorVariant), ConstructorVariant (InfixConstructor, NormalConstructor, RecordConstructor), DatatypeInfo (datatypeCons, datatypeInstTypes, datatypeName, datatypeVariant, datatypeVars), DatatypeVariant (Newtype), TypeSubstitution (applySubstitution), reifyDatatype) +import Language.Haskell.TH.Datatype.TyVarBndr (TyVarBndrUnit) +import Language.Haskell.TH.Syntax + ( BangType, + Body (NormalB), + Clause (Clause), + Dec (FunD, TySynInstD, ValD), + Exp (AppTypeE), + TyLit (StrTyLit), + TySynEqn (TySynEqn), + Type (ArrowT, EqualityT, LitT, SigT), + VarBangType, + mkNameS, + ) +import qualified Language.Haskell.TH + +data DatatypeFields + = NormalFields [Type] + | RecordFields [(String, Type)] + deriving (Show) + +data DownhillRecord = DownhillRecord + { ddtTypeConName :: Name, + ddtDataConName :: Name, + ddtFieldTypes :: [Type], + ddtFieldNames :: Maybe [String], + ddtTypeVars :: [TyVarBndrUnit], + ddtFieldCount :: Int, + ddtVariant :: DatatypeVariant + } + deriving (Show) + +data RecordNamer = RecordNamer + { typeConNamer :: String -> String, + dataConNamer :: String -> String, + fieldNamer :: String -> String + } + +data RecordTranstorm = RecordTranstorm RecordNamer (Type -> Type) + +data AffineSpaceOptions + = -- | Generate AffineSpace instance + MakeAffineSpace + | -- | Don't generate AffineSpace instance + NoAffineSpace + | -- | Generate AffineSpace instance if @optExcludeFields@ is empty + AutoAffineSpace + +data BVarOptions = BVarOptions + { optTangNamer :: RecordNamer, + optGradNamer :: RecordNamer, + optMetricNamer :: RecordNamer, + optBuilderNamer :: RecordNamer, + optAffineSpace :: AffineSpaceOptions, + -- | List of fields that take no part in differentiation + optExcludeFields :: [String] + } + +pattern ConP :: Name -> [Pat] -> Pat +#if MIN_VERSION_template_haskell(2,18,0) +pattern ConP x y = Language.Haskell.TH.ConP x [] y +#else +pattern ConP x y = Language.Haskell.TH.ConP x y +#endif + +defaultTangRecordNamer :: RecordNamer +defaultTangRecordNamer = + RecordNamer + { typeConNamer = (++ "Tang"), + dataConNamer = (++ "Tang"), + fieldNamer = id + } + +defaultGradRecordNamer :: RecordNamer +defaultGradRecordNamer = + RecordNamer + { typeConNamer = (++ "Grad"), + dataConNamer = (++ "Grad"), + fieldNamer = id + } + +defaultMetricRecordNamer :: RecordNamer +defaultMetricRecordNamer = + RecordNamer + { typeConNamer = (++ "Metric"), + dataConNamer = (++ "Metric"), + fieldNamer = id + } + +defaultBuilderRecordNamer :: RecordNamer +defaultBuilderRecordNamer = + RecordNamer + { typeConNamer = (++ "Builder"), + dataConNamer = (++ "Builder"), + fieldNamer = id + } + +defaultBVarOptions :: BVarOptions +defaultBVarOptions = + BVarOptions + { optTangNamer = defaultTangRecordNamer, + optGradNamer = defaultGradRecordNamer, + optMetricNamer = defaultMetricRecordNamer, + optBuilderNamer = defaultBuilderRecordNamer, + optAffineSpace = AutoAffineSpace, + optExcludeFields = [] + } + +mkConstructor :: DownhillRecord -> Con +mkConstructor record = + case ddtFieldNames record of + Nothing -> + NormalC newConstrName (map mkType (ddtFieldTypes record)) + Just names -> + RecC newConstrName (zipWith mkRecType names (ddtFieldTypes record)) + where + newConstrName :: Name + newConstrName = ddtDataConName record + mkRecType :: String -> Type -> VarBangType + mkRecType name type_ = + ( mkNameS name, + Bang NoSourceUnpackedness NoSourceStrictness, + type_ + ) + mkType :: Type -> BangType + mkType type_ = + ( Bang NoSourceUnpackedness NoSourceStrictness, + type_ + ) + +parseGradConstructor :: Name -> DatatypeInfo -> ConstructorInfo -> [TyVarBndrUnit] -> Q DownhillRecord +parseGradConstructor tyName dinfo cinfo typevars = do + let types = constructorFields cinfo + n = length types + (fieldTypes, fieldNames) <- case constructorVariant cinfo of + NormalConstructor -> return (types, Nothing) + InfixConstructor -> return (types, Nothing) + RecordConstructor fieldNames -> do + return (types, Just (nameBase <$> fieldNames)) + return + DownhillRecord + { ddtTypeConName = tyName, + ddtDataConName = constructorName cinfo, + ddtTypeVars = typevars, + ddtFieldCount = n, + ddtFieldTypes = fieldTypes, + ddtFieldNames = fieldNames, + ddtVariant = datatypeVariant dinfo + } + +parseDownhillRecord :: Name -> DatatypeInfo -> Q (DownhillRecord, ConstructorInfo) +parseDownhillRecord recordName record' = do + let name = datatypeName record' + let typevars = datatypeVars record' + constructors' = datatypeCons record' + constr' <- case constructors' of + [] -> fail (show recordName <> " has no data constructors") + [constr''] -> return constr'' + _ -> fail (show recordName <> " has multiple data constructors") + + r <- parseGradConstructor name record' constr' typevars + return (r, constr') + +elementwiseOp :: DownhillRecord -> Name -> Q Dec +elementwiseOp record = elementwiseOp' record record record + +elementwiseOp' :: DownhillRecord -> DownhillRecord -> DownhillRecord -> Name -> Q Dec +elementwiseOp' leftRecord rightRecord resRecord func = do + let n = ddtFieldCount resRecord + --dataConName :: Name + --dataConName = ddtDataConName record + xs <- replicateM n (newName "x") + ys <- replicateM n (newName "y") + let fieldOp :: Name -> Name -> Exp + fieldOp x y = InfixE (Just (VarE x)) (VarE func) (Just (VarE y)) + resultFields :: [Exp] + resultFields = zipWith fieldOp xs ys + leftPat = ConP (ddtDataConName leftRecord) (map VarP xs) + rightPat = ConP (ddtDataConName rightRecord) (map VarP ys) + rhs :: Exp + rhs = foldl AppE (ConE (ddtDataConName resRecord)) resultFields + dec = + FunD + func + [ Clause + [leftPat, rightPat] + (NormalB rhs) + [] + ] + return dec + +elementwiseValue :: DownhillRecord -> Name -> Q Dec +elementwiseValue record func = do + let n = ddtFieldCount record + dataConName :: Name + dataConName = ddtDataConName record + rhs :: Exp + rhs = foldl AppE (ConE dataConName) (replicate n (VarE 'zeroV)) + dec = ValD (VarP func) (NormalB rhs) [] + return dec + +elementwiseFunc :: DownhillRecord -> Name -> Q Dec +elementwiseFunc record func = do + let n = ddtFieldCount record + dataConName :: Name + dataConName = ddtDataConName record + rhsConName = ddtDataConName record + xs <- case ddtFieldNames record of + Nothing -> replicateM n (newName "x") + Just names -> traverse newName names + let fieldOp :: Name -> Exp + fieldOp = AppE (VarE func) . VarE + resultFields :: [Exp] + resultFields = map fieldOp xs + leftPat = ConP dataConName (map VarP xs) + rhs :: Exp + rhs = foldl AppE (ConE rhsConName) resultFields + dec = + FunD + func + [ Clause + [leftPat] + (NormalB rhs) + [] + ] + return dec + +mkClassInstance :: Name -> Cxt -> DownhillRecord -> [Type] -> [Dec] -> Q [Dec] +mkClassInstance className cxt record instVars decs = do + let recordType = ConT (ddtTypeConName record) + ihead = AppT (ConT className) (foldl AppT recordType instVars) + return [InstanceD Nothing cxt ihead decs] + +mkSemigroupInstance :: Cxt -> DownhillRecord -> [Type] -> Q [Dec] +mkSemigroupInstance cxt record instVars = do + dec <- elementwiseOp record '(<>) + mkClassInstance ''Semigroup cxt record instVars [dec] + +mkAdditiveGroupInstance :: Cxt -> DownhillRecord -> [Type] -> Q [Dec] +mkAdditiveGroupInstance cxt record instVars = do + zeroVDec <- elementwiseValue record 'zeroV + negateDec <- elementwiseFunc record 'negateV + plusDec <- elementwiseOp record '(^+^) + minusDec <- elementwiseOp record '(^-^) + let decs = + [ zeroVDec, + negateDec, + plusDec, + minusDec + ] + mkClassInstance ''AdditiveGroup cxt record instVars decs + +mkVectorSpaceInstance :: DownhillRecord -> Type -> Cxt -> [Type] -> Q [Dec] +mkVectorSpaceInstance record scalarType cxt instVars = do + let n = ddtFieldCount record + dataConName :: Name + dataConName = ddtDataConName record + xs <- case ddtFieldNames record of + Nothing -> replicateM n (newName "x") + Just names -> traverse newName names + + lhsName <- newName "s" + let rightPat = ConP (ddtDataConName record) (map VarP xs) + recordType = foldl AppT (ConT (ddtTypeConName record)) instVars + mulField :: Name -> Exp + mulField y = InfixE (Just (VarE lhsName)) (VarE '(*^)) (Just (VarE y)) + rhsMulV :: Exp + rhsMulV = foldl AppE (ConE dataConName) (map mulField xs) + let vmulDec = + FunD + '(*^) + [ Clause + [VarP lhsName, rightPat] + (NormalB rhsMulV) + [] + ] + scalarTypeDec = + TySynInstD + ( TySynEqn + Nothing + (AppT (ConT ''Scalar) recordType) + scalarType + ) + decs = [scalarTypeDec, vmulDec] + mkClassInstance ''VectorSpace cxt record instVars decs + +mkBasicVectorInstance :: DownhillRecord -> BVarOptions -> Cxt -> [Type] -> Q [Dec] +mkBasicVectorInstance vectorRecord options cxt instVars = do + sumBuilderDec <- mkSumBuilder + mkClassInstance ''BasicVector cxt vectorRecord instVars [vecbuilderDec, sumBuilderDec] + where + n = ddtFieldCount vectorRecord + builderRecord = renameDownhillRecord (builderTransform options) vectorRecord + + -- not an elementiseOp, because right hand side is wrapped in Maybe + mkSumBuilder :: Q Dec + mkSumBuilder = do + builders <- replicateM n (newName "x") + let pat :: Pat + pat = ConP (ddtDataConName builderRecord) (map VarP builders) + rhs :: Exp + rhs = + foldl + AppE + (ConE (ddtDataConName vectorRecord)) + [AppE (VarE 'sumBuilder) (VarE x) | x <- builders] + return $ + FunD + 'sumBuilder + [ Clause [ConP 'Nothing []] (NormalB (VarE 'zeroV)) [], + Clause [ConP 'Just [pat]] (NormalB rhs) [] + ] + + vecbuilderDec = + TySynInstD + ( TySynEqn + Nothing + (AppT (ConT ''VecBuilder) vectorType) + (AppT (ConT ''Maybe) builderType) + ) + where + vectorType = foldl AppT (ConT (ddtTypeConName vectorRecord)) instVars + builderType = foldl AppT (ConT (ddtTypeConName builderRecord)) instVars + +sumVExpr :: [Exp] -> Exp +sumVExpr = \case + [] -> VarE 'zeroV + exps -> foldl1 (zipExpInfix '(^+^)) exps + where + zipExpInfix :: Name -> Exp -> Exp -> Exp + zipExpInfix f x y = InfixE (Just x) (VarE f) (Just y) + +mkDualInstance :: + DownhillRecord -> + DownhillRecord -> + Type -> + Cxt -> + [Type] -> + Q [Dec] +mkDualInstance tangRecord gradRecord scalarType cxt instVars = do + when (ddtFieldCount tangRecord /= ddtFieldCount gradRecord) $ + fail "mkDualInstance: ddtFieldCount tangRecord /= ddtFieldCount gradRecord" + scalarTypeName <- newName "s" + mkClassDec (VarT scalarTypeName) + where + n = ddtFieldCount tangRecord + + -- instance (cxt, AdditiveGroup s, s ~ scalarType) => AdditiveGroup (Record a1 … an) where + -- … + mkClassDec :: Type -> Q [Dec] + mkClassDec scalarVar = do + evalGradDec <- mkEvalGradDec + return [InstanceD Nothing (cxt ++ newConstraints) ihead [evalGradDec]] + where + -- Dual s (RecordTang a1 … an) (RecordGrad a1 … an) + ihead :: Type + ihead = ConT ''Dual `AppT` scalarVar `AppT` vecType `AppT` gradType + where + vecType = foldl AppT (ConT $ ddtTypeConName tangRecord) instVars + gradType = foldl AppT (ConT $ ddtTypeConName gradRecord) instVars + newConstraints :: Cxt + newConstraints = + [ -- AdditiveGroup s + AppT (ConT ''AdditiveGroup) scalarVar, + -- s ~ scalarType + AppT (AppT EqualityT scalarVar) scalarType + ] + + -- evalGrad (RecordGrad x1 … xn) (RecordTang y1 … yn) = evalGrad x1 y1 ^+^ … ^+^ evalGrad xn yn + mkEvalGradDec :: Q Dec + mkEvalGradDec = do + xs <- replicateM n (newName "x") + ys <- replicateM n (newName "y") + let leftPat = ConP (ddtDataConName gradRecord) (map VarP xs) + rightPat = ConP (ddtDataConName tangRecord) (map VarP ys) + -- terms = [evalGrad x1 y1, …, evalGrad xn yn] + terms :: [Exp] + terms = zipWith evalGradExp xs ys + where + evalGradExp :: Name -> Name -> Exp + evalGradExp x y = VarE 'evalGrad `AppE` VarE x `AppE` VarE y + rhs = sumVExpr terms + return $ + FunD + 'evalGrad + [ Clause + [leftPat, rightPat] + (NormalB rhs) + [] + ] + +mkMetricInstance :: + DownhillRecord -> + DownhillRecord -> + DownhillRecord -> + Type -> + Cxt -> + [Type] -> + Q [Dec] +mkMetricInstance metricRecord tangRecord gradRecord scalarType cxt instVars = do + scalarTypeName <- newName "s" + mkClassDec (VarT scalarTypeName) + where + -- instance (ctx, s ~ scalarType) => MetricTensor s (RecordMetric a1 … an) where + -- … + mkClassDec :: Type -> Q [Dec] + mkClassDec scalarVar = do + let newConstraints = + [ -- s ~ scalarType + AppT (AppT EqualityT scalarVar) scalarType + ] + -- MetricTensor s (RecordMetric a1 … an) + ihead = ConT ''MetricTensor `AppT` metricType + evalMetricDec <- mkEvalMetric + sqrNormDec <- mkSqrNorm + return + [ InstanceD + Nothing + (cxt ++ newConstraints) + ihead + [vectypeDec, covectorTypeDec, evalMetricDec, sqrNormDec] + ] + where + vectorType :: Type + vectorType = foldl AppT (ConT $ ddtTypeConName tangRecord) instVars + covectorType :: Type + covectorType = foldl AppT (ConT $ ddtTypeConName gradRecord) instVars + metricType :: Type + metricType = foldl AppT (ConT $ ddtTypeConName metricRecord) instVars + -- type MtVector (RecordMetric a1 … an) = RecordTang a1 … an + vectypeDec = + TySynInstD + ( TySynEqn + Nothing + (AppT (ConT ''MtVector) metricType) + vectorType + ) + -- type MtCovector (RecordMetric a1 … an) = RecordGrad a1 … an + covectorTypeDec = + TySynInstD + ( TySynEqn + Nothing + (AppT (ConT ''MtCovector) metricType) + covectorType + ) + + mkEvalMetric :: Q Dec + mkEvalMetric = do + let n = ddtFieldCount metricRecord + xs <- replicateM n (newName "m") + ys <- replicateM n (newName "dv") + let leftPat, rightPat :: Pat + leftPat = ConP (ddtDataConName metricRecord) (map VarP xs) + rightPat = ConP (ddtDataConName gradRecord) (map VarP ys) + terms :: [Exp] + terms = zipWith evalGradExp xs ys + where + evalGradExp :: Name -> Name -> Exp + evalGradExp x y = VarE 'evalMetric `AppE` VarE x `AppE` VarE y + rhs = + foldl + AppE + (ConE (ddtDataConName tangRecord)) + terms + return $ + FunD + 'evalMetric + [ Clause + [leftPat, rightPat] + (NormalB rhs) + [] + ] + + mkSqrNorm :: Q Dec + mkSqrNorm = do + let n = ddtFieldCount metricRecord + xs <- replicateM n (newName "m") + ys <- replicateM n (newName "dv") + let leftPat, rightPat :: Pat + leftPat = ConP (ddtDataConName metricRecord) (map VarP xs) + rightPat = ConP (ddtDataConName gradRecord) (map VarP ys) + terms :: [Exp] + terms = zipWith evalSqrtNorm xs ys + where + evalSqrtNorm :: Name -> Name -> Exp + evalSqrtNorm x y = VarE 'sqrNorm `AppE` VarE x `AppE` VarE y + rhs = sumVExpr terms + return $ + FunD + 'sqrNorm + [ Clause + [leftPat, rightPat] + (NormalB rhs) + [] + ] + +mkRecord :: DownhillRecord -> Q [Dec] +mkRecord record = do + let newConstr = mkConstructor record + let newRecordName = ddtTypeConName record + let dataType = case ddtVariant record of + Newtype -> NewtypeD [] newRecordName (ddtTypeVars record) Nothing newConstr [] + _ -> DataD [] newRecordName (ddtTypeVars record) Nothing [newConstr] [] + return [dataType] + +renameTypeS :: (String -> String) -> Name -> Name +renameTypeS f = mkNameS . f . nameBase + +data FieldInfo = FieldInfo + { fiName :: String, + fiIndex :: Int, + fiType :: Type + } + +mkGetField :: + DownhillRecord -> + DownhillRecord -> + Cxt -> + [Type] -> + FieldInfo -> + Q [Dec] +mkGetField pointRecord gradBuilderRecord cxt instVars field = do + rName <- newName "r" + xName <- newName "x" + dxName <- newName "dx" + goName <- newName "go" + dxdaName <- newName "dx_da" + let rhsFieldList :: [Exp] + rhsFieldList = + replicate (fiIndex field) (VarE 'mempty) + ++ [VarE dxdaName] + ++ replicate (n - fiIndex field - 1) (VarE 'mempty) + -- rhs = MyRecordGradBuilder mempty … mempty dx_da_a6SX mempty … mempty + rhs :: Exp + rhs = foldl AppE (ConE (ddtDataConName gradBuilderRecord)) rhsFieldList + return + [ InstanceD + Nothing + cxt + ( AppT + ( AppT + (AppT (ConT ''HasField) (LitT (StrTyLit (fiName field)))) + (AppT (AppT (ConT ''BVar) (VarT rName)) pointType) + ) + (AppT (AppT (ConT ''BVar) (VarT rName)) (fiType field)) + ) + [ FunD + 'getField + [ Clause + [ConP 'BVar [VarP xName, VarP dxName]] + ( NormalB + ( AppE + ( AppE + (ConE 'BVar) + (AppE (AppTypeE (VarE 'getField) (LitT (StrTyLit (fiName field)))) (VarE xName)) + ) + (AppE (AppE (VarE 'lift1_sparse) (VarE goName)) (VarE dxName)) + ) + ) + [ SigD + goName + ( AppT + ( AppT + ArrowT + ( ConT ''VecBuilder + `AppT` AppT (ConT ''Grad) (fiType field) + ) + ) + (ConT ''Maybe `AppT` gradBuilderType) + ), + FunD + goName + [ Clause + [VarP dxdaName] + ( NormalB + ( AppE + (ConE 'Just) + rhs + ) + ) + [] + ] + ] + ] + ] + ] + where + n = ddtFieldCount pointRecord + applyVars :: Type -> Type + applyVars x = foldl AppT x instVars + pointType :: Type + pointType = applyVars (ConT $ ddtTypeConName pointRecord) + gradBuilderType = applyVars (ConT $ ddtTypeConName gradBuilderRecord) + +renameDownhillRecord :: RecordTranstorm -> DownhillRecord -> DownhillRecord +renameDownhillRecord (RecordTranstorm namer typeFun) record = + DownhillRecord + { ddtTypeConName = renameTypeS (typeConNamer namer) (ddtTypeConName record), + ddtDataConName = renameTypeS (dataConNamer namer) (ddtDataConName record), + ddtTypeVars = ddtTypeVars record, + ddtFieldCount = ddtFieldCount record, + ddtFieldTypes = typeFun <$> ddtFieldTypes record, + ddtFieldNames = fmap (fmap (fieldNamer namer)) (ddtFieldNames record), + ddtVariant = ddtVariant record + } + +builderTransform :: BVarOptions -> RecordTranstorm +builderTransform options = RecordTranstorm (optBuilderNamer options) (AppT (ConT ''VecBuilder)) + +tangTransform :: BVarOptions -> RecordTranstorm +tangTransform options = RecordTranstorm (optTangNamer options) (AppT (ConT ''Tang)) + +gradTransform :: BVarOptions -> RecordTranstorm +gradTransform options = RecordTranstorm (optGradNamer options) (AppT (ConT ''Grad)) + +metricTransform :: BVarOptions -> RecordTranstorm +metricTransform options = RecordTranstorm (optMetricNamer options) (AppT (ConT ''Metric)) + +mkVec :: Cxt -> [Type] -> Type -> DownhillRecord -> BVarOptions -> Q [Dec] +mkVec cxt instVars scalarType vectorType options = do + let builderType = renameDownhillRecord (builderTransform options) vectorType + tangDec <- mkRecord vectorType + tangBuilderDec <- mkRecord builderType + tangSemigroup <- mkSemigroupInstance cxt builderType instVars + tangInst <- mkBasicVectorInstance vectorType options cxt instVars + additiveTang <- mkAdditiveGroupInstance cxt vectorType instVars + vspaceTang <- mkVectorSpaceInstance vectorType scalarType cxt instVars + return + ( concat + [ tangDec, + tangBuilderDec, + tangInst, + tangSemigroup, + additiveTang, + vspaceTang + ] + ) + +mkDVar'' :: + Cxt -> + DownhillRecord -> + BVarOptions -> + Type -> + [Type] -> + ConstructorInfo -> + Q [Dec] +mkDVar'' cxt pointRecord options scalarType instVars substitutedCInfo = do + let tangRecord = renameDownhillRecord (tangTransform options) pointRecord + gradRecord = renameDownhillRecord (gradTransform options) pointRecord + metricRecord = renameDownhillRecord (metricTransform options) pointRecord + + tangDecs <- mkVec cxt instVars scalarType tangRecord options + gradDecs <- mkVec cxt instVars scalarType gradRecord options + + metricDec <- mkRecord metricRecord + additiveMetric <- mkAdditiveGroupInstance cxt metricRecord instVars + vspaceMetric <- mkVectorSpaceInstance metricRecord scalarType cxt instVars + dualInstance <- mkDualInstance tangRecord gradRecord scalarType cxt instVars + metricInstance <- mkMetricInstance metricRecord tangRecord gradRecord scalarType cxt instVars + let needAffineSpace = case optAffineSpace options of + MakeAffineSpace -> True + NoAffineSpace -> False + AutoAffineSpace -> null (optExcludeFields options) + + affineSpaceInstance <- + if needAffineSpace + then mkAffineSpaceInstance cxt pointRecord tangRecord instVars + else return [] + + hasFieldInstance <- case ddtFieldNames pointRecord of + Nothing -> return [] + Just names -> + let info :: Int -> String -> Type -> FieldInfo + info index name = FieldInfo name index + substitutedFields = constructorFields substitutedCInfo + fields :: [FieldInfo] + fields = zipWith3 info [0 ..] names substitutedFields + in concat + <$> traverse + ( mkGetField + pointRecord + ( renameDownhillRecord (builderTransform options) gradRecord + ) + cxt + instVars + ) + fields + + let decs = + [ tangDecs, + gradDecs, + additiveMetric, + vspaceMetric, + dualInstance, + metricDec, + metricInstance, + hasFieldInstance, + affineSpaceInstance + ] + return (concat decs) + +parseRecordType :: Type -> [Type] -> Q (Name, [Type]) +parseRecordType type_ vars = case type_ of + AppT inner typeVar -> parseRecordType inner (typeVar : vars) + ConT recordName -> return (recordName, vars) + _ -> fail "Expected (T a1 ... an) in constraint" + +mkAffineSpaceInstance :: Cxt -> DownhillRecord -> DownhillRecord -> [Type] -> Q [Dec] +mkAffineSpaceInstance cxt recordPoint recordTang instVars = do + plusDec <- elementwiseOp' recordPoint recordTang recordPoint '(.+^) + minusDec <- elementwiseOp' recordPoint recordPoint recordTang '(.-.) + let recordTypePoint = foldl AppT (ConT (ddtTypeConName recordPoint)) instVars + recordTypeTang = foldl AppT (ConT (ddtTypeConName recordTang)) instVars + diffTypeDec = + TySynInstD + ( TySynEqn + Nothing + (AppT (ConT ''Diff) recordTypePoint) + recordTypeTang + ) + let decs = + [ plusDec, + minusDec, + diffTypeDec + ] + mkClassInstance ''AffineSpace cxt recordPoint instVars decs + +filterFields :: forall m. MonadFail m => BVarOptions -> DownhillRecord -> m DownhillRecord +filterFields options record = + case optExcludeFields options of + [] -> return record + _ -> do + fieldList <- case ddtFieldNames record of + Just fields -> return fields + Nothing -> fail (nameBase (ddtTypeConName record) ++ " is not a records, can't exclude fields") + doFilterFields fieldList + where + doFilterFields fieldList = do + traverse_ check (optExcludeFields options) + return + record + { ddtFieldTypes = go (ddtFieldTypes record), + ddtFieldNames = go <$> ddtFieldNames record, + ddtFieldCount = goN (ddtFieldCount record) + } + where + check :: String -> m () + check name + | name `elem` fieldList = return () + | otherwise = fail ("Field " ++ name ++ " is not a member of " ++ nameBase (ddtTypeConName record)) + excludeZipList :: [x -> Maybe x] + excludeZipList = filterField <$> fieldList + where + filterField :: String -> x -> Maybe x + filterField fieldName x + | fieldName `elem` optExcludeFields options = Nothing + | otherwise = Just x + go :: [a] -> [a] + go = catMaybes . zipWith ($) excludeZipList + goN :: Int -> Int + goN n = length . go $ replicate n () + +mkDVarC1 :: BVarOptions -> Dec -> Q [Dec] +mkDVarC1 options = \case + InstanceD mayOverlap cxt type_ decs -> do + case mayOverlap of + Just _ -> fail "Overlapping instances not implemented" + _ -> return () + case type_ of + AppT (ConT hasgradCtx) recordInConstraintType -> do + when (hasgradCtx /= ''HasGrad) $ + fail $ "Constraint must be `HasGrad`, got " ++ show hasgradCtx + (recordName, instVars) <- parseRecordType recordInConstraintType [] + record' <- reifyDatatype recordName + + (fullParsedRecord, cinfo) <- parseDownhillRecord recordName record' + parsedRecord <- filterFields options fullParsedRecord + recordTypeVarNames <- do + let getName x = case x of + SigT (VarT y) _ -> return y + _ -> fail "Type variable is not VarT" + traverse getName (datatypeInstTypes record') + -- We have two sets of type variables: one in record definition (as in `data MyRecord a b c = ...`) + -- and another one in instance head (`instance HasGrad (MyRecord a' b' c')). We need + -- those from instance head for HasField instances. + let substPairs = zip recordTypeVarNames instVars + substitutedRecord = applySubstitution (Map.fromList substPairs) cinfo + + scalarType <- case decs of + [] -> fail "`HasGrad` instance has no declarations" + [dec1] -> case dec1 of + TySynInstD (TySynEqn _ (AppT (ConT scalarName) _) scalarType) -> do + when (scalarName /= ''MScalar) $ + fail ("Expected `Scalar` equation, got " ++ show scalarName) + return scalarType + _ -> fail "HasGrad instance must contain `Scalar ... = ...` declaration" + _ -> fail "`HasGrad` has multiple declarations" + + dvar <- mkDVar'' cxt parsedRecord options scalarType instVars substitutedRecord + + let tangName = ddtTypeConName (renameDownhillRecord (tangTransform options) parsedRecord) + gradName = ddtTypeConName (renameDownhillRecord (gradTransform options) parsedRecord) + metricName = ddtTypeConName (renameDownhillRecord (metricTransform options) parsedRecord) + tangTypeDec = + TySynInstD + ( TySynEqn + Nothing + (AppT (ConT ''Tang) recordInConstraintType) + (foldl AppT (ConT tangName) instVars) + ) + gradTypeDec = + TySynInstD + ( TySynEqn + Nothing + (AppT (ConT ''Grad) recordInConstraintType) + (foldl AppT (ConT gradName) instVars) + ) + metricTypeDec = + TySynInstD + ( TySynEqn + Nothing + (AppT (ConT ''Metric) recordInConstraintType) + (foldl AppT (ConT metricName) instVars) + ) + + hasgradInstance = + InstanceD + Nothing + cxt + type_ + ( decs + ++ [ tangTypeDec, + gradTypeDec, + metricTypeDec + ] + ) + return $ dvar ++ [hasgradInstance] + _ -> fail "Instance head is not a constraint" + _ -> fail "Expected instance declaration" + +-- | Generates @HasGrad@ instance, along with @Tang@ and @Grad@ types, +-- @VecBuilder@ types and all other instances needed for @HasGrad@. +mkHasGradInstances :: BVarOptions -> Q [Dec] -> Q [Dec] +mkHasGradInstances options decs = concat <$> (traverse (mkDVarC1 options) =<< decs) diff --git a/test/DownhillTest/Point.hs b/test/DownhillTest/Point.hs new file mode 100644 --- /dev/null +++ b/test/DownhillTest/Point.hs @@ -0,0 +1,7 @@ +module DownhillTest.Point where + +data Vector a = Vector { vectorX :: a, vectorY :: a } +data Point a = Point { pointX :: a, pointY :: a} + +--vectorX' :: BVar (Vector a) da dv -> BVar a da dv +--vectorX' = diff --git a/test/DownhillTest/TH.hs b/test/DownhillTest/TH.hs new file mode 100644 --- /dev/null +++ b/test/DownhillTest/TH.hs @@ -0,0 +1,102 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE DuplicateRecordFields #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module DownhillTest.TH (thTest) where + +import Data.AffineSpace (AffineSpace (..)) +import Downhill.Grad (HasGrad (MScalar, Tang)) +import Downhill.TH (BVarOptions (..), RecordNamer (..), mkHasGradInstances) +import Test.Tasty (TestTree, testGroup) +import DownhillTest.TestTHOptions (defaultDVarOptions) + +{-# ANN module "HLint: ignore Use newtype instead of data" #-} +newtype MyRecord1 = MyRecord1 Float + +data MyRecord2 = MyRecord2 Float + +mkHasGradInstances + defaultDVarOptions + [d| + instance HasGrad MyRecord1 where + type MScalar MyRecord1 = Float + |] + +mkHasGradInstances + defaultDVarOptions + [d| + instance HasGrad MyRecord2 where + type MScalar MyRecord2 = Float + |] + +data MyRecord3 = MyRecord3 + +mkHasGradInstances + defaultDVarOptions + [d| + instance HasGrad MyRecord3 where + type MScalar MyRecord3 = () + |] + +data MyRecord4 a = MyRecord4 a + +mkHasGradInstances + defaultDVarOptions + [d| + instance (AffineSpace a, HasGrad a, Diff a ~ Tang a) => HasGrad (MyRecord4 a) where + type MScalar (MyRecord4 a) = MScalar a + |] + +data MyRecord5 a b = MyRecord5 a b + +mkHasGradInstances + defaultDVarOptions + [d| + instance + ( AffineSpace a, + AffineSpace b, + HasGrad a, + HasGrad b, + MScalar a ~ MScalar b, + Diff a ~ Tang a, + Diff b ~ Tang b + ) => + HasGrad (MyRecord5 a b) + where + type MScalar (MyRecord5 a b) = MScalar a + |] + +data MyRecord6 a b = MyRecord6 a b + +mkHasGradInstances + defaultDVarOptions + [d| + instance + ( AffineSpace a, + HasGrad a, + MScalar a ~ Float, + Diff a ~ Tang a + ) => + HasGrad (MyRecord6 a Float) + where + type MScalar (MyRecord6 a Float) = Float + |] + +data MyRecord7 a = MyRecord7 + { myField7 :: a + , myLabel7 :: String + } + +mkHasGradInstances + defaultDVarOptions {optExcludeFields = ["myLabel7"]} + [d| + instance HasGrad a => HasGrad (MyRecord7 a) where + type MScalar (MyRecord7 a) = MScalar a + |] + +thTest :: TestTree +thTest = testGroup "Template Haskell" [] -- just test if it compiles... diff --git a/test/DownhillTest/TestTHOptions.hs b/test/DownhillTest/TestTHOptions.hs new file mode 100644 --- /dev/null +++ b/test/DownhillTest/TestTHOptions.hs @@ -0,0 +1,46 @@ +module DownhillTest.TestTHOptions(defaultDVarOptions) where +import Downhill.TH ( mkHasGradInstances, RecordNamer(..), BVarOptions(..), AffineSpaceOptions (AutoAffineSpace)) + +defaultTangRecordNamer :: RecordNamer +defaultTangRecordNamer = + RecordNamer + { typeConNamer = (++ "TangT"), + dataConNamer = (++ "TangD"), + fieldNamer = id + } + +defaultGradRecordNamer :: RecordNamer +defaultGradRecordNamer = + RecordNamer + { typeConNamer = (++ "GradT"), + dataConNamer = (++ "GradD"), + fieldNamer = id + } + +defaultMetricRecordNamer :: RecordNamer +defaultMetricRecordNamer = + RecordNamer + { typeConNamer = (++ "MetricT"), + dataConNamer = (++ "MetricD"), + fieldNamer = id + } + +defaultBuilderRecordNamer :: RecordNamer +defaultBuilderRecordNamer = + RecordNamer + { typeConNamer = (++ "BuilderT"), + dataConNamer = (++ "BuilderD"), + fieldNamer = id + } + +defaultDVarOptions :: BVarOptions +defaultDVarOptions = + BVarOptions + { optTangNamer = defaultTangRecordNamer, + optGradNamer = defaultGradRecordNamer, + optMetricNamer = defaultMetricRecordNamer, + optBuilderNamer = defaultBuilderRecordNamer, + optAffineSpace = AutoAffineSpace, + optExcludeFields = [] + } + diff --git a/test/DownhillTest/Traversable.hs b/test/DownhillTest/Traversable.hs new file mode 100644 --- /dev/null +++ b/test/DownhillTest/Traversable.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE UndecidableInstances #-} + +module DownhillTest.Traversable(recordTest) where + +import Downhill.BVar.Traversable (TraversableVar (TraversableVar), backpropTraversable, splitTraversable) +import Downhill.BVar (BVar (BVar), backprop, var) +import Downhill.Grad (HasGrad (Grad)) +import Test.Tasty (TestTree) +import Test.Tasty.HUnit (testCase, (@?=)) + +data MyRecord a = MyRecord + { memberPair :: (a, a), + memberList :: [a] + } + deriving (Eq, Functor, Foldable, Traversable, Show) + +deriving via (TraversableVar MyRecord a) instance HasGrad a => HasGrad (MyRecord a) + +test_r :: MyRecord Integer +test_r = MyRecord (10, 11) [12, 13, 14] + +expectedResult :: MyRecord (Integer, Integer) +expectedResult = + MyRecord + ((10, 2), (11, 3)) + [(12, 5), (13, 5), (14, 5)] + +test_fun :: Num a => MyRecord a -> a +test_fun (MyRecord (x, y) zs) = 2 * x + 3 * y + 5 * sum zs + +type MyGrad a = Grad (MyRecord a) + +actualResult :: MyRecord (Integer, Integer) +actualResult = backpropTraversable 1 (,) test_fun test_r + where + x :: BVar (MyGrad Integer) (MyRecord Integer) + x = var test_r + x' :: MyRecord (BVar (MyGrad Integer) Integer) + x' = splitTraversable x + y :: BVar (MyGrad Integer) Integer + y = test_fun x' + +recordTest :: TestTree +recordTest = testCase "Traverse record" (actualResult @?= expectedResult) + +main :: IO () +main = return () diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,25 @@ +import Downhill.BVar(bvarValue) +import Test.Tasty (defaultMain, testGroup, TestTree) +import Test.Tasty.HUnit (Assertion, testCase, (@?=)) + +import qualified Test.Tasty as Tasty +import Downhill.BVar.Num (NumBVar(..), backpropNum, constant, var, numbvarValue, AsNum) +import DownhillTest.Traversable(recordTest) +import DownhillTest.TH (thTest) + +basicTests = testGroup "Basic tests" + [ testCase "Derivative of constant == 0" testConstant + , testCase "Derivative of identity == 1" testIdentity + , testCase "Derivative of simple polynomial" testPoly + ] + where testConstant = backpropNum (constant 3 :: NumBVar Integer) @?= 0 + testIdentity = backpropNum (var 3 :: NumBVar Integer) @?= 1 + testPoly = + let x = var 5 :: NumBVar Integer + y = 3*x :: NumBVar Integer + in backpropNum ((2+3*x) * (5+7*x)) @?= 29 + 42 * numbvarValue x + +tests :: TestTree +tests = testGroup "Tests" [basicTests, recordTest, thTest] + +main = defaultMain tests