diff --git a/Build.hs b/Build.hs new file mode 100644 --- /dev/null +++ b/Build.hs @@ -0,0 +1,59 @@ +#!/usr/bin/env stack +-- stack --install-ghc runghc --package shake + +import Development.Shake +import Development.Shake.FilePath +import System.Directory + +opts = shakeOptions { shakeFiles = ".shake" + , shakeVersion = "1.0" + , shakeVerbosity = Normal + , shakeThreads = 0 + } + +data Doc = Lab + +main :: IO () +main = getDirectoryFilesIO "samples" ["/*.lhs"] >>= \allSamps -> + getDirectoryFilesIO "src" ["//*.hs"] >>= \allSrc -> + shakeArgs opts $ do + + want ["all"] + + "all" ~> + need ["pdf", "md", "haddocks", "gentags"] + + "pdf" ~> + need (map (\f -> "renders" takeFileName f -<.> "pdf") allSamps) + + "md" ~> + need (map (\f -> "renders" takeFileName f -<.> "md") allSamps) + + "haddocks" ~> + cmd "jle-git-haddocks" + + "gentags" ~> + need ["tags", "TAGS"] + + ["renders/*.pdf", "renders/*.md"] |%> \f -> do + let src = "samples" takeFileName f -<.> "lhs" + need [src] + liftIO $ createDirectoryIfMissing True "renders" + cmd "pandoc" "-V geometry:margin=1in" + "-V fontfamily:palatino,cmtt" + "-V links-as-notes" + "-sS" + "--highlight-style tango" + "--reference-links" + "--reference-location block" + "-o" f + src + + ["tags","TAGS"] &%> \_ -> do + need (("src" ) <$> allSrc) + cmd "hasktags" "src/" + + "clean" ~> do + unit $ cmd "stack clean" + removeFilesAfter ".shake" ["//*"] + diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,14 @@ +Changelog +========= + +Version 0.0.1.0 +--------------- + + + +* Initial pre-release, as a request for comments. API is in a usable form + and everything is fully documented, but there are definitely some things + left to be done. (See [README.md][readme-0.0.1.0]) + + [readme-0.0.1.0]: https://github.com/mstksg/backprop/tree/v0.0.1.0#readme + diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Justin Le (c) 2017 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Justin Le nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,163 @@ +backprop +======== + +[![Build Status](https://travis-ci.org/mstksg/backprop.svg?branch=master)](https://travis-ci.org/mstksg/backprop) + +[**Literate Haskell Tutorial/Demo on MNIST data set**][mnist-lhs] (and [PDF +rendering][mnist-pdf]) + +Automatic *heterogeneous* back-propagation that can be used either *implicitly* +(in the style of the [ad][] library) or using *explicit* graphs built in +monadic style. Implements reverse-mode automatic differentiation. Differs +from [ad][] by offering full heterogeneity -- each intermediate step and the +resulting value can have different types. Mostly intended for usage with +tensor manipulation libraries to implement automatic back-propagation for +gradient descent and other optimization techniques. + +[ad]: http://hackage.haskell.org/package/ad + +Documentation is currently rendered [on github pages][docs]! + +[docs]: https://mstksg.github.io/backprop + +MNIST Digit Classifier Example +------------------------------ + +Tutorial and example on training on the MNIST data set [available here as a +literate haskell file][mnist-lhs], or [rendered here as a PDF][mnist-pdf]! +**Read this first!** + +[mnist-lhs]: https://github.com/mstksg/backprop/blob/master/samples/MNIST.lhs +[mnist-pdf]: https://github.com/mstksg/backprop/blob/master/renders/MNIST.pdf + + +Brief example +------------- + +The quick example below describes the running of a neural network with one +hidden layer to calculate its squared error with respect to target `targ`, +which is parameterized by two weight matrices and two bias vectors. +Vector/matrix types are from the *hmatrix* package. + +~~~haskell +logistic :: Floating a => a -> a +logistic x = 1 / (1 + exp (-x)) + +matVec + :: (KnownNat m, KnownNat n) + => Op '[ L m n, R n ] (R m) + +neuralNetImplicit + :: (KnownNat m, KnownNat n, KnownNat o) + => R m + -> BPOpI s '[ L n m, R n, L o n, R o ] (R o) +neuralNetImplicit inp = \(w1 :< b1 :< w2 :< b2 :< Ø) -> + let z = logistic (liftB2 matVec w1 x + b1) + in logistic (liftB2 matVec w2 z + b2) + where + x = constRef inp + +neuralNetExplicit + :: (KnownNat m, KnownNat n, KnownNat o) + => R m + -> BPOp s '[ L n m, R n, L o n, R o ] (R o) +neuralNetExplicit inp = withInps $ \(w1 :< b1 :< w2 :< b2 :< Ø) -> do + y1 <- matVec ~$ (w1 :< x1 :< Ø) + let x2 = logistic (y1 + b1) + y2 <- matVec ~$ (w2 :< x2 :< Ø) + return $ logistic (y2 + b2) + where + x1 = constVar inp +~~~ + +Now `neuralNetExplicit` and `neuralNetImplicit` can be "run" with the input +vectors and parameters (a `L n m`, `R n`, `L o n`, and `R o`) and calculate the +output of the neural net. + +~~~haskell +runNet + :: (KnownNat m, KnownNat n, KnownNat o) + => R m + -> Tuple '[ L n m, R n, L o n, R o ] + -> R o +runNet inp = evalBPOp (neuralNetExplicit inp) +~~~ + +But, in defining `neuralNet`, we also generated a graph that *backprop* can +use to do back-propagation, too! + +~~~haskell +dot :: KnownNat n + => Op '[ R n , R n ] Double + +netGrad + :: forall m n o. (KnownNat m, KnownNat n, KnownNat o) + => R m + -> R o + -> Tuple '[ L n m, R n, L o n, R o ] + -> Tuple '[ L n m, R n, L o n, R o ] +netGrad inp targ params = gradBPOp opError params + where + -- calculate squared error, in *explicit* style + opError :: BPOp s '[ L n m, R n, L o n, R o ] Double + opError = do + res <- neuralNetExplicit inp + err <- bindRef (res - t) + dot ~$ (err :< err :< Ø) + where + t = constRef targ +~~~ + +The result is the gradient of the input tuple's components, with respect +to the `Double` result of `opError` (the squared error). We can then use +this gradient to do gradient descent. + +For a more fleshed out example, see the [MNIST tutorial][mnist-lhs] (also +[rendered as a pdf][mnist-pdf]) + +Todo +---- + +1. Actual profiling and benchmarking, to gauge how much overhead this library + adds over "manual" back-propagation. + + Ideally this can be brought down to 0? + +2. Some simple performance and API tweaks that are probably possible now and + would clearly benefit: (if you want to contribute) + + a. Providing optimized `Num`/`Fractional`/`Floating` instances for `BVal` + by supplying known gradients directly instead of relying on *ad*. + + b. Switch from `ST s` to `IO`, and use `unsafePerformIO` to automatically + bind `BVal`s (like *ad* does) when using `liftB`. This might remove + some overhead during graph building, and, from an API standpoint, + remove the need for explicit binding. + + c. Switch from `STRef`s/`IORef`s to `Array`. (This one I'm unclear if it + would help any) + +3. Benchmark against competing back-propagation libraries like *ad*, and + auto-differentiating tensor libraries like *[grenade][]* + + [grenade]: https://github.com/HuwCampbell/grenade + +4. Explore opportunities for parallelization. There are some naive ways of + directly parallelizing right now, but potential overhead should be + investigated. + +5. Some open questions: + + a. Is it possible to offer pattern matching on sum types/with different + constructors for implicit-graph backprop? It's possible for + explicit-graph versions already, with `choicesVar`, but not yet with + the implicit-graph interface. Could be similar to an "Applicative vs. + Monad" issue where you can only have pre-determined fixed computation + paths when using `Applicative`, but I'm not sure. Still, it would be + nice, because if this was possible, we could possibly do away with + explicit-graph mode completely. + + b. Though we already have sum type support with explicit-graph mode, we + can't support GADTs yet. It'd be nice to see if this is possible, + because a lot of dependently typed neural network stuff is made much + simpler with GADTs. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/backprop.cabal b/backprop.cabal new file mode 100644 --- /dev/null +++ b/backprop.cabal @@ -0,0 +1,102 @@ +name: backprop +version: 0.0.1.0 +synopsis: Heterogeneous, type-safe automatic backpropagation in Haskell +description: See +homepage: https://github.com/mstksg/backprop +license: BSD3 +license-file: LICENSE +author: Justin Le +maintainer: justin@jle.im +copyright: (c) Justin Le 2017 +category: Web +build-type: Simple +extra-source-files: README.md + CHANGELOG.md + Build.hs + renders/MNIST.md + renders/MNIST.pdf + renders/NeuralTest.md + renders/NeuralTest.pdf +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: Numeric.Backprop + Numeric.Backprop.Implicit + Numeric.Backprop.Iso + Numeric.Backprop.Mono + Numeric.Backprop.Mono.Implicit + Numeric.Backprop.Op + Numeric.Backprop.Op.Mono + other-modules: Numeric.Backprop.Internal + Numeric.Backprop.Internal.Helper + Data.Type.Util + build-depends: base >= 4.7 && < 5 + , ad + , generics-sop + , microlens + , microlens-mtl + , microlens-th + , mtl + , profunctors + , reflection + , tagged + , transformers-base + , type-combinators + default-language: Haskell2010 + ghc-options: -Wall + +executable backprop-monotest + hs-source-dirs: samples + main-is: MonoTest.hs + ghc-options: -threaded -rtsopts -with-rtsopts=-N -Wall -O2 + build-depends: base + , backprop + default-language: Haskell2010 + +executable backprop-neuraltest + hs-source-dirs: samples + main-is: NeuralTest.lhs + ghc-options: -threaded -rtsopts -with-rtsopts=-N -Wall -O2 + build-depends: base + , ad + , backprop + , generics-sop + , hmatrix >= 0.18 + , mwc-random + , primitive + , singletons + , type-combinators + default-language: Haskell2010 + +executable backprop-mnist + hs-source-dirs: samples + main-is: MNIST.lhs + ghc-options: -threaded -rtsopts -with-rtsopts=-N -Wall -O2 + build-depends: base + , backprop + , bifunctors + , deepseq + , finite-typelits + , generics-sop + , hmatrix >= 0.18 + , mnist-idx + , mwc-random + , split + , time + , transformers + , vector + default-language: Haskell2010 + +-- test-suite backprop-test +-- type: exitcode-stdio-1.0 +-- hs-source-dirs: test +-- main-is: Spec.hs +-- build-depends: base +-- , backprop +-- ghc-options: -threaded -rtsopts -with-rtsopts=-N +-- default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/mstksg/backprop diff --git a/renders/MNIST.md b/renders/MNIST.md new file mode 100644 --- /dev/null +++ b/renders/MNIST.md @@ -0,0 +1,554 @@ +--- +author: +- Justin Le +fontfamily: 'palatino,cmtt' +geometry: margin=1in +links-as-notes: true +title: Learning MNIST with Neural Networks with backprop library +--- + +The *backprop* library performs back-propagation over a *hetereogeneous* +system of relationships. It offers both an implicit (*[ad]*-like) and +explicit graph building usage style. Let’s use it to build neural +networks and learn mnist! + + [ad]: http://hackage.haskell.org/package/ad + +Repository source is [on github], and so are the [rendered docs]. + + [on github]: https://github.com/mstksg/backprop + [rendered docs]: https://mstksg.github.io/backprop + +If you’re reading this as a literate haskell file, you should know that +a [rendered pdf version is available on github.]. If you are reading +this as a pdf file, you should know that a [literate haskell version +that you can run] is also available on github! + + [rendered pdf version is available on github.]: https://github.com/mstksg/backprop/blob/master/renders/MNIST.pdf + [literate haskell version that you can run]: https://github.com/mstksg/backprop/blob/master/samples/MNIST.lhs + +``` {.sourceCode .literate .haskell} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-} +{-# OPTIONS_GHC -fno-warn-unused-top-binds #-} + +import Control.DeepSeq +import Control.Exception +import Control.Monad +import Control.Monad.IO.Class +import Control.Monad.Trans.Maybe +import Control.Monad.Trans.State +import Data.Bitraversable +import Data.Foldable +import Data.IDX +import Data.List.Split +import Data.Maybe +import Data.Time.Clock +import Data.Traversable +import Data.Tuple +import GHC.Generics (Generic) +import GHC.TypeLits +import Numeric.Backprop +import Numeric.LinearAlgebra.Static hiding (dot) +import Text.Printf +import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Unboxed as VU +import qualified Generics.SOP as SOP +import qualified Numeric.LinearAlgebra as HM +import qualified System.Random.MWC as MWC +import qualified System.Random.MWC.Distributions as MWC +``` + +Types +===== + +For the most part, we’re going to be using the great *[hmatrix]* library +and its vector and matrix types. It offers a type `L m n` for +$m \times n$ matrices, and a type `R n` for an $n$ vector. + + [hmatrix]: http://hackage.haskell.org/package/hmatrix + +First thing’s first: let’s define our neural networks as simple +containers of parameters (weight matrices and bias vectors). + +First, a type for layers: + +``` {.sourceCode .literate .haskell} +data Layer i o = + Layer { _lWeights :: !(L o i) + , _lBiases :: !(R o) + } + deriving (Show, Generic) + +instance SOP.Generic (Layer i o) +instance NFData (Layer i o) +``` + +And a type for a simple feed-forward network with two hidden layers: + +``` {.sourceCode .literate .haskell} +data Network i h1 h2 o = + Net { _nLayer1 :: !(Layer i h1) + , _nLayer2 :: !(Layer h1 h2) + , _nLayer3 :: !(Layer h2 o) + } + deriving (Show, Generic) + +instance SOP.Generic (Network i h1 h2 o) +instance NFData (Network i h1 h2 o) +``` + +These are pretty straightforward container types…pretty much exactly the +type you’d make to represent these networks! Note that, following true +Haskell form, we separate out logic from data. This should be all we +need. + +We derive an instance of `SOP.Generic` from the *[generics-sop]* +package, which *backprop* uses to propagate derivatives on values inside +product types. + + [generics-sop]: http://hackage.haskell.org/package/generics-sop + +Instances +--------- + +Things are much simplier if we had `Num` and `Fractional` instances for +everything, so let’s just go ahead and define that now, as well. Just a +little bit of boilerplate. + +``` {.sourceCode .literate .haskell} +instance (KnownNat i, KnownNat o) => Num (Layer i o) where + Layer w1 b1 + Layer w2 b2 = Layer (w1 + w2) (b1 + b2) + Layer w1 b1 - Layer w2 b2 = Layer (w1 - w2) (b1 - b2) + Layer w1 b1 * Layer w2 b2 = Layer (w1 * w2) (b1 * b2) + abs (Layer w b) = Layer (abs w) (abs b) + signum (Layer w b) = Layer (signum w) (signum b) + negate (Layer w b) = Layer (negate w) (negate b) + fromInteger x = Layer (fromInteger x) (fromInteger x) + +instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Num (Network i h1 h2 o) where + Net a b c + Net d e f = Net (a + d) (b + e) (c + f) + Net a b c - Net d e f = Net (a - d) (b - e) (c - f) + Net a b c * Net d e f = Net (a * d) (b * e) (c * f) + abs (Net a b c) = Net (abs a) (abs b) (abs c) + signum (Net a b c) = Net (signum a) (signum b) (signum c) + negate (Net a b c) = Net (negate a) (negate b) (negate c) + fromInteger x = Net (fromInteger x) (fromInteger x) (fromInteger x) + +instance (KnownNat i, KnownNat o) => Fractional (Layer i o) where + Layer w1 b1 / Layer w2 b2 = Layer (w1 / w2) (b1 / b2) + recip (Layer w b) = Layer (recip w) (recip b) + fromRational x = Layer (fromRational x) (fromRational x) + +instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Fractional (Network i h1 h2 o) where + Net a b c / Net d e f = Net (a / d) (b / e) (c / f) + recip (Net a b c) = Net (recip a) (recip b) (recip c) + fromRational x = Net (fromRational x) (fromRational x) (fromRational x) +``` + +`KnownNat` comes from *base*; it’s a typeclass that *hmatrix* uses to +refer to the numbers in its type and use it to go about its normal +hmatrixy business. + +Ops +=== + +Now, *backprop* does require *primitive* differentiable operations on +our relevant types to be defined. *backprop* uses these primitive `Op`s +to tie everything together. Ideally we’d import these from a library +that implements these for you, and the end-user never has to make `Op` +primitives. + +But in this case, I’m going to put the definitions here to show that +there isn’t any magic going on. If you’re curious, refer to +[documentation for `Op`] for more details on how `Op` is implemented and +how this works. + + [documentation for `Op`]: https://mstksg.github.io/backprop/Numeric-Backprop-Op.html + +First, matrix-vector multiplication primitive, giving an explicit +gradient function. + +``` {.sourceCode .literate .haskell} +matVec + :: (KnownNat m, KnownNat n) + => Op '[ L m n, R n ] (R m) +matVec = op2' $ \m v -> + ( m #> v, \(fromMaybe 1 -> g) -> + (g `outer` v, tr m #> g) + ) +``` + +Dot products would be nice too. + +``` {.sourceCode .literate .haskell} +dot :: KnownNat n + => Op '[ R n, R n ] Double +dot = op2' $ \x y -> + ( x <.> y, \case Nothing -> (y, x) + Just g -> (konst g * y, x * konst g) + ) +``` + +Also a “scaling” function, scales a vector by a given factor. + +``` {.sourceCode .literate .haskell} +scale + :: KnownNat n + => Op '[ Double, R n ] (R n) +scale = op2' $ \a x -> + ( konst a * x + , \case Nothing -> (HM.sumElements (extract x ), konst a ) + Just g -> (HM.sumElements (extract (x * g)), konst a * g) + ) +``` + +Finally, an operation to sum all of the items in the vector. + +``` {.sourceCode .literate .haskell} +vsum + :: KnownNat n + => Op '[ R n ] Double +vsum = op1' $ \x -> (HM.sumElements (extract x), maybe 1 konst) +``` + +And why not, here’s the [logistic function], which we’ll use as an +activation function for internal layers. We don’t need to define this as +an `Op` up-front right now, because the library can automatically +promote any numeric polymorphic function (an `a -> a` or `a -> a -> a`, +etc.) to an `Op` anyways. + + [logistic function]: https://en.wikipedia.org/wiki/Logistic_function + +``` {.sourceCode .literate .haskell} +logistic :: Floating a => a -> a +logistic x = 1 / (1 + exp (-x)) +``` + +Running our Network +=================== + +Now that we have our primitives in place, let’s actually write a +function to run our network! + +``` {.sourceCode .literate .haskell} +runLayer + :: (KnownNat i, KnownNat o) + => BPOp s '[ R i, Layer i o ] (R o) +runLayer = withInps $ \(x :< l :< Ø) -> do + w :< b :< Ø <- gTuple #<~ l + y <- matVec ~$ (w :< x :< Ø) + return $ y + b +``` + +A `BPOp s '[ R i, Layer i o ] (R o)` is a backpropagatable function that +produces an `R o` (a vector with `o` elements, from the *[hmatrix]* +library) given an input environment of an `R i` (the “input” of the +layer) and a layer. + + [hmatrix]: http://hackage.haskell.org/package/hmatrix + +We use `withInps` to bring the environment into scope as a bunch of +`BVar`s. `x` is a `BVar` containing the input vector, and `l` is a +`BVar` containing the layer. + +The first thing we do is split out the parts of the layer so we can work +with the internal matrices. We can use `#<~` to “split out” the +components of a `BVar`, splitting on `gTuple` (which uses `GHC.Generics` +to automatically figure out how to split up a product type). + +Then we apply `matVec` (our primitive `Op` that does matrix-vector +multiplication) to `w` and `x`, and then the result is that added to the +bias vector `b`. + +We can write the `runNetwork` function pretty much the same way. + +``` {.sourceCode .literate .haskell} +runNetwork + :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) + => BPOp s '[ R i, Network i h1 h2 o ] (R o) +runNetwork = withInps $ \(x :< n :< Ø) -> do + l1 :< l2 :< l3 :< Ø <- gTuple #<~ n + y <- runLayer -$ (x :< l1 :< Ø) + z <- runLayer -$ (logistic y :< l2 :< Ø) + r <- runLayer -$ (logistic z :< l3 :< Ø) + softmax -$ (r :< Ø) + where + softmax :: KnownNat n => BPOp s '[ R n ] (R n) + softmax = withInps $ \(x :< Ø) -> do + expX <- bindVar (exp x) + totX <- vsum ~$ (expX :< Ø) + scale ~$ (1/totX :< expX :< Ø) +``` + +After splitting out the layers in the input `Network`, we run each layer +successively using our previously defined `runLayer`, giving inputs +using `-$`. We can directly apply `logistic` to `BVar`s. At the end, we +run a [softmax function] because MNIST is a classification challenge. +The softmax is done by applying $e^x$ for every item in the input +vector, and dividing each element by the total. + + [softmax function]: https://en.wikipedia.org/wiki/Softmax_function + +The Magic +--------- + +What did we just define? Well, with a `BPOp s rs a`, we can *run* it and +get the output: + +``` {.sourceCode .literate .haskell} +runNetOnInp + :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) + => Network i h1 h2 o + -> R i + -> R o +runNetOnInp n x = evalBPOp runNetwork (x ::< n ::< Ø) +``` + +But, the magic part is that we can also get the gradient! + +``` {.sourceCode .literate .haskell} +gradNet + :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) + => Network i h1 h2 o + -> R i + -> Network i h1 h2 o +gradNet n x = case gradBPOp runNetwork (x ::< n ::< Ø) of + _gradX ::< gradN ::< Ø -> gradN +``` + +This gives the gradient of all of the parameters in the matrices and +vectors inside the `Network`, which we can use to “train”! + +Training +======== + +Now for the real work. To train a network, we can do gradient descent +based on the gradient of some type of *error function* with respect to +the network parameters. Let’s use the [cross entropy], which is popular +for classification problems. + + [cross entropy]: https://en.wikipedia.org/wiki/Cross_entropy + +``` {.sourceCode .literate .haskell} +crossEntropy + :: KnownNat n + => R n + -> BPOpI s '[ R n ] Double +crossEntropy targ (r :< Ø) = negate (dot .$ (log r :< t :< Ø)) + where + t = constVar targ +``` + +Given a target vector and a `BVar` referring to the result of the +network, we can directly apply: + +$$ +H(\mathbf{r}, \mathbf{t}) = - (log(\mathbf{r}) \cdot \mathbf{t}) +$$ + +Just for fun, I implemented `crossEntropy` in “implicit-graph” mode, so +you don’t see any binds or returns. + +Now, a function to make one gradient descent step based on an input +vector and a target, using `gradBPOp`: + +``` {.sourceCode .literate .haskell} +trainStep + :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) + => Double + -> R i + -> R o + -> Network i h1 h2 o + -> Network i h1 h2 o +trainStep r !x !t !n = case gradBPOp o (x ::< n ::< Ø) of + _ ::< gN ::< Ø -> + n - (realToFrac r * gN) + where + o :: BPOp s '[ R i, Network i h1 h2 o ] Double + o = do + y <- runNetwork + implicitly (crossEntropy t) -$ (y :< Ø) +``` + +A convenient wrapper for training over all of the observations in a +list: + +``` {.sourceCode .literate .haskell} +trainList + :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) + => Double + -> [(R i, R o)] + -> Network i h1 h2 o + -> Network i h1 h2 o +trainList r = flip $ foldl' (\n (x,y) -> trainStep r x y n) +``` + +Pulling it all together +======================= + +`testNet` will be a quick way to test our net by computing the +percentage of correct guesses: (mostly using *hmatrix* stuff) + +``` {.sourceCode .literate .haskell} +testNet + :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) + => [(R i, R o)] + -> Network i h1 h2 o + -> Double +testNet xs n = sum (map (uncurry test) xs) / fromIntegral (length xs) + where + test :: R i -> R o -> Double + test x (extract->t) + | HM.maxIndex t == HM.maxIndex (extract r) = 1 + | otherwise = 0 + where + r :: R o + r = evalBPOp runNetwork (x ::< n ::< Ø) +``` + +And now, a main loop! + +If you are following along at home, download the [mnist data set files] +and uncompress them into the folder `data`, and everything should work +fine. + + [mnist data set files]: http://yann.lecun.com/exdb/mnist/ + +``` {.sourceCode .literate .haskell} +main :: IO () +main = MWC.withSystemRandom $ \g -> do + Just train <- loadMNIST "data/train-images-idx3-ubyte" "data/train-labels-idx1-ubyte" + Just test <- loadMNIST "data/t10k-images-idx3-ubyte" "data/t10k-labels-idx1-ubyte" + putStrLn "Loaded data." + net0 <- MWC.uniformR @(Network 784 300 100 9) (-1, 1) g + flip evalStateT net0 . forM_ [1..] $ \e -> do + train' <- liftIO . fmap V.toList $ MWC.uniformShuffle (V.fromList train) g + liftIO $ printf "[Epoch %d]\n" (e :: Int) + + forM_ ([1..] `zip` chunksOf batch train') $ \(b, chnk) -> StateT $ \n0 -> do + printf "(Batch %d)\n" (b :: Int) + + t0 <- getCurrentTime + n' <- evaluate . force $ trainList rate chnk n0 + t1 <- getCurrentTime + printf "Trained on %d points in %s.\n" batch (show (t1 `diffUTCTime` t0)) + + let trainScore = testNet chnk n' + testScore = testNet test n' + printf "Training error: %.2f%%\n" ((1 - trainScore) * 100) + printf "Validation error: %.2f%%\n" ((1 - testScore) * 100) + + return ((), n') + where + rate = 0.02 + batch = 5000 +``` + +Each iteration of the loop: + +1. Shuffles the training set +2. Splits it into chunks of `batch` size +3. Uses `trainList` to train over the batch +4. Computes the score based on `testNet` based on the training set and + the test set +5. Prints out the results + +And, that’s really it! + +Result +------ + +I haven’t put much into optimizing the library yet, but the network +(with hidden layer sizes 300 and 100) seems to take 25s on my computer +to finish a batch of 5000 training points. It’s slow (five minutes per +60000 point epooch), but it’s a first unoptimized run and a proof of +concept! It’s my goal to get this down to a point where the result has +the same performance characteristics as the actual backend (*hmatrix*), +and so overhead is 0. + +Main takeaways +============== + +Most of the actual heavy lifting/logic actually came from the *hmatrix* +library itself. We just created simple types to wrap up our bare +matrices. + +Basically, all that *backprop* did was give you an API to define *how to +run* a neural net — how to *run* a net based on a `Network` and `R i` +input you were given. The goal of the library is to let you write down +how to run things in as natural way as possible. + +And then, after things are run, we can just get the gradient and roll +from there! + +Because the heavy lifting is done by the data types themselves, we can +presumably plug in *any* type and any tensor/numerical backend, and reap +the benefits of those libraries’ optimizations and parallelizations. +*Any* type can be backpropagated! :D + +What now? +--------- + +Check out the docs for the [Numeric.Backprop] module for a more detailed +picture of what’s going on, or find more examples at the [github repo]! + + [Numeric.Backprop]: https://mstksg.github.io/backprop/Numeric-Backprop.html + [github repo]: https://github.com/mstksg/backprop + +Boring stuff +============ + +Here is a small wrapper function over the [mnist-idx] library loading +the contents of the idx files into *hmatrix* vectors: + + [mnist-idx]: http://hackage.haskell.org/package/mnist-idx + +``` {.sourceCode .literate .haskell} +loadMNIST + :: FilePath + -> FilePath + -> IO (Maybe [(R 784, R 9)]) +loadMNIST fpI fpL = runMaybeT $ do + i <- MaybeT $ decodeIDXFile fpI + l <- MaybeT $ decodeIDXLabelsFile fpL + d <- MaybeT . return $ labeledIntData l i + r <- MaybeT . return $ for d (bitraverse mkImage mkLabel . swap) + liftIO . evaluate $ force r + where + mkImage :: VU.Vector Int -> Maybe (R 784) + mkImage = create . VG.convert . VG.map (\i -> fromIntegral i / 255) + mkLabel :: Int -> Maybe (R 9) + mkLabel n = create $ HM.build 9 (\i -> if round i == n then 1 else 0) +``` + +And here are instances to generating random +vectors/matrices/layers/networks, used for the initialization step. + +``` {.sourceCode .literate .haskell} +instance KnownNat n => MWC.Variate (R n) where + uniform g = randomVector <$> MWC.uniform g <*> pure Uniform + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g + +instance (KnownNat m, KnownNat n) => MWC.Variate (L m n) where + uniform g = uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1 + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g + +instance (KnownNat i, KnownNat o) => MWC.Variate (Layer i o) where + uniform g = Layer <$> MWC.uniform g <*> MWC.uniform g + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g + +instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => MWC.Variate (Network i h1 h2 o) where + uniform g = Net <$> MWC.uniform g <*> MWC.uniform g <*> MWC.uniform g + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g +``` diff --git a/renders/MNIST.pdf b/renders/MNIST.pdf new file mode 100644 Binary files /dev/null and b/renders/MNIST.pdf differ diff --git a/renders/NeuralTest.md b/renders/NeuralTest.md new file mode 100644 --- /dev/null +++ b/renders/NeuralTest.md @@ -0,0 +1,447 @@ +--- +author: +- Justin Le +fontfamily: 'palatino,cmtt' +geometry: margin=1in +links-as-notes: true +title: Neural networks with backprop library +--- + +The *backprop* library performs back-propagation over a *hetereogeneous* +system of relationships. It offers both an implicit ([ad]-like) and +explicit graph building usage style. Let’s use it to build neural +networks! + + [ad]: http://hackage.haskell.org/package/ad + +Repository source is [on github], and so are the [rendered unstable +docs]. + + [on github]: https://github.com/mstksg/backprop + [rendered unstable docs]: https://mstksg.github.io/backprop + +``` {.sourceCode .literate .haskell} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -fno-warn-unused-top-binds #-} + +import Data.Functor +import Data.Kind +import Data.Maybe +import Data.Singletons +import Data.Singletons.Prelude +import Data.Singletons.TypeLits +import Data.Type.Combinator +import Data.Type.Product +import GHC.Generics (Generic) +import Numeric.Backprop +import Numeric.Backprop.Iso +import Numeric.LinearAlgebra.Static hiding (dot) +import System.Random.MWC +import qualified Generics.SOP as SOP +``` + +Ops +=== + +First, we define values of `Op` for the operations we want to do. `Op`s +are bundles of functions packaged with their hetereogeneous gradients. +For simple numeric functions, *backprop* can derive `Op`s automatically. +But for matrix operations, we have to derive them ourselves. + +The types help us with matching up the dimensions, but we still need to +be careful that our gradients are calculated correctly. + +`L` and `R` are matrix and vector types from the great *hmatrix* +library. + +First, matrix-vector multiplication: + +``` {.sourceCode .literate .haskell} +matVec + :: (KnownNat m, KnownNat n) + => Op '[ L m n, R n ] (R m) +matVec = op2' $ \m v -> ( m #> v + , \(fromMaybe 1 -> g) -> + (g `outer` v, tr m #> g) + ) +``` + +Now, dot products: + +``` {.sourceCode .literate .haskell} +dot :: KnownNat n + => Op '[ R n, R n ] Double +dot = op2' $ \x y -> ( x <.> y + , \case Nothing -> (y, x) + Just g -> (konst g * y, x * konst g) + ) +``` + +Polymorphic functions can be easily turned into `Op`s with `op1`/`op2` +etc., but they can also be run directly on graph nodes. + +``` {.sourceCode .literate .haskell} +logistic :: Floating a => a -> a +logistic x = 1 / (1 + exp (-x)) +``` + +A Simple Complete Example +========================= + +At this point, we already have enough to train a simple +single-hidden-layer neural network: + +``` {.sourceCode .literate .haskell} +simpleOp + :: (KnownNat m, KnownNat n, KnownNat o) + => R m + -> BPOpI s '[ L n m, R n, L o n, R o ] (R o) +simpleOp inp = \(w1 :< b1 :< w2 :< b2 :< Ø) -> + let z = logistic $ liftB2 matVec w1 x + b1 + in logistic $ liftB2 matVec w2 z + b2 + where + x = constVar inp +``` + +Here, `simpleOp` is defined in implicit (non-monadic) style, given a +tuple of inputs and returning outputs. Now `simpleOp` can be “run” with +the input vectors and parameters (a `L n m`, `R n`, `L o n`, and `R o`) +and calculate the output of the neural net. + +``` {.sourceCode .literate .haskell} +runSimple + :: (KnownNat m, KnownNat n, KnownNat o) + => R m + -> Tuple '[ L n m, R n, L o n, R o ] + -> R o +runSimple inp = evalBPOp (implicitly $ simpleOp inp) +``` + +Alternatively, we can define `simpleOp` in explicit monadic style, were +we specify our graph nodes explicitly. The results should be the same. + +``` {.sourceCode .literate .haskell} +simpleOpExplicit + :: (KnownNat m, KnownNat n, KnownNat o) + => R m + -> BPOp s '[ L n m, R n, L o n, R o ] (R o) +simpleOpExplicit inp = withInps $ \(w1 :< b1 :< w2 :< b2 :< Ø) -> do + -- First layer + y1 <- matVec ~$ (w1 :< x1 :< Ø) + let x2 = logistic (y1 + b1) + -- Second layer + y2 <- matVec ~$ (w2 :< x2 :< Ø) + return $ logistic (y2 + b2) + where + x1 = constVar inp +``` + +Now, for the magic of *backprop*: the library can now take advantage of +the implicit (or explicit) graph and use it to do back-propagation, too! + +``` {.sourceCode .literate .haskell} +simpleGrad + :: forall m n o. (KnownNat m, KnownNat n, KnownNat o) + => R m + -> R o + -> Tuple '[ L n m, R n, L o n, R o ] + -> Tuple '[ L n m, R n, L o n, R o ] +simpleGrad inp targ params = gradBPOp opError params + where + opError :: BPOp s '[ L n m, R n, L o n, R o ] Double + opError = do + res <- implicitly $ simpleOp inp + -- we explicitly bind err to prevent recomputation + err <- bindVar $ res - t + dot ~$ (err :< err :< Ø) + where + t = constVar targ +``` + +The result is the gradient of the input tuple’s components, with respect +to the `Double` result of `opError` (the squared error). We can then use +this gradient to do gradient descent. + +With Parameter Containers +========================= + +This method doesn’t quite scale, because we might want to make networks +with multiple layers and parameterize networks by layers. Let’s make +some basic container data types to help us organize our types, including +a recursive `Network` type that lets us chain multiple layers. + +``` {.sourceCode .literate .haskell} +data Layer :: Nat -> Nat -> Type where + Layer :: { _lWeights :: L m n + , _lBiases :: R m + } + -> Layer n m + deriving (Show, Generic) + + +data Network :: Nat -> [Nat] -> Nat -> Type where + NØ :: !(Layer a b) -> Network a '[] b + (:&) :: !(Layer a b) -> Network b bs c -> Network a (b ': bs) c +``` + +A `Layer n m` is a layer taking an n-vector and returning an m-vector. A +`Network a '[b, c, d] e` would be a Network that takes in an a-vector +and outputs an e-vector, with hidden layers of sizes b, c, and d. + +Isomorphisms +------------ + +The *backprop* library lets you apply operations on “parts” of data +types (like on the weights and biases of a `Layer`) by using `Iso`’s +(isomorphisms), like the ones from the *lens* library. The library +doesn’t depend on lens, but it can use the `Iso`s from the library and +also custom-defined ones. + +First, we can auto-generate isomorphisms using the *generics-sop* +library: + +``` {.sourceCode .literate .haskell} +instance SOP.Generic (Layer n m) +``` + +And then can create isomorphisms by hand for the two `Network` +constructors: + +``` {.sourceCode .literate .haskell} +netExternal :: Iso' (Network a '[] b) (Tuple '[Layer a b]) +netExternal = iso (\case NØ x -> x ::< Ø) + (\case I x :< Ø -> NØ x ) + +netInternal :: Iso' (Network a (b ': bs) c) (Tuple '[Layer a b, Network b bs c]) +netInternal = iso (\case x :& xs -> x ::< xs ::< Ø) + (\case I x :< I xs :< Ø -> x :& xs ) +``` + +An `Iso' a (Tuple as)` means that an `a` can really just be seen as a +tuple of `as`. + +Running a network +================= + +Now, we can write the `BPOp` that reprenents running the network and +getting a result. We pass in a `Sing bs` (a singleton list of the hidden +layer sizes) so that we can “pattern match” on the list and handle the +different network constructors differently. + +``` {.sourceCode .literate .haskell} +netOp + :: forall s a bs c. (KnownNat a, KnownNat c) + => Sing bs + -> BPOp s '[ R a, Network a bs c ] (R c) +netOp sbs = go sbs + where + go :: forall d es. KnownNat d + => Sing es + -> BPOp s '[ R d, Network d es c ] (R c) + go = \case + SNil -> withInps $ \(x :< n :< Ø) -> do + -- peek into the NØ using netExternal iso + l :< Ø <- netExternal #<~ n + -- run the 'layerOp' BP, with x and l as inputs + bpOp layerOp ~$ (x :< l :< Ø) + SNat `SCons` ses -> withInps $ \(x :< n :< Ø) -> withSingI ses $ do + -- peek into the (:&) using the netInternal iso + l :< n' :< Ø <- netInternal #<~ n + -- run the 'layerOp' BP, with x and l as inputs + z <- bpOp layerOp ~$ (x :< l :< Ø) + -- run the 'go ses' BP, with z and n as inputs + bpOp (go ses) ~$ (z :< n' :< Ø) + layerOp + :: forall d e. (KnownNat d, KnownNat e) + => BPOp s '[ R d, Layer d e ] (R e) + layerOp = withInps $ \(x :< l :< Ø) -> do + -- peek into the layer using the gTuple iso, auto-generated with SOP.Generic + w :< b :< Ø <- gTuple #<~ l + y <- matVec ~$ (w :< x :< Ø) + return $ logistic (y + b) +``` + +There’s some singletons work going on here, but it’s fairly standard +singletons stuff. Most of the complexity here is from the static typing +in our neural network type, and *not* from *backprop*. + +From *backprop* specifically, the only elements are `#<~` lets you +“split” an input ref with the given iso, and `bpOp`, which converts a +`BPOp` into an `Op` that you can bind with `~$`. + +Note that this library doesn’t support truly pattern matching on GADTs, +and that we had to pass in `Sing bs` as a reference to the structure of +our networks. + +Gradient Descent +---------------- + +Now we can do simple gradient descent. Defining an error function: + +``` {.sourceCode .literate .haskell} +errOp + :: KnownNat m + => R m + -> BVar s rs (R m) + -> BPOp s rs Double +errOp targ r = do + err <- bindVar $ r - t + dot ~$ (err :< err :< Ø) + where + t = constVar targ +``` + +And now, we can use `backprop` to generate the gradient, and shift the +`Network`! Things are made a bit cleaner from the fact that +`Network a bs c` has a `Num` instance, so we can use `(-)` and `(*)` +etc. + +``` {.sourceCode .literate .haskell} +train + :: (KnownNat a, SingI bs, KnownNat c) + => Double + -> R a + -> R c + -> Network a bs c + -> Network a bs c +train r x t n = case backprop (errOp t =<< netOp sing) (x ::< n ::< Ø) of + (_, _ :< I g :< Ø) -> n - (realToFrac r * g) +``` + +(`(::<)` is cons and `Ø` is nil for tuples.) + +Main +==== + +`main`, which will train on sample data sets, is still in progress! +Right now it just generates a random network using the *mwc-random* +library and prints each internal layer. + +``` {.sourceCode .literate .haskell} +main :: IO () +main = withSystemRandom $ \g -> do + n <- uniform @(Network 4 '[3,2] 1) g + void $ traverseNetwork sing (\l -> l <$ print l) n +``` + +Appendix: Boilerplate +===================== + +And now for some typeclass instances and boilerplates unrelated to the +*backprop* library that makes our custom types easier to use. + +``` {.sourceCode .literate .haskell} +instance KnownNat n => Variate (R n) where + uniform g = randomVector <$> uniform g <*> pure Uniform + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> uniform g + +instance (KnownNat m, KnownNat n) => Variate (L m n) where + uniform g = uniformSample <$> uniform g <*> pure 0 <*> pure 1 + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> uniform g + +instance (KnownNat n, KnownNat m) => Variate (Layer n m) where + uniform g = subtract 1 . (* 2) <$> (Layer <$> uniform g <*> uniform g) + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> uniform g + +instance (KnownNat m, KnownNat n) => Num (Layer n m) where + Layer w1 b1 + Layer w2 b2 = Layer (w1 + w2) (b1 + b2) + Layer w1 b1 - Layer w2 b2 = Layer (w1 - w2) (b1 - b2) + Layer w1 b1 * Layer w2 b2 = Layer (w1 * w2) (b1 * b2) + abs (Layer w b) = Layer (abs w) (abs b) + signum (Layer w b) = Layer (signum w) (signum b) + negate (Layer w b) = Layer (negate w) (negate b) + fromInteger x = Layer (fromInteger x) (fromInteger x) + +instance (KnownNat m, KnownNat n) => Fractional (Layer n m) where + Layer w1 b1 / Layer w2 b2 = Layer (w1 / w2) (b1 / b2) + recip (Layer w b) = Layer (recip w) (recip b) + fromRational x = Layer (fromRational x) (fromRational x) + +instance (KnownNat a, SingI bs, KnownNat c) => Variate (Network a bs c) where + uniform g = genNet sing (uniform g) + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> uniform g + +genNet + :: forall f a bs c. (Applicative f, KnownNat a, KnownNat c) + => Sing bs + -> (forall d e. (KnownNat d, KnownNat e) => f (Layer d e)) + -> f (Network a bs c) +genNet sbs f = go sbs + where + go :: forall d es. KnownNat d => Sing es -> f (Network d es c) + go = \case + SNil -> NØ <$> f + SNat `SCons` ses -> (:&) <$> f <*> go ses + +mapNetwork0 + :: forall a bs c. (KnownNat a, KnownNat c) + => Sing bs + -> (forall d e. (KnownNat d, KnownNat e) => Layer d e) + -> Network a bs c +mapNetwork0 sbs f = getI $ genNet sbs (I f) + +traverseNetwork + :: forall a bs c f. (KnownNat a, KnownNat c, Applicative f) + => Sing bs + -> (forall d e. (KnownNat d, KnownNat e) => Layer d e -> f (Layer d e)) + -> Network a bs c + -> f (Network a bs c) +traverseNetwork sbs f = go sbs + where + go :: forall d es. KnownNat d => Sing es -> Network d es c -> f (Network d es c) + go = \case + SNil -> \case + NØ x -> NØ <$> f x + SNat `SCons` ses -> \case + x :& xs -> (:&) <$> f x <*> go ses xs + +mapNetwork1 + :: forall a bs c. (KnownNat a, KnownNat c) + => Sing bs + -> (forall d e. (KnownNat d, KnownNat e) => Layer d e -> Layer d e) + -> Network a bs c + -> Network a bs c +mapNetwork1 sbs f = getI . traverseNetwork sbs (I . f) + +mapNetwork2 + :: forall a bs c. (KnownNat a, KnownNat c) + => Sing bs + -> (forall d e. (KnownNat d, KnownNat e) => Layer d e -> Layer d e -> Layer d e) + -> Network a bs c + -> Network a bs c + -> Network a bs c +mapNetwork2 sbs f = go sbs + where + go :: forall d es. KnownNat d => Sing es -> Network d es c -> Network d es c -> Network d es c + go = \case + SNil -> \case + NØ x -> \case + NØ y -> NØ (f x y) + SNat `SCons` ses -> \case + x :& xs -> \case + y :& ys -> f x y :& go ses xs ys + +instance (KnownNat a, SingI bs, KnownNat c) => Num (Network a bs c) where + (+) = mapNetwork2 sing (+) + (-) = mapNetwork2 sing (-) + (*) = mapNetwork2 sing (*) + negate = mapNetwork1 sing negate + abs = mapNetwork1 sing abs + signum = mapNetwork1 sing signum + fromInteger x = mapNetwork0 sing (fromInteger x) + +instance (KnownNat a, SingI bs, KnownNat c) => Fractional (Network a bs c) where + (/) = mapNetwork2 sing (/) + recip = mapNetwork1 sing recip + fromRational x = mapNetwork0 sing (fromRational x) +``` diff --git a/renders/NeuralTest.pdf b/renders/NeuralTest.pdf new file mode 100644 Binary files /dev/null and b/renders/NeuralTest.pdf differ diff --git a/samples/MNIST.lhs b/samples/MNIST.lhs new file mode 100644 --- /dev/null +++ b/samples/MNIST.lhs @@ -0,0 +1,506 @@ +% Learning MNIST with Neural Networks with backprop library +% Justin Le + +The *backprop* library performs back-propagation over a *hetereogeneous* +system of relationships. It offers both an implicit (*[ad][]*-like) and explicit graph +building usage style. Let's use it to build neural networks and learn +mnist! + +[ad]: http://hackage.haskell.org/package/ad + +Repository source is [on github][repo], and so are the [rendered +docs][docs]. + +[repo]: https://github.com/mstksg/backprop +[docs]: https://mstksg.github.io/backprop + +If you're reading this as a literate haskell file, you should know that a +[rendered pdf version is available on github.][rendered]. If you are reading +this as a pdf file, you should know that a [literate haskell version that +you can run][lhs] is also available on github! + +[rendered]: https://github.com/mstksg/backprop/blob/master/renders/MNIST.pdf +[lhs]: https://github.com/mstksg/backprop/blob/master/samples/MNIST.lhs + + +> {-# LANGUAGE BangPatterns #-} +> {-# LANGUAGE DataKinds #-} +> {-# LANGUAGE DeriveGeneric #-} +> {-# LANGUAGE GADTs #-} +> {-# LANGUAGE LambdaCase #-} +> {-# LANGUAGE ScopedTypeVariables #-} +> {-# LANGUAGE TupleSections #-} +> {-# LANGUAGE TypeApplications #-} +> {-# LANGUAGE ViewPatterns #-} +> {-# OPTIONS_GHC -fno-warn-orphans #-} +> {-# OPTIONS_GHC -fno-warn-incomplete-patterns #-} +> {-# OPTIONS_GHC -fno-warn-unused-top-binds #-} +> +> import Control.DeepSeq +> import Control.Exception +> import Control.Monad +> import Control.Monad.IO.Class +> import Control.Monad.Trans.Maybe +> import Control.Monad.Trans.State +> import Data.Bitraversable +> import Data.Foldable +> import Data.IDX +> import Data.List.Split +> import Data.Maybe +> import Data.Time.Clock +> import Data.Traversable +> import Data.Tuple +> import GHC.Generics (Generic) +> import GHC.TypeLits +> import Numeric.Backprop +> import Numeric.LinearAlgebra.Static hiding (dot) +> import Text.Printf +> import qualified Data.Vector as V +> import qualified Data.Vector.Generic as VG +> import qualified Data.Vector.Unboxed as VU +> import qualified Generics.SOP as SOP +> import qualified Numeric.LinearAlgebra as HM +> import qualified System.Random.MWC as MWC +> import qualified System.Random.MWC.Distributions as MWC + +Types +===== + +For the most part, we're going to be using the great *[hmatrix][]* library +and its vector and matrix types. It offers a type `L m n` for $m \times n$ +matrices, and a type `R n` for an $n$ vector. + +[hmatrix]: http://hackage.haskell.org/package/hmatrix + +First thing's first: let's define our neural networks as simple containers +of parameters (weight matrices and bias vectors). + +First, a type for layers: + +> data Layer i o = +> Layer { _lWeights :: !(L o i) +> , _lBiases :: !(R o) +> } +> deriving (Show, Generic) +> +> instance SOP.Generic (Layer i o) +> instance NFData (Layer i o) + +And a type for a simple feed-forward network with two hidden layers: + +> data Network i h1 h2 o = +> Net { _nLayer1 :: !(Layer i h1) +> , _nLayer2 :: !(Layer h1 h2) +> , _nLayer3 :: !(Layer h2 o) +> } +> deriving (Show, Generic) +> +> instance SOP.Generic (Network i h1 h2 o) +> instance NFData (Network i h1 h2 o) + +These are pretty straightforward container types...pretty much exactly the +type you'd make to represent these networks! Note that, following true +Haskell form, we separate out logic from data. This should be all we need. + +We derive an instance of `SOP.Generic` from the *[generics-sop][]* package, +which *backprop* uses to propagate derivatives on values inside product +types. + +[generics-sop]: http://hackage.haskell.org/package/generics-sop + +Instances +--------- + +Things are much simplier if we had `Num` and `Fractional` instances for +everything, so let's just go ahead and define that now, as well. Just a +little bit of boilerplate. + +> instance (KnownNat i, KnownNat o) => Num (Layer i o) where +> Layer w1 b1 + Layer w2 b2 = Layer (w1 + w2) (b1 + b2) +> Layer w1 b1 - Layer w2 b2 = Layer (w1 - w2) (b1 - b2) +> Layer w1 b1 * Layer w2 b2 = Layer (w1 * w2) (b1 * b2) +> abs (Layer w b) = Layer (abs w) (abs b) +> signum (Layer w b) = Layer (signum w) (signum b) +> negate (Layer w b) = Layer (negate w) (negate b) +> fromInteger x = Layer (fromInteger x) (fromInteger x) +> +> instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Num (Network i h1 h2 o) where +> Net a b c + Net d e f = Net (a + d) (b + e) (c + f) +> Net a b c - Net d e f = Net (a - d) (b - e) (c - f) +> Net a b c * Net d e f = Net (a * d) (b * e) (c * f) +> abs (Net a b c) = Net (abs a) (abs b) (abs c) +> signum (Net a b c) = Net (signum a) (signum b) (signum c) +> negate (Net a b c) = Net (negate a) (negate b) (negate c) +> fromInteger x = Net (fromInteger x) (fromInteger x) (fromInteger x) +> +> instance (KnownNat i, KnownNat o) => Fractional (Layer i o) where +> Layer w1 b1 / Layer w2 b2 = Layer (w1 / w2) (b1 / b2) +> recip (Layer w b) = Layer (recip w) (recip b) +> fromRational x = Layer (fromRational x) (fromRational x) +> +> instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Fractional (Network i h1 h2 o) where +> Net a b c / Net d e f = Net (a / d) (b / e) (c / f) +> recip (Net a b c) = Net (recip a) (recip b) (recip c) +> fromRational x = Net (fromRational x) (fromRational x) (fromRational x) + +`KnownNat` comes from *base*; it's a typeclass that *hmatrix* uses to refer +to the numbers in its type and use it to go about its normal hmatrixy +business. + +Ops +=== + +Now, *backprop* does require *primitive* differentiable operations on our +relevant types to be defined. *backprop* uses these primitive `Op`s to tie +everything together. Ideally we'd import these from a library that +implements these for you, and the end-user never has to make `Op` +primitives. + +But in this case, I'm going to put the definitions here to show that there +isn't any magic going on. If you're curious, refer to [documentation for +`Op`][opdoc] for more details on how `Op` is implemented and how this +works. + +[opdoc]: https://mstksg.github.io/backprop/Numeric-Backprop-Op.html + +First, matrix-vector multiplication primitive, giving an explicit gradient +function. + +> matVec +> :: (KnownNat m, KnownNat n) +> => Op '[ L m n, R n ] (R m) +> matVec = op2' $ \m v -> +> ( m #> v, \(fromMaybe 1 -> g) -> +> (g `outer` v, tr m #> g) +> ) + +Dot products would be nice too. + +> dot :: KnownNat n +> => Op '[ R n, R n ] Double +> dot = op2' $ \x y -> +> ( x <.> y, \case Nothing -> (y, x) +> Just g -> (konst g * y, x * konst g) +> ) + +Also a "scaling" function, scales a vector by a given factor. + +> scale +> :: KnownNat n +> => Op '[ Double, R n ] (R n) +> scale = op2' $ \a x -> +> ( konst a * x +> , \case Nothing -> (HM.sumElements (extract x ), konst a ) +> Just g -> (HM.sumElements (extract (x * g)), konst a * g) +> ) + +Finally, an operation to sum all of the items in the vector. + +> vsum +> :: KnownNat n +> => Op '[ R n ] Double +> vsum = op1' $ \x -> (HM.sumElements (extract x), maybe 1 konst) + +And why not, here's the [logistic function][], which we'll use as an +activation function for internal layers. We don't need to define this as +an `Op` up-front right now, because the library can automatically promote +any numeric polymorphic function (an `a -> a` or `a -> a -> a`, etc.) to an +`Op` anyways. + +[logistic function]: https://en.wikipedia.org/wiki/Logistic_function + +> logistic :: Floating a => a -> a +> logistic x = 1 / (1 + exp (-x)) + +Running our Network +=================== + +Now that we have our primitives in place, let's actually write a function +to run our network! + +> runLayer +> :: (KnownNat i, KnownNat o) +> => BPOp s '[ R i, Layer i o ] (R o) +> runLayer = withInps $ \(x :< l :< Ø) -> do +> w :< b :< Ø <- gTuple #<~ l +> y <- matVec ~$ (w :< x :< Ø) +> return $ y + b + +A `BPOp s '[ R i, Layer i o ] (R o)` is a backpropagatable function that +produces an `R o` (a vector with `o` elements, from the *[hmatrix][]* +library) given an input environment of an `R i` (the "input" of the layer) +and a layer. + +We use `withInps` to bring the environment into scope as a bunch of +`BVar`s. `x` is a `BVar` containing the input vector, and `l` is a `BVar` +containing the layer. + +The first thing we do is split out the parts of the layer so we can work +with the internal matrices. We can use `#<~` to "split out" the components +of a `BVar`, splitting on `gTuple` (which uses `GHC.Generics` to +automatically figure out how to split up a product type). + +Then we apply `matVec` (our primitive `Op` that does matrix-vector +multiplication) to `w` and `x`, and then the result is that added to the +bias vector `b`. + +We can write the `runNetwork` function pretty much the same way. + +> runNetwork +> :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) +> => BPOp s '[ R i, Network i h1 h2 o ] (R o) +> runNetwork = withInps $ \(x :< n :< Ø) -> do +> l1 :< l2 :< l3 :< Ø <- gTuple #<~ n +> y <- runLayer -$ (x :< l1 :< Ø) +> z <- runLayer -$ (logistic y :< l2 :< Ø) +> r <- runLayer -$ (logistic z :< l3 :< Ø) +> softmax -$ (r :< Ø) +> where +> softmax :: KnownNat n => BPOp s '[ R n ] (R n) +> softmax = withInps $ \(x :< Ø) -> do +> expX <- bindVar (exp x) +> totX <- vsum ~$ (expX :< Ø) +> scale ~$ (1/totX :< expX :< Ø) + + +After splitting out the layers in the input `Network`, we run each layer +successively using our previously defined `runLayer`, giving inputs using +`-$`. We can directly apply `logistic` to `BVar`s. At the end, we run a +[softmax function][] because MNIST is a classification challenge. The softmax +is done by applying $e^x$ for every item in the input vector, and dividing +each element by the total. + +[softmax function]: https://en.wikipedia.org/wiki/Softmax_function + + +The Magic +--------- + +What did we just define? Well, with a `BPOp s rs a`, we can *run* it and +get the output: + +> runNetOnInp +> :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) +> => Network i h1 h2 o +> -> R i +> -> R o +> runNetOnInp n x = evalBPOp runNetwork (x ::< n ::< Ø) + +But, the magic part is that we can also get the gradient! + +> gradNet +> :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) +> => Network i h1 h2 o +> -> R i +> -> Network i h1 h2 o +> gradNet n x = case gradBPOp runNetwork (x ::< n ::< Ø) of +> _gradX ::< gradN ::< Ø -> gradN + +This gives the gradient of all of the parameters in the matrices and +vectors inside the `Network`, which we can use to "train"! + +Training +======== + +Now for the real work. To train a network, we can do gradient descent +based on the gradient of some type of *error function* with respect to the +network parameters. Let's use the [cross entropy][], which is popular for +classification problems. + +[cross entropy]: https://en.wikipedia.org/wiki/Cross_entropy + +> crossEntropy +> :: KnownNat n +> => R n +> -> BPOpI s '[ R n ] Double +> crossEntropy targ (r :< Ø) = negate (dot .$ (log r :< t :< Ø)) +> where +> t = constVar targ + +Given a target vector and a `BVar` referring to the result of the network, +we can directly apply: + +$$ +H(\mathbf{r}, \mathbf{t}) = - (log(\mathbf{r}) \cdot \mathbf{t}) +$$ + +Just for fun, I implemented `crossEntropy` in "implicit-graph" mode, so you +don't see any binds or returns. + +Now, a function to make one gradient descent step based on an input vector +and a target, using `gradBPOp`: + +> trainStep +> :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) +> => Double +> -> R i +> -> R o +> -> Network i h1 h2 o +> -> Network i h1 h2 o +> trainStep r !x !t !n = case gradBPOp o (x ::< n ::< Ø) of +> _ ::< gN ::< Ø -> +> n - (realToFrac r * gN) +> where +> o :: BPOp s '[ R i, Network i h1 h2 o ] Double +> o = do +> y <- runNetwork +> implicitly (crossEntropy t) -$ (y :< Ø) + +A convenient wrapper for training over all of the observations in a list: + +> trainList +> :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) +> => Double +> -> [(R i, R o)] +> -> Network i h1 h2 o +> -> Network i h1 h2 o +> trainList r = flip $ foldl' (\n (x,y) -> trainStep r x y n) + +Pulling it all together +======================= + +`testNet` will be a quick way to test our net by computing the percentage +of correct guesses: (mostly using *hmatrix* stuff) + +> testNet +> :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) +> => [(R i, R o)] +> -> Network i h1 h2 o +> -> Double +> testNet xs n = sum (map (uncurry test) xs) / fromIntegral (length xs) +> where +> test :: R i -> R o -> Double +> test x (extract->t) +> | HM.maxIndex t == HM.maxIndex (extract r) = 1 +> | otherwise = 0 +> where +> r :: R o +> r = evalBPOp runNetwork (x ::< n ::< Ø) + +And now, a main loop! + +If you are following along at home, download the [mnist data set +files][mnist] and uncompress them into the folder `data`, and everything +should work fine. + +[mnist]: http://yann.lecun.com/exdb/mnist/ + +> main :: IO () +> main = MWC.withSystemRandom $ \g -> do +> Just train <- loadMNIST "data/train-images-idx3-ubyte" "data/train-labels-idx1-ubyte" +> Just test <- loadMNIST "data/t10k-images-idx3-ubyte" "data/t10k-labels-idx1-ubyte" +> putStrLn "Loaded data." +> net0 <- MWC.uniformR @(Network 784 300 100 9) (-1, 1) g +> flip evalStateT net0 . forM_ [1..] $ \e -> do +> train' <- liftIO . fmap V.toList $ MWC.uniformShuffle (V.fromList train) g +> liftIO $ printf "[Epoch %d]\n" (e :: Int) +> +> forM_ ([1..] `zip` chunksOf batch train') $ \(b, chnk) -> StateT $ \n0 -> do +> printf "(Batch %d)\n" (b :: Int) +> +> t0 <- getCurrentTime +> n' <- evaluate . force $ trainList rate chnk n0 +> t1 <- getCurrentTime +> printf "Trained on %d points in %s.\n" batch (show (t1 `diffUTCTime` t0)) +> +> let trainScore = testNet chnk n' +> testScore = testNet test n' +> printf "Training error: %.2f%%\n" ((1 - trainScore) * 100) +> printf "Validation error: %.2f%%\n" ((1 - testScore) * 100) +> +> return ((), n') +> where +> rate = 0.02 +> batch = 5000 + +Each iteration of the loop: + +1. Shuffles the training set +2. Splits it into chunks of `batch` size +3. Uses `trainList` to train over the batch +4. Computes the score based on `testNet` based on the training set and the + test set +5. Prints out the results + +And, that's really it! + +Result +------ + +I haven't put much into optimizing the library yet, but the network (with +hidden layer sizes 300 and 100) seems to take 25s on my computer to finish +a batch of 5000 training points. It's slow (five minutes per 60000 point +epooch), but it's a first unoptimized run and a proof of concept! It's my +goal to get this down to a point where the result has the same performance +characteristics as the actual backend (*hmatrix*), and so overhead is 0. + +Main takeaways +============== + +Most of the actual heavy lifting/logic actually came from the *hmatrix* +library itself. We just created simple types to wrap up our bare matrices. + +Basically, all that *backprop* did was give you an API to define *how to +run* a neural net --- how to *run* a net based on a `Network` and `R i` input +you were given. The goal of the library is to let you write down how to +run things in as natural way as possible. + +And then, after things are run, we can just get the gradient and roll from +there! + +Because the heavy lifting is done by the data types themselves, we can +presumably plug in *any* type and any tensor/numerical backend, and reap +the benefits of those libraries' optimizations and parallelizations. *Any* +type can be backpropagated! :D + +What now? +--------- + +Check out the docs for the [Numeric.Backprop][] module for a more detailed +picture of what's going on, or find more examples at the [github repo][repo]! + +[Numeric.Backprop]: https://mstksg.github.io/backprop/Numeric-Backprop.html + +Boring stuff +============ + +Here is a small wrapper function over the [mnist-idx][] library loading the +contents of the idx files into *hmatrix* vectors: + +[mnist-idx]: http://hackage.haskell.org/package/mnist-idx + +> loadMNIST +> :: FilePath +> -> FilePath +> -> IO (Maybe [(R 784, R 9)]) +> loadMNIST fpI fpL = runMaybeT $ do +> i <- MaybeT $ decodeIDXFile fpI +> l <- MaybeT $ decodeIDXLabelsFile fpL +> d <- MaybeT . return $ labeledIntData l i +> r <- MaybeT . return $ for d (bitraverse mkImage mkLabel . swap) +> liftIO . evaluate $ force r +> where +> mkImage :: VU.Vector Int -> Maybe (R 784) +> mkImage = create . VG.convert . VG.map (\i -> fromIntegral i / 255) +> mkLabel :: Int -> Maybe (R 9) +> mkLabel n = create $ HM.build 9 (\i -> if round i == n then 1 else 0) + +And here are instances to generating random +vectors/matrices/layers/networks, used for the initialization step. + +> instance KnownNat n => MWC.Variate (R n) where +> uniform g = randomVector <$> MWC.uniform g <*> pure Uniform +> uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g +> +> instance (KnownNat m, KnownNat n) => MWC.Variate (L m n) where +> uniform g = uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1 +> uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g +> +> instance (KnownNat i, KnownNat o) => MWC.Variate (Layer i o) where +> uniform g = Layer <$> MWC.uniform g <*> MWC.uniform g +> uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g +> +> instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => MWC.Variate (Network i h1 h2 o) where +> uniform g = Net <$> MWC.uniform g <*> MWC.uniform g <*> MWC.uniform g +> uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g +> diff --git a/samples/MonoTest.hs b/samples/MonoTest.hs new file mode 100644 --- /dev/null +++ b/samples/MonoTest.hs @@ -0,0 +1,18 @@ +{-# LANGUAGE GADTs #-} + +import Numeric.Backprop.Mono + +testImplicit :: BPOp s N3 Double Double +testImplicit = implicitly $ \(x :* y :* z :* ØV) -> + ((x * y) + y) * z + +testExplicit :: BPOp s N3 Double Double +testExplicit = withInps $ \(x :* y :* z :* ØV) -> do + xy <- op2 (*) ~$ (x :* y :* ØV) + xyy <- op2 (+) ~$ (xy :* y :* ØV) + op2 (*) ~$ (xyy :* z :* ØV) + +main :: IO () +main = do + print $ backprop testImplicit (2 :+ 3 :+ 4 :+ ØV) + print $ backprop testExplicit (2 :+ 3 :+ 4 :+ ØV) diff --git a/samples/NeuralTest.lhs b/samples/NeuralTest.lhs new file mode 100644 --- /dev/null +++ b/samples/NeuralTest.lhs @@ -0,0 +1,405 @@ +% Neural networks with backprop library +% Justin Le + +The *backprop* library performs back-propagation over a *hetereogeneous* +system of relationships. It offers both an implicit ([ad][]-like) and explicit graph +building usage style. Let's use it to build neural networks! + +[ad]: http://hackage.haskell.org/package/ad + +Repository source is [on github][repo], and so are the [rendered unstable +docs][docs]. + +[repo]: https://github.com/mstksg/backprop +[docs]: https://mstksg.github.io/backprop + +> {-# LANGUAGE DeriveGeneric #-} +> {-# LANGUAGE GADTs #-} +> {-# LANGUAGE LambdaCase #-} +> {-# LANGUAGE RankNTypes #-} +> {-# LANGUAGE ScopedTypeVariables #-} +> {-# LANGUAGE StandaloneDeriving #-} +> {-# LANGUAGE TypeApplications #-} +> {-# LANGUAGE TypeInType #-} +> {-# LANGUAGE TypeOperators #-} +> {-# LANGUAGE ViewPatterns #-} +> {-# OPTIONS_GHC -fno-warn-orphans #-} +> {-# OPTIONS_GHC -fno-warn-unused-top-binds #-} +> +> import Data.Functor +> import Data.Kind +> import Data.Maybe +> import Data.Singletons +> import Data.Singletons.Prelude +> import Data.Singletons.TypeLits +> import Data.Type.Combinator +> import Data.Type.Product +> import GHC.Generics (Generic) +> import Numeric.Backprop +> import Numeric.Backprop.Iso +> import Numeric.LinearAlgebra.Static hiding (dot) +> import System.Random.MWC +> import qualified Generics.SOP as SOP + +Ops +=== + +First, we define values of `Op` for the operations we want to do. `Op`s +are bundles of functions packaged with their hetereogeneous gradients. For +simple numeric functions, *backprop* can derive `Op`s automatically. But +for matrix operations, we have to derive them ourselves. + +The types help us with matching up the dimensions, but we still need to be +careful that our gradients are calculated correctly. + +`L` and `R` are matrix and vector types from the great *hmatrix* library. + +First, matrix-vector multiplication: + +> matVec +> :: (KnownNat m, KnownNat n) +> => Op '[ L m n, R n ] (R m) +> matVec = op2' $ \m v -> ( m #> v +> , \(fromMaybe 1 -> g) -> +> (g `outer` v, tr m #> g) +> ) + +Now, dot products: + +> dot :: KnownNat n +> => Op '[ R n, R n ] Double +> dot = op2' $ \x y -> ( x <.> y +> , \case Nothing -> (y, x) +> Just g -> (konst g * y, x * konst g) +> ) + +Polymorphic functions can be easily turned into `Op`s with `op1`/`op2` +etc., but they can also be run directly on graph nodes. + +> logistic :: Floating a => a -> a +> logistic x = 1 / (1 + exp (-x)) + +A Simple Complete Example +========================= + +At this point, we already have enough to train a simple single-hidden-layer +neural network: + +> simpleOp +> :: (KnownNat m, KnownNat n, KnownNat o) +> => R m +> -> BPOpI s '[ L n m, R n, L o n, R o ] (R o) +> simpleOp inp = \(w1 :< b1 :< w2 :< b2 :< Ø) -> +> let z = logistic $ liftB2 matVec w1 x + b1 +> in logistic $ liftB2 matVec w2 z + b2 +> where +> x = constVar inp + +Here, `simpleOp` is defined in implicit (non-monadic) style, given a tuple +of inputs and returning outputs. Now `simpleOp` can be "run" with the +input vectors and parameters (a `L n m`, `R n`, `L o n`, and `R o`) and +calculate the output of the neural net. + +> runSimple +> :: (KnownNat m, KnownNat n, KnownNat o) +> => R m +> -> Tuple '[ L n m, R n, L o n, R o ] +> -> R o +> runSimple inp = evalBPOp (implicitly $ simpleOp inp) + +Alternatively, we can define `simpleOp` in explicit monadic style, were we +specify our graph nodes explicitly. The results should be the same. + +> simpleOpExplicit +> :: (KnownNat m, KnownNat n, KnownNat o) +> => R m +> -> BPOp s '[ L n m, R n, L o n, R o ] (R o) +> simpleOpExplicit inp = withInps $ \(w1 :< b1 :< w2 :< b2 :< Ø) -> do +> -- First layer +> y1 <- matVec ~$ (w1 :< x1 :< Ø) +> let x2 = logistic (y1 + b1) +> -- Second layer +> y2 <- matVec ~$ (w2 :< x2 :< Ø) +> return $ logistic (y2 + b2) +> where +> x1 = constVar inp + +Now, for the magic of *backprop*: the library can now take advantage of +the implicit (or explicit) graph and use it to do back-propagation, too! + +> simpleGrad +> :: forall m n o. (KnownNat m, KnownNat n, KnownNat o) +> => R m +> -> R o +> -> Tuple '[ L n m, R n, L o n, R o ] +> -> Tuple '[ L n m, R n, L o n, R o ] +> simpleGrad inp targ params = gradBPOp opError params +> where +> opError :: BPOp s '[ L n m, R n, L o n, R o ] Double +> opError = do +> res <- implicitly $ simpleOp inp +> -- we explicitly bind err to prevent recomputation +> err <- bindVar $ res - t +> dot ~$ (err :< err :< Ø) +> where +> t = constVar targ + +The result is the gradient of the input tuple's components, with respect +to the `Double` result of `opError` (the squared error). We can then use +this gradient to do gradient descent. + +With Parameter Containers +========================= + +This method doesn't quite scale, because we might want to make networks +with multiple layers and parameterize networks by layers. Let's make some +basic container data types to help us organize our types, including a +recursive `Network` type that lets us chain multiple layers. + +> data Layer :: Nat -> Nat -> Type where +> Layer :: { _lWeights :: L m n +> , _lBiases :: R m +> } +> -> Layer n m +> deriving (Show, Generic) +> +> +> data Network :: Nat -> [Nat] -> Nat -> Type where +> NØ :: !(Layer a b) -> Network a '[] b +> (:&) :: !(Layer a b) -> Network b bs c -> Network a (b ': bs) c + +A `Layer n m` is a layer taking an n-vector and returning an m-vector. A +`Network a '[b, c, d] e` would be a Network that takes in an a-vector and +outputs an e-vector, with hidden layers of sizes b, c, and d. + +Isomorphisms +------------ + +The *backprop* library lets you apply operations on "parts" of data types +(like on the weights and biases of a `Layer`) by using `Iso`'s +(isomorphisms), like the ones from the *lens* library. The library doesn't +depend on lens, but it can use the `Iso`s from the library and also +custom-defined ones. + +First, we can auto-generate isomorphisms using the *generics-sop* library: + +> instance SOP.Generic (Layer n m) + +And then can create isomorphisms by hand for the two `Network` +constructors: + +> netExternal :: Iso' (Network a '[] b) (Tuple '[Layer a b]) +> netExternal = iso (\case NØ x -> x ::< Ø) +> (\case I x :< Ø -> NØ x ) +> +> netInternal :: Iso' (Network a (b ': bs) c) (Tuple '[Layer a b, Network b bs c]) +> netInternal = iso (\case x :& xs -> x ::< xs ::< Ø) +> (\case I x :< I xs :< Ø -> x :& xs ) + +An `Iso' a (Tuple as)` means that an `a` can really just be seen as a tuple +of `as`. + +Running a network +================= + +Now, we can write the `BPOp` that reprenents running the network and +getting a result. We pass in a `Sing bs` (a singleton list of the hidden +layer sizes) so that we can "pattern match" on the list and handle the +different network constructors differently. + +> netOp +> :: forall s a bs c. (KnownNat a, KnownNat c) +> => Sing bs +> -> BPOp s '[ R a, Network a bs c ] (R c) +> netOp sbs = go sbs +> where +> go :: forall d es. KnownNat d +> => Sing es +> -> BPOp s '[ R d, Network d es c ] (R c) +> go = \case +> SNil -> withInps $ \(x :< n :< Ø) -> do +> -- peek into the NØ using netExternal iso +> l :< Ø <- netExternal #<~ n +> -- run the 'layerOp' BP, with x and l as inputs +> bpOp layerOp ~$ (x :< l :< Ø) +> SNat `SCons` ses -> withInps $ \(x :< n :< Ø) -> withSingI ses $ do +> -- peek into the (:&) using the netInternal iso +> l :< n' :< Ø <- netInternal #<~ n +> -- run the 'layerOp' BP, with x and l as inputs +> z <- bpOp layerOp ~$ (x :< l :< Ø) +> -- run the 'go ses' BP, with z and n as inputs +> bpOp (go ses) ~$ (z :< n' :< Ø) +> layerOp +> :: forall d e. (KnownNat d, KnownNat e) +> => BPOp s '[ R d, Layer d e ] (R e) +> layerOp = withInps $ \(x :< l :< Ø) -> do +> -- peek into the layer using the gTuple iso, auto-generated with SOP.Generic +> w :< b :< Ø <- gTuple #<~ l +> y <- matVec ~$ (w :< x :< Ø) +> return $ logistic (y + b) + +There's some singletons work going on here, but it's fairly standard +singletons stuff. Most of the complexity here is from the static typing in +our neural network type, and *not* from *backprop*. + +From *backprop* specifically, the only elements are `#<~` lets you "split" an +input ref with the given iso, and `bpOp`, which converts a `BPOp` into an `Op` +that you can bind with `~$`. + +Note that this library doesn't support truly pattern matching on GADTs, and +that we had to pass in `Sing bs` as a reference to the structure of our +networks. + +Gradient Descent +---------------- + +Now we can do simple gradient descent. Defining an error function: + +> errOp +> :: KnownNat m +> => R m +> -> BVar s rs (R m) +> -> BPOp s rs Double +> errOp targ r = do +> err <- bindVar $ r - t +> dot ~$ (err :< err :< Ø) +> where +> t = constVar targ + +And now, we can use `backprop` to generate the gradient, and shift the +`Network`! Things are made a bit cleaner from the fact that `Network a bs c` +has a `Num` instance, so we can use `(-)` and `(*)` etc. + +> train +> :: (KnownNat a, SingI bs, KnownNat c) +> => Double +> -> R a +> -> R c +> -> Network a bs c +> -> Network a bs c +> train r x t n = case backprop (errOp t =<< netOp sing) (x ::< n ::< Ø) of +> (_, _ :< I g :< Ø) -> n - (realToFrac r * g) + +(`(::<)` is cons and `Ø` is nil for tuples.) + +Main +==== + +`main`, which will train on sample data sets, is still in progress! Right +now it just generates a random network using the *mwc-random* library and +prints each internal layer. + +> main :: IO () +> main = withSystemRandom $ \g -> do +> n <- uniform @(Network 4 '[3,2] 1) g +> void $ traverseNetwork sing (\l -> l <$ print l) n + +Appendix: Boilerplate +===================== + +And now for some typeclass instances and boilerplates unrelated to the +*backprop* library that makes our custom types easier to use. + +> instance KnownNat n => Variate (R n) where +> uniform g = randomVector <$> uniform g <*> pure Uniform +> uniformR (l, h) g = (\x -> x * (h - l) + l) <$> uniform g +> +> instance (KnownNat m, KnownNat n) => Variate (L m n) where +> uniform g = uniformSample <$> uniform g <*> pure 0 <*> pure 1 +> uniformR (l, h) g = (\x -> x * (h - l) + l) <$> uniform g +> +> instance (KnownNat n, KnownNat m) => Variate (Layer n m) where +> uniform g = subtract 1 . (* 2) <$> (Layer <$> uniform g <*> uniform g) +> uniformR (l, h) g = (\x -> x * (h - l) + l) <$> uniform g +> +> instance (KnownNat m, KnownNat n) => Num (Layer n m) where +> Layer w1 b1 + Layer w2 b2 = Layer (w1 + w2) (b1 + b2) +> Layer w1 b1 - Layer w2 b2 = Layer (w1 - w2) (b1 - b2) +> Layer w1 b1 * Layer w2 b2 = Layer (w1 * w2) (b1 * b2) +> abs (Layer w b) = Layer (abs w) (abs b) +> signum (Layer w b) = Layer (signum w) (signum b) +> negate (Layer w b) = Layer (negate w) (negate b) +> fromInteger x = Layer (fromInteger x) (fromInteger x) +> +> instance (KnownNat m, KnownNat n) => Fractional (Layer n m) where +> Layer w1 b1 / Layer w2 b2 = Layer (w1 / w2) (b1 / b2) +> recip (Layer w b) = Layer (recip w) (recip b) +> fromRational x = Layer (fromRational x) (fromRational x) +> +> instance (KnownNat a, SingI bs, KnownNat c) => Variate (Network a bs c) where +> uniform g = genNet sing (uniform g) +> uniformR (l, h) g = (\x -> x * (h - l) + l) <$> uniform g +> +> genNet +> :: forall f a bs c. (Applicative f, KnownNat a, KnownNat c) +> => Sing bs +> -> (forall d e. (KnownNat d, KnownNat e) => f (Layer d e)) +> -> f (Network a bs c) +> genNet sbs f = go sbs +> where +> go :: forall d es. KnownNat d => Sing es -> f (Network d es c) +> go = \case +> SNil -> NØ <$> f +> SNat `SCons` ses -> (:&) <$> f <*> go ses +> +> mapNetwork0 +> :: forall a bs c. (KnownNat a, KnownNat c) +> => Sing bs +> -> (forall d e. (KnownNat d, KnownNat e) => Layer d e) +> -> Network a bs c +> mapNetwork0 sbs f = getI $ genNet sbs (I f) +> +> traverseNetwork +> :: forall a bs c f. (KnownNat a, KnownNat c, Applicative f) +> => Sing bs +> -> (forall d e. (KnownNat d, KnownNat e) => Layer d e -> f (Layer d e)) +> -> Network a bs c +> -> f (Network a bs c) +> traverseNetwork sbs f = go sbs +> where +> go :: forall d es. KnownNat d => Sing es -> Network d es c -> f (Network d es c) +> go = \case +> SNil -> \case +> NØ x -> NØ <$> f x +> SNat `SCons` ses -> \case +> x :& xs -> (:&) <$> f x <*> go ses xs +> +> mapNetwork1 +> :: forall a bs c. (KnownNat a, KnownNat c) +> => Sing bs +> -> (forall d e. (KnownNat d, KnownNat e) => Layer d e -> Layer d e) +> -> Network a bs c +> -> Network a bs c +> mapNetwork1 sbs f = getI . traverseNetwork sbs (I . f) +> +> mapNetwork2 +> :: forall a bs c. (KnownNat a, KnownNat c) +> => Sing bs +> -> (forall d e. (KnownNat d, KnownNat e) => Layer d e -> Layer d e -> Layer d e) +> -> Network a bs c +> -> Network a bs c +> -> Network a bs c +> mapNetwork2 sbs f = go sbs +> where +> go :: forall d es. KnownNat d => Sing es -> Network d es c -> Network d es c -> Network d es c +> go = \case +> SNil -> \case +> NØ x -> \case +> NØ y -> NØ (f x y) +> SNat `SCons` ses -> \case +> x :& xs -> \case +> y :& ys -> f x y :& go ses xs ys +> +> instance (KnownNat a, SingI bs, KnownNat c) => Num (Network a bs c) where +> (+) = mapNetwork2 sing (+) +> (-) = mapNetwork2 sing (-) +> (*) = mapNetwork2 sing (*) +> negate = mapNetwork1 sing negate +> abs = mapNetwork1 sing abs +> signum = mapNetwork1 sing signum +> fromInteger x = mapNetwork0 sing (fromInteger x) +> +> instance (KnownNat a, SingI bs, KnownNat c) => Fractional (Network a bs c) where +> (/) = mapNetwork2 sing (/) +> recip = mapNetwork1 sing recip +> fromRational x = mapNetwork0 sing (fromRational x) diff --git a/src/Data/Type/Util.hs b/src/Data/Type/Util.hs new file mode 100644 --- /dev/null +++ b/src/Data/Type/Util.hs @@ -0,0 +1,188 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE EmptyCase #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} + +module Data.Type.Util where + +import Control.Applicative +import Data.Bifunctor +import Data.Kind +import Data.Monoid hiding (Sum) +import Data.Type.Conjunction +import Data.Type.Fin +import Data.Type.Index +import Data.Type.Length +import Data.Type.Nat +import Data.Type.Product +import Data.Type.Sum +import Data.Type.Vector +import Lens.Micro +import Type.Class.Higher +import Type.Class.Known +import Type.Class.Witness +import Type.Family.List +import Type.Family.Nat + +-- | @'Replicate' n a@ is a list of @a@s repeated @n@ times. +-- +-- >>> :kind! Replicate N3 Int +-- '[Int, Int, Int] +-- >>> :kind! Replicate N5 Double +-- '[Double, Double, Double, Double, Double] +type family Replicate (n :: N) (a :: k) = (as :: [k]) | as -> n where + Replicate 'Z a = '[] + Replicate ('S n) a = a ': Replicate n a + +vecToProd + :: VecT n f a + -> Prod f (Replicate n a) +vecToProd = \case + ØV -> Ø + x :* xs -> x :< vecToProd xs + +prodToVec' + :: Nat n + -> Prod f (Replicate n a) + -> VecT n f a +prodToVec' = \case + Z_ -> \case + Ø -> ØV + S_ n -> \case + x :< xs -> x :* prodToVec' n xs + +prodAlong + :: VecT n f b + -> Prod f (Replicate n a) + -> VecT n f a +prodAlong = \case + ØV -> \case + Ø -> ØV + _ :* v -> \case + x :< xs -> x :* prodAlong v xs + +finIndex + :: Fin n + -> Index (Replicate n a) a +finIndex = \case + FZ -> IZ + FS f -> IS (finIndex f) + +traverse1_ + :: (Applicative h, Traversable1 t) + => (forall a. f a -> h ()) + -> t f b + -> h () +traverse1_ f = ($ pure ()) + . appEndo + . getConst + . foldMap1 (\y -> Const (Endo (f y *>))) + +itraverse1_ + :: (Applicative h, IxFoldable1 i t) + => (forall a. i b a -> f a -> h ()) + -> t f b + -> h () +itraverse1_ f = ($ pure ()) + . appEndo + . getConst + . ifoldMap1 (\i y -> Const (Endo (f i y *>))) + +for1 + :: (Applicative h, Traversable1 t) + => t f b + -> (forall a. f a -> h (g a)) + -> h (t g b) +for1 x f = traverse1 f x + +for1_ + :: (Applicative h, Traversable1 t) + => t f b + -> (forall a. f a -> h ()) + -> h () +for1_ x f = traverse1_ f x + +ifor1 + :: (Applicative h, IxTraversable1 i t) + => t f b + -> (forall a. i b a -> f a -> h (g a)) + -> h (t g b) +ifor1 x f = itraverse1 f x + +ifor1_ + :: (Applicative h, IxFoldable1 i t) + => t f b + -> (forall a. i b a -> f a -> h ()) + -> h () +ifor1_ x f = itraverse1_ f x + +zipP + :: Prod f as + -> Prod g as + -> Prod (f :&: g) as +zipP = \case + Ø -> \case + Ø -> Ø + x :< xs -> \case + y :< ys -> x :&: y :< zipP xs ys + +unzipP + :: Prod (f :&: g) as + -> (Prod f as, Prod g as) +unzipP = \case + Ø -> (Ø, Ø) + (x :&: y) :< zs -> bimap (x :<) (y :<) (unzipP zs) + +indexP :: Index as a -> Lens' (Prod g as) (g a) +indexP = \case + IZ -> \f -> \case + x :< xs -> (:< xs) <$> f x + IS i -> \f -> \case + x :< xs -> (x :<) <$> indexP i f xs + +reIndex + :: forall k (f :: k -> Type) (as :: [k]) (a :: k). () + => Index as a + -> Index (f <$> as) (f a) +reIndex = undefined + +prodLength + :: Prod f as + -> Length as +prodLength = \case + Ø -> LZ + _ :< xs -> LS (prodLength xs) + +withEvery + :: forall c f as. (Known Length as, Every c as) + => (forall a. c a => f a) + -> Prod f as +withEvery = withEvery' @c known + +withEvery' + :: forall c f as. Every c as + => Length as + -> (forall a. c a => f a) + -> Prod f as +withEvery' l x = map1 ((// x) . every @_ @c) (indices' l) + +tagSum + :: Prod f as + -> Sum g as + -> Sum (f :&: g) as +tagSum = \case + Ø -> \case + x :< xs -> \case + InL y -> InL (x :&: y) + InR ys -> InR (tagSum xs ys) diff --git a/src/Numeric/Backprop.hs b/src/Numeric/Backprop.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop.hs @@ -0,0 +1,1597 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} + +-- | +-- Module : Numeric.Backprop +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- +-- Provides the 'BP' monad and the 'BVar' type; after manipulating 'BVar's +-- (inputs to your function) to produce a result, the library tracks internal data +-- dependencies, which are used to perform back-propagation (reverse-mode +-- automatic differentiation) to calculate the gradient of the output with +-- respect to the inputs. +-- +-- Similar to automatic differentiation from the /ad/ library and +-- "Numeric.AD.Mode.Reverse", except for a few key differences: +-- +-- 1. Most importantly, this library implements /heterogeneous/ +-- back-propagation, so you can manipulate values of different types (like +-- different matrix and vector types, and product and sum types). This is +-- essential for things like back-propagation for neural networks. +-- +-- 2. This module allows you to /explicitly/ build your data dependency +-- graph if you wish, which allows the library to perform optimizations and +-- reduce extra allocation, which may or may not provide advantages over +-- "Numeric.AD.Mode.Reverse"'s 'System.IO.Unsafe.unsafePerformIO'-based +-- implicit graph building. +-- +-- See the for more information +-- and links to demonstrations and tutorials. If you want to plunge right +-- in, you can also look directly at the main types, 'BP', 'BPOp', 'BVar', +-- 'Op', and the main functions, 'backprop' and 'opVar'. +-- +-- + +module Numeric.Backprop ( + -- * Types + -- ** Backprop types + BP, BPOp, BPOpI, BVar, Op, OpB + -- ** Tuple types#prod# + -- $prod + , Prod(..), Tuple, I(..) + -- * BP + -- ** Backprop + , backprop, evalBPOp, gradBPOp + , backprop', gradBPOp' + -- ** Utility combinators + , withInps, implicitly + , withInps', implicitly' + -- * Vars + , constVar + , inpVar, inpVars + , bpOp + , bindVar + , inpVars' + , bpOp' + , bindVar' + -- ** From Ops + , opVar, (~$) + , opVar1, opVar2, opVar3 + , (-$) + , opVar' + , opVar1', opVar2', opVar3' + -- ** Var manipulation + -- *** As parts + , partsVar, (#<~), withParts + , splitVars, gSplit, gTuple + , partsVar', withParts' + , splitVars', gSplit' + -- *** As sums + , choicesVar, (?<~), withChoices + , choicesVar', withChoices' + -- $sum + , Sum(..) + -- *** As sums of products + , sopVar, gSplits, gSOP + , sopVar', gSplits' + -- ** Combining + , liftB, (.$), liftB1, liftB2, liftB3 + -- * Op + , op1, op2, op3, opN, composeOp, composeOp1, (~.) + , op1', op2', op3' + -- * Utility + , pattern (:>), only, head' + , pattern (::<), only_ + , Summer(..), Unity(..) + , summers, unities + , summers', unities' + ) where + +import Control.Monad.Base +import Control.Monad.Reader +import Control.Monad.ST +import Control.Monad.State +import Data.Maybe +import Data.Monoid ((<>)) +import Data.STRef +import Data.Type.Combinator +import Data.Type.Conjunction +import Data.Type.Index +import Data.Type.Length +import Data.Type.Product +import Data.Type.Sum hiding (index) +import Data.Type.Util +import Lens.Micro hiding (ix) +import Lens.Micro.Mtl hiding (view) +import Numeric.Backprop.Internal +import Numeric.Backprop.Iso +import Numeric.Backprop.Op +import Type.Class.Higher +import Type.Class.Known +import Type.Class.Witness +import qualified Generics.SOP as SOP + +-- $prod +-- +-- 'Prod' is a heterogeneous list/tuple type, which allows you to tuple +-- together multiple values of different types and operate on them +-- generically. +-- +-- A @'Prod' f '[a, b, c]@ contains an @f a@, an @f b@, and an @f c@, and +-- is constructed by consing them together with ':<' (using 'Ø' as nil): +-- +-- @ +-- 'I' "hello" ':<' I True :< I 7.8 :< Ø :: 'Prod' 'I' '[String, Bool, Double] +-- 'C' "hello" :< C "world" :< C "ok" :< Ø :: 'Prod' ('C' String) '[a, b, c] +-- 'Proxy' :< Proxy :< Proxy :< Ø :: 'Prod' 'Proxy' '[a, b, c] +-- @ +-- +-- ('I' is the identity functor, and 'C' is the constant functor) +-- +-- So, in general: +-- +-- @ +-- x :: f a +-- y :: f b +-- z :: f c +-- x :< y :< z :< Ø :: Prod f '[a, b, c] +-- @ +-- +-- If you're having problems typing 'Ø', you can use 'only': +-- +-- @ +-- only z :: Prod f '[c] +-- x :< y :< only z :: Prod f '[a, b, c] +-- @ +-- +-- 'Tuple' is provided as a convenient type synonym for 'Prod' 'I', and has +-- a convenient pattern synonym '::<' (and 'only_'), which can also be used +-- for pattern matching: +-- +-- @ +-- x :: a +-- y :: b +-- z :: c +-- +-- 'only_' z :: 'Tuple' '[c] +-- x '::<' y ::< z ::< Ø :: 'Tuple' '[a, b, c] +-- x ::< y ::< only_ z :: 'Tuple' '[a, b, c] +-- @ + + +-- $sum +-- +-- #sum# +-- +-- Like the 'Prod' type (see mini-tutorial at "Numeric.Backprop#prod"), the +-- 'Sum' type lets you make arbitrary sum types over different types and +-- work with them generically. +-- +-- A @'Sum' f '[a, b, c]@ contains /either/ an @f a@, an @f b@, /or/ an @f +-- c@, and is constructed with the constructors 'InL' and 'InR', which are +-- analogous to 'Left' and 'Right'. +-- +-- For a value of type @'Sum' f '[Int, Bool, String]@, there are three +-- constructors: +-- +-- @ +-- 'InL' :: f Int -> 'Sum' f '[Int, Bool, String] +-- InL . InR :: f Bool -> Sum f '[Int, Bool, String] +-- InL . InR . InR :: f String -> Sum f '[Int, Bool, String] +-- @ +-- +-- Each 'InR' "pushes deeper" into the 'Sum'. +-- +-- Likewise, if you have a value of type @'Sum' f '[Int, Bool, String]@, +-- you can see which constructor it was made (and what type it contains) +-- with by pattern matching: +-- +-- @ +-- foo :: 'Sum' f '[Int, Bool, String] +-- +-- case foo of +-- 'InL' i -> -- foo contains an "f Int" +-- 'InR' (InL b) -> -- foo contains an "f Bool" +-- InR (InR (InL s)) -> -- foo contains an "f String" +-- @ + + + +-- | A handy type synonym representing a 'BP' action that returns a 'BVar'. +-- This is handy because this is the form of 'BP' actions that +-- 'backprop' and 'gradBPOp' (etc.) expects. +-- +-- A value of type: +-- +-- @ +-- 'BPOp' s rs a +-- @ +-- +-- is an action that takes an input environment of @rs@ and produces +-- a 'BVar' containing a value of type @a@. Because it returns a 'BVar', +-- the library can track the data dependencies between the 'BVar' and the +-- input environment and perform back-propagation. +-- +-- See documentation for 'BP' for an explanation of the phantom type +-- parameter @s@. +type BPOp s rs a = BP s rs (BVar s rs a) + +-- | An "implicit" operation on 'BVar's that can be backpropagated. +-- A value of type: +-- +-- @ +-- 'BPOpI' s rs a +-- @ +-- +-- takes a bunch of 'BVar's containg @rs@ and uses them to (purely) produce +-- a 'BVar' containing an @a@. +-- +-- @ +-- foo :: BPOpI s '[ Double, Double ] Double +-- foo (x :< y :< Ø) = x + sqrt y +-- @ +-- +-- If you are exclusively doing implicit back-propagation by combining +-- 'BVar's and using 'BPOpI's, you are probably better off just importing +-- "Numeric.Backprop.Implicit", which provides better tools. This type +-- synonym exists in "Numeric.Backprop" just for the 'implicitly' function, +-- which can convert "implicit" backprop functions like a @'BPOpI' s rs a@ +-- into an "explicit" graph backprop function, a @'BPOp' s rs a@. +type BPOpI s rs a = Prod (BVar s rs) rs -> BVar s rs a + + +-- | A version of 'opVar' taking an explicit 'Summer', so can be used on +-- values of types that aren't instances of 'Num'. +opVar' + :: forall s rs as a. () + => Summer a + -> OpB s as a + -> Prod (BVar s rs) as + -> BP s rs (BVar s rs a) +opVar' s o i = do + xs <- traverse1 (fmap I . BP . resolveVar) i + (res, gf) <- BP . liftBase $ runOpM' o xs + let bp = BPN { _bpnOut = only $ FRInternal [] + , _bpnRes = only_ res + , _bpnGradFunc = gf . head' + , _bpnGradCache = Nothing + , _bpnSummer = only s + } + r <- BP . liftBase $ newSTRef bp + itraverse1_ (registerVar . flip IRNode r) i + return (BVNode IZ r) + +-- | A version of 'splitVars' taking explicit 'Summer's and 'Unity's, so it +-- can be run with types that aren't instances of 'Num'. +splitVars' + :: forall s rs as. () + => Prod Summer as + -> Prod Unity as + -> BVar s rs (Tuple as) + -> BP s rs (Prod (BVar s rs) as) +splitVars' ss us = partsVar' ss us id + +-- | Split out a 'BVar' of a tuple into a tuple ('Prod') of 'BVar's. +-- +-- @ +-- -- the environment is a single Int-Bool tuple, tup +-- stuff :: 'BP' s '[ Tuple '[Int, Bool] ] a +-- stuff = 'withInps' $ \\(tup :< Ø) -\> do +-- i :< b :< Ø <- 'splitVars' tup +-- -- now, i is a 'BVar' pointing to the 'Int' inside tup +-- -- and b is a 'BVar' pointing to the 'Bool' inside tup +-- -- you can do stuff with the i and b here +-- @ +-- +-- Note that +-- +-- @ +-- 'splitVars' = 'partsVar' 'id' +-- @ +splitVars + :: forall s rs as. (Every Num as, Known Length as) + => BVar s rs (Tuple as) + -> BP s rs (Prod (BVar s rs) as) +splitVars = partsVar id + +-- | A version of 'partsVar' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +partsVar' + :: forall s rs bs b. () + => Prod Summer bs + -> Prod Unity bs + -> Iso' b (Tuple bs) + -> BVar s rs b + -> BP s rs (Prod (BVar s rs) bs) +partsVar' ss us i = + fmap (view sum1) . sopVar' (only ss) (only us) (i . resum1) + +-- | Use an 'Iso' (or compatible 'Control.Lens.Iso.Iso' from the lens +-- library) to "pull out" the parts of a data type and work with each part +-- as a 'BVar'. +-- +-- If there is an isomorphism between a @b@ and a @'Tuple' as@ (that is, if +-- an @a@ is just a container for a bunch of @as@), then it lets you break +-- out the @as@ inside and work with those. +-- +-- @ +-- data Foo = F Int Bool +-- +-- fooIso :: 'Iso'' Foo (Tuple '[Int, Bool]) +-- fooIso = 'iso' (\\(F i b) -\> i ::\< b ::\< Ø) +-- (\\(i ::\< b ::\< Ø) -\> F i b ) +-- +-- 'partsVar' fooIso :: 'BVar' rs Foo -> 'BP' s rs ('Prod' ('BVar' s rs) '[Int, Bool]) +-- +-- stuff :: 'BP' s '[Foo] a +-- stuff = 'withInps' $ \\(foo :< Ø) -\> do +-- i :< b :< Ø <- partsVar fooIso foo +-- -- now, i is a 'BVar' pointing to the 'Int' inside foo +-- -- and b is a 'BVar' pointing to the 'Bool' inside foo +-- -- you can do stuff with the i and b here +-- @ +-- +-- You can use this to pass in product types as the environment to a 'BP', +-- and then break out the type into its constituent products. +-- +-- Note that for a type like @Foo@, @fooIso@ can be generated automatically +-- with 'GHC.Generics.Generic' from "GHC.Generics" and +-- 'Generics.SOP.Generic' from "Generics.SOP" and /generics-sop/, using the +-- 'gTuple' iso. See 'gSplit' for more information. +-- +-- Also, if you are literally passing a tuple (like +-- @'BP' s '[Tuple '[Int, Bool]@) then you can give in the identity +-- isomorphism ('id') or use 'splitVars'. +partsVar + :: forall s rs bs b. (Every Num bs, Known Length bs) + => Iso' b (Tuple bs) + -> BVar s rs b + -> BP s rs (Prod (BVar s rs) bs) +partsVar = partsVar' summers unities + +-- | A useful infix alias for 'partsVar'. +-- +-- Building on the example from 'partsVar': +-- +-- @ +-- data Foo = F Int Bool +-- +-- fooIso :: 'Iso'' Foo (Tuple '[Int, Bool]) +-- fooIso = 'iso' (\\(F i b) -\> i ::\< b ::\< Ø) +-- (\\(i ::\< b ::\< Ø) -\> F i b ) +-- +-- stuff :: 'BP' s '[Foo] a +-- stuff = 'withInps' $ \\(foo :< Ø) -\> do +-- i :< b :< Ø <- fooIso '#<~' foo +-- -- now, i is a 'BVar' pointing to the 'Int' inside foo +-- -- and b is a 'BVar' pointing to the 'Bool' inside foo +-- -- you can do stuff with the i and b here +-- @ +-- +-- See 'gSplit' for an example usage of splitting up an arbitrary product +-- type (like @Foo@) using "GHC.Geneics" and "Generics.SOP". +infixr 1 #<~ +(#<~) + :: (Every Num bs, Known Length bs) + => Iso' b (Tuple bs) + -> BVar s rs b + -> BP s rs (Prod (BVar s rs) bs) +(#<~) = partsVar + +-- | A version of 'withParts' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +withParts' + :: Prod Summer bs + -> Prod Unity bs + -> Iso' b (Tuple bs) + -> BVar s rs b + -> (Prod (BVar s rs) bs -> BP s rs a) + -> BP s rs a +withParts' ss us i r f = do + p <- partsVar' ss us i r + f p + +-- | A continuation-based version of 'partsVar'. Instead of binding the +-- parts and using it in the rest of the block, provide a continuation to +-- handle do stuff with the parts inside. +-- +-- Building on the example from 'partsVar': +-- +-- @ +-- data Foo = F Int Bool +-- +-- fooIso :: 'Iso'' Foo (Tuple '[Int, Bool]) +-- fooIso = 'iso' (\\(F i b) -\> i ::\< b ::\< Ø) +-- (\\(i ::\< b ::\< Ø) -\> F i b ) +-- +-- stuff :: 'BP' s '[Foo] a +-- stuff = 'withInps' $ \\(foo :< Ø) -\> do +-- 'withParts' fooIso foo $ \\(i :< b :< Ø) -\> do +-- -- now, i is a 'BVar' pointing to the 'Int' inside foo +-- -- and b is a 'BVar' pointing to the 'Bool' inside foo +-- -- you can do stuff with the i and b here +-- @ +-- +-- Useful so that you can work with the internal parts of the data type +-- in a closure, so the parts don't leak out to the rest of your 'BP'. +-- But, mostly just a stylistic choice. +withParts + :: (Every Num bs, Known Length bs) + => Iso' b (Tuple bs) + -> BVar s rs b + -> (Prod (BVar s rs) bs -> BP s rs a) + -> BP s rs a +withParts i r f = do + p <- partsVar i r + f p + +-- | A version of 'gSplit' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +gSplit' + :: (SOP.Generic b, SOP.Code b ~ '[bs]) + => Prod Summer bs + -> Prod Unity bs + -> BVar s rs b + -> BP s rs (Prod (BVar s rs) bs) +gSplit' ss us = partsVar' ss us gTuple + +-- | Using 'GHC.Generics.Generic' from "GHC.Generics" and +-- 'Generics.SOP.Generic' from "Generics.SOP", /split/ a 'BVar' containing +-- a product type into a tuple ('Prod') of 'BVar's pointing to each value +-- inside. +-- +-- Building on the example from 'partsVar': +-- +-- @ +-- import qualified Generics.SOP as SOP +-- +-- data Foo = F Int Bool +-- deriving Generic +-- +-- instance SOP.Generic Foo +-- +-- 'gSplit' :: 'BVar' rs Foo -> 'BP' s rs ('Prod' ('BVar' s rs) '[Int, Bool]) +-- +-- stuff :: 'BP' s '[Foo] a +-- stuff = 'withInps' $ \\(foo :< Ø) -\> do +-- i :< b :< Ø <- 'gSplit' foo +-- -- now, i is a 'BVar' pointing to the 'Int' inside foo +-- -- and b is a 'BVar' pointing to the 'Bool' inside foo +-- -- you can do stuff with the i and b here +-- @ +-- +-- Because @Foo@ is a straight up product type, 'gSplit' can use +-- "GHC.Generics" and take out the items inside. +-- +-- Note that because +-- +-- @ +-- 'gSplit' = 'splitVars' 'gTuple' +-- @ +-- +-- Then, you can also use 'gTuple' with '#<~': +-- +-- @ +-- stuff :: 'BP' s '[Foo] a +-- stuff = 'withInps' $ \\(foo :< Ø) -\> do +-- i :< b :< Ø <- 'gTuple' '#<~' foo +-- -- now, i is a 'BVar' pointing to the 'Int' inside foo +-- -- and b is a 'BVar' pointing to the 'Bool' inside foo +-- -- you can do stuff with the i and b here +-- @ +-- +gSplit + :: (Every Num bs, Known Length bs, SOP.Generic b, SOP.Code b ~ '[bs]) + => BVar s rs b + -> BP s rs (Prod (BVar s rs) bs) +gSplit = gSplit' summers unities + +-- | A version of 'choicesVar' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +choicesVar' + :: forall s rs bs b. () + => Prod Summer bs + -> Prod Unity bs + -> Iso' b (Sum I bs) + -> BVar s rs b + -> BP s rs (Sum (BVar s rs) bs) +choicesVar' ss us i r = do + x <- BP $ resolveVar r + let xs :: Sum I bs + xs = view i x + ifor1 ((ss `zipP` us) `tagSum` xs) $ \ix ((s :&: u) :&: I (y :: c)) -> do + let bp :: BPNode s rs '[b] '[c] + bp = BPN { _bpnOut = only $ FRInternal [] + , _bpnRes = only_ y + , _bpnGradFunc = return . only_ . review i + . injectSum ix + . maybe (I (getUnity u)) I + . head' + , _bpnGradCache = Nothing + , _bpnSummer = only s + } + r' <- BP . liftBase $ newSTRef bp + registerVar (IRNode IZ r') r + return $ BVNode IZ r' +-- TODO: cannot implement via sopVar? oh well. + +-- | Use an 'Iso' (or compatible 'Control.Lens.Iso.Iso' from the lens +-- library) to "pull out" the different constructors of a sum type and +-- return a (choice) sum of 'BVar's that you can pattern match on. +-- +-- If there is an isomorphism between a @b@ and a @'Sum' 'I' as@ (that is, +-- if an @a@ is just a sum type for every type in @as@), then it lets you +-- /branch/ on which constructor is used inside the @b@. +-- +-- Essentially implements pattern matching on 'BVar' values. +-- +-- @ +-- data Bar = A Int | B Bool | C String +-- +-- barIso :: 'Iso'' Bar ('Sum' I '[Int, Bool, String]) +-- barIso = 'iso' (\\case A i -> 'InL' (I i) +-- B b -> 'InR' ('InL' (I b)) +-- C s -> 'InR' ('InR' ('InL' (I s)) +-- ) +-- (\\case 'InL' (I i) -> A i +-- 'InR' ('InL' (I b)) -> B b +-- 'InR' ('InR' ('InL' (I s))) -> C s +-- ) +-- +-- choicesVar barIso :: BVar rs Bar -> BP s rs (Sum I (BVar s rs) '[Int, Bool, String]) +-- +-- stuff :: 'BP' s '[Bar] a +-- stuff = 'withInps' $ \\(bar :< Ø) -\> do +-- c <- 'choicesVar' barIso bar +-- case c of +-- 'InL' i -> do +-- -- in this branch, bar was made with the A constructor +-- -- i is the Int inside it +-- 'InR' ('InL' b) -> do +-- -- in this branch, bar was made with the B constructor +-- -- b is the Bool inside it +-- 'InR' ('InR' ('InL' s)) -> do +-- -- in this branch, bar was made with the B constructor +-- -- s is the String inside it +-- @ +-- +-- You can use this to pass in sum types as the environment to a 'BP', and +-- then branch on which constructor the value was made with. +-- +-- See "Numeric.Backprop#sum" for a mini-tutorial on 'Sum'. +choicesVar + :: forall s rs bs b. (Every Num bs, Known Length bs) + => Iso' b (Sum I bs) + -> BVar s rs b + -> BP s rs (Sum (BVar s rs) bs) +choicesVar = choicesVar' summers unities + +-- | A version of 'withChoices' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +withChoices' + :: forall s rs bs b a. () + => Prod Summer bs + -> Prod Unity bs + -> Iso' b (Sum I bs) + -> BVar s rs b + -> (Sum (BVar s rs) bs -> BP s rs a) + -> BP s rs a +withChoices' ss us i r f = do + c <- choicesVar' ss us i r + f c + +-- | A continuation-based version of 'choicesVar'. Instead of binding the +-- parts and using it in the rest of the block, provide a continuation that +-- will handle every possible constructor/case of the type of the value the +-- 'BVar' points to. +-- +-- Building on the example from 'choicesVar': +-- +-- @ +-- data Bar = A Int | B Bool | C String +-- +-- barIso :: 'Iso'' Bar ('Sum' I '[Int, Bool, String]) +-- barIso = 'iso' (\\case A i -> 'InL' (I i) +-- B b -> 'InR' ('InL' (I b)) +-- C s -> 'InR' ('InR' ('InL' (I s)) +-- ) +-- (\\case 'InL' (I i) -> A i +-- 'InR' ('InL' (I b)) -> B b +-- 'InR' ('InR' ('InL' (I s))) -> C s +-- ) +-- +-- 'choicesVar' barIso :: BVar rs Bar -> BP s rs (Sum I (BVar s rs) '[Int, Bool, String]) +-- +-- stuff :: 'BP' s '[Bar] a +-- stuff = 'withInps' $ \\(bar :< Ø) -\> do +-- 'withChoices' barIso bar $ \case +-- 'InL' i -> do +-- -- in this branch, bar was made with the A constructor +-- -- i is the Int inside it +-- 'InR' ('InL' b) -> do +-- -- in this branch, bar was made with the B constructor +-- -- b is the Bool inside it +-- 'InR' ('InR' ('InL' s)) -> do +-- -- in this branch, bar was made with the B constructor +-- -- s is the String inside it +-- @ +-- +-- Nicer than 'choicesVar' directly, because you don't have to give the +-- result a superfluous name before pattern matching on it. You can just +-- directly pattern match in the lambda, so there's a lot less syntactical +-- noise. +withChoices + :: forall s rs bs b a. (Every Num bs, Known Length bs) + => Iso' b (Sum I bs) + -> BVar s rs b + -> (Sum (BVar s rs) bs -> BP s rs a) + -> BP s rs a +withChoices i r f = do + c <- choicesVar i r + f c + +-- | A useful infix alias for 'choicesVar'. +-- +-- Building on the example from 'choicesVar': +-- +-- @ +-- data Bar = A Int | B Bool | C String +-- +-- barIso :: 'Iso'' Bar ('Sum' I '[Int, Bool, String]) +-- barIso = 'iso' (\\case A i -> 'InL' (I i) +-- B b -> 'InR' ('InL' (I b)) +-- C s -> 'InR' ('InR' ('InL' (I s)) +-- ) +-- (\\case 'InL' (I i) -> A i +-- 'InR' ('InL' (I b)) -> B b +-- 'InR' ('InR' ('InL' (I s))) -> C s +-- ) +-- +-- stuff :: 'BP' s '[Bar] a +-- stuff = 'withInps' $ \\(bar :< Ø) -\> do +-- c <- barIso '?<~' bar +-- case c of +-- 'InL' i -> do +-- -- in this branch, bar was made with the A constructor +-- -- i is the Int inside it +-- 'InR' ('InL' b) -> do +-- -- in this branch, bar was made with the B constructor +-- -- b is the Bool inside it +-- 'InR' ('InR' ('InL' s)) -> do +-- -- in this branch, bar was made with the B constructor +-- -- s is the String inside it +-- @ +infixr 1 ?<~ +(?<~) + :: (Every Num bs, Known Length bs) + => Iso' b (Sum I bs) + -> BVar s rs b + -> BP s rs (Sum (BVar s rs) bs) +(?<~) = choicesVar + +-- | A version of 'sopVar' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +sopVar' + :: forall s rs bss b. () + => Prod (Prod Summer) bss + -> Prod (Prod Unity) bss + -> Iso' b (Sum Tuple bss) + -> BVar s rs b + -> BP s rs (Sum (Prod (BVar s rs)) bss) +sopVar' sss uss i r = do + x <- BP $ resolveVar r + let xs :: Sum Tuple bss + xs = view i x + ifor1 ((sss `zipP` uss) `tagSum` xs) $ \ix ((ss :&: us) :&: (ys :: Tuple bs)) -> do + let bp :: BPNode s rs '[b] bs + bp = BPN { _bpnOut = map1 (const (FRInternal [])) ys + , _bpnRes = ys + , _bpnGradFunc = return . only_ + . review i . injectSum ix + . map1 (uncurryFan $ \u -> + maybe (I (getUnity u)) I + ) + . zipP us + , _bpnGradCache = Nothing + , _bpnSummer = ss + } + r' <- BP . liftBase $ newSTRef bp + registerVar (IRNode IZ r') r + return $ imap1 (\ix' _ -> BVNode ix' r') ys + +-- | A combination of 'partsVar' and 'choicesVar', that lets you split +-- a type into a sum of products. Using an 'Iso' (or compatible +-- 'Control.Lens.Iso.Iso' from the lens library), you can pull out a type +-- that is a sum of products into a sum of products of 'BVar's. +-- +-- Implements branching on the constructors of a value that a 'BVar' +-- contains, and also splitting out the different items inside each +-- constructor. +-- +-- @ +-- data Baz = A Int Bool +-- | B String Double +-- +-- +-- bazIso :: 'Iso'' Baz ('Sum' 'Tuple' '[ '[Int, Bool], '[String, Double] ]) +-- bazIso = 'iso' (\\case A i b -> 'InL' (I (i ::< b ::< Ø)) +-- B s d -> 'InR' ('InL' (I (s ::< d ::< Ø))) +-- ) +-- (\\case 'InL' (I (i ::< b ::< Ø)) -> A i b +-- 'InR' ('InL' (I (s ::< d ::< Ø))) -> B s d +-- ) +-- +-- 'sopVar' bazIso :: 'BVar' rs Baz -> 'BP' s rs ('Sum' ('Prod' ('BVar' s rs)) '[ '[Int, Bool], '[String, Double] ]) +-- +-- stuff :: 'BP' s '[Baz] a +-- stuff = 'withInps' $ \\(baz :< Ø) -\> do +-- c <- 'sopVar' barIso baz +-- case c of +-- 'InL' (i :< b :< Ø) -> do +-- -- in this branch, baz was made with the A constructor +-- -- i and b are the Int and Bool inside it +-- 'InR' ('InL' (s :< d :< Ø)) -> do +-- -- in this branch, baz was made with the B constructor +-- -- s and d are the String and Double inside it +-- @ +-- +-- Essentially exists to implement "pattern matching" on multiple +-- constructors and fields for the value inside a 'BVar'. +-- +-- Note that for a type like @Baz@, @bazIso@ can be generated automatically +-- with 'GHC.Generics.Generic' from "GHC.Generics" and +-- 'Generics.SOP.Generic' from "Generics.SOP" and /generics-sop/, with +-- 'gSOP'. See 'gSplits' for more information. +-- +-- See "Numeric.Backprop#sum" for a mini-tutorial on 'Sum'. +sopVar + :: forall s rs bss b. (Known Length bss, Every (Every Num ∧ Known Length) bss) + => Iso' b (Sum Tuple bss) + -> BVar s rs b + -> BP s rs (Sum (Prod (BVar s rs)) bss) +sopVar = sopVar' (withEvery @(Every Num ∧ Known Length) summers) + (withEvery @(Every Num ∧ Known Length) unities) + +-- | A version of 'gSplits' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +gSplits' + :: forall s rs b. SOP.Generic b + => Prod (Prod Summer) (SOP.Code b) + -> Prod (Prod Unity) (SOP.Code b) + -> BVar s rs b + -> BP s rs (Sum (Prod (BVar s rs)) (SOP.Code b)) +gSplits' sss uss = sopVar' sss uss gSOP + +-- | Using 'GHC.Generics.Generic' from "GHC.Generics" and +-- 'Generics.SOP.Generic' from "Generics.SOP", /split/ a 'BVar' containing +-- a sum of products (any simple ADT, essentialy) into a 'Sum' of each +-- constructor, each containing a tuple ('Prod') of 'BVar's pointing to +-- each value inside. +-- +-- Building on the example from 'sopVar': +-- +-- @ +-- import qualified Generics.SOP as SOP +-- +-- data Baz = A Int Bool +-- | B String Double +-- deriving Generic +-- +-- instance SOP.Generic Baz +-- +-- 'gSplits' :: 'BVar' rs Baz -> 'BP' s rs ('Sum' ('Prod' ('BVar' s rs)) '[ '[Int, Bool], '[String, Double] ]) +-- +-- stuff :: 'BP' s '[Baz] a +-- stuff = 'withInps' $ \\(baz :< Ø) -\> do +-- c <- gSplits baz +-- case c of +-- 'InL' (i :< b :< Ø) -> do +-- -- in this branch, baz was made with the A constructor +-- -- i and b are the Int and Bool inside it +-- 'InR' ('InL' (s :< d :< Ø)) -> do +-- -- in this branch, baz was made with the B constructor +-- -- s and d are the String and Double inside it +-- @ +-- +-- Because @Foo@ is a straight up sum-of-products type, 'gSplits' can use +-- "GHC.Generics" and take out the items inside. +-- +-- Note: +-- +-- @ +-- 'gSplit' = 'splitVars' 'gSOP' +-- @ +-- +-- See "Numeric.Backprop#sum" for a mini-tutorial on 'Sum'. +gSplits + :: forall s rs b. + ( SOP.Generic b + , Known Length (SOP.Code b) + , Every (Every Num ∧ Known Length) (SOP.Code b) + ) + => BVar s rs b + -> BP s rs (Sum (Prod (BVar s rs)) (SOP.Code b)) +gSplits = sopVar gSOP + + +resolveVar + :: (MonadReader (Tuple rs) m, MonadBase (ST s) m) + => BVar s rs a + -> m a +resolveVar = \case + BVNode ix r -> getI . index ix . _bpnRes <$> liftBase (readSTRef r) + BVInp ix -> getI . index ix <$> ask + BVConst x -> return x + BVOp rs o -> do + xs <- traverse1 (fmap I . resolveVar) rs + liftBase $ runOpM o xs + +registerVar + :: forall s rs a. () + => BPInpRef s rs a + -> BVar s rs a + -> BP s rs () +registerVar bpir = \case + BVNode ix' r' -> BP . liftBase . modifySTRef r' $ + over (bpnOut . indexP ix' . _FRInternal) (bpir :) + BVInp ix' -> BP $ modifying (bpsSources . indexP ix' . _FRInternal) (bpir :) + BVConst _ -> return () + -- This independently makes a new BPPipe for every usage site of the + -- BVOp, so it's a bit inefficient. + BVOp (rs :: Prod (BVar s rs) ds) (o :: OpM (ST s) ds a) -> do + xs :: Tuple ds <- traverse1 (fmap I . BP . resolveVar) rs + (res, gF) <- BP . liftBase $ runOpM' o xs + let bpp :: BPPipe s rs ds '[a] + bpp = BPP { _bppOut = only bpir + , _bppRes = only_ res + , _bppGradFunc = gF . Just . getI . head' + , _bppGradCache = Nothing + } + r' <- BP . liftBase $ newSTRef bpp + ifor1_ rs $ \ix' (bpr :: BVar s rs d) -> + registerVar (IRPipe ix' r') bpr + +-- | Apply an 'OpB' to a 'Prod' (tupling) of 'BVar's. +-- +-- If you had an @'OpB' s '[a, b, c] d@, this function will expect a 3-Prod +-- of a @'BVar' s rs a@, a @'BVar' s rs b@, and a @'BVar' s rs c@, and the +-- result will be a @'BVar' s rs d@: +-- +-- @ +-- myOp :: 'OpB' s '[a, b, c] d +-- x :: 'BVar' s rs a +-- y :: 'BVar' s rs b +-- z :: 'BVar' s rs c +-- +-- x :< y :< z :< Ø :: 'Prod' ('BVar' s rs) '[a, b, c] +-- 'opVar' myOp (x :< y :< z :< Ø) :: 'BP' s rs ('BVar' s rs d) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can provide any 'Op' +-- here, as well (like those created by 'op1', 'op2', 'constOp', 'op0' +-- etc.) +-- +-- 'opVar' has an infix alias, '~$', so the above example can also be +-- written as: +-- +-- @ +-- myOp '~$' (x :< y :< z :< Ø) :: 'BP' s rs ('BVar' s rs d) +-- @ +-- +-- to let you pretend that you're applying the 'myOp' function to three +-- inputs. +-- +-- Also note the relation between 'opVar' and 'liftB' and 'bindVar': +-- +-- @ +-- 'opVar' o xs = 'bindVar' ('liftB' o xs) +-- @ +-- +-- 'opVar' can be thought of as a "binding" version of 'liftB'. +opVar + :: Num a + => OpB s as a + -> Prod (BVar s rs) as + -> BP s rs (BVar s rs a) +opVar = opVar' known + +-- | Infix synonym for 'opVar', which lets you pretend that you're applying +-- 'OpB's as if they were functions: +-- +-- @ +-- myOp :: 'OpB' s '[a, b, c] d +-- x :: 'BVar' s rs a +-- y :: 'BVar' s rs b +-- z :: 'BVar' s rs c +-- +-- x :< y :< z :< Ø :: 'Prod' ('BVar' s rs) '[a, b, c] +-- myOp '~$' (x :< y :< z :< Ø) :: 'BP' s rs ('BVar' s rs d) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in any 'Op' +-- here, as well (like those created by 'op1', 'op2', 'constOp', 'op0' +-- etc.) +-- +-- '~$' can also be thought of as a "binding" version of '.$': +-- +-- @ +-- o '~$' xs = 'bindVar' (o '.$' xs) +-- @ +-- +infixr 5 ~$ +(~$) + :: Num a + => OpB s as a + -> Prod (BVar s rs) as + -> BP s rs (BVar s rs a) +(~$) = opVar + +-- | Lets you treat a @'BPOp' s as b@ as an @'Op' as b@, and "apply" +-- arguments to it just like you would with an 'Op' and '~$' / 'opVar'. +-- +-- Basically a convenient wrapper over 'bpOp' and '~$': +-- +-- @ +-- o '-$' xs = bpOp o '~$' xs +-- @ +-- +-- So for a @'BPOp' s as b@, you can "plug in" 'BVar's to @as@, and get +-- a @b@ as a result. +-- +-- Useful for running a @'BPOp' s as b@ that you got from a different function, and +-- "plugging in" its @as@ inputs with 'BVar's from your current +-- environment. +infixr 5 -$ +(-$) + :: (Every Num as, Known Length as, Num a) + => BPOp s as a + -> Prod (BVar s rs) as + -> BPOp s rs a +o -$ xs = bpOp o ~$ xs + +-- | Create a 'BVar' that represents just a specific value, that doesn't +-- depend on any other 'BVar's. +constVar :: a -> BVar s rs a +constVar = BVConst + +-- | A version of 'opVar1' taking an explicit 'Summer', so can be used on +-- values of types that aren't instances of 'Num'. +opVar1' + :: Summer b + -> OpB s '[a] b + -> BVar s rs a + -> BP s rs (BVar s rs b) +opVar1' s o = opVar' s o . only + +-- | Convenient wrapper over 'opVar' that takes an 'OpB' with one argument +-- and a single 'BVar' argument. Lets you not have to type out the entire +-- 'Prod'. +-- +-- @ +-- 'opVar1' o x = 'opVar' o (x ':<' 'Ø') +-- +-- myOp :: 'Op' '[a] b +-- x :: 'BVar' s rs a +-- +-- 'opVar1' myOp x :: 'BP' s rs ('BVar' s rs b) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op1') as well. +opVar1 + :: Num b + => OpB s '[a] b + -> BVar s rs a + -> BP s rs (BVar s rs b) +opVar1 = opVar1' known + +-- | A version of 'opVar2' taking an explicit 'Summer', so can be used on +-- values of types that aren't instances of 'Num'. +opVar2' + :: Summer c + -> OpB s '[a,b] c + -> BVar s rs a + -> BVar s rs b + -> BP s rs (BVar s rs c) +opVar2' s o rx ry = opVar' s o (rx :< ry :< Ø) + +-- | Convenient wrapper over 'opVar' that takes an 'OpB' with two arguments +-- and two 'BVar' arguments. Lets you not have to type out the entire +-- 'Prod'. +-- +-- @ +-- 'opVar2' o x y = 'opVar' o (x ':<' y ':<' 'Ø') +-- +-- myOp :: 'Op' '[a, b] c +-- x :: 'BVar' s rs a +-- y :: 'BVar' s rs b +-- +-- 'opVar2' myOp x y :: 'BP' s rs ('BVar' s rs c) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op2') as well. +opVar2 + :: Num c + => OpB s '[a,b] c + -> BVar s rs a + -> BVar s rs b + -> BP s rs (BVar s rs c) +opVar2 = opVar2' known + +-- | A version of 'opVar3' taking an explicit 'Summer', so can be used on +-- values of types that aren't instances of 'Num'. +opVar3' + :: Summer d + -> OpB s '[a,b,c] d + -> BVar s rs a + -> BVar s rs b + -> BVar s rs c + -> BP s rs (BVar s rs d) +opVar3' s o rx ry rz = opVar' s o (rx :< ry :< rz :< Ø) + +-- | Convenient wrapper over 'opVar' that takes an 'OpB' with three arguments +-- and three 'BVar' arguments. Lets you not have to type out the entire +-- 'Prod'. +-- +-- @ +-- 'opVar3' o x y z = 'opVar' o (x ':<' y ':<' z ':<' 'Ø') +-- +-- myOp :: 'Op' '[a, b, c] d +-- x :: 'BVar' s rs a +-- y :: 'BVar' s rs b +-- z :: 'BVar' s rs c +-- +-- 'opVar3' myOp x y z :: 'BP' s rs ('BVar' s rs d) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op3') as well. +opVar3 + :: Num d + => OpB s '[a,b,c] d + -> BVar s rs a + -> BVar s rs b + -> BVar s rs c + -> BP s rs (BVar s rs d) +opVar3 = opVar3' known + +-- | A version of 'bindVar' that requires an explicit 'Summer', so that you +-- can use it on values whose types aren't instances of 'Num'. +bindVar' + :: Summer a + -> BVar s rs a + -> BP s rs (BVar s rs a) +bindVar' s r = case r of + BVNode _ _ -> return r + BVInp _ -> return r + BVConst _ -> return r + BVOp rs o -> opVar' s o rs + +-- | Concretizes a delayed 'BVar'. If you build up a 'BVar' using numeric +-- functions like '+' or '*' or using 'liftB', it'll defer the evaluation, +-- and all of its usage sites will create a separate graph node. +-- +-- Use 'bindVar' if you ever intend to use a 'BVar' in more than one +-- location. +-- +-- @ +-- -- bad +-- errSquared :: Num a => 'BP' s '[a, a] a +-- errSquared = 'withInp' $ \\(r :< t :< Ø) -\> do +-- let err = r - t +-- 'return' (err * err) -- err is used twice! +-- +-- -- good +-- errSquared :: Num a => 'BP' s '[a, a] a +-- errSquared = 'withInps' $ \\(r :< t :< Ø) -\> do +-- let err = r - t +-- e <- 'bindVar' err -- force e, so that it's safe to use twice! +-- 'return' (e * e) +-- +-- -- better +-- errSquared :: Num a => 'BP' s '[a, a] a +-- errSquared = 'withInps' $ \\(r :< t :< Ø) -\> do +-- let err = r - t +-- e <- 'bindVar' err +-- 'bindVar' (e * e) -- result is forced so user doesn't have to worry +-- @ +-- +-- Note the relation to 'opVar' / '~$' / 'liftB' / '.$': +-- +-- @ +-- 'opVar' o xs = 'bindVar' ('liftB' o xs) +-- o '~$' xs = 'bindVar' (o '.$' xs) +-- 'op2' (*) '~$' (x :< y :< Ø) = 'bindVar' (x * y) +-- @ +-- +-- So you can avoid 'bindVar' altogether if you use the explicitly binding +-- '~$' and 'opVar' etc. +-- +-- Note that 'bindVar' on 'BVar's that are already forced is a no-op. +bindVar + :: Num a + => BVar s rs a + -> BP s rs (BVar s rs a) +bindVar = bindVar' known + + + +backwardPass + :: forall s rs a. () + => BPInpRef s rs a + -> ST s a +backwardPass = \case + IRNode ix r' -> getI . index ix <$> pullNode r' + IRPipe ix r' -> getI . index ix <$> pullPipe r' + IRConst g -> return g + where + pullNode + :: forall as bs. () + => STRef s (BPNode s rs as bs) + -> ST s (Tuple as) + pullNode r = caching bpnGradCache r $ \BPN{..} -> do + totdervs <- for1 (_bpnSummer `zipP` _bpnOut) $ \case + s :&: FRInternal rs -> Just . runSummer s + <$> traverse backwardPass rs + _ :&: FRTerminal g -> return g + g <- _bpnGradFunc totdervs + return g + pullPipe + :: forall as bs. () + => STRef s (BPPipe s rs as bs) + -> ST s (Tuple as) + pullPipe r = caching bppGradCache r $ \BPP{..} -> + _bppGradFunc =<< traverse1 (fmap I . backwardPass) _bppOut + +-- | A version of 'backprop' taking explicit 'Summer's and 'Unity's, so it +-- can be run with types that aren't instances of 'Num'. +backprop' + :: Prod Summer rs + -> Prod Unity rs + -> (forall s. BPOp s rs a) + -> Tuple rs + -> (a, Tuple rs) +backprop' ss us bp env = runST $ do + (res, gFunc) <- backpropWith ss us bp env + grad <- gFunc Nothing + return (res, grad) + +-- | Perform back-propagation on the given 'BPOp'. Returns the result of +-- the operation it represents, as well as the gradient of the result with +-- respect to its inputs. See module header for "Numeric.Backprop" and +-- package documentation for examples and usages. +backprop + :: forall rs a. Every Num rs + => (forall s. BPOp s rs a) + -> Tuple rs + -> (a, Tuple rs) +backprop bp xs = backprop' (summers' l) (unities' l) bp xs + where + l :: Length rs + l = prodLength xs + +-- | 'bpOp', but taking explicit 'Summer's and 'Unity's, for the situation +-- where the @rs@ are not instance of 'Num'. +bpOp' + :: Prod Summer rs + -> Prod Unity rs + -> BPOp s rs a + -> OpB s rs a +bpOp' ss us bp = OpM $ backpropWith ss us bp + +-- | Turn a 'BPOp' into an 'OpB'. Basically converts a 'BP' taking an @rs@ +-- and producing an @a@ into an 'Op' taking an @rs@ and returning an @a@, +-- with all of the powers and utility of an 'Op', including all of its +-- gradient-finding glory. +-- +-- Really just reveals the fact that any @'BPOp' s rs a@ is itself an 'Op', +-- an @'OpB' s rs a@, which makes it a differentiable function. +-- +-- Handy because an 'OpB' can be used with almost all of +-- the 'Op'-related functions in this moduel, including 'opVar', '~$', etc. +bpOp + :: (Every Num rs, Known Length rs) + => BPOp s rs a + -> OpB s rs a +bpOp = bpOp' summers unities + +-- | Simply run the 'BPOp' on an input tuple, getting the result without +-- bothering with the gradient or with back-propagation. +evalBPOp + :: (forall s. BPOp s rs a) -- ^ 'BPOp' to run + -> Tuple rs -- ^ input + -> a -- ^ output +evalBPOp bp env = runST $ do + r <- evalStateT (runReaderT (bpST bp) env) + (BPS (map1 (\_ -> FRInternal []) env)) + runReaderT (resolveVar r) env + +-- | A version of 'gradBPOp' taking explicit 'Summer's and 'Unity's, so it +-- can be run with types that aren't instances of 'Num'. +gradBPOp' + :: Prod Summer rs + -> Prod Unity rs + -> (forall s. BPOp s rs a) -- ^ 'BPOp' to differentiate' + -> Tuple rs -- ^ input + -> Tuple rs -- ^ gradient +gradBPOp' ss us bp = snd . backprop' ss us bp + +-- | Run the 'BPOp' on an input tuple and return the gradient of the result +-- with respect to the input tuple. +gradBPOp + :: Every Num rs + => (forall s. BPOp s rs a) -- ^ 'BPOp' to differentiate + -> Tuple rs -- ^ input + -> Tuple rs -- ^ gradient +gradBPOp bp = snd . backprop bp + + +closeOff + :: (MonadReader (Tuple rs) m, MonadState (BPState s rs) m, MonadBase (ST s) m) + => Bool + -> Maybe a + -> BVar s rs a + -> m () +closeOff isTerminal gOut = \case + BVNode ix sr -> liftBase $ modifySTRef sr (over (bpnOut . indexP ix) (<> fr)) + BVInp ix' -> modifying (bpsSources . indexP ix') (<> fr) + BVConst _ -> return () + BVOp rs o -> do + xs <- traverse1 (fmap I . resolveVar) rs + gs <- liftBase $ gradOpWithM' o xs gOut + for1_ (gs `zipP` rs) $ \(I g :&: r) -> + closeOff False (Just g) r + where + fr | isTerminal = FRTerminal gOut + | otherwise = FRInternal (IRConst <$> maybeToList gOut) + +backpropWith + :: Prod Summer rs + -> Prod Unity rs + -> BPOp s rs a + -> Tuple rs + -> ST s (a, Maybe a -> ST s (Tuple rs)) +backpropWith ss us bp env = do + (r, bps0) <- runStateT (runReaderT (bpST bp) env) + (BPS (map1 (\_ -> FRInternal []) env)) + res <- runReaderT (resolveVar r) env + let gradFunc gradOut = do + BPS{..} <- execStateT (runReaderT (closeOff True gradOut r) env) bps0 + for1 (ss `zipP` us `zipP` _bpsSources) $ \((s :&: u) :&: rs) -> do + I <$> case rs of + FRInternal rs' -> runSummer s <$> traverse backwardPass rs' + FRTerminal g -> return $ fromMaybe (getUnity u) g + return (res, gradFunc) + +-- | A version of 'implicitly' taking explicit 'Length' and an explicit +-- 'Summer', indicating the number of inputs required and their types, and +-- also allowing it to work on types that aren't instances of 'Num'. +-- +-- Requiring an explicit 'Length' is mostly useful for rare "extremely +-- polymorphic" situations, where GHC can't infer the type and length of +-- the list of inputs. If you ever actually explicitly write down @rs@ as +-- a list of types, you should be able to just use 'implicitly'. +implicitly' + :: Length rs + -> Summer a + -> BPOpI s rs a + -> BPOp s rs a +implicitly' l s f = withInps' l (bindVar' s . f) + +-- | Convert a 'BPOpI' into a 'BPOp'. That is, convert a function on +-- a bundle of 'BVar's (generating an implicit graph) into a fully fledged +-- 'BPOp' that you can run 'backprop' on. See 'BPOpI' for more +-- information. +-- +-- If you are going to write exclusively using implicit 'BVar' operations, +-- it might be more convenient to use "Numeric.Backprop.Implicit" instead, +-- which is geared around that use case. +implicitly + :: (Known Length rs, Num a) + => BPOpI s rs a + -> BPOp s rs a +implicitly = implicitly' known known + +-- | Create a 'BVar' given an index into the input environment. For an +-- example, +-- +-- @ +-- 'inpVar' 'IZ' +-- @ +-- +-- would refer to the /first/ input variable (the 'Int' in a +-- @'BP' s '[Int, Bool]@), and +-- +-- @ +-- 'inpVar' ('IS' 'IZ') +-- @ +-- +-- Would refer to the /second/ input variable (the 'Bool' in a +-- @'BP' s '[Int, Bool]@) +-- +-- Typically, there shouldn't be any reason to use 'inpVar' directly. It's +-- cleaner to get all of your input 'BVar's together using 'withInps' or +-- 'inpVars'. +inpVar + :: Index rs a + -> BVar s rs a +inpVar = BVInp + +-- | Get a 'Prod' (tupling) of 'BVar's for all of the input environment +-- (@rs@) of the @'BP' s rs@ +-- +-- For example, if your 'BP' has an 'Int' and 'Double' in its input +-- environment (a @'BP' s '[Int, Double]@), this would return a 'BVar' +-- pointing to the 'Int' and a 'BVar' pointing to the 'Double'. +-- +-- @ +-- case ('inpVars' :: 'Prod' ('BVar' s '[Int, Double]) '[Int, Double]) of +-- x :\< y :\< Ø -\> do +-- -- the first item, x, is a var to the input 'Int' +-- -- x :: 'BVar' s '[Int, Double] Int +-- -- the second item, y, is a var to the input 'Double' +-- -- y :: 'BVar' s '[Int, Double] Double +-- @ +inpVars + :: Known Length rs + => Prod (BVar s rs) rs +inpVars = inpVars' known + +-- | A version of 'inpVars' taking explicit 'Length', indicating the +-- number of inputs required and their types. +-- +-- Mostly useful for rare "extremely polymorphic" situations, where GHC +-- can't infer the type and length of the list of inputs. If you ever +-- actually explicitly write down @rs@ as a list of types, you should be +-- able to just use 'inpVars'. +inpVars' + :: Length rs + -> Prod (BVar s rs) rs +inpVars' = map1 inpVar . indices' + +-- | A version of 'withInps' taking explicit 'Length', indicating the +-- number of inputs required and their types. +-- +-- Mostly useful for rare "extremely polymorphic" situations, where GHC +-- can't infer the type and length of the list of inputs. If you ever +-- actually explicitly write down @rs@ as a list of types, you should be +-- able to just use 'withInps'. +withInps' + :: Length rs + -> (Prod (BVar s rs) rs -> BP s rs a) + -> BP s rs a +withInps' l f = f (inpVars' l) + +-- | Runs a continuation on a 'Prod' of all of the input 'BVar's. +-- +-- Handy for bringing the environment into scope and doing stuff with it: +-- +-- @ +-- foo :: 'BPOp' '[Double, Int] a +-- foo = 'withInps' $ \\(x :< y :< Ø) -\> do +-- -- do stuff with inputs +-- @ +-- +-- Looks kinda like @foo (x :< y :< Ø) = -- ...@, don't it? +-- +-- Note that the above is the same as +-- +-- @ +-- foo :: 'BPOp' '[Double, Int] a +-- foo = do +-- case 'inpVars' of +-- x :< y :< Ø -> do +-- -- do stuff with inputs +-- @ +-- +-- But just a little nicer! +withInps + :: Known Length rs + => (Prod (BVar s rs) rs -> BP s rs a) + -> BP s rs a +withInps = withInps' known + +-- | Apply 'OpB' over a 'Prod' of 'BVar's, as inputs. Provides +-- "implicit-graph" back-propagation, with deferred evaluation. +-- +-- If you had an @'OpB' s '[a, b, c] d@, this function will expect a 3-Prod +-- of a @'BVar' s rs a@, a @'BVar' s rs b@, and a @'BVar' s rs c@, and the +-- result will be a @'BVar' s rs d@: +-- +-- @ +-- myOp :: 'OpB' s '[a, b, c] d +-- x :: 'BVar' s rs a +-- y :: 'BVar' s rs b +-- z :: 'BVar' s rs c +-- +-- x :< y :< z :< Ø :: 'Prod' ('BVar' s rs) '[a, b, c] +-- 'liftB' myOp (x :< y :< z :< Ø) :: 'BVar' s rs d +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can provide any 'Op' +-- here, as well (like those created by 'op1', 'op2', 'constOp', 'op0' +-- etc.) +-- +-- 'liftB' has an infix alias, '.$', so the above example can also be +-- written as: +-- +-- @ +-- myOp '.$' (x :< y :< z :< Ø) :: 'BVar' s rs d +-- @ +-- +-- to let you pretend that you're applying the 'myOp' function to three +-- inputs. +-- +-- The result is a new /deferred/ 'BVar'. This should be fine in most +-- cases, unless you use the result in more than one location. This will +-- cause evaluation to be duplicated and multiple redundant graph nodes to +-- be created. If you need to use it in two locations, you should use +-- 'opVar' instead of 'liftB', or use 'bindVar': +-- +-- @ +-- 'opVar' o xs = 'bindVar' ('liftB' o xs) +-- @ +-- +-- 'liftB' can be thought of as a "deferred evaluation" version of 'opVar'. +liftB + :: OpB s as a + -> Prod (BVar s rs) as + -> BVar s rs a +liftB = flip BVOp + + +-- | Infix synonym for 'liftB', which lets you pretend that you're applying +-- 'OpB's as if they were functions: +-- +-- @ +-- myOp :: 'OpB' s '[a, b, c] d +-- x :: 'BVar' s rs a +-- y :: 'BVar' s rs b +-- z :: 'BVar' s rs c +-- +-- x :< y :< z :< Ø :: 'Prod' ('BVar' s rs) '[a, b, c] +-- myOp '.$' (x :< y :< z :< Ø) :: 'BVar' s rs d +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in any 'Op' +-- here, as well (like those created by 'op1', 'op2', 'constOp', 'op0' +-- etc.) +-- +-- See the documentation for 'liftB' for all the caveats of this usage. +-- +-- '.$' can also be thought of as a "deferred evaluation" version of '~$': +-- +-- @ +-- o '~$' xs = 'bindVar' (o '.$' xs) +-- @ +-- +infixr 5 .$ +(.$) + :: OpB s as a + -> Prod (BVar s rs) as + -> BVar s rs a +(.$) = liftB + + +-- | Convenient wrapper over 'liftB' that takes an 'OpB' with one argument +-- and a single 'BVar' argument. Lets you not have to type out the entire +-- 'Prod'. +-- +-- @ +-- 'liftB1' o x = 'liftB' o (x ':<' 'Ø') +-- +-- myOp :: 'Op' '[a] b +-- x :: 'BVar' s rs a +-- +-- 'liftB1' myOp x :: 'BVar' s rs b +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op1') as well. +-- +-- See the documentation for 'liftB' for caveats and potential problematic +-- situations with this. +liftB1 + :: OpB s '[a] b + -> BVar s rs a + -> BVar s rs b +liftB1 o = liftB o . only + +-- | Convenient wrapper over 'liftB' that takes an 'OpB' with two arguments +-- and two 'BVar' arguments. Lets you not have to type out the entire +-- 'Prod'. +-- +-- @ +-- 'liftB2' o x y = 'liftB' o (x ':<' y ':<' 'Ø') +-- +-- myOp :: 'Op' '[a, b] c +-- x :: 'BVar' s rs a +-- y :: 'BVar' s rs b +-- +-- 'liftB2' myOp x y :: 'BVar' s rs c +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op2') as well. +-- +-- See the documentation for 'liftB' for caveats and potential problematic +-- situations with this. +liftB2 + :: OpB s '[a,b] c + -> BVar s rs a + -> BVar s rs b + -> BVar s rs c +liftB2 o x y = liftB o (x :< y :< Ø) + +-- | Convenient wrapper over 'liftB' that takes an 'OpB' with three arguments +-- and three 'BVar' arguments. Lets you not have to type out the entire +-- 'Prod'. +-- +-- @ +-- 'liftB3' o x y z = 'liftB' o (x ':<' y ':<' z ':<' 'Ø') +-- +-- myOp :: 'Op' '[a, b, c] d +-- x :: 'BVar' s rs a +-- y :: 'BVar' s rs b +-- z :: 'BVar' s rs c +-- +-- 'liftB3' myOp x y z :: 'BVar' s rs d +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op3') as well. +-- +-- See the documentation for 'liftB' for caveats and potential problematic +-- situations with this. +liftB3 + :: OpB s '[a,b,c] d + -> BVar s rs a + -> BVar s rs b + -> BVar s rs c + -> BVar s rs d +liftB3 o x y z = liftB o (x :< y :< z :< Ø) + + + + + + + + + + + +-- | Apply a function to the contents of an STRef, and cache the results +-- using the given lens. If already calculated, simply returned the cached +-- result. +caching + :: Lens' a (Maybe b) + -> STRef s a + -> (a -> ST s b) + -> ST s b +caching l r f = do + x <- readSTRef r + let y = view l x + case y of + Just z -> + return z + Nothing -> do + z <- f x + modifySTRef r (set l (Just z)) + return z + diff --git a/src/Numeric/Backprop/Implicit.hs b/src/Numeric/Backprop/Implicit.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop/Implicit.hs @@ -0,0 +1,382 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +-- | +-- Module : Numeric.Backprop.Implicit +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Offers full functionality for implicit-graph back-propagation. The +-- intended usage is to write a 'BPOp', which is a normal Haskell +-- function from 'BVar's to a result 'BVar'. These 'BVar's can be +-- manipulated using their 'Num' \/ 'Fractional' \/ 'Floating' instances. +-- +-- The library can then perform back-propagation on the function (using +-- 'backprop' or 'grad') by using an implicitly built graph. +-- +-- This should actually be powerful enough for most use cases, but falls +-- short for a couple of situations: +-- +-- 1. If the result of a function on 'BVar's is used twice +-- (like @z@ in @let z = x * y in z + z@), this will allocate a new +-- redundant graph node for every usage site of @z@. You can explicitly +-- /force/ @z@, but only using an explicit graph description using +-- "Numeric.Backprop". +-- +-- 2. This can't handle sum types, like "Numeric.Backprop" can. You can +-- never pattern match on the constructors of a value inside a 'BVar'. I'm +-- not sure if this is a fundamental limitation (I suspect it might be) or +-- if I just can't figure out how to implement it. Suggestions welcome! +-- +-- As a comparison, this module offers functionality and an API very +-- similar to "Numeric.AD.Mode.Reverse" from the /ad/ library, except for +-- the fact that it can handle /heterogeneous/ values. +-- + + +module Numeric.Backprop.Implicit ( + -- * Types + -- ** Backprop types + BPOp, BVar, Op, OpB + -- ** Tuple types + -- | See "Numeric.Backprop#prod" for a mini-tutorial on 'Prod' and + -- 'Tuple' + , Prod(..), Tuple, I(..) + -- * back-propagation + , backprop, grad, eval + , backprop', grad' + -- * Var manipulation + , BP.constVar, BP.liftB, (BP..$), BP.liftB1, BP.liftB2, BP.liftB3 + -- ** As Parts + , partsVar, withParts + , splitVars, gSplit, gTuple + , partsVar', withParts' + , splitVars', gSplit' + -- * Op + , BP.op1, BP.op2, BP.op3, BP.opN + , BP.op1', BP.op2', BP.op3' + -- * Utility + , pattern (:>), only, head' + , pattern (::<), only_ + , Summer(..), Unity(..) + , summers, unities + , summers', unities' + ) where + +import Data.Type.Combinator +import Data.Type.Index +import Data.Type.Length +import Data.Type.Product +import Data.Type.Util +import Lens.Micro hiding (ix) +import Lens.Micro.Extras +import Numeric.Backprop.Internal +import Numeric.Backprop.Iso +import Numeric.Backprop.Op +import Type.Class.Higher +import Type.Class.Known +import qualified Generics.SOP as SOP +import qualified Numeric.Backprop as BP + +-- | An operation on 'BVar's that can be backpropagated. A value of type: +-- +-- @ +-- 'BPOp' rs a +-- @ +-- +-- takes a bunch of 'BVar's containg @rs@ and uses them to (purely) produce +-- a 'BVar' containing an @a@. +-- +-- @ +-- foo :: 'BPOp' '[ Double, Double ] Double +-- foo (x ':<' y ':<' 'Ø') = x + sqrt y +-- @ +-- +-- 'BPOp' here is related to 'Numeric.Backprop.BPOpI' from the normal +-- explicit-graph backprop module "Numeric.Backprop". +type BPOp rs a = forall s. Prod (BVar s rs) rs -> BVar s rs a + +-- | A version of 'backprop' taking explicit 'Summer's and 'Unity's, so it +-- can be run with types that aren't instances of 'Num'. +backprop' + :: Prod Summer rs + -> Prod Unity rs + -> BPOp rs a + -> Tuple rs + -> (a, Tuple rs) +backprop' ss us f = BP.backprop' ss us $ BP.withInps' (prodLength ss) (return . f) + +-- | Run back-propagation on a 'BPOp' function, getting both the result and +-- the gradient of the result with respect to the inputs. +-- +-- @ +-- foo :: 'BPOp' '[Double, Double] Double +-- foo (x :< y :< Ø) = +-- let z = x * sqrt y +-- in z + x ** y +-- @ +-- +-- >>> 'backprop' foo (2 ::< 3 ::< Ø) +-- (11.46, 13.73 ::< 6.12 ::< Ø) +backprop + :: (Known Length rs, Every Num rs, Num a) + => BPOp rs a + -> Tuple rs + -> (a, Tuple rs) +backprop f = BP.backprop $ BP.implicitly f + +-- | A version of 'grad' taking explicit 'Summer's and 'Unity's, so it +-- can be run with types that aren't instances of 'Num'. +grad' + :: Prod Summer rs + -> Prod Unity rs + -> BPOp rs a + -> Tuple rs + -> Tuple rs +grad' ss us f = snd . backprop' ss us f + +-- | Run the 'BPOp' on an input tuple and return the gradient of the result +-- with respect to the input tuple. +-- +-- @ +-- foo :: 'BPOp' '[Double, Double] Double +-- foo (x :< y :< Ø) = +-- let z = x * sqrt y +-- in z + x ** y +-- @ +-- +-- >>> grad foo (2 ::< 3 ::< Ø) +-- 13.73 ::< 6.12 ::< Ø +grad + :: (Known Length rs, Every Num rs, Num a) + => BPOp rs a + -> Tuple rs + -> Tuple rs +grad f = snd . backprop f + +-- | Simply run the 'BPOp' on an input tuple, getting the result without +-- bothering with the gradient or with back-propagation. +-- +-- @ +-- foo :: 'BPOp' '[Double, Double] Double +-- foo (x :< y :< Ø) = +-- let z = x * sqrt y +-- in z + x ** y +-- @ +-- +-- >>> eval foo (2 ::< 3 ::< Ø) +-- 11.46 +eval + :: (Known Length rs, Every Num rs, Num a) + => BPOp rs a + -> Tuple rs + -> a +eval f = BP.evalBPOp $ BP.implicitly f + +-- | A version of 'partsVar' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +partsVar' + :: forall s rs bs a. () + => Prod Summer bs + -> Prod Unity bs + -> Iso' a (Tuple bs) + -> BVar s rs a + -> Prod (BVar s rs) bs +partsVar' ss us i r = imap1 (\ix u -> BP.liftB1 (BP.op1' (f ix u)) r) us + where + f :: Index bs b + -> Unity b + -> a + -> (b, Maybe b -> a) + f ix u x = ( getI . index ix . view i $ x + , review i + . flip (set (indexP ix)) zeroes + . maybe (I (getUnity u)) I + ) + zeroes :: Tuple bs + zeroes = map1 (\s -> I $ runSummer s []) ss + +-- | Use an 'Iso' (or compatible 'Control.Lens.Iso.Iso' from the lens +-- library) to "pull out" the parts of a data type and work with each part +-- as a 'BVar'. +-- +-- If there is an isomorphism between a @b@ and a @'Tuple' as@ (that is, if +-- an @a@ is just a container for a bunch of @as@), then it lets you break +-- out the @as@ inside and work with those. +-- +-- @ +-- data Foo = F Int Bool +-- +-- fooIso :: 'Iso'' Foo (Tuple '[Int, Bool]) +-- fooIso = 'iso' (\\(F i b) -\> i ::\< b ::\< Ø) +-- (\\(i ::\< b ::\< Ø) -\> F i b ) +-- +-- 'partsVar' fooIso :: 'BVar' rs Foo -> 'Prod' ('BVar' s rs) '[Int, Bool] +-- +-- stuff :: 'BPOp' s '[Foo] a +-- stuff (foo :< Ø) = +-- case 'partsVar' fooIso foo of +-- i :< b :< Ø -> +-- -- now, i is a 'BVar' pointing to the 'Int' inside foo +-- -- and b is a 'BVar' pointing to the 'Bool' inside foo +-- -- you can do stuff with the i and b here +-- @ +-- +-- You can use this to pass in product types as the environment to a 'BP', +-- and then break out the type into its constituent products. +-- +-- Note that for a type like @Foo@, @fooIso@ can be generated automatically +-- with 'GHC.Generics.Generic' from "GHC.Generics" and +-- 'Generics.SOP.Generic' from "Generics.SOP" and /generics-sop/, using the +-- 'gTuple' iso. See 'gSplit' for more information. +-- +-- Also, if you are literally passing a tuple (like +-- @'BP' s '[Tuple '[Int, Bool]@) then you can give in the identity +-- isomorphism ('id') or use 'splitVars'. +-- +-- At the moment, this implicit 'partsVar' is less efficient than the +-- explicit 'Numeric.Backprop.partsVar', but this might change in the +-- future. +partsVar + :: forall s rs bs a. (Known Length bs, Every Num bs) + => Iso' a (Tuple bs) + -> BVar s rs a + -> Prod (BVar s rs) bs +partsVar = partsVar' summers unities + +-- | A version of 'withParts' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +withParts' + :: forall s rs bs a r. () + => Prod Summer bs + -> Prod Unity bs + -> Iso' a (Tuple bs) + -> BVar s rs a + -> (Prod (BVar s rs) bs -> r) + -> r +withParts' ss us i r f = f (partsVar' ss us i r) + +-- | A continuation-based version of 'partsVar'. Instead of binding the +-- parts and using it in the rest of the block, provide a continuation to +-- handle do stuff with the parts inside. +-- +-- Building on the example from 'partsVar': +-- +-- @ +-- data Foo = F Int Bool +-- +-- fooIso :: 'Iso'' Foo (Tuple '[Int, Bool]) +-- fooIso = 'iso' (\\(F i b) -\> i ::\< b ::\< Ø) +-- (\\(i ::\< b ::\< Ø) -\> F i b ) +-- +-- stuff :: 'BPOp' s '[Foo] a +-- stuff (foo :< Ø) = 'withParts' fooIso foo $ \\case +-- i :\< b :< Ø -\> +-- -- now, i is a 'BVar' pointing to the 'Int' inside foo +-- -- and b is a 'BVar' pointing to the 'Bool' inside foo +-- -- you can do stuff with the i and b here +-- @ +-- +-- Mostly just a stylistic alternative to 'partsVar'. +withParts + :: forall s rs bs a r. (Known Length bs, Every Num bs) + => Iso' a (Tuple bs) + -> BVar s rs a + -> (Prod (BVar s rs) bs -> r) + -> r +withParts i r f = f (partsVar i r) + +-- | A version of 'splitVars' taking explicit 'Summer's and 'Unity's, so it +-- can be run with types that aren't instances of 'Num'. +splitVars' + :: forall s rs as. () + => Prod Summer as + -> Prod Unity as + -> BVar s rs (Tuple as) + -> Prod (BVar s rs) as +splitVars' ss us = partsVar' ss us id + +-- | Split out a 'BVar' of a tuple into a tuple ('Prod') of 'BVar's. +-- +-- @ +-- -- the environment is a single Int-Bool tuple, tup +-- stuff :: 'BPOp' s '[ Tuple '[Int, Bool] ] a +-- stuff (tup :< Ø) = +-- case 'splitVar' tup of +-- i :< b :< Ø <- 'splitVars' tup +-- -- now, i is a 'BVar' pointing to the 'Int' inside tup +-- -- and b is a 'BVar' pointing to the 'Bool' inside tup +-- -- you can do stuff with the i and b here +-- @ +-- +-- Note that +-- +-- @ +-- 'splitVars' = 'partsVar' 'id' +-- @ +splitVars + :: forall s rs as. (Known Length as, Every Num as) + => BVar s rs (Tuple as) + -> Prod (BVar s rs) as +splitVars = partsVar id + +-- | A version of 'gSplit' taking explicit 'Summer's and 'Unity's, so it +-- can be run with internal types that aren't instances of 'Num'. +gSplit' + :: forall s rs as a. (SOP.Generic a, SOP.Code a ~ '[as]) + => Prod Summer as + -> Prod Unity as + -> BVar s rs a + -> Prod (BVar s rs) as +gSplit' ss us = partsVar' ss us gTuple + +-- | Using 'GHC.Generics.Generic' from "GHC.Generics" and +-- 'Generics.SOP.Generic' from "Generics.SOP", /split/ a 'BVar' containing +-- a product type into a tuple ('Prod') of 'BVar's pointing to each value +-- inside. +-- +-- Building on the example from 'partsVar': +-- +-- @ +-- import qualified Generics.SOP as SOP +-- +-- data Foo = F Int Bool +-- deriving Generic +-- +-- instance SOP.Generic Foo +-- +-- 'gSplit' :: 'BVar' rs Foo -> 'Prod' ('BVar' s rs) '[Int, Bool] +-- +-- stuff :: 'BPOp' s '[Foo] a +-- stuff (foo :< Ø) = +-- case 'gSplit' foo of +-- i :< b :< Ø -> +-- -- now, i is a 'BVar' pointing to the 'Int' inside foo +-- -- and b is a 'BVar' pointing to the 'Bool' inside foo +-- -- you can do stuff with the i and b here +-- @ +-- +-- Because @Foo@ is a straight up product type, 'gSplit' can use +-- "GHC.Generics" and take out the items inside. +-- +-- Note that +-- +-- @ +-- 'gSplit' = 'splitVars' 'gTuple' +-- @ +gSplit + :: forall s rs as a. (SOP.Generic a, SOP.Code a ~ '[as], Known Length as, Every Num as) + => BVar s rs a + -> Prod (BVar s rs) as +gSplit = partsVar gTuple + +-- TODO: figure out how to split sums diff --git a/src/Numeric/Backprop/Internal.hs b/src/Numeric/Backprop/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop/Internal.hs @@ -0,0 +1,294 @@ +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : Numeric.Backprop.Internal +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Provides the types and instances used for the graph +-- building/back-propagation for the library. + +module Numeric.Backprop.Internal + ( Summer(..), summers, summers' + , Unity(..), unities, unities' + , OpB + , BPState(..), bpsSources + , BP(..) + , BPInpRef(..) + , BPNode(..), bpnOut, bpnRes, bpnGradFunc, bpnGradCache, bpnSummer + , BPPipe(..), bppOut, bppRes, bppGradFunc, bppGradCache + , BVar(..) + , ForwardRefs(..), _FRInternal + ) where + +import Control.Monad.Reader +import Control.Monad.ST +import Control.Monad.State +import Data.Kind +import Data.STRef +import Data.Type.Index +import Data.Type.Product +import Lens.Micro hiding (ix) +import Lens.Micro.TH +import Numeric.Backprop.Internal.Helper +import Numeric.Backprop.Op + +-- | A subclass of 'OpM' (and superclass of 'Op'), representing 'Op's that +-- the /backprop/ library uses to perform backpropation. +-- +-- An +-- +-- @ +-- 'OpB' s rs a +-- @ +-- +-- represents a differentiable function that takes a tuple of @rs@ and +-- produces an a @a@, which can be run on @'BVar' s@s and also inside @'BP' +-- s@s. For example, an @'OpB' s '[ Int, Double ] Bool@ takes an 'Int' and +-- a 'Double' and produces a 'Bool', and does it in a differentiable way. +-- +-- 'OpB' is a /superset/ of 'Op', so, if you see any function +-- that expects an 'OpB' (like 'Numeric.Backprop.opVar'' and +-- 'Numeric.Backprop.~$', for example), you can give them an 'Op', as well. +-- +-- You can think of 'OpB' as a superclass/parent class of 'Op' in this +-- sense, and of 'Op' as a subclass of 'OpB'. +type OpB s as a = OpM (ST s) as a + +-- | Reference to /usage sites/ for a given entity, used to get partial or +-- total derivatives. +data ForwardRefs s rs a + -- | A list of 'BPInpRef's pointing to places that use the entity, to + -- provide partial derivatives. + = FRInternal ![BPInpRef s rs a] + -- | The entity is the terminal result of a BP, so its total derivative + -- is fixed. + | FRTerminal !(Maybe a) + +-- | Combines two 'FRInternal' lists. If either input is an 'FRTerminal', +-- then throws away the other result and keeps the new terminal forced +-- total derivative. (Biases to the left) +instance Monoid (ForwardRefs s rs a) where + mempty = FRInternal [] + mappend = \case + FRInternal rs -> \case + FRInternal rs' -> FRInternal (rs ++ rs') + t@(FRTerminal _) -> t + FRTerminal _ -> id + +-- | The "state" of a 'BP' action, which keeps track of what nodes, if any, +-- refer to any of the inputs. +data BPState :: Type -> [Type] -> Type where + BPS :: { _bpsSources :: !(Prod (ForwardRefs s rs) rs) + } + -> BPState s rs + +-- | A Monad allowing you to explicitly build hetereogeneous data +-- dependency graphs and that the library can perform back-propagation on. +-- +-- A @'BP' s rs a@ is a 'BP' action that uses an environment of @rs@ +-- returning a @a@. When "run", it will compute a gradient that is a tuple +-- of @rs@. (The phantom parameter @s@ is used to ensure that any 'BVar's +-- aren't leaked out of the monad) +-- +-- Note that you can only "run" a @'BP' s rs@ that produces a 'BVar' -- +-- that is, things of the form +-- +-- @ +-- 'BP' s rs ('BVar' s rs a) +-- @ +-- +-- The above is a 'BP' action that returns a 'BVar' containing an @a@. +-- When this is run, it'll produce a result of type @a@ and a gradient of +-- that is a tuple of @rs@. (This form has a type synonym, +-- 'Numeric.Backprop.BPOp', for convenience) +-- +-- For example, a @'BP' s '[ Int, Double, Double ]@ is a monad that +-- represents a computation with an 'Int', 'Double', and 'Double' as +-- inputs. And, if you ran a +-- +-- @ +-- 'BP' s '[ Int, Double, Double ] ('BVar' s '[ Int, Double, Double ] Double) +-- @ +-- +-- Or, using the 'BPOp' type synonym: +-- +-- @ +-- 'Numeric.Backprop.BPOp' s '[ Int, Double, Double ] Double +-- @ +-- +-- with 'Numeric.Backprop.backprop' or 'Numeric.Backprop.gradBPOp', it'll +-- return a gradient on the inputs ('Int', 'Double', and 'Double') and +-- produce a value of type 'Double'. +-- +-- Now, one powerful thing about this type is that a 'BP' is itself an +-- 'Op' (or more precisely, an 'Numeric.Backprop.OpB', which is a subtype of +-- 'OpM'). So, once you create your fancy 'BP' computation, you can +-- transform it into an 'OpM' using 'Numeric.Backprop.bpOp'. +newtype BP s rs a = BP { bpST :: ReaderT (Tuple rs) (StateT (BPState s rs) (ST s)) a } + deriving ( Functor + , Applicative + , Monad + ) + +-- | The basic unit of manipulation inside 'BP' (or inside an +-- implicit-graph backprop function). Instead of directly working with +-- values, you work with 'BVar's contating those values. When you work +-- with a 'BVar', the /backprop/ library can keep track of what values +-- refer to which other values, and so can perform back-propagation to +-- compute gradients. +-- +-- A @'BVar' s rs a@ refers to a value of type @a@, with an environment +-- of values of the types @rs@. The phantom parameter @s@ is used to +-- ensure that stray 'BVar's don't leak outside of the backprop process. +-- +-- (That is, if you're using implicit backprop, it ensures that you interact +-- with 'BVar's in a polymorphic way. And, if you're using explicit +-- backprop, it ensures that a @'BVar' s rs a@ never leaves the @'BP' s rs@ +-- that it was created in.) +-- +-- 'BVar's have 'Num', 'Fractional', 'Floating', etc. instances, so they +-- can be manipulated using polymorphic functions and numeric functions in +-- Haskell. You can add them, subtract them, etc., in "implicit" backprop +-- style. +-- +-- (However, note that if you directly manipulate 'BVar's using those +-- instances or using 'Numeric.Backprop.liftB', it delays evaluation, so every usage site +-- has to re-compute the result/create a new node. If you want to re-use +-- a 'BVar' you created using '+' or '-' or 'Numeric.Backprop.liftB', use +-- 'Numeric.Backprop.bindVar' to force it first. See documentation for +-- 'Numeric.Backprop.bindVar' for more details.) +data BVar :: Type -> [Type] -> Type -> Type where + -- | A BVar referring to a 'BPNode' + BVNode :: !(Index bs a) + -> !(STRef s (BPNode s rs as bs)) + -> BVar s rs a + -- | A BVar referring to an environment input variable + BVInp :: !(Index rs a) + -> BVar s rs a + -- | A constant BVar that refers to a specific Haskell value + BVConst :: !a + -> BVar s rs a + -- | A BVar that combines several other BVars using a function (an + -- 'Op'). Essentially a branch of a tree. + BVOp :: !(Prod (BVar s rs) as) + -> !(OpB s as a) + -> BVar s rs a + +-- | Used exclusively by 'ForwardRefs' to specify "where" and "how" to look +-- for partial derivatives at usage sites of a given entity. +data BPInpRef :: Type -> [Type] -> Type -> Type where + -- | The entity is used in a 'BPNode', and as an Nth input + IRNode :: !(Index bs a) + -> !(STRef s (BPNode s rs bs cs)) + -> BPInpRef s rs a + -- | The entity is used in a 'BPPipe', and as an Nth input + IRPipe :: !(Index bs a) + -> !(STRef s (BPPipe s rs bs cs)) + -> BPInpRef s rs a + -- | The entity is used somehow in the terminal result of a 'BP', and + -- so therefore has a fixed partial derivative contribution. + IRConst :: !a + -> BPInpRef s rs a + +-- | A (stateful) node in the graph of operations/data dependencies in 'BP' +-- that the library uses. 'BVar's can refer to these to get results from +-- them, and 'BPInpRef's can refer to these to get partial derivatives from +-- them. +data BPNode :: Type -> [Type] -> [Type] -> [Type] -> Type where + BPN :: { _bpnOut :: !(Prod (ForwardRefs s rs) bs) + , _bpnRes :: !(Tuple bs) + , _bpnGradFunc :: !(Prod Maybe bs -> ST s (Tuple as)) + , _bpnGradCache :: !(Maybe (Tuple as)) -- nothing if is the "final output" + , _bpnSummer :: !(Prod Summer bs) + } + -> BPNode s rs as bs + +-- | Essentially a "single-usage" 'BPNode'. It's a stateful node, but only +-- ever has a single consumer (and so its total derivative comes from +-- a single partial derivative). Used when keeping track of 'BVOp's. +data BPPipe :: Type -> [Type] -> [Type] -> [Type] -> Type where + BPP :: { _bppOut :: !(Prod (BPInpRef s rs) bs) + , _bppRes :: !(Tuple bs) + , _bppGradFunc :: !(Tuple bs -> ST s (Tuple as)) + , _bppGradCache :: !(Maybe (Tuple as)) + } + -> BPPipe s rs as bs + +makeLenses ''BPState +makeLenses ''BPNode +makeLenses ''BPPipe + +-- | Traversal (fake prism) to refer to the list of internal refs if the +-- 'ForwardRef' isn't associated with a terminal entity. +_FRInternal + :: Traversal (ForwardRefs s as a) (ForwardRefs t bs a) + [BPInpRef s as a] [BPInpRef t bs a] +_FRInternal f = \case + FRInternal xs -> FRInternal <$> f xs + FRTerminal g -> pure (FRTerminal g) + + + + +-- | Note that if you use the 'Num' instance to create 'BVar's, the +-- resulting 'BVar' is deferred/delayed. At every location you use it, it +-- will be recomputed, and a separate graph node will be created. If you +-- are using a 'BVar' you made with the 'Num' instance in multiple +-- locations, use 'Numeric.Backprop.bindVar' first to force it and prevent +-- recomputation. +instance Num a => Num (BVar s rs a) where + r1 + r2 = BVOp (r1 :< r2 :< Ø) $ op2 (+) + r1 - r2 = BVOp (r1 :< r2 :< Ø) $ op2 (-) + r1 * r2 = BVOp (r1 :< r2 :< Ø) $ op2 (*) + negate r = BVOp (r :< Ø) $ op1 negate + signum r = BVOp (r :< Ø) $ op1 signum + abs r = BVOp (r :< Ø) $ op1 abs + fromInteger x = BVConst (fromInteger x) + +-- | See note for 'Num' instance. +instance Fractional a => Fractional (BVar s rs a) where + r1 / r2 = BVOp (r1 :< r2 :< Ø) $ op2 (/) + recip r = BVOp (r :< Ø) $ op1 recip + fromRational x = BVConst (fromRational x) + +-- | See note for 'Num' instance. +instance Floating a => Floating (BVar s rs a) where + pi = BVConst pi + exp r = BVOp (r :< Ø) $ op1 exp + log r = BVOp (r :< Ø) $ op1 log + sqrt r = BVOp (r :< Ø) $ op1 sqrt + r1 ** r2 = BVOp (r1 :< r2 :< Ø) $ op2 (**) + logBase r1 r2 = BVOp (r1 :< r2 :< Ø) $ op2 logBase + sin r = BVOp (r :< Ø) $ op1 sin + cos r = BVOp (r :< Ø) $ op1 cos + tan r = BVOp (r :< Ø) $ op1 tan + asin r = BVOp (r :< Ø) $ op1 asin + acos r = BVOp (r :< Ø) $ op1 acos + atan r = BVOp (r :< Ø) $ op1 atan + sinh r = BVOp (r :< Ø) $ op1 sinh + cosh r = BVOp (r :< Ø) $ op1 cosh + tanh r = BVOp (r :< Ø) $ op1 tanh + asinh r = BVOp (r :< Ø) $ op1 asinh + acosh r = BVOp (r :< Ø) $ op1 acosh + atanh r = BVOp (r :< Ø) $ op1 atanh + diff --git a/src/Numeric/Backprop/Internal/Helper.hs b/src/Numeric/Backprop/Internal/Helper.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop/Internal/Helper.hs @@ -0,0 +1,134 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} + +-- | +-- Module : Numeric.Backprop.Internal.Helper +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Provides general helper types like 'Summer' and 'Unity' that both +-- "Numeric.Backprop.Op" and "Numeric.Backprop.Internal" use. + +module Numeric.Backprop.Internal.Helper ( + -- * Summer + Summer(..), summers, nSummers', summers' + -- * Unity + , Unity(..), unities, nUnities', unities' + ) where + +import Data.Type.Index +import Data.Type.Length +import Data.Type.Nat +import Data.Type.Product +import Data.Type.Util +import Type.Class.Known + +-- | Instructions on how to "sum" a list of values of a given type. +-- Basically used as an explicit witness for a 'Num' instance. +-- +-- For most types, the only meaningful value of type @'Summer' a@ is +-- @'Summer' 'sum'@. However, using 'Summer' lets us use 'BP' with types +-- that are /not/ instances of 'Num'. Any type can be used, as long as you +-- provide a way to "sum" it! +-- +-- For most of the functions in this library, you can completely ignore +-- this, as they will be generated automatically. You only need to work +-- with this directly if you want to use custom types that /aren't/ +-- instances of 'Num' with this library. +-- +-- If 'Num a' is satisfied, one can create the canonical 'Summer' using +-- @'known' :: 'Num' a => 'Summer' a@. +newtype Summer a = Summer { runSummer :: [a] -> a } + +-- | A canonical "unity" (the multiplicative identity) for a given type. +-- Basically used as an explicit witness for a 'Num' instance. +-- +-- For most types, the only meaningful value of type @'Unity' a@ is +-- @'Unity' 1'@. However, using 'Unity' lets us use 'BP' with types +-- that are /not/ instances of 'Num'. Any type can be used, as long as you +-- provide a way to get a multiplicative identity in it! +-- +-- For most of the functions in this library, you can completely ignore +-- this, as they will be generated automatically. You only need to work +-- with this directly if you want to use custom types that /aren't/ +-- instances of 'Num' with this library. +-- +-- If 'Num a' is satisfied, one can create the canonical 'Unity' using +-- @'known' :: 'Num' a => 'Unity' a@. +newtype Unity a = Unity { getUnity :: a } + deriving (Functor, Show) + +-- | If @a@ is an instance of 'Num', then the canonical @'Summer' a@ is +-- @'Summer' 'sum'@. +instance Num a => Known Summer a where + type KnownC Summer a = Num a + known = Summer sum + +-- | If @a@ is an instance of 'Num', then the canonical @'Unity' a@ is +-- @'Unity' 1@. +instance Num a => Known Unity a where + type KnownC Unity a = Num a + known = Unity 1 + +-- | If all the types in @as@ are instances of 'Num', generate a @'Prod' +-- 'Summer' as@, or a tuple of 'Summer's for every type in @as@. +summers + :: (Every Num as, Known Length as) + => Prod Summer as +summers = summers' known + +-- | Like 'summers', but requiring an explicit witness for the number of +-- types in the list @as@. +summers' + :: Every Num as + => Length as + -> Prod Summer as +summers' l = withEvery' @Num l known + +-- | If all the types in @as@ are instances of 'Num', generate a @'Prod' +-- 'Unity' as@, or a tuple of 'Unity's for every type in @as@. +unities + :: (Every Num as, Known Length as) + => Prod Unity as +unities = unities' known + +-- | Like 'unities', but requiring an explicit witness for the number of +-- types in the list @as@. +unities' + :: Every Num as + => Length as + -> Prod Unity as +unities' l = withEvery' @Num l known + +-- | Create @n@ canonical 'Summer's of for the same type, using its 'Num' +-- instance. +nSummers' + :: forall n a. Num a + => Nat n + -> Prod Summer (Replicate n a) +nSummers' = \case + Z_ -> Ø + S_ (n :: Nat n') -> Summer sum :< nSummers' @n' @a n + +-- | Create @n@ canonical 'Unity's of for the same type, using its 'Num' +-- instance. +nUnities' + :: forall n a. Num a + => Nat n + -> Prod Unity (Replicate n a) +nUnities' = \case + Z_ -> Ø + S_ (n :: Nat n') -> Unity 1 :< nUnities' @n' @a n + diff --git a/src/Numeric/Backprop/Iso.hs b/src/Numeric/Backprop/Iso.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop/Iso.hs @@ -0,0 +1,209 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} + +-- | +-- Module : Numeric.Backprop.Iso +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- A poor substitute for the "Control.Lens.Iso" module in /lens/, providing +-- the 'Iso' type synonym and some sample useful 'Iso's for usage with +-- /backprop/, without incuring a lens dependency. +-- +-- If you also import lens, you should only use this module for the +-- 'Iso's it exports, and not import the redefined 'Iso' type synonym or +-- 'from' \/ 'iso' \/ 'review'. +-- + +module Numeric.Backprop.Iso ( + -- * Isomorphisms + Iso, Iso' + -- ** Construction and usage + , iso + , from, review, view + -- * Useful Isos + , coerced + , gTuple, gSOP + , sum1, resum1 + -- * Utility types + -- | See "Numeric.Backprop#prod" for a mini-tutorial on 'Prod' and + -- 'Tuple', and "Numeric.Backprop#sum" for a mini-tutorial on 'Sum'. + , Prod(..), Tuple, Sum(..), I(..) + ) where + +import Data.Coerce +import Data.Functor.Identity +import Data.Profunctor.Unsafe +import Data.Tagged +import Data.Type.Combinator +import Data.Type.Product +import Data.Type.Sum +import Lens.Micro.Extras +import Type.Class.Higher +import qualified Generics.SOP as SOP + +-- | A family of isomorphisms. See 'Iso''. +type Iso s t a b = forall p f. (Profunctor p, Functor f) => p a (f b) -> p s (f t) + +-- | An @'Iso'' s a@ encodes an isomorphism between an 's' and an 'a'. It +-- basically lets you go from @s -> a@ and back (from @a -> s@) while +-- preserving structure. You can basically imagine an @'Iso'' s a@ to be +-- an @(s -> a, a -> s)@ tuple. +-- +-- You can get the "forward" direction of an 'Iso'' with 'view': +-- +-- @ +-- 'view' :: Iso'' s a -> (s -> a) +-- @ +-- +-- And the "backwards" direction with 'review': +-- +-- @ +-- 'review' :: Iso'' s a -> (a -> s) +-- @ +-- +-- You can construct an 'Iso'' using 'iso', giving the forward and +-- backwards functions: +-- +-- >>> myIso :: Iso' (Identity a) a +-- myIso = iso runIdentity Identity +-- >>> view myIso (Identity "hello") +-- "hello" +-- >>> review myIso "hello" +-- Identity "hello" +-- +-- One powerful thing about 'Iso''s is that they're /composable/ using '.': +-- +-- @ +-- ('.') :: 'Iso'' c b -> 'Iso'' b a -> 'Iso'' c a +-- @ +-- +-- This is basically provided here so that this package doesn't incurr +-- a /lens/ dependecy, but if you already depend on /lens/, you should use +-- the version from "Control.Lens.Iso" instead. +type Iso' s a = Iso s s a a + +-- | Construct an 'Iso' by giving the "forward" and "backward" direction +-- functions: +-- +-- >>> myIso :: Iso' (Identity a) a +-- myIso = iso runIdentity Identity +-- >>> view myIso (Identity "hello") +-- "hello" +-- >>> review myIso "hello" +-- Identity "hello" +-- +-- This is basically provided here so that this package doesn't incurr +-- a /lens/ dependecy, but if you already depend on /lens/, you should use +-- the version from "Control.Lens.Iso" instead. +iso :: (s -> a) -> (b -> t) -> Iso s t a b +iso to_ from_ = dimap to_ (fmap from_) + +-- | Get the "reverse" direction function from an 'Iso'. +-- +-- This is basically provided here so that this package doesn't incurr +-- a /lens/ dependecy, but if you already depend on /lens/, you should use +-- the version from "Control.Lens.Review" instead. +review :: Iso s t a b -> b -> t +review i = runIdentity #. unTagged #. i .# Tagged .# Identity + +-- | A useful 'Iso' between two types with the same runtime representation. +coerced :: Coercible s a => Iso' s a +coerced = iso coerce coerce + +-- | An 'Iso' between a type that is a product type, and a tuple that +-- contains all of its components. Uses "Generics.SOP" and the +-- 'SOP.Generic' typeclass. +-- +-- >>> import qualified Generics.SOP as SOP +-- >>> data Foo = A Int Bool deriving Generic +-- >>> instance SOP.Generic Foo +-- >>> view gTuple (A 10 True) +-- 10 ::< True ::< Ø +-- >>> review gTuple (15 ::< False ::< Ø) +-- A 15 False +-- +gTuple :: (SOP.Generic a, SOP.Code a ~ '[as]) => Iso' a (Tuple as) +gTuple = gSOP . sum1 + +-- | An 'Iso' between a sum type whose constructors are products, and a sum +-- ('Sum') of products ('Tuple'). Uses "Generics.SOP" and the +-- 'SOP.Generic' typeclass. +-- +-- >>> import qualified Generics.SOP as SOP +-- >>> data Bar = A Int Bool | B String Double +-- >>> instance SOP.Generic Bar +-- >>> 'view' 'gSOP' (A 10 True) +-- 'InL' (10 ::< True ::< Ø) +-- >>> 'view' 'gSOP' (B "hello" 3.4) +-- 'InR' ('InL' ("hello" ::< 3.4 ::< Ø)) +-- >>> 'review' 'gTuple' ('InL' (15 ::< False ::< Ø)) +-- A 15 False +-- >>> 'review' 'gTuple' ('InR' ('InL' ("bye" ::< 9.8 ::< Ø))) +-- B "bye" 9.8 +gSOP :: SOP.Generic a => Iso' a (Sum Tuple (SOP.Code a)) +gSOP = sop . sopTC + . iso (map1 (map1 (I . SOP.unI))) (map1 (map1 (SOP.I . getI))) + +-- | An iso between a single-type 'Sum' and the single type. +sum1 :: Iso' (Sum f '[a]) (f a) +sum1 = iso (\case InL x -> x + InR _ -> error "inaccessible?" + ) InL + +-- | An iso between a single type and a single-type 'Sum'. +resum1 :: Iso' (f a) (Sum f '[a]) +resum1 = iso InL + (\case InL x -> x + InR _ -> error "inaccessible?" + ) + +-- | Reverse an 'Iso''. The forward function becomes the backwards +-- function, and the backwards function becomes the forward function. +-- +-- This is basically provided here so that this package doesn't incurr +-- a /lens/ dependecy, but if you already depend on /lens/, you should use +-- the version from "Control.Lens.Review" instead. +from :: Iso' s a -> Iso' a s +from i = iso (review i) (view i) + +sop :: SOP.Generic a => Iso' a (SOP.SOP SOP.I (SOP.Code a)) +sop = iso SOP.from SOP.to + +sopTC :: Iso' (SOP.SOP f as) (Sum (Prod f) as) +sopTC = iso SOP.unSOP SOP.SOP + . nsSum + . iso (map1 (view npProd)) (map1 (review npProd)) + +npProd :: Iso' (SOP.NP f as) (Prod f as) +npProd = iso to_ from_ + where + to_ :: SOP.NP f as -> Prod f as + to_ = \case + SOP.Nil -> Ø + x SOP.:* xs -> x :< to_ xs + from_ :: Prod f as -> SOP.NP f as + from_ = \case + Ø -> SOP.Nil + x :< xs -> x SOP.:* from_ xs + +nsSum :: Iso' (SOP.NS f as) (Sum f as) +nsSum = iso to_ from_ + where + to_ :: SOP.NS f as -> Sum f as + to_ = \case + SOP.Z x -> InL x + SOP.S xs -> InR (to_ xs) + from_ :: Sum f as -> SOP.NS f as + from_ = \case + InL x -> SOP.Z x + InR xs -> SOP.S (from_ xs) + diff --git a/src/Numeric/Backprop/Mono.hs b/src/Numeric/Backprop/Mono.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop/Mono.hs @@ -0,0 +1,828 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE TypeOperators #-} + +-- | +-- Module : Numeric.Backprop.Mono +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- +-- Provides a monomorphic interface to the library and to the +-- "Numeric.Backprop" module. +-- +-- They are monomorphic in the sense that all of the /inputs/ have to be of +-- the same type. So, something like +-- +-- @ +-- 'Numeric.Backprop.BP' s '[Double, Double, Double] Int +-- @ +-- +-- From "Numeric.Backprop" would, in this module, be: +-- +-- @ +-- 'BP' s 'N3' Double Int +-- @ +-- +-- Instead of dealing with 'Prod's and 'Tuple's, this module works with +-- 'VecT's and 'Vec's, respectively. These are fixed-length vectors whose +-- length are encoded in their types, constructed with ':*' (for 'VecT') or +-- ':+' (for 'Vec'). +-- +-- Most of the concepts in normal heterogeneous backprop (for +-- "Numeric.Backprop") should apply here as well, so you can look at any of +-- the tutorials or examples and repurpose them to work here. Just +-- remember to convert something like @'Numeric.Backprop.Op.Op' '[a, a] b@ +-- to @'Op' 'N2' a b@. +-- +-- As a comparison, this implements something similar in functionality to +-- "Numeric.AD" and "Numeric.AD.Mode.Reverse" from the /ad/ package, in +-- that they both offer monomorphic automatic differentiation through +-- back-propagation. This module doesn't allow the computation of jacobians +-- or generalized gradients for \(\mathbb{R}^N \rightarrow \mathbb{R}^M\) +-- functions. This module only computs gradients for \(\mathbb{R}^N +-- \rightarrow \mathbb{R}\)-like functions. This is more of a conscious +-- design decision in the API of this module rather than a fundamental +-- limitation of the implementation. +-- +-- This module also allows you to build explicit data dependency graphs so +-- the library can reduce duplication and perform optimizations, which may +-- or may not provide advantages over "Numeric.AD.Mode.Reverse"'s +-- 'System.IO.Unsafe.unsafePerformIO'-based implicit graph building. +-- + +module Numeric.Backprop.Mono ( + -- * Types + -- ** Backprop types + BP, BPOp, BPOpI, BVar + , Op, OpB + -- ** Vectors#vec# + -- $vec + , VecT(..), Vec, I(..) + -- * BP + -- ** Backprop + , backprop, evalBPOp, gradBPOp + -- ** Utility combinators + , withInps, implicitly + -- * Vars + , constVar + , inpVar, inpVars + , bpOp + , bindVar + -- ** From Ops + , opVar, (~$) + , opVar1, opVar2, opVar3 + , (-$) + -- ** Combining + , liftB, (.$), liftB1, liftB2, liftB3 + -- * Op + , op1, op2, op3, opN, composeOp, composeOp1, (~.) + , op1', op2', op3' + -- * Utility + , pattern (:+), (*:), (+:), head' + -- ** 'Nat' type synonyms + , N0, N1, N2, N3, N4, N5, N6, N7, N8, N9, N10 + ) where + +import Data.Type.Fin +import Data.Type.Nat +import Data.Type.Product hiding (head') +import Data.Type.Util +import Data.Type.Vector +import Numeric.Backprop.Internal.Helper +import Numeric.Backprop.Op.Mono +import Type.Class.Known +import qualified Numeric.Backprop as BP + +-- $vec +-- +-- A 'VecT' is a fixed-length list of a given type. It's basically the +-- "monomorphic" version of a 'Prod' (see the mini-tutorial in +-- "Numeric.Backprop#prod"). +-- +-- A @'VecT' n f a@ is a list of @n@ @f a@s, and is constructed by consing +-- them together with ':*' (using 'ØV' as nil): +-- +-- +-- @ +-- 'I' "hello" ':*' I "world" :* I "ok" :* ØV :: 'VecT' 'N3' 'I' String +-- [1,2,3] :* [4,5,6,7] :* ØV :: 'VecT' 'N2' [] Int +-- @ +-- +-- ('I' is the identity functor) +-- +-- So, in general: +-- +-- @ +-- x :: f a +-- y :: f a +-- z :: f a +-- k :: f a +-- x :* y :* z :* k :* ØV :: 'VecT' f 'N4' a +-- @ +-- +-- 'Vec' is provided as a convenient type synonym for 'VecT' 'I', and has +-- a convenient pattern synonym ':+', which can also be used for pattern +-- matching: +-- +-- @ +-- x :: a +-- y :: a +-- z :: a +-- k :: a +-- +-- x '::<' y ::< z ::< k ::< ØV :: 'Vec' 'N4' a +-- @ + +-- | A Monad allowing you to explicitly build hetereogeneous data +-- dependency graphs and that the library can perform back-propagation on. +-- +-- A @'BP' s n r a@ is a 'BP' action that uses an environment @n@ values of +-- type @r@, and returns an @a@. When "run", it will compute a gradient that +-- is a vector ('Vec') of @n@ @r@s. (The phantom parameter @s@ is used to +-- ensure that any 'BVar's aren't leaked out of the monad) +-- +-- Note that you can only "run" a @'BP' s n r@ that produces a 'BVar' -- +-- that is, things of the form +-- +-- @ +-- 'BP' s n r ('BVar' n r a) +-- @ +-- +-- The above is a 'BP' action that returns a 'BVar' containing an @a@. +-- When this is run, it'll produce a result of type @a@ and a gradient of +-- that is a vector of @n@ values of type @r@. (This form has a type +-- synonym, 'BPOp', for convenience) +-- +-- For example, @'BP' s 'N3' Double@ is a monad that represents +-- a computation with three 'Double's as inputs. And, if you ran a +-- +-- @ +-- 'BP' s 'N3' Double ('BVar' N3 Double Int) +-- @ +-- +-- Or, using the 'BPOp' type synonym: +-- +-- @ +-- 'BPOp' s 'N3' Double Int +-- @ +-- +-- with 'backprop' or 'gradBPOp', it'll return a gradient on the inputs (a +-- vector of three 'Double's) and produce a value of type 'Int'. +-- +-- Now, one powerful thing about this type is that a 'BP' is itself an +-- 'Op' (or more precisely, an 'OpM'). So, once you create your fancy 'BP' +-- computation, you can transform it into an 'OpM' using 'bpOp'. +type BP s n r = BP.BP s (Replicate n r) + +-- | The basic unit of manipulation inside 'BP' (or inside an +-- implicit-graph backprop function). Instead of directly working with +-- values, you work with 'BVar's contating those values. When you work +-- with a 'BVar', the /backprop/ library can keep track of what values +-- refer to which other values, and so can perform back-propagation to +-- compute gradients. +-- +-- A @'BVar' s n r a@ refers to a value of type @a@, with an environment +-- of @n@ values of type @r@. The phantom parameter @s@ is used to +-- ensure that stray 'BVar's don't leak outside of the backprop process. +-- +-- (That is, if you're using implicit backprop, it ensures that you interact +-- with 'BVar's in a polymorphic way. And, if you're using explicit +-- backprop, it ensures that a @'BVar' s n r a@ never leaves the @'BP' +-- s n r@ that it was created in.) +-- +-- 'BVar's have 'Num', 'Fractional', 'Floating', etc. instances, so they +-- can be manipulated using polymorphic functions and numeric functions in +-- Haskell. You can add them, subtract them, etc., in "implicit" backprop +-- style. +-- +-- (However, note that if you directly manipulate 'BVar's using those +-- instances or using 'liftB', it delays evaluation, so every usage site +-- has to re-compute the result/create a new node. If you want to re-use +-- a 'BVar' you created using '+' or '-' or 'liftB', use +-- 'bindVar' to force it first. See documentation for +-- 'bindVar' for more details.) +type BVar s n a = BP.BVar s (Replicate n a) + +-- | A handy type synonym representing a 'BP' action that returns a 'BVar'. +-- This is handy because this is the form of 'BP' actions that +-- 'backprop' and 'gradBPOp' (etc.) expects. +-- +-- A value of type: +-- +-- @ +-- 'BPOp' s n r a +-- @ +-- +-- is an action that takes an input environment of @n@ values of type @r@ +-- and produces a 'BVar' containing a value of type @a@. Because it +-- returns a 'BVar', the library can track the data dependencies between +-- the 'BVar' and the input environment and perform back-propagation. +-- +-- See documentation for 'BP' for an explanation of the phantom type +-- parameter @s@. +type BPOp s n r a = BP s n r (BVar s n r a) + +-- | An "implicit" operation on 'BVar's that can be backpropagated. +-- A value of type: +-- +-- @ +-- 'BPOpI' s n r a +-- @ +-- +-- takes a vector ('Vec') of @n@ of 'BVar's containg @r@s and uses them to (purely) +-- produce a 'BVar' containing an @a@. +-- +-- @ +-- foo :: BPOpI s 'N2' Double Double +-- foo (x :* y :* ØV) = x + sqrt y +-- @ +-- +-- If you are exclusively doing implicit back-propagation by combining +-- 'BVar's and using 'BPOpI's, you are probably better off just importing +-- "Numeric.Backprop.Mono.Implicit", which provides better tools. This +-- type synonym exists in "Numeric.Backprop.Mono" just for the 'implicitly' +-- function, which can convert "implicit" backprop functions like +-- a @'BPOpI' s rs a@ into an "explicit" graph backprop function, a @'BPOp' +-- s rs a@. +type BPOpI s n r a = VecT n (BVar s n r) r -> BVar s n r a + +-- | A subclass of 'Numeric.Backprop.Op.Mono.OpM' (and superclass of 'Op'), +-- representing 'Op's that the /backprop/ library uses to perform +-- backpropation. +-- +-- An +-- +-- @ +-- 'OpB' s n a b +-- @ +-- +-- represents a differentiable function that takes a @n@ values of type @a@ +-- produces an a @b@, which can be run on @'BVar' s@s and also inside +-- @'BP' s@s. For example, an @'OpB' s 'N2' Double Bool@ takes two 'Double's +-- and produces a 'Bool', and does it in a differentiable way. +-- +-- 'OpB' is a /superset/ of 'Op', so, if you see any function that expects +-- an 'OpB' (like 'Numeric.Backprop.opVar'' and 'Numeric.Backprop.~$', for +-- example), you can give them an 'Op', as well. +-- +-- You can think of 'OpB' as a superclass/parent class of 'Op' in this +-- sense, and of 'Op' as a subclass of 'OpB'. +type OpB s n a b = BP.OpB s (Replicate n a) b + +-- | Apply an 'OpB' to a 'VecT' (vector) of 'BVar's. +-- +-- If you had an @'OpB' s N3 a b@, this function will expect a vector of of +-- three @'BVar' s n r a@s, and the result will be a @'BVar' s n r b@: +-- +-- @ +-- myOp :: 'OpB' s N3 a b +-- x :: 'BVar' s n r a +-- y :: 'BVar' s n r a +-- z :: 'BVar' s n r a +-- +-- x ':*' y :* z :* 'ØV' :: 'VecT' N3 ('BVar' s n r) a +-- 'opVar' myOp (x :* y :* z :* ØV) :: 'BP' s n r ('BVar' s n r b) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can provide any 'Op' +-- here, as well (like those created by 'op1', 'op2', 'constOp', 'op0' +-- etc.) +-- +-- 'opVar' has an infix alias, '~$', so the above example can also be +-- written as: +-- +-- @ +-- myOp '~$' (x :* y :* z :* ØV) :: 'BP' s n r ('BVar' s n r b) +-- @ +-- +-- to let you pretend that you're applying the 'myOp' function to three +-- inputs. +-- +-- Also note the relation between 'opVar' and 'liftB' and 'bindVar': +-- +-- @ +-- 'opVar' o xs = 'bindVar' ('liftB' o xs) +-- @ +-- +-- 'opVar' can be thought of as a "binding" version of 'liftB'. +opVar + :: forall s m n r a b. Num b + => OpB s m a b + -> VecT m (BVar s n r) a + -> BP s n r (BVar s n r b) +opVar o = BP.opVar o . vecToProd + +-- | Infix synonym for 'opVar', which lets you pretend that you're applying +-- 'OpB's as if they were functions: +-- +-- @ +-- myOp :: 'OpB' s N3 a b +-- x :: 'BVar' s n r a +-- y :: 'BVar' s n r a +-- z :: 'BVar' s n r a +-- +-- x ':*' y :* z :* 'ØV' :: 'VecT' N3 ('BVar' s n r) a +-- myOp '~$' (x :* y :* z :* ØV) :: 'BP' s n r ('BVar' s n r b) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in any 'Op' +-- here, as well (like those created by 'op1', 'op2', 'constOp', 'op0' +-- etc.) +-- +-- '~$' can also be thought of as a "binding" version of '.$': +-- +-- @ +-- o '~$' xs = 'bindVar' (o '.$' xs) +-- @ +-- +infixr 5 ~$ +(~$) + :: forall s m n r a b. Num b + => OpB s m a b + -> VecT m (BVar s n r) a + -> BP s n r (BVar s n r b) +(~$) = opVar @_ @_ @_ @r + +-- | Lets you treat a @'BPOp' s n a b@ as an @'Op' n a b@, and "apply" +-- arguments to it just like you would with an 'Op' and '~$' / 'opVar'. +-- +-- Basically a convenient wrapper over 'bpOp' and '~$': +-- +-- @ +-- o '-$' xs = bpOp o '~$' xs +-- @ +-- +-- So for a @'BPOp' s n a b@, you can "plug in" 'BVar's to each @a@, and +-- get a @b@ as a result. +-- +-- Useful for running a @'BPOp' s n a b@ that you got from a different function, and +-- "plugging in" its @a@ inputs with 'BVar's from your current +-- environment. +infixr 5 -$ +(-$) + :: forall s m n r a b. (Num a, Num b, Known Nat m) + => BPOp s m a b + -> VecT m (BVar s n r) a + -> BP s n r (BVar s n r b) +o -$ xs = opVar @_ @_ @_ @r (bpOp @_ @_ @a @b o) xs + +-- | Create a 'BVar' that represents just a specific value, that doesn't +-- depend on any other 'BVar's. +constVar + :: a + -> BVar s n r a +constVar = BP.constVar + +-- | Convenient wrapper over 'opVar' that takes an 'OpB' with one argument +-- and a single 'BVar' argument. Lets you not have to type out the entire +-- 'VecT'. +-- +-- @ +-- 'opVar1' o x = 'opVar' o (x ':*' 'ØV') +-- +-- myOp :: 'Op' N2 a b +-- x :: 'BVar' s n r a +-- +-- 'opVar1' myOp x :: 'BP' s n r ('BVar' s n r b) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op1') as well. +opVar1 + :: forall s n r a b. Num b + => OpB s N1 a b + -> BVar s n r a + -> BP s n r (BVar s n r b) +opVar1 o x = opVar @_ @_ @n @r o (x :* ØV) + +-- | Convenient wrapper over 'opVar' that takes an 'OpB' with two arguments +-- and two 'BVar' arguments. Lets you not have to type out the entire +-- 'VecT'. +-- +-- @ +-- 'opVar2' o x y = 'opVar' o (x ':*' y ':*' 'ØV') +-- +-- myOp :: 'Op' N2 a b +-- x :: 'BVar' s n r a +-- y :: 'BVar' s n r b +-- +-- 'opVar2' myOp x y :: 'BP' s n r ('BVar' s n r b) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op2') as well. +opVar2 + :: forall s n r a b. Num b + => OpB s N2 a b + -> BVar s n r a + -> BVar s n r a + -> BP s n r (BVar s n r b) +opVar2 o x y = opVar @_ @_ @n @r o (x :* y :* ØV) + +-- | Convenient wrapper over 'opVar' that takes an 'OpB' with three arguments +-- and three 'BVar' arguments. Lets you not have to type out the entire +-- 'VecT'. +-- +-- @ +-- 'opVar3' o x y z = 'opVar' o (x ':*' y ':*' z ':*' 'ØV') +-- +-- myOp :: 'Op' N3 a b +-- x :: 'BVar' s n r a +-- y :: 'BVar' s n r a +-- z :: 'BVar' s n r a +-- +-- 'opVar3' myOp x y z :: 'BP' s n r ('BVar' s n r b) +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op3') as well. +opVar3 + :: forall s n r a b. Num b + => OpB s N3 a b + -> BVar s n r a + -> BVar s n r a + -> BVar s n r a + -> BP s n r (BVar s n r b) +opVar3 o x y z = opVar @_ @_ @n @r o (x :* y :* z :* ØV) + +-- | Concretizes a delayed 'BVar'. If you build up a 'BVar' using numeric +-- functions like '+' or '*' or using 'liftB', it'll defer the evaluation, +-- and all of its usage sites will create a separate graph node. +-- +-- Use 'bindVar' if you ever intend to use a 'BVar' in more than one +-- location. +-- +-- @ +-- -- bad +-- errSquared :: Num a => 'BP' s N2 a a +-- errSquared = 'withInp' $ \\(x :* y :* Ø) -\> do +-- let err = r - t +-- 'return' (err * err) -- err is used twice! +-- +-- -- good +-- errSquared :: Num a => 'BP' s N2 a a +-- errSquared = 'withInp' $ \\(x :* y :* Ø) -\> do +-- let err = r - t +-- e <- 'bindVar' err -- force e, so that it's safe to use twice! +-- 'return' (e * e) +-- +-- -- better +-- errSquared :: Num a => 'BP' s N2 a a +-- errSquared = 'withInp' $ \\(x :* y :* Ø) -\> do +-- let err = r - t +-- e <- 'bindVar' err +-- 'bindVar' (e * e) -- result is forced so user doesn't have to worry +-- @ +-- +-- Note the relation to 'opVar' / '~$' / 'liftB' / '.$': +-- +-- @ +-- 'opVar' o xs = 'bindVar' ('liftB' o xs) +-- o '~$' xs = 'bindVar' (o '.$' xs) +-- 'op2' (*) '~$' (x :< y :< Ø) = 'bindVar' (x * y) +-- @ +-- +-- So you can avoid 'bindVar' altogether if you use the explicitly binding +-- '~$' and 'opVar' etc. +-- +-- Note that 'bindVar' on 'BVar's that are already forced is a no-op. +bindVar + :: forall s n r a. Num a + => BVar s n r a + -> BP s n r (BVar s n r a) +bindVar = BP.bindVar + +-- | Perform back-propagation on the given 'BPOp'. Returns the result of +-- the operation it represents, as well as the gradient of the result with +-- respect to its inputs. See module header for "Numeric.Backprop.Mono" +-- and package documentation for examples and usages. +backprop + :: forall n r a. Num r + => (forall s. BPOp s n r a) + -> Vec n r + -> (a, Vec n r) +backprop bp i = (x, prodAlong i g) + where + (x, g) = BP.backprop' (toSummers i) (toUnities i) bp (vecToProd i) + +-- | Simply run the 'BPOp' on an input vector, getting the result without +-- bothering with the gradient or with back-propagation. +evalBPOp + :: forall n r a. () + => (forall s. BPOp s n r a) + -> Vec n r + -> a +evalBPOp bp = BP.evalBPOp bp . vecToProd + +-- | Run the 'BPOp' on an input vector and return the gradient of the result +-- with respect to the input vector +gradBPOp + :: forall n r a. Num r + => (forall s. BPOp s n r a) + -> Vec n r + -> Vec n r +gradBPOp bp = snd . backprop bp + +-- | Turn a 'BPOp' into an 'OpB'. Basically converts a 'BP' taking @n@ +-- @r@s and producing an @a@ into an 'Op' taking an @n@ @r@s and returning +-- an @a@, with all of the powers and utility of an 'Op', including all of +-- its gradient-finding glory. +-- +-- Really just reveals the fact that any @'BPOp' s rs a@ is itself an 'Op', +-- an @'OpB' s rs a@, which makes it a differentiable function. +-- +-- Handy because an 'OpB' can be used with almost all of +-- the 'Op'-related functions in this moduel, including 'opVar', '~$', etc. +bpOp + :: forall s n r a. (Num r, Known Nat n) + => BPOp s n r a + -> OpB s n r a +bpOp b = BP.bpOp' (nSummers' @n @r n) (nUnities' @n @r n) b + where + n :: Nat n + n = known + + +-- | Create a 'BVar' given an index ('Fin') into the input environment. For an +-- example, +-- +-- @ +-- 'inpVar' 'FZ' +-- @ +-- +-- would refer to the /first/ input variable, Bool]@), and +-- +-- @ +-- 'inpVar' ('FS' 'FZ') +-- @ +-- +-- Would refer to the /second/ input variable. +-- +-- Typically, there shouldn't be any reason to use 'inpVar' directly. It's +-- cleaner to get all of your input 'BVar's together using 'withInps' or +-- 'inpVars'. +inpVar + :: Fin n + -> BVar s n r r +inpVar = BP.inpVar . finIndex + +-- | Get a 'VecT' (vector) of 'BVar's for all of the input environment +-- (the @n@ @r@s) of the @'BP' s n r@ +-- +-- For example, if your 'BP' has two 'Double's inside its input +-- environment (a @'BP' s 'N2' Double@), this would return two 'BVar's, +-- pointing to each input 'Double'. +-- +-- @ +-- case ('inpVars' :: 'VecT' 'N2' ('BVar' s 'N2' Double) Double) of +-- x :* y :* ØV -> do +-- -- the first item, x, is a var to the first input +-- x :: 'BVar' s N2 Double +-- -- the second item, y, is a var to the second input +-- y :: 'BVar' s N2 Double +-- @ +inpVars + :: Known Nat n + => VecT n (BVar s n r) r +inpVars = vgen_ inpVar + +-- | Runs a continuation on a 'Vec' of all of the input 'BVar's. +-- +-- Handy for bringing the environment into scope and doing stuff with it: +-- +-- @ +-- foo :: 'BPOp' 'N2' Double Int +-- foo = 'withInps' $ \\(x :* y :* ØV) -\> do +-- -- do stuff with inputs +-- @ +-- +-- Looks kinda like @foo (x :* y *+ ØV) = -- ...@, don't it? +-- +-- Note that the above is the same as +-- +-- @ +-- foo :: 'BPOp' 'N2' Double Int +-- foo = do +-- case 'inpVars' of +-- x :* y :* ØV -> do +-- -- do stuff with inputs +-- @ +-- +-- But just a little nicer! +withInps + :: Known Nat n + => (VecT n (BVar s n r) r -> BP s n r a) + -> BP s n r a +withInps f = f inpVars + +-- | Convert a 'BPOpI' into a 'BPOp'. That is, convert a function on +-- a bundle of 'BVar's (generating an implicit graph) into a fully fledged +-- 'BPOp' that you can run 'backprop' on. See 'BPOpI' for more +-- information. +-- +-- If you are going to write exclusively using implicit 'BVar' operations, +-- it might be more convenient to use "Numeric.Backprop.Mono.Implicit" +-- instead, which is geared around that use case. +implicitly + :: Known Nat n + => BPOpI s n r a + -> BPOp s n r a +implicitly f = withInps (return . f) + +-- | Apply 'OpB' over a 'VecT' of 'BVar's, as inputs. Provides "implicit" +-- back-propagation, with deferred evaluation. +-- +-- If you had an @'OpB' s N3 a b@, this function will expect a vector of of +-- three @'BVar' s n r a@s, and the result will be a @'BVar' s n r b@: +-- +-- @ +-- myOp :: 'OpB' s N3 a b +-- x :: 'BVar' s n r a +-- y :: 'BVar' s n r a +-- z :: 'BVar' s n r a +-- +-- x ':*' y :* z :* 'ØV' :: 'VecT' N3 ('BVar' s n r) a +-- 'liftB' myOp (x :* y :* z :* ØV) :: 'BVar' s n r b +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can provide any 'Op' +-- here, as well (like those created by 'op1', 'op2', 'constOp', 'op0' +-- etc.) +-- +-- 'liftB' has an infix alias, '.$', so the above example can also be +-- written as: +-- +-- @ +-- myOp '.$' (x :* y :* z :* ØV) :: 'BVar' s n r b +-- @ +-- +-- to let you pretend that you're applying the 'myOp' function to three +-- inputs. +-- +-- The result is a new /deferred/ 'BVar'. This should be fine in most +-- cases, unless you use the result in more than one location. This will +-- cause evaluation to be duplicated and multiple redundant graph nodes to +-- be created. If you need to use it in two locations, you should use +-- 'opVar' instead of 'liftB', or use 'bindVar': +-- +-- @ +-- 'opVar' o xs = 'bindVar' ('liftB' o xs) +-- @ +-- +-- 'liftB' can be thought of as a "deferred evaluation" version of 'opVar'. +liftB + :: forall s m n a b r. () + => OpB s m a b + -> VecT m (BVar s n r) a + -> BVar s n r b +liftB o = BP.liftB o . vecToProd + +-- | Infix synonym for 'liftB', which lets you pretend that you're applying +-- 'OpB's as if they were functions: +-- +-- @ +-- myOp :: 'OpB' s N3 a b +-- x :: 'BVar' s n r a +-- y :: 'BVar' s n r a +-- z :: 'BVar' s n r a +-- +-- x ':*' y :* z :* 'ØV' :: 'VecT' N3 ('BVar' s n r) a +-- myOp '.$' (x :* y :* z :* ØV) :: 'BVar' s n r b +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in any 'Op' +-- here, as well (like those created by 'op1', 'op2', 'constOp', 'op0' +-- etc.) +-- +-- See the documentation for 'liftB' for all the caveats of this usage. +-- +-- '.$' can also be thought of as a "deferred evaluation" version of '~$': +-- +-- @ +-- o '~$' xs = 'bindVar' (o '.$' xs) +-- @ +-- +(.$) + :: forall s m n a b r. () + => OpB s m a b + -> VecT m (BVar s n r) a + -> BVar s n r b +o .$ x = liftB @_ @_ @_ @_ @_ @r o x + +-- | Convenient wrapper over 'liftB' that takes an 'OpB' with one argument +-- and a single 'BVar' argument. Lets you not have to type out the entire +-- 'VecT'. +-- +-- @ +-- 'liftB1' o x = 'liftB' o (x ':*' 'ØV') +-- +-- myOp :: 'Op' N2 a b +-- x :: 'BVar' s n r a +-- +-- 'liftB1' myOp x :: 'BVar' s n r b +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op1') as well. +-- +-- See the documentation for 'liftB' for caveats and potential problematic +-- situations with this. +liftB1 + :: OpB s N1 a a + -> BVar s n r a + -> BVar s n r a +liftB1 = BP.liftB1 + +-- | Convenient wrapper over 'liftB' that takes an 'OpB' with two arguments +-- and two 'BVar' arguments. Lets you not have to type out the entire +-- 'VecT'. +-- +-- @ +-- 'liftB2' o x y = 'liftB' o (x ':*' y ':*' 'ØV') +-- +-- myOp :: 'Op' N2 a b +-- x :: 'BVar' s n r a +-- y :: 'BVar' s n r b +-- +-- 'liftB2' myOp x y :: 'BVar' s n r b +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op2') as well. +-- +-- See the documentation for 'liftB' for caveats and potential problematic +-- situations with this. +liftB2 + :: OpB s N2 a a + -> BVar s n r a + -> BVar s n r a + -> BVar s n r a +liftB2 = BP.liftB2 + +-- | Convenient wrapper over 'liftB' that takes an 'OpB' with three arguments +-- and three 'BVar' arguments. Lets you not have to type out the entire +-- 'Prod'. +-- +-- @ +-- 'liftB3' o x y z = 'liftB' o (x ':*' y ':*' z ':*' 'ØV') +-- +-- myOp :: 'Op' N3 a b +-- x :: 'BVar' s n r a +-- y :: 'BVar' s n r b +-- z :: 'BVar' s n r b +-- +-- 'liftB3' myOp x y z :: 'BVar' s n r b +-- @ +-- +-- Note that 'OpB' is a superclass of 'Op', so you can pass in an 'Op' here +-- (like one made with 'op3') as well. +-- +-- See the documentation for 'liftB' for caveats and potential problematic +-- situations with this. +liftB3 + :: OpB s N3 a a + -> BVar s n r a + -> BVar s n r a + -> BVar s n r a + -> BVar s n r a +liftB3 = BP.liftB3 + + + + + + + + +toSummers + :: Num a + => VecT n f a + -> Prod BP.Summer (Replicate n a) +toSummers = \case + ØV -> Ø + _ :* xs -> BP.Summer sum :< toSummers xs + +toUnities + :: Num a + => VecT n f a + -> Prod BP.Unity (Replicate n a) +toUnities = \case + ØV -> Ø + _ :* xs -> BP.Unity 1 :< toUnities xs + diff --git a/src/Numeric/Backprop/Mono/Implicit.hs b/src/Numeric/Backprop/Mono/Implicit.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop/Mono/Implicit.hs @@ -0,0 +1,147 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} + +-- | +-- Module : Numeric.Backprop.Mono.Implicit +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Offers full functionality for implicit-graph back-propagation with +-- monomorphic inputs. The intended usage is to write a 'BPOp', which is +-- a normal Haskell function from 'BVar's to a result 'BVar'. These 'BVar's +-- can be manipulated using their 'Num' / 'Fractional' / 'Floating' +-- instances. +-- +-- The library can then perform back-propagation on the function (using +-- 'backprop' or 'grad') by using an implicitly built graph. +-- +-- This is an "implicit-only" version of "Numeric.Backprop.Mono", and +-- a monomorphic version of "Numeric.Backprop.Implicit", monomorphic in the +-- sense that all of the inputs are of the same type. +-- +-- Like for "Numeric.Backprop.Implicit", this should actually be powerful +-- enough for most use cases, but falls short because without explicit +-- graph capabilities, recomputation can sometimes be inevitable. If the +-- result of a function on 'BVar's is used twice (like @z@ in @let +-- z = x * y in z + z@), this will allocate a new redundant graph node for +-- every usage site of @z@. You can explicitly /force/ @z@, but only using +-- an explicit graph description using "Numeric.Backprop.Mono". +-- +-- Like "Numeric.Backprop.Implicit", this can't handle sum types, but +-- neither can "Numeric.Backprop.Mono", so no loss here :) +-- +-- This module implements pretty much the same functionality as +-- "Numeric.AD" and "Numeric.AD.Mode.Reverse" from the /ad/ package, +-- because it uses the same implicit-graph back-propagation method. It +-- can't compute jacobians/generalized gradients, however. This isn't +-- a fundamental limitation of the implementaiton, though, but rather just +-- a conscious design decision for this module's API. +-- + + +module Numeric.Backprop.Mono.Implicit ( + -- * Types + -- ** Backprop types + BVar, BPOp, Op, BP.OpB + -- ** Vectors + -- | See "Numeric.Backprop.Mono#vec" for a mini-tutorial on 'VecT' and + -- 'Vec' + , VecT(..), Vec, I(..) + -- * back-propagation + , backprop, grad, eval + -- * Var manipulation + , constVar, liftB, (.$), liftB1, liftB2, liftB3 + -- * Op + , op1, op2, op3, opN + -- * Utility + , pattern (:+), (*:), (+:), head' + -- ** 'Nat' type synonyms + , N0, N1, N2, N3, N4, N5, N6, N7, N8, N9, N10 + ) where + +import Data.Type.Nat +import Data.Type.Vector +import Numeric.Backprop.Mono hiding (backprop, BPOp) +import Type.Class.Known +import qualified Numeric.Backprop.Mono as BP + +-- | An operation on 'BVar's that can be backpropagated. A value of type: +-- +-- @ +-- 'BPOp' n r a +-- @ +-- +-- takes a vector ('VecT') of 'BVar's containg @n@ @r@s and uses them to +-- (purely) produce a 'BVar' containing an @a@. +-- +-- @ +-- foo :: 'BPOp' 'N2' Double Double +-- foo (x ':*' y ':*' 'ØV') = x + sqrt y +-- @ +-- +-- 'BPOp' here is related to 'Numeric.Backprop.Mono.BPOpI' from the normal +-- explicit-graph backprop module "Numeric.Backprop.Mono". +type BPOp n a b = forall s. VecT n (BVar s n a) a -> BVar s n a b + +-- | Run back-propagation on a 'BPOp' function, getting both the result and +-- the gradient of the result with respect to the inputs. +-- +-- @ +-- foo :: 'BPOp' 'N2' Double Double +-- foo (x :* y :* ØV) = +-- let z = x * sqrt y +-- in z + x ** y +-- @ +-- +-- >>> 'backprop' foo (2 :+ 3 :+ ØV) +-- (11.46, 13.73 :+ 6.12 :+ ØV) +backprop + :: forall n a b. (Num a, Known Nat n) + => BPOp n a b + -> Vec n a + -> (b, Vec n a) +backprop f = BP.backprop $ BP.withInps (return . f) + +-- | Run the 'BPOp' on an input tuple and return the gradient of the result +-- with respect to the input tuple. +-- +-- @ +-- foo :: 'BPOp' 'N2' Double Double +-- foo (x :* y :* ØV) = +-- let z = x * sqrt y +-- in z + x ** y +-- @ +-- +-- >>> 'grad' foo (2 :+ 3 :+ ØV) +-- 13.73 :+ 6.12 :+ ØV +grad + :: forall n a b. (Num a, Known Nat n) + => BPOp n a b + -> Vec n a + -> Vec n a +grad f = snd . backprop f + +-- | Simply run the 'BPOp' on an input tuple, getting the result without +-- bothering with the gradient or with back-propagation. +-- +-- @ +-- foo :: 'BPOp' 'N2' Double Double +-- foo (x :* y :* ØV) = +-- let z = x * sqrt y +-- in z + x ** y +-- @ +-- +-- >>> 'eval' foo (2 :+ 3 :+ ØV) +-- 11.46 +eval + :: forall n a b. (Num a, Known Nat n) + => BPOp n a b + -> Vec n a + -> b +eval f = fst . backprop f + diff --git a/src/Numeric/Backprop/Op.hs b/src/Numeric/Backprop/Op.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop/Op.hs @@ -0,0 +1,710 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Numeric.Backprop.Op +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Provides the 'Op' (and 'OpM') type and combinators, which represent +-- differentiable functions/operations on values, and are used by the +-- library to perform back-propagation. +-- +-- Note that 'Op' is a /subset/ or /subtype/ of 'OpM', and so, any function +-- that expects an @'OpM' m as a@ (or an @'Numeric.Backprop.OpB' s as a@) +-- can be given an @'Op' as a@ and it'll work just fine. +-- + +module Numeric.Backprop.Op ( + -- * Implementation + -- $opdoc + -- * Types + -- ** Op and Synonyms + Op, pattern Op, OpM(..) + -- ** Tuple Types + -- | See "Numeric.Backprop#prod" for a mini-tutorial on 'Prod' and + -- 'Tuple' + , Prod(..), Tuple, I(..) + -- * Running + -- ** Pure + , runOp, gradOp, gradOp', gradOpWith, gradOpWith', runOp' + -- ** Monadic + , runOpM, gradOpM, gradOpM', gradOpWithM, gradOpWithM', runOpM' + -- * Manipulation + , composeOp, composeOp1, (~.) + , composeOp', composeOp1' + -- * Creation + , op0, opConst + , opConst' + -- ** Automatic creation using the /ad/ library + , op1, op2, op3, opN + , Replicate + -- ** Giving gradients directly + , op1', op2', op3' + -- ** From Isomorphisms + , opCoerce, opTup, opIso + , opCoerce', opTup', opIso' + -- * Utility + , pattern (:>), only, head' + , pattern (::<), only_ + ) where + +import Data.Bifunctor +import Data.Coerce +import Data.Maybe +import Data.Reflection (Reifies) +import Data.Type.Combinator +import Data.Type.Conjunction +import Data.Type.Index +import Data.Type.Length +import Data.Type.Nat +import Data.Type.Product +import Data.Type.Util +import Data.Type.Vector hiding (head') +import Lens.Micro.Extras +import Numeric.AD +import Numeric.AD.Internal.Reverse (Reverse, Tape) +import Numeric.AD.Mode.Forward hiding (grad') +import Numeric.Backprop.Internal.Helper +import Numeric.Backprop.Iso +import Type.Class.Higher +import Type.Class.Known +import Type.Class.Witness + +-- instead of Tuple as, Prod Diff as, where Diff can be a value, or zero, +-- or one? + +-- $opdoc +-- 'Op's contain information on a function as well as its gradient, but +-- provides that information in a way that allows them to be "chained". +-- +-- For example, for a function +-- +-- \[ +-- f : \mathbb{R}^n \rightarrow \mathbb{R} +-- \] +-- +-- We might want to apply a function \(g\) to the result we get, to get +-- our "final" result: +-- +-- \[ +-- \eqalign{ +-- y &= f(\mathbf{x})\cr +-- z &= g(y) +-- } +-- \] +-- +-- Now, we might want the gradient \(\nabla z\) with respect to +-- \(\mathbf{x}\), or \(\nabla_\mathbf{x} z\). Explicitly, this is: +-- +-- \[ +-- \nabla_\mathbf{x} z = \left< \frac{\partial z}{\partial x_1}, \frac{\partial z}{\partial x_2}, \ldots \right> +-- \] +-- +-- We can compute that by multiplying the total derivative of \(z\) with +-- respect to \(y\) (that is, \(\frac{dz}{dy}\)) with the gradient of +-- \(f\)) itself: +-- +-- \[ +-- \eqalign{ +-- \nabla_\mathbf{x} z &= \frac{dz}{dy} \left< \frac{\partial y}{\partial x_1}, \frac{\partial y}{\partial x_2}, \ldots \right>\cr +-- \nabla_\mathbf{x} z &= \frac{dz}{dy} \nabla_\mathbf{x} y +-- } +-- \] +-- +-- So, to create an @'Op' as a@ with the 'Op' constructor (or an 'OpM' with the +-- 'OpM' constructor), you give a function that returns a tuple, +-- containing: +-- +-- 1. An @a@: The result of the function +-- 2. An @Maybe a -> Tuple as@: A function that, when given +-- \(\frac{dz}{dy}\) (in a 'Just'), returns the total gradient +-- \(\nabla_z \mathbf{x}\). If the function is given is given +-- 'Nothing', then \(\frac{dz}{dy}\) should be taken to be 1. In other +-- words, you would simply need to return \(\nabla_y \mathbf{x}\), +-- unchanged. That is, an input of 'Nothing' indicates that the "final +-- result" is just simply \(f(\mathbf{x})\), and not some +-- \(g(f(\mathbf{x}))\). +-- +-- This is done so that 'Op's can easily be "chained" together, one after +-- the other. If you have an 'Op' for \(f\) and an 'Op' for \(g\), you can +-- compute the gradient of \(f\) knowing that the result target is +-- \(g \circ f\). +-- +-- Note that end users should probably never be required to construct an +-- 'Op' or 'OpM' explicitly this way. Instead, libraries should provide +-- carefuly pre-constructed ones, or provide ways to generate them +-- automatically (like 'op1', 'op2', and 'op3' here). + +-- | An @'OpM' m as a@ represents a /differentiable/ (monadic) function +-- from @as@ to @a@, in the context of a 'Monad' @m@. +-- +-- For example, an +-- +-- @ +-- 'OpM' IO '[Int, Bool] Double +-- @ +-- +-- would be a function that takes an 'Int' and a 'Bool' and returns +-- a 'Double' (in 'IO'). It can be differentiated to give a /gradient/ of +-- an 'Int' and a 'Bool' (also in 'IO') if given the total derivative for +-- the @Double@. +-- +-- Note that an 'OpM' is a /superclass/ of 'Op', so any function that +-- expects an @'OpM' m as a@ can also accept an @'Op' as a@. +-- +-- See 'runOpM', 'gradOpM', and 'gradOpWithM' for examples on how to run +-- it. +newtype OpM m as a = + -- | Construct an 'OpM' by giving a (monadic) function creating the + -- result, and also a continuation on how to create the gradient, given + -- the total derivative of @a@. + -- + -- See the module documentation for "Numeric.Backprop.Op" for more + -- details on the function that this constructor and 'Op' expect. + OpM (Tuple as -> m (a, Maybe a -> m (Tuple as))) + +-- | An @'Op' as a@ describes a differentiable function from @as@ to @a@. +-- +-- For example, a value of type +-- +-- @ +-- 'Op' '[Int, Bool] Double +-- @ +-- +-- is a function from an 'Int' and a 'Bool', returning a 'Double'. It can +-- be differentiated to give a /gradient/ of an 'Int' and a 'Bool' if given +-- a total derivative for the @Double@. If we call 'Bool' \(2\), then, +-- mathematically, it is akin to a: +-- +-- \[ +-- f : \mathbb{Z} \times 2 \rightarrow \mathbb{R} +-- \] +-- +-- See 'runOp', 'gradOp', and 'gradOpWith' for examples on how to run it, +-- and 'Op' for instructions on creating it. +-- +-- This type is abstracted over using the pattern synonym with constructor +-- 'Op', so you can create one from scratch with it. However, it's +-- simplest to create it using 'op2'', 'op1'', 'op2'', and 'op3'' helper +-- smart constructors And, if your function is a numeric function, they +-- can even be created automatically using 'op1', 'op2', 'op3', and 'opN' +-- with a little help from "Numeric.AD" from the /ad/ library. +-- +-- Note that this type is a /subset/ or /subtype/ of 'OpM' (and also of +-- 'Numeric.Backprop.OpB'). So, if a function ever expects an @'OpM' m as +-- a@ (or a 'Numeric.Backprop.OpB'), you can always provide an @'Op' as a@ +-- instead. +-- +-- Many functions in this library will expect an @'OpM' m as a@ (or +-- an @'Numeric.Backprop.OpB' s as a@), and in all of these cases, you can +-- provide an @'Op' as a@. +type Op as a = forall m. Monad m => OpM m as a + +-- | Helper wrapper used for the implementation of 'composeOp'. +newtype OpCont m as a = OC { runOpCont :: Maybe a -> m (Tuple as) } + +-- | Construct an 'Op' by giving a function creating the result, and also +-- a continuation on how to create the gradient, given the total derivative +-- of @a@. +-- +-- See the module documentation for "Numeric.Backprop.Op" for more details +-- on the function that this constructor and 'OpM' expect. +pattern Op :: (Tuple as -> (a, Maybe a -> Tuple as)) -> Op as a +pattern Op runOp' <- OpM (\f -> (second . fmap) getI . getI . f -> runOp') + where + Op f = OpM (pure . (second . fmap) pure . f) + +-- | A combination of 'runOpM' and 'gradOpWithM''. Given an 'OpM' and +-- inputs, returns the result of the 'OpM' and a continuation that gives +-- its gradient. +-- +-- The continuation takes the total derivative of the result as input. See +-- documenation for 'gradOpWithM'' and module documentation for +-- "Numeric.Backprop.Op" for more information. +runOpM' + :: OpM m as a -- ^ 'OpM' to run + -> Tuple as -- ^ Inputs + -> m (a, Maybe a -> m (Tuple as)) -- ^ Result, and continuation to + -- get the gradient +runOpM' (OpM f) = f + +-- | A combination of 'runOp' and 'gradOpWith''. Given an 'Op' and inputs, +-- returns the result of the 'Op' and a continuation that gives its +-- gradient. +-- +-- The continuation takes the total derivative of the result as input. See +-- documenation for 'gradOpWith'' and module documentation for +-- "Numeric.Backprop.Op" for more information. +runOp' + :: Op as a -- ^ 'Op' to run + -> Tuple as -- ^ Inputs + -> (a, Maybe a -> Tuple as) -- ^ Result, and continuation to get + -- the gradient +runOp' o = (second . fmap) getI . getI . runOpM' o + +-- | 'composeOp', but taking explicit 'Summer's, for the situation where +-- the @as@ are not instance of 'Num'. +composeOp' + :: Monad m + => Prod Summer as -- ^ Explicit 'Summer's + -> Prod (OpM m as) bs -- ^ 'Prod' of 'OpM's taking @as@ and returning + -- different @b@ in @bs@ + -> OpM m bs c -- ^ 'OpM' taking eac of the @bs@ from the + -- input 'Prod'. + -> OpM m as c -- ^ Composed 'OpM' +composeOp' ss os o = OpM $ \xs -> do + (ys, conts) <- fmap unzipP + . traverse1 (fmap (\(x, c) -> I x :&: OC c) . flip runOpM' xs) + $ os + (z, gFz) <- runOpM' o ys + let gFunc g0 = do + g1 <- gFz g0 + g2s <- sequenceA + . toList (\(oc :&: I g) -> runOpCont oc (Just g)) + $ conts `zipP` g1 + return $ map1 (\(s :&: gs) -> I (runSummer s gs)) + . zipP ss + . foldr (\x -> map1 (uncurryFan (\(I y) -> (y:))) . zipP x) + (map1 (const []) ss) + $ g2s + return (z, gFunc) + +-- | Compose 'OpM's together, similar to '.'. But, because all 'OpM's are +-- \(\mathbb{R}^N \rightarrow \mathbb{R}\), this is more like 'sequence' +-- for functions, or @liftAN@. +-- +-- That is, given an @'OpM' m as b1@, an @'OpM' m as b2@, and an @'OpM' +-- m as b3@, it can compose them with an @'OpM' m '[b1,b2,b3] c@ to create +-- an @'OpM' m as c@. +composeOp + :: (Monad m, Known Length as, Every Num as) + => Prod (OpM m as) bs -- ^ 'Prod' of 'OpM's taking @as@ and returning + -- different @b@ in @bs@ + -> OpM m bs c -- ^ 'OpM' taking eac of the @bs@ from the + -- input 'Prod'. + -> OpM m as c -- ^ Composed 'OpM' +composeOp = composeOp' summers + +-- | 'composeOp1', but taking explicit 'Summer's, for the situation where +-- the @as@ are not instance of 'Num'. +composeOp1' + :: Monad m + => Prod Summer as + -> OpM m as b + -> OpM m '[b] c + -> OpM m as c +composeOp1' ss = composeOp' ss . only + +-- | Convenient wrappver over 'composeOp' for the case where the second +-- function only takes one input, so the two 'OpM's can be directly piped +-- together, like for '.'. +composeOp1 + :: (Monad m, Known Length as, Every Num as) + => OpM m as b + -> OpM m '[b] c + -> OpM m as c +composeOp1 = composeOp . only + +-- | Convenient infix synonym for (flipped) 'composeOp1'. Meant to be used +-- just like '.': +-- +-- @ +-- 'op1' negate :: 'Op' '[a] a +-- 'op2' (+) :: Op '[a,a] a +-- +-- op1 negate '~.' op2 (+) :: Op '[a, a] a +-- @ +infixr 9 ~. +(~.) + :: (Monad m, Known Length as, Every Num as) + => OpM m '[b] c + -> OpM m as b + -> OpM m as c +(~.) = flip composeOp1 + + +-- | Run the function that an 'Op' encodes, to get the result. +-- +-- >>> runOp (op2 (*)) (3 ::< 5 ::< Ø) +-- 15 +runOp :: Op as a -> Tuple as -> a +runOp o = fst . runOp' o + +-- | Run the function that an 'Op' encodes, to get the resulting output and +-- also its gradient with respect to the inputs. +-- +-- >>> gradOpM' (op2 (*)) (3 ::< 5 ::< Ø) :: IO (Int, Tuple '[Int, Int]) +-- (15, 5 ::< 3 ::< Ø) +gradOp' :: Op as a -> Tuple as -> (a, Tuple as) +gradOp' o = second ($ Nothing) . runOp' o + +-- | The monadic version of 'runOp', for 'OpM's. +-- +-- >>> runOpM (op2 (*)) (3 ::< 5 ::< Ø) :: IO Int +-- 15 +runOpM :: Functor m => OpM m as a -> Tuple as -> m a +runOpM o = fmap fst . runOpM' o + +-- | The monadic version of 'gradOp'', for 'OpM's. +gradOpM' :: Monad m => OpM m as a -> Tuple as -> m (a, Tuple as) +gradOpM' o x = do + (y, gF) <- runOpM' o x + g <- gF Nothing + return (y, g) + +-- | A combination of 'gradOp' and 'gradOpWith'. The third argument is +-- (optionally) the total derivative the result. Give 'Nothing' and it is +-- assumed that the result is the final result (and the total derivative is +-- 1), and this behaves the same as 'gradOp'. Give @'Just' d@ and it uses +-- the @d@ as the total derivative of the result, and this behaves like +-- 'gradOpWith'. +-- +-- See 'gradOp' and the module documentaiton for "Numeric.Backprop.Op" for +-- more information. +gradOpWith' + :: Op as a -- ^ 'Op' to run + -> Tuple as -- ^ Inputs to run it with + -> Maybe a -- ^ If 'Just', taken as the total derivative of the + -- result. If 'Nothing', assumes that the result is + -- the final result. + -> Tuple as -- ^ The gradient +gradOpWith' o = snd . runOp' o + +-- | The monadic version of 'gradOpWith'', for 'OpM's. +gradOpWithM' + :: Monad m + => OpM m as a -- ^ 'OpM' to run + -> Tuple as -- ^ Inputs to run it with + -> Maybe a -- ^ If 'Just', taken as the total derivative of the + -- result. If 'Nothing', assumes that the result is + -- the final result. + -> m (Tuple as) -- ^ The gradient +gradOpWithM' o xs g = do + (_, f) <- runOpM' o xs + f g + +-- | Run the function that an 'Op' encodes, and get the gradient of +-- a "final result" with respect to the inputs, given the total derivative +-- of the output with the final result. +-- +-- See 'gradOp' and the module documentaiton for "Numeric.Backprop.Op" for +-- more information. +gradOpWith + :: Op as a -- ^ 'Op' to run + -> Tuple as -- ^ Inputs to run it with + -> a -- ^ The total derivative of the result + -> Tuple as -- ^ The gradient +gradOpWith o i = gradOpWith' o i . Just + +-- | The monadic version of 'gradOpWith', for 'OpM's. +gradOpWithM + :: Monad m + => OpM m as a -- ^ 'OpM' to run + -> Tuple as -- ^ Inputs to run it with + -> a -- ^ The total derivative of the result + -> m (Tuple as) -- ^ the gradient +gradOpWithM o i = gradOpWithM' o i . Just + +-- | Run the function that an 'Op' encodes, and get the gradient of the +-- output with respect to the inputs. +-- +-- >>> gradOp (op2 (*)) (3 ::< 5 ::< Ø) +-- 5 ::< 3 ::< Ø +-- -- the gradient of x*y is (y, x) +gradOp :: Op as a -> Tuple as -> Tuple as +gradOp o i = gradOpWith' o i Nothing + +-- | The monadic version of 'gradOp', for 'OpM's. +gradOpM :: Monad m => OpM m as a -> Tuple as -> m (Tuple as) +gradOpM o i = do + (_, gF) <- runOpM' o i + gF Nothing + +-- | A version of 'opCoerce' that takes an explicit 'Unity', so can be run +-- on values that aren't 'Num' instances. +opCoerce' :: Coercible a b => Unity a -> Op '[a] b +opCoerce' u = opIso' u coerced + +-- | An 'Op' that coerces an item into another item whose type has the same +-- runtime representation. Requires the input to be an instance of 'Num'. +-- +-- >>> gradOp' opCoerce (Identity 5) :: (Int, Identity Int) +-- (5, Identity 1) +-- +-- @ +-- 'opCoerce' = 'opIso' 'coerced' +-- @ +opCoerce :: (Coercible a b, Num a) => Op '[a] b +opCoerce = opIso coerced + +-- | A version of 'opTup' that takes explicit 'Unity's, so can be run on +-- values of types that aren't 'Num' instances. +opTup' + :: Prod Unity as + -> Op as (Tuple as) +opTup' u = Op $ \xs -> (xs, fromMaybe (map1 (I . getUnity) u)) + +-- | An 'Op' that takes @as@ and returns exactly the input tuple. +-- +-- >>> gradOp' opTup (1 ::< 2 ::< 3 ::< Ø) +-- (1 ::< 2 ::< 3 ::< Ø, 1 ::< 1 ::< 1 ::< Ø) +opTup + :: (Every Num as, Known Length as) + => Op as (Tuple as) +opTup = opTup' (map1 ((// known) . every @_ @Num) indices) + +-- | A version of 'opIso' that takes an explicit 'Unity', so can be run on +-- values of types that aren't 'Num' instances. +opIso' :: Unity a -> Iso' a b -> Op '[ a ] b +opIso' u i = op1' $ \x -> (view i x, maybe (getUnity u) (review i)) + +-- | An 'Op' that runs the input value through the isomorphism encoded in +-- the 'Iso'. Requires the input to be an instance of 'Num'. +-- +-- Warning: This is unsafe! It assumes that the isomorphisms themselves +-- have derivative 1, so will break for things like +-- 'Numeric.Lens.exponentiating'. Basically, don't use this for any +-- "numeric" isomorphisms. +opIso :: Num a => Iso' a b -> Op '[ a ] b +opIso = opIso' known + +-- | A version of 'opConst' that takes explicit 'Summer's, so can be run on +-- values of types that aren't 'Num' instances. +opConst' :: Prod Summer as -> a -> Op as a +opConst' ss x = Op $ \_ -> + (x , const $ map1 (\s -> I $ runSummer s []) ss) + +-- | An 'Op' that ignores all of its inputs and returns a given constant +-- value. +-- +-- >>> gradOp' (opConst 10) (1 ::< 2 ::< 3 ::< Ø) +-- (10, 0 ::< 0 ::< 0 ::< Ø) +opConst :: (Every Num as, Known Length as) => a -> Op as a +opConst = opConst' summers + +-- | Create an 'Op' that takes no inputs and always returns the given +-- value. +-- +-- There is no gradient, of course (using 'gradOp' will give you an empty +-- tuple), because there is no input to have a gradient of. +-- +-- >>> gradOp' (op0 10) Ø +-- (10, Ø) +-- +-- For a constant 'Op' that takes input and ignores it, see 'opConst' and +-- 'opConst''. +-- +-- Note that because this returns an 'Op', it can be used with any function +-- that expects an 'OpM' or 'Numeric.Backprop.OpB', as well. +op0 :: a -> Op '[] a +op0 x = Op $ \case + Ø -> (x, const Ø) + +-- | Create an 'Op' of a function taking one input, by giving its explicit +-- derivative. The function should return a tuple containing the result of +-- the function, and also a function taking the derivative of the result +-- and return the derivative of the input. +-- +-- If we have +-- +-- \[ +-- \eqalign{ +-- f &: \mathbb{R} \rightarrow \mathbb{R}\cr +-- y &= f(x)\cr +-- z &= g(y) +-- } +-- \] +-- +-- Then the derivative \( \frac{dz}{dx} \), it would be: +-- +-- \[ +-- \frac{dz}{dx} = \frac{dz}{dy} \frac{dy}{dx} +-- \] +-- +-- If our 'Op' represents \(f\), then the second item in the resulting +-- tuple should be a function that takes \(\frac{dz}{dy}\) and returns +-- \(\frac{dz}{dx}\). +-- +-- If the input is 'Nothing', then \(\frac{dz}{dy}\) should be taken to be +-- \(1\). +-- +-- As an example, here is an 'Op' that squares its input: +-- +-- @ +-- square :: Num a => 'Op' '[a] a +-- square = 'op1'' $ \\x -> (x*x, \\case Nothing -> 2 * x +-- Just d -> 2 * d * x +-- ) +-- @ +-- +-- Remember that, generally, end users shouldn't directly construct 'Op's; +-- they should be provided by libraries or generated automatically. +-- +-- For numeric functions, single-input 'Op's can be generated automatically +-- using 'op1'. +op1' + :: (a -> (b, Maybe b -> a)) + -> Op '[a] b +op1' f = Op $ \case + I x :< Ø -> + let (y, dx) = f x + in (y, only_ . dx) + +-- | Create an 'Op' of a function taking two inputs, by giving its explicit +-- gradient. The function should return a tuple containing the result of +-- the function, and also a function taking the derivative of the result +-- and return the derivative of the input. +-- +-- If we have +-- +-- \[ +-- \eqalign{ +-- f &: \mathbb{R}^2 \rightarrow \mathbb{R}\cr +-- z &= f(x, y)\cr +-- k &= g(z) +-- } +-- \] +-- +-- Then the gradient \( \left< \frac{\partial k}{\partial x}, \frac{\partial k}{\partial y} \right> \) +-- would be: +-- +-- \[ +-- \left< \frac{\partial k}{\partial x}, \frac{\partial k}{\partial y} \right> = +-- \left< \frac{dk}{dz} \frac{\partial z}{dx}, \frac{dk}{dz} \frac{\partial z}{dy} \right> +-- \] +-- +-- If our 'Op' represents \(f\), then the second item in the resulting +-- tuple should be a function that takes \(\frac{dk}{dz}\) and returns +-- \( \left< \frac{\partial k}{dx}, \frac{\partial k}{dx} \right> \). +-- +-- If the input is 'Nothing', then \(\frac{dk}{dz}\) should be taken to be +-- \(1\). +-- +-- As an example, here is an 'Op' that multiplies its inputs: +-- +-- @ +-- mul :: Num a => 'Op' '[a, a] a +-- mul = 'op2'' $ \\x y -> (x*y, \\case Nothing -> (y , x ) +-- Just d -> (d*y, x*d) +-- ) +-- @ +-- +-- Remember that, generally, end users shouldn't directly construct 'Op's; +-- they should be provided by libraries or generated automatically. +-- +-- For numeric functions, two-input 'Op's can be generated automatically +-- using 'op2'. +op2' + :: (a -> b -> (c, Maybe c -> (a, b))) + -> Op '[a,b] c +op2' f = Op $ \case + I x :< I y :< Ø -> + let (z, dxdy) = f x y + in (z, (\(dx,dy) -> dx ::< dy ::< Ø) . dxdy) + +-- | Create an 'Op' of a function taking three inputs, by giving its explicit +-- gradient. See documentation for 'op2'' for more details. +op3' + :: (a -> b -> c -> (d, Maybe d -> (a, b, c))) + -> Op '[a,b,c] d +op3' f = Op $ \case + I x :< I y :< I z :< Ø -> + let (q, dxdydz) = f x y z + in (q, (\(dx, dy, dz) -> dx ::< dy ::< dz ::< Ø) . dxdydz) + +-- | Automatically create an 'Op' of a numerical function taking one +-- argument. Uses 'Numeric.AD.diff', and so can take any numerical +-- function polymorphic over the standard numeric types. +-- +-- >>> gradOp' (op1 (recip . negate)) (5 ::< Ø) +-- (-0.2, 0.04 ::< Ø) +op1 :: Num a + => (forall s. AD s (Forward a) -> AD s (Forward a)) + -> Op '[a] a +op1 f = op1' $ \x -> + let (z, dx) = diff' f x + in (z, maybe dx (* dx)) + +-- | Automatically create an 'Op' of a numerical function taking two +-- arguments. Uses 'Numeric.AD.grad', and so can take any numerical function +-- polymorphic over the standard numeric types. +-- +-- >>> gradOp' (op2 (\x y -> x * sqrt y)) (3 ::< 4 ::< Ø) +-- (6.0, 2.0 ::< 0.75 ::< Ø) +op2 :: Num a + => (forall s. Reifies s Tape => Reverse s a -> Reverse s a -> Reverse s a) + -> Op '[a,a] a +op2 f = opN $ \case I x :* I y :* ØV -> f x y + +-- | Automatically create an 'Op' of a numerical function taking three +-- arguments. Uses 'Numeric.AD.grad', and so can take any numerical function +-- polymorphic over the standard numeric types. +-- +-- >>> gradOp' (op3 (\x y z -> (x * sqrt y)**z)) (3 ::< 4 ::< 2 ::< Ø) +-- (36.0, 24.0 ::< 9.0 ::< 64.503 ::< Ø) +op3 :: Num a + => (forall s. Reifies s Tape => Reverse s a -> Reverse s a -> Reverse s a -> Reverse s a) + -> Op '[a,a,a] a +op3 f = opN $ \case I x :* I y :* I z :* ØV -> f x y z + +-- | Automatically create an 'Op' of a numerical function taking multiple +-- arguments. Uses 'Numeric.AD.grad', and so can take any numerical +-- function polymorphic over the standard numeric types. +-- +-- >>> gradOp' (opN (\(x :+ y :+ Ø) -> x * sqrt y)) (3 ::< 4 ::< Ø) +-- (6.0, 2.0 ::< 0.75 ::< Ø) +opN :: (Num a, Known Nat n) + => (forall s. Reifies s Tape => Vec n (Reverse s a) -> Reverse s a) + -> Op (Replicate n a) a +opN f = Op $ \xs -> + let (y, dxs) = grad' f (prodToVec' known xs) + in (y, vecToProd . maybe dxs (\q -> (q *) <$> dxs)) + +instance (Monad m, Known Length as, Every Num as, Num a) => Num (OpM m as a) where + o1 + o2 = composeOp (o1 :< o2 :< Ø) $ op2 (+) + o1 - o2 = composeOp (o1 :< o2 :< Ø) $ op2 (-) + o1 * o2 = composeOp (o1 :< o2 :< Ø) $ op2 (*) + negate o = composeOp (o :< Ø) $ op1 negate + signum o = composeOp (o :< Ø) $ op1 signum + abs o = composeOp (o :< Ø) $ op1 abs + fromInteger x = opConst (fromInteger x) + +instance (Monad m, Known Length as, Every Fractional as, Every Num as, Fractional a) => Fractional (OpM m as a) where + o1 / o2 = composeOp (o1 :< o2 :< Ø) $ op2 (/) + recip o = composeOp (o :< Ø) $ op1 recip + fromRational x = opConst (fromRational x) + +instance (Monad m, Known Length as, Every Floating as, Every Fractional as, Every Num as, Floating a) => Floating (OpM m as a) where + pi = opConst pi + exp o = composeOp (o :< Ø) $ op1 exp + log o = composeOp (o :< Ø) $ op1 log + sqrt o = composeOp (o :< Ø) $ op1 sqrt + o1 ** o2 = composeOp (o1 :< o2 :< Ø) $ op2 (**) + logBase o1 o2 = composeOp (o1 :< o2 :< Ø) $ op2 logBase + sin o = composeOp (o :< Ø) $ op1 sin + cos o = composeOp (o :< Ø) $ op1 cos + tan o = composeOp (o :< Ø) $ op1 tan + asin o = composeOp (o :< Ø) $ op1 asin + acos o = composeOp (o :< Ø) $ op1 acos + atan o = composeOp (o :< Ø) $ op1 atan + sinh o = composeOp (o :< Ø) $ op1 sinh + cosh o = composeOp (o :< Ø) $ op1 cosh + asinh o = composeOp (o :< Ø) $ op1 asinh + acosh o = composeOp (o :< Ø) $ op1 acosh + atanh o = composeOp (o :< Ø) $ op1 atanh + diff --git a/src/Numeric/Backprop/Op/Mono.hs b/src/Numeric/Backprop/Op/Mono.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Backprop/Op/Mono.hs @@ -0,0 +1,484 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Numeric.Backprop.Op.Mono +-- Copyright : (c) Justin Le 2017 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Provides monomorphic versions of the types and combinators in +-- "Numeric.Backprop.Op", for usage with "Numeric.Backprop.Mono" and +-- "Numeric.Backprop.Mono.Implicit". +-- +-- They are monomorphic in the sense that all of the /inputs/ have to be of +-- the same type. So, something like +-- +-- @ +-- 'Numeric.Backprop.Op' '[Double, Double, Double] Int +-- @ +-- +-- From "Numeric.Backprop" would, in this module, be: +-- +-- @ +-- 'Op' 'N3' Double Int +-- @ +-- +-- See the module header for "Numeric.Backprop.Op" for more explicitly +-- details on how to encode an 'Op' and how they are implemented. For the +-- most part, the same principles will apply. +-- +-- Note that 'Op' is a /subset/ or /subtype/ of 'OpM', and so, any function +-- that expects an @'OpM' m as a@ (or an @'Numeric.Backprop.Mono.OpB' s as a@) +-- can be given an @'Op' as a@ and it'll work just fine. +-- + +module Numeric.Backprop.Op.Mono ( + -- * Types + -- ** Op and synonyms + Op, pattern Op, OpM, pattern OpM + -- ** Vector types + -- | See "Numeric.Backprop.Mono#vec" for a mini-tutorial on 'VecT' and + -- 'Vec' + , VecT(..), Vec, I(..) + -- * Running + -- ** Pure + , runOp, gradOp, gradOp', gradOpWith, gradOpWith', runOp' + -- ** Monadic + , runOpM, gradOpM, gradOpM', gradOpWithM, gradOpWithM', runOpM' + -- * Creation + , op0, opConst, composeOp, composeOp1, (~.) + -- ** Automatic creation using the /ad/ library + , op1, op2, op3, opN + , Replicate + -- ** Giving gradients directly + , op1', op2', op3' + -- * Utility + , pattern (:+), (*:), (+:), head' + -- ** 'Nat' type synonyms + , N0, N1, N2, N3, N4, N5, N6, N7, N8, N9, N10 + ) where + +import Data.Bifunctor +import Data.Reflection (Reifies) +import Data.Type.Combinator +import Data.Type.Nat +import Data.Type.Util +import Data.Type.Vector +import Numeric.AD.Internal.Reverse (Reverse, Tape) +import Numeric.AD.Mode.Forward (AD, Forward) +import Type.Class.Known +import Type.Family.Nat +import qualified Numeric.Backprop.Internal.Helper as BP +import qualified Numeric.Backprop.Op as BP + +-- | An @'Op' n a b@ describes a differentiable function from @n@ values of +-- type @a@ to a value of type @b@. +-- +-- For example, a value of type +-- +-- @ +-- 'Op' 'N2' Int Double +-- @ +-- +-- is a function that takes two 'Int's and returns a 'Double'. +-- It can be differentiated to give a /gradient/ of two 'Int's, if given +-- a total derivative for the 'Double'. Mathematically, it is akin to a: +-- +-- \[ +-- f : \mathbb{Z}^2 \rightarrow \mathbb{R} +-- \] +-- +-- See 'runOp', 'gradOp', and 'gradOpWith' for examples on how to run it, +-- and 'Op' for instructions on creating it. +-- +-- This type is abstracted over using the pattern synonym with constructor +-- 'Op', so you can create one from scratch with it. However, it's +-- simplest to create it using 'op2'', 'op1'', 'op2'', and 'op3'' helper +-- smart constructors And, if your function is a numeric function, they +-- can even be created automatically using 'op1', 'op2', 'op3', and 'opN' +-- with a little help from "Numeric.AD" from the /ad/ library. +-- +-- Note that this type is a /subset/ or /subtype/ of 'OpM' (and also of +-- 'Numeric.Backprop.Mono.OpB'). So, if a function ever expects an @'OpM' +-- m as a@ (or a 'Numeric.Backprop.Mono.OpB'), you can always provide an +-- @'Op' as a@ instead. +-- +-- Many functions in this library will expect an @'OpM' m as a@ (or +-- an @'Numeric.Backprop.Mono.OpB' s as a@), and in all of these cases, you can +-- provide an @'Op' as a@. +type Op n a b = BP.Op (Replicate n a) b + +-- | An @'OpM' m n a b@ represents a differentiable (monadic) function from +-- @n@ values of type @a@ to a value of type @b@. +-- +-- For example, an +-- +-- @ +-- 'OpM' IO 'N2' Int Double +-- @ +-- +-- would be a function that takes two 'Int's and returns a 'Double' (in +-- 'IO'). It can be differentiated to give a /gradient/ of the two input +-- 'Int's (also in 'IO') if given the total derivative for @a@. +-- +-- Note that an 'OpM' is a /superclass/ of 'Op', so any function that +-- expects an @'OpM' m as a@ can also accept an @'Op' as a@. +-- +-- See 'runOpM', 'gradOpM', and 'gradOpWithM' for examples on how to run +-- it. +type OpM m n a = BP.OpM m (Replicate n a) + +-- | Construct an 'Op' by giving a function creating the result, and also +-- a continuation on how to create the gradient, given the total derivative +-- of @a@. +-- +-- See the module documentation for "Numeric.Backprop.Op" for more details +-- on the function that this constructor and 'OpM' expect. +pattern Op :: Known Nat n => (Vec n a -> (b, Maybe b -> Vec n a)) -> Op n a b +pattern Op runOp' <- BP.Op (\f xs -> (second . fmap) (prodAlong xs) + . f + . vecToProd + $ xs + -> runOp' + ) + where + Op f = BP.Op (\xs -> (second . fmap) vecToProd . f . prodToVec' known $ xs) + +-- | Construct an 'OpM' by giving a (monadic) function creating the result, +-- and also a continuation on how to create the gradient, given the total +-- derivative of @a@. +-- +-- See the module documentation for "Numeric.Backprop.Op" for more details +-- on the function that this constructor and 'Op' expect. +pattern OpM :: (Known Nat n, Functor m) => (Vec n a -> m (b, Maybe b -> m (Vec n a))) -> OpM m n a b +pattern OpM runOpM' <- BP.OpM (\f xs -> (fmap . second . fmap . fmap) (prodAlong xs) + . f + . vecToProd + $ xs + -> runOpM' + ) + where + OpM f = BP.OpM (\xs -> (fmap . second . fmap . fmap) vecToProd . f . prodToVec' known $ xs) + +-- | Create an 'Op' that takes no inputs and always returns the given +-- value. +-- +-- There is no gradient, of course (using 'gradOp' will give you an empty +-- vector), because there is no input to have a gradient of. +-- +-- >>> gradOp' (op0 10) ØV +-- (10, ØV) +-- +-- For a constant 'Op' that takes input and ignores it, see 'opConst'. +-- +-- Note that because this returns an 'Op', it can be used with any function +-- that expects an 'OpM' or 'Numeric.Backprop.Mono.OpB', as well. +op0 :: a -> Op N0 b a +op0 x = BP.op0 x + +-- | An 'Op' that ignores all of its inputs and returns a given constant +-- value. +-- +-- >>> gradOp' (opConst 10) (1 :+ 2 :+ 3 :+ ØV) +-- (10, 0 :+ 0 :+ 0 :+ ØV) +opConst :: forall n a b. (Known Nat n, Num b) => a -> Op n b a +opConst x = BP.opConst' (BP.nSummers' @n @b known) x + +-- | Automatically create an 'Op' of a numerical function taking one +-- argument. Uses 'Numeric.AD.diff', and so can take any numerical +-- function polymorphic over the standard numeric types. +-- +-- >>> gradOp' (op1 (recip . negate)) (5 :+ ØV) +-- (-0.2, 0.04 :+ ØV) +op1 :: Num a + => (forall s. AD s (Forward a) -> AD s (Forward a)) + -> Op N1 a a +op1 f = BP.op1 f + +-- | Automatically create an 'Op' of a numerical function taking two +-- arguments. Uses 'Numeric.AD.grad', and so can take any numerical function +-- polymorphic over the standard numeric types. +-- +-- >>> gradOp' (op2 (\x y -> x * sqrt y)) (3 :+ 4 :+ ØV) +-- (6.0, 2.0 :+ 0.75 :+ ØV) +op2 :: Num a + => (forall s. Reifies s Tape => Reverse s a -> Reverse s a -> Reverse s a) + -> Op N2 a a +op2 = BP.op2 + +-- | Automatically create an 'Op' of a numerical function taking three +-- arguments. Uses 'Numeric.AD.grad', and so can take any numerical function +-- polymorphic over the standard numeric types. +-- +-- >>> gradOp' (op3 (\x y z -> (x * sqrt y)**z)) (3 :+ 4 :+ 2 :+ ØV) +-- (36.0, 24.0 :+ 9.0 :+ 64.503 :+ ØV) +op3 :: Num a + => (forall s. Reifies s Tape => Reverse s a -> Reverse s a -> Reverse s a -> Reverse s a) + -> Op N3 a a +op3 = BP.op3 + +-- | Automatically create an 'Op' of a numerical function taking multiple +-- arguments. Uses 'Numeric.AD.grad', and so can take any numerical +-- function polymorphic over the standard numeric types. +-- +-- >>> gradOp' (opN (\(x :+ y :+ Ø) -> x * sqrt y)) (3 :+ 4 :+ ØV) +-- (6.0, 2.0 :+ 0.75 :+ ØV) +opN :: (Num a, Known Nat n) + => (forall s. Reifies s Tape => Vec n (Reverse s a) -> Reverse s a) + -> Op n a a +opN = BP.opN + +-- | Create an 'Op' of a function taking one input, by giving its explicit +-- derivative. The function should return a tuple containing the result of +-- the function, and also a function taking the derivative of the result +-- and return the derivative of the input. +-- +-- If we have +-- +-- \[ +-- \eqalign{ +-- f &: \mathbb{R} \rightarrow \mathbb{R}\cr +-- y &= f(x)\cr +-- z &= g(y) +-- } +-- \] +-- +-- Then the derivative \( \frac{dz}{dx} \), it would be: +-- +-- \[ +-- \frac{dz}{dx} = \frac{dz}{dy} \frac{dy}{dx} +-- \] +-- +-- If our 'Op' represents \(f\), then the second item in the resulting +-- tuple should be a function that takes \(\frac{dz}{dy}\) and returns +-- \(\frac{dz}{dx}\). +-- +-- If the input is 'Nothing', then \(\frac{dz}{dy}\) should be taken to be +-- \(1\). +-- +-- As an example, here is an 'Op' that squares its input: +-- +-- @ +-- square :: Num a => 'Op' 'N1' a a +-- square = 'op1'' $ \\x -> (x*x, \\case Nothing -> 2 * x +-- Just d -> 2 * d * x +-- ) +-- @ +-- +-- Remember that, generally, end users shouldn't directly construct 'Op's; +-- they should be provided by libraries or generated automatically. +-- +-- For numeric functions, single-input 'Op's can be generated automatically +-- using 'op1'. +op1' + :: (a -> (b, Maybe b -> a)) + -> Op N1 a b +op1' = BP.op1' + +-- | Create an 'Op' of a function taking two inputs, by giving its explicit +-- gradient. The function should return a tuple containing the result of +-- the function, and also a function taking the derivative of the result +-- and return the derivative of the input. +-- +-- If we have +-- +-- \[ +-- \eqalign{ +-- f &: \mathbb{R}^2 \rightarrow \mathbb{R}\cr +-- z &= f(x, y)\cr +-- k &= g(z) +-- } +-- \] +-- +-- Then the gradient \( \left< \frac{\partial k}{\partial x}, \frac{\partial k}{\partial y} \right> \) +-- would be: +-- +-- \[ +-- \left< \frac{\partial k}{\partial x}, \frac{\partial k}{\partial y} \right> = +-- \left< \frac{dk}{dz} \frac{\partial z}{dx}, \frac{dk}{dz} \frac{\partial z}{dy} \right> +-- \] +-- +-- If our 'Op' represents \(f\), then the second item in the resulting +-- tuple should be a function that takes \(\frac{dk}{dz}\) and returns +-- \( \left< \frac{\partial k}{dx}, \frac{\partial k}{dx} \right> \). +-- +-- If the input is 'Nothing', then \(\frac{dk}{dz}\) should be taken to be +-- \(1\). +-- +-- As an example, here is an 'Op' that multiplies its inputs: +-- +-- @ +-- mul :: Num a => 'Op' 'N2' a a +-- mul = 'op2'' $ \\x y -> (x*y, \\case Nothing -> (y , x ) +-- Just d -> (d*y, x*d) +-- ) +-- @ +-- +-- Remember that, generally, end users shouldn't directly construct 'Op's; +-- they should be provided by libraries or generated automatically. +-- +-- For numeric functions, two-input 'Op's can be generated automatically +-- using 'op2'. +op2' + :: (a -> a -> (b, Maybe b -> (a, a))) + -> Op N2 a b +op2' = BP.op2' + +-- | Create an 'Op' of a function taking three inputs, by giving its explicit +-- gradient. See documentation for 'op2'' for more details. +op3' + :: (a -> a -> a -> (b, Maybe b -> (a, a, a))) + -> Op N3 a b +op3' = BP.op3' + +-- | A combination of 'runOp' and 'gradOpWith''. Given an 'Op' and inputs, +-- returns the result of the 'Op' and a continuation that gives its +-- gradient. +-- +-- The continuation takes the total derivative of the result as input. See +-- documenation for 'gradOpWith'' and module documentation for +-- "Numeric.Backprop.Op" for more information. +runOp' :: Op n a b -> Vec n a -> (b, Maybe b -> Vec n a) +runOp' o xs = (second . fmap) (prodAlong xs) + . BP.runOp' o + . vecToProd + $ xs + +-- | Run the function that an 'Op' encodes, to get the result. +-- +-- >>> runOp (op2 (*)) (3 :+ 5 :+ Ø) +-- 15 +runOp :: Op n a b -> Vec n a -> b +runOp o = fst . runOp' o + +-- | A combination of 'gradOp' and 'gradOpWith'. The third argument is +-- (optionally) the total derivative the result. Give 'Nothing' and it is +-- assumed that the result is the final result (and the total derivative is +-- 1), and this behaves the same as 'gradOp'. Give @'Just' d@ and it uses +-- the @d@ as the total derivative of the result, and this behaves like +-- 'gradOpWith'. +-- +-- See 'gradOp' and the module documentaiton for "Numeric.Backprop.Op" for +-- more information. +gradOpWith' :: Op n a b -> Vec n a -> Maybe b -> Vec n a +gradOpWith' o = snd . runOp' o + +-- | Run the function that an 'Op' encodes, and get the gradient of +-- a "final result" with respect to the inputs, given the total derivative +-- of the output with the final result. +-- +-- See 'gradOp' and the module documentaiton for "Numeric.Backprop.Op" for +-- more information. +gradOpWith :: Op n a b -> Vec n a -> b -> Vec n a +gradOpWith o i = gradOpWith' o i . Just + +-- | Run the function that an 'Op' encodes, and get the gradient of the +-- output with respect to the inputs. +-- +-- >>> gradOp (op2 (*)) (3 :+ 5 :+ ØV) +-- 5 :+ 3 :+ ØV +-- -- the gradient of x*y is (y, x) +gradOp :: Op n a b -> Vec n a -> Vec n a +gradOp o i = gradOpWith' o i Nothing + +-- | Run the function that an 'Op' encodes, to get the resulting output and +-- also its gradient with respect to the inputs. +-- +-- >>> gradOpM' (op2 (*)) (3 :+ 5 :+ ØV) :: IO (Int, Vec N2 Int) +-- (15, 5 :+ 3 :+ ØV) +gradOp' :: Op n a b -> Vec n a -> (b, Vec n a) +gradOp' o = second ($ Nothing) . runOp' o + +-- | The monadic version of 'runOp', for 'OpM's. +-- +-- >>> runOpM (op2 (*)) (3 :+ 5 :+ ØV) :: IO Int +-- 15 +runOpM' :: Functor m => OpM m n a b -> Vec n a -> m (b, Maybe b -> m (Vec n a)) +runOpM' o xs = (fmap . second . fmap . fmap) (prodAlong xs) + . BP.runOpM' o + . vecToProd + $ xs + +-- | The monadic version of 'runOp', for 'OpM's. +-- +-- >>> runOpM (op2 (*)) (3 :+ 5 :+ ØV) :: IO Int +-- 15 +runOpM :: Functor m => OpM m n a b -> Vec n a -> m b +runOpM o = fmap fst . runOpM' o + +-- | The monadic version of 'gradOp', for 'OpM's. +gradOpM :: Monad m => OpM m n a b -> Vec n a -> m (Vec n a) +gradOpM o i = do + (_, gF) <- runOpM' o i + gF Nothing + +-- | The monadic version of 'gradOp'', for 'OpM's. +gradOpM' :: Monad m => OpM m n a b -> Vec n a -> m (b, Vec n a) +gradOpM' o i = do + (x, gF) <- runOpM' o i + g <- gF Nothing + return (x, g) + +-- | The monadic version of 'gradOpWith'', for 'OpM's. +gradOpWithM' :: Monad m => OpM m n a b -> Vec n a -> Maybe b -> m (Vec n a) +gradOpWithM' o i d = do + (_, gF) <- runOpM' o i + gF d + +-- | The monadic version of 'gradOpWith', for 'OpM's. +gradOpWithM :: Monad m => OpM m n a b -> Vec n a -> b -> m (Vec n a) +gradOpWithM o i d = do + (_, gF) <- runOpM' o i + gF (Just d) + +-- | Compose 'OpM's together, similar to '.'. But, because all 'OpM's are +-- \(\mathbb{R}^N \rightarrow \mathbb{R}\), this is more like 'sequence' +-- for functions, or @liftAN@. +-- +-- That is, given an @o@ of @'OpM' m n a b@s, it can compose them with an +-- @'OpM' m o b c@ to create an @'OpM' m o a c@. +composeOp + :: forall m n o a b c. (Monad m, Num a, Known Nat n) + => VecT o (OpM m n a) b + -> OpM m o b c + -> OpM m n a c +composeOp v o = BP.composeOp' (BP.nSummers' @n @a known) (vecToProd v) o + +-- | Convenient wrappver over 'composeOp' for the case where the second +-- function only takes one input, so the two 'OpM's can be directly piped +-- together, like for '.'. +composeOp1 + :: forall m n a b c. (Monad m, Num a, Known Nat n) + => OpM m n a b + -> OpM m N1 b c + -> OpM m n a c +composeOp1 v o = composeOp @_ @_ @_ @a (v :* ØV) o + +-- | Convenient infix synonym for (flipped) 'composeOp1'. Meant to be used +-- just like '.': +-- +-- @ +-- 'op1' negate :: 'Op' '[a] a +-- 'op2' (+) :: Op '[a,a] a +-- +-- op1 negate '~.' op2 (+) :: Op '[a, a] a +-- @ +infixr 9 ~. +(~.) + :: forall m n a b c. (Monad m, Num a, Known Nat n) + => OpM m N1 b c + -> OpM m n a b + -> OpM m n a c +f ~. g = composeOp1 @_ @_ @a g f