diff --git a/Build.hs b/Build.hs --- a/Build.hs +++ b/Build.hs @@ -1,164 +1,182 @@ #!/usr/bin/env stack -- stack --install-ghc runghc --package shake-0.16.4 --stack-yaml stack.yaml -import Development.Shake -import Development.Shake.FilePath -import System.Directory +import Development.Shake +import Development.Shake.FilePath +import System.Directory -opts = shakeOptions { shakeFiles = ".shake" - , shakeVersion = "1.0" - , shakeVerbosity = Normal - , shakeThreads = 1 - } +opts = + shakeOptions + { shakeFiles = ".shake" + , shakeVersion = "1.0" + , shakeVerbosity = Normal + , shakeThreads = 1 + } data Doc = Lab main :: IO () -main = getDirectoryFilesIO "samples" ["/*.lhs", "/*.hs"] >>= \allSamps -> - getDirectoryFilesIO "src" ["//*.hs"] >>= \allSrc -> - shakeArgs opts $ do - - want ["all"] - - "all" ~> - need ["pdf", "md", "gentags", "install", "exe"] +main = + getDirectoryFilesIO "samples" ["/*.lhs", "/*.hs"] >>= \allSamps -> + getDirectoryFilesIO "src" ["//*.hs"] >>= \allSrc -> + shakeArgs opts $ do + want ["all"] - "pdf" ~> - need [ "renders" takeFileName f -<.> ".pdf" - | f <- allSamps, takeExtension f == ".lhs" - ] + "all" + ~> need ["pdf", "md", "gentags", "install", "exe"] - "md" ~> - need [ "renders" takeFileName f -<.> ".md" - | f <- allSamps, takeExtension f == ".lhs" - ] + "pdf" + ~> need + [ "renders" takeFileName f -<.> ".pdf" + | f <- allSamps + , takeExtension f == ".lhs" + ] - "exe" ~> - need (map (\f -> "samples-exe" dropExtension f) allSamps) + "md" + ~> need + [ "renders" takeFileName f -<.> ".md" + | f <- allSamps + , takeExtension f == ".lhs" + ] - "haddocks" ~> do - need $ ("src" ) <$> allSrc - cmd "jle-git-haddocks" + "exe" + ~> need (map (\f -> "samples-exe" dropExtension f) allSamps) - "install" ~> do - need $ ("src" ) <$> allSrc - cmd "stack install" + "haddocks" ~> do + need $ ("src" ) <$> allSrc + cmd "jle-git-haddocks" - "install-profile" ~> do - need $ ("src" ) <$> allSrc - cmd "stack install --profile" + "install" ~> do + need $ ("src" ) <$> allSrc + cmd "stack install" - "gentags" ~> - need ["tags", "TAGS"] + "install-profile" ~> do + need $ ("src" ) <$> allSrc + cmd "stack install --profile" - ["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" - "-s" - "--highlight-style tango" - "--reference-links" - "--reference-location block" - "-o" f - src + "gentags" + ~> need ["tags", "TAGS"] - "samples-exe/*" %> \f -> do - need ["install"] - [src] <- getDirectoryFiles "samples" $ (takeFileName f <.>) <$> ["hs","lhs"] - liftIO $ do - createDirectoryIfMissing True "samples-exe" - createDirectoryIfMissing True ".build" - removeFilesAfter "samples" ["/*.o"] - cmd "stack ghc" - "--stack-yaml stack.yaml" - "--package finite-typelits" - "--package hmatrix-backprop" - "--package hmatrix-vector-sized" - "--package microlens-th" - "--package mnist-idx" - "--package mwc-random" - "--package one-liner" - "--package one-liner-instances" - "--package random" - "--package singletons" - "--package split" - "--package vector-sized" - "--" - ("samples" src) - "-o" f - "-hidir .build" - "-Wall" - "-O2" + ["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" + "-s" + "--highlight-style tango" + "--reference-links" + "--reference-location block" + "-o" + f + src - "profile" ~> do - need $ do - s <- ["manual","bp-lens","bp-hkd","hybrid"] - e <- ["prof.html","svg"] - return $ "bench-prof/bench-" ++ s <.> e + "samples-exe/*" %> \f -> do + need ["install"] + [src] <- getDirectoryFiles "samples" $ (takeFileName f <.>) <$> ["hs", "lhs"] + liftIO $ do + createDirectoryIfMissing True "samples-exe" + createDirectoryIfMissing True ".build" + removeFilesAfter "samples" ["/*.o"] + cmd + "stack ghc" + "--stack-yaml stack.yaml" + "--package finite-typelits" + "--package hmatrix-backprop" + "--package hmatrix-vector-sized" + "--package microlens-th" + "--package mnist-idx" + "--package mwc-random" + "--package one-liner" + "--package one-liner-instances" + "--package random" + "--package singletons" + "--package split" + "--package vector-sized" + "--" + ("samples" src) + "-o" + f + "-hidir .build" + "-Wall" + "-O2" - "bench-prof/bench" %> \f -> do - let src = "bench" takeFileName f <.> ".hs" - need ["install-profile", src] - unit $ cmd "stack install" - "--profile" - "--stack-yaml stack.yaml" - [ "lens" - , "hmatrix" - , "one-liner-instances" - , "split" - , "criterion" - ] - unit $ cmd "stack ghc" - "--profile" - "--stack-yaml stack.yaml" - src - "--" - "-o" f - "-hidir .build" - "-O2" - "-prof" - "-fexternal-interpreter" + "profile" ~> do + need $ do + s <- ["manual", "bp-lens", "bp-hkd", "hybrid"] + e <- ["prof.html", "svg"] + return $ "bench-prof/bench-" ++ s <.> e - "bench-prof/bench-*.prof" %> \f -> do - need ["bench-prof/bench"] - let b = drop 6 $ takeBaseName f - unit $ cmd "./bench-prof/bench" - ("gradient/" ++ b) - "+RTS" - "-p" - cmd "mv" "bench.prof" f + "bench-prof/bench" %> \f -> do + let src = "bench" takeFileName f <.> ".hs" + need ["install-profile", src] + unit $ + cmd + "stack install" + "--profile" + "--stack-yaml stack.yaml" + [ "lens" + , "hmatrix" + , "one-liner-instances" + , "split" + , "criterion" + ] + unit $ + cmd + "stack ghc" + "--profile" + "--stack-yaml stack.yaml" + src + "--" + "-o" + f + "-hidir .build" + "-O2" + "-prof" + "-fexternal-interpreter" - "**/*.prof.html" %> \f -> do - let src = f -<.> "" - need [src] - cmd "profiteur" src + "bench-prof/bench-*.prof" %> \f -> do + need ["bench-prof/bench"] + let b = drop 6 $ takeBaseName f + unit $ + cmd + "./bench-prof/bench" + ("gradient/" ++ b) + "+RTS" + "-p" + cmd "mv" "bench.prof" f - "**/*.prof.folded" %> \f -> do - let src = f -<.> "" - need [src] - Stdout out <- cmd "cat" [src] - cmd (Stdin out) - (FileStdout f) - "ghc-prof-flamegraph" + "**/*.prof.html" %> \f -> do + let src = f -<.> "" + need [src] + cmd "profiteur" src - "bench-prof/*.svg" %> \f -> do - let src = f -<.> "prof.folded" - need [src] - cmd (FileStdout f) - "flamegraph.pl" - "--width 2000" - src + "**/*.prof.folded" %> \f -> do + let src = f -<.> "" + need [src] + Stdout out <- cmd "cat" [src] + cmd + (Stdin out) + (FileStdout f) + "ghc-prof-flamegraph" - ["tags","TAGS"] &%> \_ -> do - need (("src" ) <$> allSrc) - cmd "hasktags" "src/" + "bench-prof/*.svg" %> \f -> do + let src = f -<.> "prof.folded" + need [src] + cmd + (FileStdout f) + "flamegraph.pl" + "--width 2000" + src - "clean" ~> do - unit $ cmd "stack clean" - removeFilesAfter ".shake" ["//*"] - removeFilesAfter ".build" ["//*"] + ["tags", "TAGS"] &%> \_ -> do + need (("src" ) <$> allSrc) + cmd "hasktags" "src/" + "clean" ~> do + unit $ cmd "stack clean" + removeFilesAfter ".shake" ["//*"] + removeFilesAfter ".build" ["//*"] diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,15 @@ Changelog ========= +Version 0.2.7.0 +--------------- + +*June 3, 2025* + + + +* Export `Numeric.Backprop.Internal`. + Version 0.2.6.5 --------------- diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -235,7 +235,7 @@ 3. There is a link between *backprop* and deep learning/neural network libraries like *[tensorflow][]*, *[caffe][]*, and *[theano][]*, which all - all support some form of heterogeneous automatic differentiation. Haskell + support some form of heterogeneous automatic differentiation. Haskell libraries doing similar things include *[grenade][]*. These are all frameworks for working with neural networks or other diff --git a/Setup.hs b/Setup.hs --- a/Setup.hs +++ b/Setup.hs @@ -1,2 +1,3 @@ import Distribution.Simple + main = defaultMain diff --git a/backprop.cabal b/backprop.cabal --- a/backprop.cabal +++ b/backprop.cabal @@ -1,113 +1,111 @@ -cabal-version: 1.12 - --- This file has been generated from package.yaml by hpack version 0.35.2. --- --- see: https://github.com/sol/hpack --- --- hash: b1229ff8088fce05c68ad20c4c853349f1b30d0f9e19aa5b2c594a833b05790a +cabal-version: 1.12 +name: backprop +version: 0.2.7.0 +synopsis: Heterogeneous automatic differentation +description: + Write your functions to compute your result, and the library will + automatically generate functions to compute your gradient. + . + Implements heterogeneous reverse-mode automatic differentiation, commonly + known as "backpropagation". + . + See for official introduction and documentation. -name: backprop -version: 0.2.6.5 -synopsis: Heterogeneous automatic differentation -description: Write your functions to compute your result, and the library will - automatically generate functions to compute your gradient. - . - Implements heterogeneous reverse-mode automatic differentiation, commonly - known as "backpropagation". - . - See for official introduction and documentation. -category: Math -homepage: https://backprop.jle.im -bug-reports: https://github.com/mstksg/backprop/issues -author: Justin Le -maintainer: justin@jle.im -copyright: (c) Justin Le 2018 -license: BSD3 -license-file: LICENSE -tested-with: - GHC >= 8.4 -build-type: Simple +category: Math +homepage: https://backprop.jle.im +bug-reports: https://github.com/mstksg/backprop/issues +author: Justin Le +maintainer: justin@jle.im +copyright: (c) Justin Le 2018 +license: BSD3 +license-file: LICENSE +build-type: Simple +tested-with: GHC >=8.4 extra-source-files: - README.md - CHANGELOG.md - Build.hs - samples/backprop-mnist.lhs - samples/extensible-neural.lhs - renders/backprop-mnist.pdf - renders/extensible-neural.pdf - renders/backprop-mnist.md - renders/extensible-neural.md - doc/01-getting-started.md - doc/02-a-detailed-look.md - doc/03-manipulating-bvars.md - doc/04-the-backprop-typeclass.md - doc/05-applications.md - doc/06-manual-gradients.md - doc/07-performance.md - doc/08-equipping-your-library.md - doc/09-comparisons.md - doc/index.md + Build.hs + CHANGELOG.md + doc/01-getting-started.md + doc/02-a-detailed-look.md + doc/03-manipulating-bvars.md + doc/04-the-backprop-typeclass.md + doc/05-applications.md + doc/06-manual-gradients.md + doc/07-performance.md + doc/08-equipping-your-library.md + doc/09-comparisons.md + doc/index.md + README.md + renders/backprop-mnist.md + renders/backprop-mnist.pdf + renders/extensible-neural.md + renders/extensible-neural.pdf + samples/backprop-mnist.lhs + samples/extensible-neural.lhs source-repository head - type: git + type: git location: https://github.com/mstksg/backprop flag vinyl_0_14 - manual: False + manual: False default: True library exposed-modules: - Numeric.Backprop - Numeric.Backprop.Class - Numeric.Backprop.Explicit - Numeric.Backprop.Num - Numeric.Backprop.Op - Prelude.Backprop - Prelude.Backprop.Explicit - Prelude.Backprop.Num - other-modules: - Numeric.Backprop.Internal - Data.Type.Util - hs-source-dirs: - src - ghc-options: -Wall -Wcompat -Wincomplete-record-updates -Wredundant-constraints + Numeric.Backprop + Numeric.Backprop.Class + Numeric.Backprop.Explicit + Numeric.Backprop.Internal + Numeric.Backprop.Num + Numeric.Backprop.Op + Prelude.Backprop + Prelude.Backprop.Explicit + Prelude.Backprop.Num + + other-modules: Data.Type.Util + hs-source-dirs: src + ghc-options: + -Wall -Wcompat -Wincomplete-record-updates -Wredundant-constraints + -Wunused-packages + build-depends: - base >=4.7 && <5 + base >=4.7 && <5 , containers , deepseq , microlens - , primitive , reflection , transformers , vector - , vinyl >=0.9.1 + , vinyl >=0.9.1 + default-language: Haskell2010 + if flag(vinyl_0_14) - build-depends: - vinyl >=0.14.2 + build-depends: vinyl >=0.14.2 + else - build-depends: - vinyl <0.14 + build-depends: vinyl <0.14 benchmark backprop-mnist-bench - type: exitcode-stdio-1.0 - main-is: bench.hs - other-modules: - Paths_backprop - hs-source-dirs: - bench - ghc-options: -Wall -Wcompat -Wincomplete-record-updates -Wredundant-constraints -threaded -rtsopts -with-rtsopts=-N -O2 + type: exitcode-stdio-1.0 + main-is: bench.hs + other-modules: Paths_backprop + hs-source-dirs: bench + ghc-options: + -Wall -Wcompat -Wincomplete-record-updates -Wredundant-constraints + -threaded -rtsopts -with-rtsopts=-N -O2 -Wunused-packages + build-depends: backprop - , base >=4.7 && <5 + , base >=4.7 && <5 , criterion , deepseq , directory - , hmatrix >=0.18 + , hmatrix >=0.18 , microlens , microlens-th , mwc-random , time , vector + default-language: Haskell2010 diff --git a/bench/bench.hs b/bench/bench.hs --- a/bench/bench.hs +++ b/bench/bench.hs @@ -1,46 +1,46 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeSynonymInstances #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ViewPatterns #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -import Control.DeepSeq -import Criterion.Main -import Criterion.Types -import Data.Char -import Data.Functor.Identity -import Data.Time -import GHC.Generics (Generic) -import GHC.TypeLits -import Lens.Micro -import Lens.Micro.TH -import Numeric.Backprop -import Numeric.Backprop.Class -import Numeric.LinearAlgebra.Static -import System.Directory -import qualified Data.Vector as V -import qualified Numeric.LinearAlgebra as HM -import qualified System.Random.MWC as MWC +import Control.DeepSeq +import Criterion.Main +import Criterion.Types +import Data.Char +import Data.Functor.Identity +import Data.Time +import qualified Data.Vector as V +import GHC.Generics (Generic) +import GHC.TypeLits +import Lens.Micro +import Lens.Micro.TH +import Numeric.Backprop +import Numeric.Backprop.Class +import qualified Numeric.LinearAlgebra as HM +import Numeric.LinearAlgebra.Static +import System.Directory +import qualified System.Random.MWC as MWC type family HKD f a where - HKD Identity a = a - HKD f a = f a + HKD Identity a = a + HKD f a = f a -data Layer' i o f = - Layer { _lWeights :: !(HKD f (L o i)) - , _lBiases :: !(HKD f (R o)) - } +data Layer' i o f + = Layer + { _lWeights :: !(HKD f (L o i)) + , _lBiases :: !(HKD f (R o)) + } deriving (Generic) type Layer i o = Layer' i o Identity @@ -50,11 +50,12 @@ makeLenses ''Layer' -data Network' i h1 h2 o f = - Net { _nLayer1 :: !(HKD f (Layer i h1)) - , _nLayer2 :: !(HKD f (Layer h1 h2)) - , _nLayer3 :: !(HKD f (Layer h2 o )) - } +data Network' i h1 h2 o f + = Net + { _nLayer1 :: !(HKD f (Layer i h1)) + , _nLayer2 :: !(HKD f (Layer h1 h2)) + , _nLayer3 :: !(HKD f (Layer h2 o)) + } deriving (Generic) type Network i h1 h2 o = Network' i h1 h2 o Identity @@ -66,60 +67,65 @@ main :: IO () main = do - g <- MWC.initialize - . V.fromList - . map (fromIntegral . ord) - $ "hello world" - test0 <- MWC.uniformR @(R 784, R 10) ((0,0),(1,1)) g - net0 <- MWC.uniformR @(Network 784 300 100 10) (-0.5, 0.5) g - t <- getZonedTime - let tstr = formatTime defaultTimeLocale "%Y%m%d-%H%M%S" t - createDirectoryIfMissing True "bench-results" - defaultMainWith defaultConfig - { reportFile = Just $ "bench-results/mnist-bench_" ++ tstr ++ ".html" - , timeLimit = 10 - } [ - bgroup "gradient" - [ let runTest x y = gradNetManual x y net0 - in bench "manual" $ nf (uncurry runTest) test0 - , let runTest x y = gradBP (netErr x y) net0 - in bench "bp-lens" $ nf (uncurry runTest) test0 - , let runTest x y = gradBP (netErrHKD x y) net0 - in bench "bp-hkd" $ nf (uncurry runTest) test0 - , let runTest x y = gradBP (\n' -> netErrHybrid n' y x) net0 - in bench "hybrid" $ nf (uncurry runTest) test0 - ] - , bgroup "descent" - [ let runTest x y = trainStepManual 0.02 x y net0 - in bench "manual" $ nf (uncurry runTest) test0 - , let runTest x y = trainStep 0.02 x y net0 - in bench "bp-lens" $ nf (uncurry runTest) test0 - , let runTest x y = trainStepHKD 0.02 x y net0 - in bench "bp-hkd" $ nf (uncurry runTest) test0 - , let runTest x y = trainStepHybrid 0.02 x y net0 - in bench "hybrid" $ nf (uncurry runTest) test0 - ] - , bgroup "run" - [ let runTest = runNetManual net0 - in bench "manual" $ nf runTest (fst test0) - , let runTest x = evalBP (`runNetwork` x) net0 - in bench "bp-lens" $ nf runTest (fst test0) - , let runTest x = evalBP (`runNetworkHKD` x) net0 - in bench "bp-hkd" $ nf runTest (fst test0) - , let runTest x = evalBP (`runNetHybrid` x) net0 - in bench "hybrid" $ nf runTest (fst test0) - ] - ] + g <- + MWC.initialize + . V.fromList + . map (fromIntegral . ord) + $ "hello world" + test0 <- MWC.uniformR @(R 784, R 10) ((0, 0), (1, 1)) g + net0 <- MWC.uniformR @(Network 784 300 100 10) (-0.5, 0.5) g + t <- getZonedTime + let tstr = formatTime defaultTimeLocale "%Y%m%d-%H%M%S" t + createDirectoryIfMissing True "bench-results" + defaultMainWith + defaultConfig + { reportFile = Just $ "bench-results/mnist-bench_" ++ tstr ++ ".html" + , timeLimit = 10 + } + [ bgroup + "gradient" + [ let runTest x y = gradNetManual x y net0 + in bench "manual" $ nf (uncurry runTest) test0 + , let runTest x y = gradBP (netErr x y) net0 + in bench "bp-lens" $ nf (uncurry runTest) test0 + , let runTest x y = gradBP (netErrHKD x y) net0 + in bench "bp-hkd" $ nf (uncurry runTest) test0 + , let runTest x y = gradBP (\n' -> netErrHybrid n' y x) net0 + in bench "hybrid" $ nf (uncurry runTest) test0 + ] + , bgroup + "descent" + [ let runTest x y = trainStepManual 0.02 x y net0 + in bench "manual" $ nf (uncurry runTest) test0 + , let runTest x y = trainStep 0.02 x y net0 + in bench "bp-lens" $ nf (uncurry runTest) test0 + , let runTest x y = trainStepHKD 0.02 x y net0 + in bench "bp-hkd" $ nf (uncurry runTest) test0 + , let runTest x y = trainStepHybrid 0.02 x y net0 + in bench "hybrid" $ nf (uncurry runTest) test0 + ] + , bgroup + "run" + [ let runTest = runNetManual net0 + in bench "manual" $ nf runTest (fst test0) + , let runTest x = evalBP (`runNetwork` x) net0 + in bench "bp-lens" $ nf runTest (fst test0) + , let runTest x = evalBP (`runNetworkHKD` x) net0 + in bench "bp-hkd" $ nf runTest (fst test0) + , let runTest x = evalBP (`runNetHybrid` x) net0 + in bench "hybrid" $ nf runTest (fst test0) + ] + ] -- ------------------------------ -- - "Backprop" Lens Mode - -- ------------------------------ -runLayer - :: (KnownNat i, KnownNat o, Reifies s W) - => BVar s (Layer i o) - -> BVar s (R i) - -> BVar s (R o) +runLayer :: + (KnownNat i, KnownNat o, Reifies s W) => + BVar s (Layer i o) -> + BVar s (R i) -> + BVar s (R o) runLayer l x = (l ^^. lWeights) #>! x + (l ^^. lBiases) {-# INLINE runLayer #-} @@ -129,44 +135,46 @@ expx = exp x {-# INLINE softMax #-} -runNetwork - :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) - => BVar s (Network i h1 h2 o) - -> R i - -> BVar s (R o) -runNetwork n = softMax - . runLayer (n ^^. nLayer3) - . logistic - . runLayer (n ^^. nLayer2) - . logistic - . runLayer (n ^^. nLayer1) - . auto +runNetwork :: + (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) => + BVar s (Network i h1 h2 o) -> + R i -> + BVar s (R o) +runNetwork n = + softMax + . runLayer (n ^^. nLayer3) + . logistic + . runLayer (n ^^. nLayer2) + . logistic + . runLayer (n ^^. nLayer1) + . auto {-# INLINE runNetwork #-} -crossEntropy - :: (KnownNat n, Reifies s W) - => R n - -> BVar s (R n) - -> BVar s Double +crossEntropy :: + (KnownNat n, Reifies s W) => + R n -> + BVar s (R n) -> + BVar s Double crossEntropy t r = negate $ log r <.>! auto t {-# INLINE crossEntropy #-} -netErr - :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) - => R i - -> R o - -> BVar s (Network i h1 h2 o) - -> BVar s Double +netErr :: + (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) => + R i -> + R o -> + BVar s (Network i h1 h2 o) -> + BVar s Double netErr x t n = crossEntropy t (runNetwork n x) {-# INLINE netErr #-} -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 :: + 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 = n - realToFrac r * gradBP (netErr x t) n {-# INLINE trainStep #-} @@ -174,44 +182,46 @@ -- - "Backprop" HKD Mode - -- ------------------------------ -runLayerHKD - :: (KnownNat i, KnownNat o, Reifies s W) - => BVar s (Layer i o) - -> BVar s (R i) - -> BVar s (R o) -runLayerHKD (splitBV->Layer w b) x = w #>! x + b +runLayerHKD :: + (KnownNat i, KnownNat o, Reifies s W) => + BVar s (Layer i o) -> + BVar s (R i) -> + BVar s (R o) +runLayerHKD (splitBV -> Layer w b) x = w #>! x + b {-# INLINE runLayerHKD #-} -runNetworkHKD - :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) - => BVar s (Network i h1 h2 o) - -> R i - -> BVar s (R o) -runNetworkHKD (splitBV->Net l1 l2 l3) = softMax - . runLayerHKD l3 - . logistic - . runLayerHKD l2 - . logistic - . runLayerHKD l1 - . auto +runNetworkHKD :: + (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) => + BVar s (Network i h1 h2 o) -> + R i -> + BVar s (R o) +runNetworkHKD (splitBV -> Net l1 l2 l3) = + softMax + . runLayerHKD l3 + . logistic + . runLayerHKD l2 + . logistic + . runLayerHKD l1 + . auto {-# INLINE runNetworkHKD #-} -netErrHKD - :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) - => R i - -> R o - -> BVar s (Network i h1 h2 o) - -> BVar s Double +netErrHKD :: + (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) => + R i -> + R o -> + BVar s (Network i h1 h2 o) -> + BVar s Double netErrHKD x t n = crossEntropy t (runNetworkHKD n x) {-# INLINE netErrHKD #-} -trainStepHKD - :: 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 +trainStepHKD :: + 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 trainStepHKD r !x !t !n = n - realToFrac r * gradBP (netErrHKD x t) n {-# INLINE trainStepHKD #-} @@ -219,11 +229,11 @@ -- - "Manual" Mode - -- ------------------------------ -runLayerManual - :: (KnownNat i, KnownNat o) - => Layer i o - -> R i - -> R o +runLayerManual :: + (KnownNat i, KnownNat o) => + Layer i o -> + R i -> + R o runLayerManual l x = (l ^. lWeights) #> x + (l ^. lBiases) {-# INLINE runLayerManual #-} @@ -233,69 +243,72 @@ expx = exp x {-# INLINE softMaxManual #-} -runNetManual - :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) - => Network i h1 h2 o - -> R i - -> R o -runNetManual n = softMaxManual - . runLayerManual (n ^. nLayer3) - . logistic - . runLayerManual (n ^. nLayer2) - . logistic - . runLayerManual (n ^. nLayer1) +runNetManual :: + (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => + Network i h1 h2 o -> + R i -> + R o +runNetManual n = + softMaxManual + . runLayerManual (n ^. nLayer3) + . logistic + . runLayerManual (n ^. nLayer2) + . logistic + . runLayerManual (n ^. nLayer1) {-# INLINE runNetManual #-} -gradNetManual - :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) - => R i - -> R o - -> Network i h1 h2 o - -> Network i h1 h2 o +gradNetManual :: + forall i h1 h2 o. + (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => + R i -> + R o -> + Network i h1 h2 o -> + Network i h1 h2 o gradNetManual x t (Net (Layer w1 b1) (Layer w2 b2) (Layer w3 b3)) = - let y1 = w1 #> x - z1 = y1 + b1 - x2 = logistic z1 - y2 = w2 #> x2 - z2 = y2 + b2 - x3 = logistic z2 - y3 = w3 #> x3 - z3 = y3 + b3 - o0 = exp z3 - o1 = HM.sumElements (extract o0) - o2 = o0 / konst o1 - -- o3 = - (log o2 <.> t) - dEdO3 = 1 - dEdO2 = dEdO3 * (- t / o2) - dEdO1 = - (dEdO2 <.> o0) / (o1 ** 2) - dEdO0 = konst dEdO1 + dEdO2 / konst o1 - dEdZ3 = dEdO0 * o0 - dEdY3 = dEdZ3 - dEdX3 = tr w3 #> dEdY3 - dEdZ2 = dEdX3 * (x3 * (1 - x3)) - dEdY2 = dEdZ2 - dEdX2 = tr w2 #> dEdY2 - dEdZ1 = dEdX2 * (x2 * (1 - x2)) - dEdY1 = dEdZ1 - dEdB3 = dEdZ3 - dEdW3 = dEdY3 `outer` x3 - dEdB2 = dEdZ2 - dEdW2 = dEdY2 `outer` x2 - dEdB1 = dEdZ1 - dEdW1 = dEdY1 `outer` x - in Net (Layer dEdW1 dEdB1) (Layer dEdW2 dEdB2) (Layer dEdW3 dEdB3) + let y1 = w1 #> x + z1 = y1 + b1 + x2 = logistic z1 + y2 = w2 #> x2 + z2 = y2 + b2 + x3 = logistic z2 + y3 = w3 #> x3 + z3 = y3 + b3 + o0 = exp z3 + o1 = HM.sumElements (extract o0) + o2 = o0 / konst o1 + -- o3 = - (log o2 <.> t) + dEdO3 = 1 + dEdO2 = - (dEdO3 * t / o2) + dEdO1 = -((dEdO2 <.> o0) / (o1 ** 2)) + dEdO0 = konst dEdO1 + dEdO2 / konst o1 + dEdZ3 = dEdO0 * o0 + dEdY3 = dEdZ3 + dEdX3 = tr w3 #> dEdY3 + dEdZ2 = dEdX3 * (x3 * (1 - x3)) + dEdY2 = dEdZ2 + dEdX2 = tr w2 #> dEdY2 + dEdZ1 = dEdX2 * (x2 * (1 - x2)) + dEdY1 = dEdZ1 + dEdB3 = dEdZ3 + dEdW3 = dEdY3 `outer` x3 + dEdB2 = dEdZ2 + dEdW2 = dEdY2 `outer` x2 + dEdB1 = dEdZ1 + dEdW1 = dEdY1 `outer` x + in Net (Layer dEdW1 dEdB1) (Layer dEdW2 dEdB2) (Layer dEdW3 dEdB3) {-# INLINE gradNetManual #-} -trainStepManual - :: 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 +trainStepManual :: + 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 trainStepManual r !x !t !n = - let gN = gradNetManual x t n - in n - (realToFrac r * gN) + let gN = gradNetManual x t n + in n - (realToFrac r * gN) -- ------------------------------ -- - "Hybrid" Mode - @@ -303,84 +316,87 @@ layerOp :: (KnownNat i, KnownNat o) => Op '[Layer i o, R i] (R o) layerOp = op2 $ \(Layer w b) x -> - ( w #> x + b - , \g -> (Layer (g `outer` x) g, tr w #> g) - ) + ( w #> x + b + , \g -> (Layer (g `outer` x) g, tr w #> g) + ) {-# INLINE layerOp #-} -logisticOp - :: Floating a - => Op '[a] a +logisticOp :: + Floating a => + Op '[a] a logisticOp = op1 $ \x -> - let lx = logistic x - in (lx, \g -> lx * (1 - lx) * g) + let lx = logistic x + in (lx, \g -> lx * (1 - lx) * g) {-# INLINE logisticOp #-} -softMaxOp - :: KnownNat n - => Op '[R n] (R n) +softMaxOp :: + KnownNat n => + Op '[R n] (R n) softMaxOp = op1 $ \x -> - let expx = exp x - tot = sumElements expx - invtot = 1 / tot - res = konst invtot * expx - in ( res - , \g -> res - konst (invtot ** 2) * exp (2 * x) * g - ) + let expx = exp x + tot = sumElements expx + invtot = 1 / tot + res = konst invtot * expx + in ( res + , \g -> res - konst (invtot ** 2) * exp (2 * x) * g + ) {-# INLINE softMaxOp #-} -softMaxCrossEntropyOp - :: KnownNat n - => R n - -> Op '[R n] Double +softMaxCrossEntropyOp :: + KnownNat n => + R n -> + Op '[R n] Double softMaxCrossEntropyOp targ = op1 $ \x -> - let expx = exp x - sm = konst (1 / sumElements expx) * expx - ce = negate $ log sm <.> targ - in ( ce - , \g -> (sm - targ) * konst g - ) + let expx = exp x + sm = konst (1 / sumElements expx) * expx + ce = negate $ log sm <.> targ + in ( ce + , \g -> (sm - targ) * konst g + ) {-# INLINE softMaxCrossEntropyOp #-} -runNetHybrid - :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) - => BVar s (Network i h1 h2 o) - -> R i - -> BVar s (R o) -runNetHybrid n = liftOp1 softMaxOp - . liftOp2 layerOp (n ^^. nLayer3) - . liftOp1 logisticOp - . liftOp2 layerOp (n ^^. nLayer2) - . liftOp1 logisticOp - . liftOp2 layerOp (n ^^. nLayer1) - . auto +runNetHybrid :: + (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) => + BVar s (Network i h1 h2 o) -> + R i -> + BVar s (R o) +runNetHybrid n = + liftOp1 softMaxOp + . liftOp2 layerOp (n ^^. nLayer3) + . liftOp1 logisticOp + . liftOp2 layerOp (n ^^. nLayer2) + . liftOp1 logisticOp + . liftOp2 layerOp (n ^^. nLayer1) + . auto {-# INLINE runNetHybrid #-} -netErrHybrid - :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) - => BVar s (Network i h1 h2 o) - -> R o - -> R i - -> BVar s Double -netErrHybrid n t = liftOp1 (softMaxCrossEntropyOp t) - . liftOp2 layerOp (n ^^. nLayer3) - . liftOp1 logisticOp - . liftOp2 layerOp (n ^^. nLayer2) - . liftOp1 logisticOp - . liftOp2 layerOp (n ^^. nLayer1) - . auto +netErrHybrid :: + (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) => + BVar s (Network i h1 h2 o) -> + R o -> + R i -> + BVar s Double +netErrHybrid n t = + liftOp1 (softMaxCrossEntropyOp t) + . liftOp2 layerOp (n ^^. nLayer3) + . liftOp1 logisticOp + . liftOp2 layerOp (n ^^. nLayer2) + . liftOp1 logisticOp + . liftOp2 layerOp (n ^^. nLayer1) + . auto {-# INLINE netErrHybrid #-} -trainStepHybrid - :: 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 +trainStepHybrid :: + 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 trainStepHybrid r !x !t !n = - let gN = gradBP (\n' -> netErrHybrid n' t x) n - in n - (realToFrac r * gN) + let gN = gradBP (\n' -> netErrHybrid n' t x) n + in n - (realToFrac r * gN) {-# INLINE trainStepHybrid #-} -- ------------------------------ @@ -388,30 +404,31 @@ -- ------------------------------ infixr 8 #>! -(#>!) - :: (KnownNat m, KnownNat n, Reifies s W) - => BVar s (L m n) - -> BVar s (R n) - -> BVar s (R m) +(#>!) :: + (KnownNat m, KnownNat n, Reifies s W) => + BVar s (L m n) -> + BVar s (R n) -> + BVar s (R m) (#>!) = liftOp2 . op2 $ \m v -> - ( m #> v, \g -> (g `outer` v, tr m #> g) ) + (m #> v, \g -> (g `outer` v, tr m #> g)) {-# INLINE (#>!) #-} infixr 8 <.>! -(<.>!) - :: (KnownNat n, Reifies s W) - => BVar s (R n) - -> BVar s (R n) - -> BVar s Double +(<.>!) :: + (KnownNat n, Reifies s W) => + BVar s (R n) -> + BVar s (R n) -> + BVar s Double (<.>!) = liftOp2 . op2 $ \x y -> - ( x <.> y, \g -> (konst g * y, x * konst g) + ( x <.> y + , \g -> (konst g * y, x * konst g) ) {-# INLINE (<.>!) #-} -konst' - :: (KnownNat n, Reifies s W) - => BVar s Double - -> BVar s (R n) +konst' :: + (KnownNat n, Reifies s W) => + BVar s Double -> + BVar s (R n) konst' = liftOp1 . op1 $ \c -> (konst c, HM.sumElements . extract) {-# INLINE konst' #-} @@ -419,10 +436,10 @@ sumElements = HM.sumElements . extract {-# INLINE sumElements #-} -sumElements' - :: (KnownNat n, Reifies s W) - => BVar s (R n) - -> BVar s Double +sumElements' :: + (KnownNat n, Reifies s W) => + BVar s (R n) -> + BVar s Double sumElements' = liftOp1 . op1 $ \x -> (sumElements x, konst) {-# INLINE sumElements' #-} @@ -435,58 +452,58 @@ -- ------------------------------ 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) + 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) + 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) + 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) + 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) 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 + 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 + 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 + 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 + uniform g = Net <$> MWC.uniform g <*> MWC.uniform g <*> MWC.uniform g + uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g instance Backprop (R n) where - zero = zeroNum - add = addNum - one = oneNum + zero = zeroNum + add = addNum + one = oneNum instance (KnownNat n, KnownNat m) => Backprop (L m n) where - zero = zeroNum - add = addNum - one = oneNum + zero = zeroNum + add = addNum + one = oneNum instance (KnownNat i, KnownNat o) => Backprop (Layer i o) instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Backprop (Network i h1 h2 o) diff --git a/src/Data/Type/Util.hs b/src/Data/Type/Util.hs --- a/src/Data/Type/Util.hs +++ b/src/Data/Type/Util.hs @@ -1,64 +1,68 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TupleSections #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeFamilyDependencies #-} -{-# LANGUAGE TypeInType #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} module Data.Type.Util ( - runzipWith - , rzipWithM_ - , Replicate - , VecT(.., (:+)) - , vmap - , withVec - , vecToRec - , fillRec - , zipVecList - , splitRec - , p1, p2, s1, s2 - ) where + runzipWith, + rzipWithM_, + Replicate, + VecT (.., (:+)), + vmap, + withVec, + vecToRec, + fillRec, + zipVecList, + splitRec, + p1, + p2, + s1, + s2, +) where -import Data.Bifunctor -import Data.Functor.Identity -import Data.Kind -import Data.Proxy -import Data.Vinyl.Core -import Data.Vinyl.TypeLevel -import GHC.Generics -import Lens.Micro +import Data.Bifunctor +import Data.Functor.Identity +import Data.Kind +import Data.Proxy +import Data.Vinyl.Core +import Data.Vinyl.TypeLevel +import GHC.Generics +import Lens.Micro -runzipWith - :: forall f g h. () - => (forall x. f x -> (g x, h x)) - -> (forall xs. Rec f xs -> (Rec g xs, Rec h xs)) +runzipWith :: + forall f g h. + () => + (forall x. f x -> (g x, h x)) -> + (forall xs. Rec f xs -> (Rec g xs, Rec h xs)) runzipWith f = go where go :: forall ys. Rec f ys -> (Rec g ys, Rec h ys) go = \case - RNil -> (RNil, RNil) - x :& xs -> let (y , z ) = f x - (ys, zs) = go xs - in (y :& ys, z :& zs) + RNil -> (RNil, RNil) + x :& xs -> + let (y, z) = f x + (ys, zs) = go xs + in (y :& ys, z :& zs) {-# INLINE runzipWith #-} data VecT :: Nat -> (k -> Type) -> k -> Type where - VNil :: VecT 'Z f a - (:*) :: !(f a) -> VecT n f a -> VecT ('S n) f a + VNil :: VecT 'Z f a + (:*) :: !(f a) -> VecT n f a -> VecT ('S n) f a pattern (:+) :: a -> VecT n Identity a -> VecT ('S n) Identity a pattern x :+ xs = Identity x :* xs -vmap - :: forall n f g a. () - => (f a -> g a) -> VecT n f a -> VecT n g a +vmap :: + forall n f g a. + () => + (f a -> g a) -> VecT n f a -> VecT n g a vmap f = go where go :: VecT m f a -> VecT m g a @@ -67,49 +71,51 @@ x :* xs -> f x :* go xs {-# INLINE vmap #-} -withVec - :: [f a] - -> (forall n. VecT n f a -> r) - -> r +withVec :: + [f a] -> + (forall n. VecT n f a -> r) -> + r withVec = \case - [] -> \f -> f VNil - x:xs -> \f -> withVec xs (f . (x :*)) + [] -> \f -> f VNil + x : xs -> \f -> withVec xs (f . (x :*)) {-# INLINE withVec #-} type family Replicate (n :: Nat) (a :: k) = (as :: [k]) | as -> n where - Replicate 'Z a = '[] - Replicate ('S n) a = a ': Replicate n a + Replicate 'Z a = '[] + Replicate ('S n) a = a ': Replicate n a -vecToRec - :: VecT n f a - -> Rec f (Replicate n a) +vecToRec :: + VecT n f a -> + Rec f (Replicate n a) vecToRec = \case - VNil -> RNil - x :* xs -> x :& vecToRec xs + VNil -> RNil + x :* xs -> x :& vecToRec xs {-# INLINE vecToRec #-} -fillRec - :: forall f g as c. () - => (forall a. f a -> c -> g a) - -> Rec f as - -> [c] - -> Maybe (Rec g as) +fillRec :: + forall f g as c. + () => + (forall a. f a -> c -> g a) -> + Rec f as -> + [c] -> + Maybe (Rec g as) fillRec f = go where go :: Rec f bs -> [c] -> Maybe (Rec g bs) go = \case RNil -> \_ -> Just RNil x :& xs -> \case - [] -> Nothing - y:ys -> (f x y :&) <$> go xs ys + [] -> Nothing + y : ys -> (f x y :&) <$> go xs ys {-# INLINE fillRec #-} -rzipWithM_ - :: forall h f g as. Applicative h - => (forall a. f a -> g a -> h ()) - -> Rec f as - -> Rec g as - -> h () +rzipWithM_ :: + forall h f g as. + Applicative h => + (forall a. f a -> g a -> h ()) -> + Rec f as -> + Rec g as -> + h () rzipWithM_ f = go where go :: forall bs. Rec f bs -> Rec g bs -> h () @@ -120,26 +126,28 @@ y :& ys -> f x y *> go xs ys {-# INLINE rzipWithM_ #-} -zipVecList - :: forall a b c f g n. () - => (f a -> Maybe b -> g c) - -> VecT n f a - -> [b] - -> VecT n g c +zipVecList :: + forall a b c f g n. + () => + (f a -> Maybe b -> g c) -> + VecT n f a -> + [b] -> + VecT n g c zipVecList f = go where go :: VecT m f a -> [b] -> VecT m g c go = \case VNil -> const VNil x :* xs -> \case - [] -> f x Nothing :* go xs [] - y:ys -> f x (Just y) :* go xs ys + [] -> f x Nothing :* go xs [] + y : ys -> f x (Just y) :* go xs ys {-# INLINE zipVecList #-} -splitRec - :: forall f as bs. RecApplicative as - => Rec f (as ++ bs) - -> (Rec f as, Rec f bs) +splitRec :: + forall f as bs. + RecApplicative as => + Rec f (as ++ bs) -> + (Rec f as, Rec f bs) splitRec = go (rpure Proxy) where go :: Rec Proxy as' -> Rec f (as' ++ bs) -> (Rec f as', Rec f bs) diff --git a/src/Numeric/Backprop.hs b/src/Numeric/Backprop.hs --- a/src/Numeric/Backprop.hs +++ b/src/Numeric/Backprop.hs @@ -1,10 +1,10 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE DataKinds #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Numeric.Backprop @@ -66,61 +66,110 @@ -- 'zero', 'add', and 'one' explicitly, which can be useful when attempting -- to avoid orphan instances or when mixing both 'Backprop' and 'Num' -- styles. --- - module Numeric.Backprop ( - -- * Types - BVar, W, Backprop(..), ABP(..), NumBP(..) - -- * Running - , backprop, E.evalBP, gradBP, backpropWith - -- ** Multiple inputs - , backprop2, E.evalBP2, gradBP2, backpropWith2 - , backpropN, E.evalBPN, gradBPN, backpropWithN - -- * Manipulating 'BVar' - , E.evalBP0 - , E.constVar, E.auto, E.coerceVar - , (^^.), (.~~), (%~~), (^^?), (^^..), (^^?!) - , viewVar, setVar, overVar - , sequenceVar, collectVar - , previewVar, toListOfVar - , pattern T2, pattern T3 - -- ** With Isomorphisms - , isoVar, isoVar2, isoVar3, isoVarN - -- ** With 'Op's#liftops# - -- $liftops - , liftOp - , liftOp1, liftOp2, liftOp3 - -- ** Generics#hkd# - -- $hkd - , splitBV - , joinBV - , pattern BV - , E.BVGroup - -- * 'Op' - , Op(..) - -- ** Creation - , op0, opConst, idOp - , bpOp - -- *** Giving gradients directly - , op1, op2, op3 - -- *** From Isomorphisms - , opCoerce, opTup, opIso, opIsoN, opLens - -- *** No gradients - , noGrad1, noGrad - -- * Utility - , Reifies - ) where + -- * Types + BVar, + W, + Backprop (..), + ABP (..), + NumBP (..), -import Data.Functor.Identity -import Data.Maybe -import Data.Reflection -import Data.Vinyl -import GHC.Generics -import Lens.Micro -import Numeric.Backprop.Class -import Numeric.Backprop.Explicit (BVar, W) -import Numeric.Backprop.Op + -- * Running + backprop, + E.evalBP, + gradBP, + backpropWith, + + -- ** Multiple inputs + backprop2, + E.evalBP2, + gradBP2, + backpropWith2, + backpropN, + E.evalBPN, + gradBPN, + backpropWithN, + + -- * Manipulating 'BVar' + E.evalBP0, + E.constVar, + E.auto, + E.coerceVar, + (^^.), + (.~~), + (%~~), + (^^?), + (^^..), + (^^?!), + viewVar, + setVar, + overVar, + sequenceVar, + collectVar, + previewVar, + toListOfVar, + pattern T2, + pattern T3, + + -- ** With Isomorphisms + isoVar, + isoVar2, + isoVar3, + isoVarN, + + -- ** With 'Op's#liftops# + -- $liftops + liftOp, + liftOp1, + liftOp2, + liftOp3, + + -- ** Generics#hkd# + -- $hkd + splitBV, + joinBV, + pattern BV, + E.BVGroup, + + -- * 'Op' + Op (..), + + -- ** Creation + op0, + opConst, + idOp, + bpOp, + + -- *** Giving gradients directly + op1, + op2, + op3, + + -- *** From Isomorphisms + opCoerce, + opTup, + opIso, + opIsoN, + opLens, + + -- *** No gradients + noGrad1, + noGrad, + + -- * Utility + Reifies, +) where + +import Data.Functor.Identity +import Data.Maybe +import Data.Reflection +import Data.Vinyl +import GHC.Generics +import Lens.Micro +import Numeric.Backprop.Class +import Numeric.Backprop.Explicit (BVar, W) import qualified Numeric.Backprop.Explicit as E +import Numeric.Backprop.Op -- $liftops -- @@ -186,11 +235,11 @@ -- If you stick to /concerete/, monomorphic usage of this (with specific -- types, typed into source code, known at compile-time), then -- @'RPureConstrained' 'Backprop' as@ should be fulfilled automatically. -backpropN - :: (RPureConstrained Backprop as, Backprop b) - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> (b, Rec Identity as) +backpropN :: + (RPureConstrained Backprop as, Backprop b) => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + (b, Rec Identity as) backpropN = E.backpropN E.zeroFuncs E.oneFunc {-# INLINE backpropN #-} @@ -201,11 +250,11 @@ -- Note that argument order changed in v0.2.4. -- -- @since 0.2.0.0 -backpropWithN - :: RPureConstrained Backprop as - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> (b, b -> Rec Identity as) +backpropWithN :: + RPureConstrained Backprop as => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + (b, b -> Rec Identity as) backpropWithN = E.backpropWithN E.zeroFuncs {-# INLINE backpropWithN #-} @@ -216,11 +265,11 @@ -- that 'BVar's do not leak out of the context (similar to how it is used -- in "Control.Monad.ST"), and also as a reference to an ephemeral Wengert -- tape used to track the graph of references. -backprop - :: (Backprop a, Backprop b) - => (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> (b, a) +backprop :: + (Backprop a, Backprop b) => + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + (b, a) backprop = E.backprop E.zeroFunc E.oneFunc {-# INLINE backprop #-} @@ -241,11 +290,11 @@ -- Note that argument order changed in v0.2.4 -- -- @since 0.2.0.0 -backpropWith - :: Backprop a - => (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> (b, b -> a) +backpropWith :: + Backprop a => + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + (b, b -> a) backpropWith = E.backpropWith E.zeroFunc {-# INLINE backpropWith #-} @@ -263,21 +312,21 @@ -- -- If you want to provide an explicit "final gradient" for the end, see -- 'backpropWith'. -gradBP - :: (Backprop a, Backprop b) - => (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> a +gradBP :: + (Backprop a, Backprop b) => + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + a gradBP = E.gradBP E.zeroFunc E.oneFunc {-# INLINE gradBP #-} -- | 'gradBP' generalized to multiple inputs of different types. See -- documentation for 'backpropN' for more details. -gradBPN - :: (RPureConstrained Backprop as, Backprop b) - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> Rec Identity as +gradBPN :: + (RPureConstrained Backprop as, Backprop b) => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + Rec Identity as gradBPN = E.gradBPN E.zeroFuncs E.oneFunc {-# INLINE gradBPN #-} @@ -288,12 +337,12 @@ -- However, this could potentially be more performant. -- -- For 3 and more arguments, consider using 'backpropN'. -backprop2 - :: (Backprop a, Backprop b, Backprop c) - => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (c, (a, b)) +backprop2 :: + (Backprop a, Backprop b, Backprop c) => + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + (c, (a, b)) backprop2 = E.backprop2 E.zeroFunc E.zeroFunc E.oneFunc {-# INLINE backprop2 #-} @@ -304,22 +353,22 @@ -- Note that argument order changed in v0.2.4 -- -- @since 0.2.0.0 -backpropWith2 - :: (Backprop a, Backprop b) - => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (c, c -> (a, b)) +backpropWith2 :: + (Backprop a, Backprop b) => + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + (c, c -> (a, b)) backpropWith2 = E.backpropWith2 E.zeroFunc E.zeroFunc {-# INLINE backpropWith2 #-} -- | 'gradBP' for a two-argument function. See 'backprop2' for notes. -gradBP2 - :: (Backprop a, Backprop b, Backprop c) - => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (a, b) +gradBP2 :: + (Backprop a, Backprop b, Backprop c) => + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + (a, b) gradBP2 = E.gradBP2 E.zeroFunc E.zeroFunc E.oneFunc {-# INLINE gradBP2 #-} @@ -332,10 +381,10 @@ -- 'liftOp' . 'bpOp' = 'id' -- 'bpOp' . 'liftOp' = 'id' -- @ -bpOp - :: RPureConstrained Backprop as - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Op as b +bpOp :: + RPureConstrained Backprop as => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Op as b bpOp = E.bpOp E.zeroFuncs {-# INLINE bpOp #-} @@ -374,13 +423,14 @@ -- -- __WARNING__: Do not use with any lenses that operate "numerically" on -- the contents (like 'multiplying'). --- -(^^.) - :: forall b a s. (Backprop b, Backprop a, Reifies s W) - => BVar s b - -> Lens' b a - -> BVar s a +(^^.) :: + forall b a s. + (Backprop b, Backprop a, Reifies s W) => + BVar s b -> + Lens' b a -> + BVar s a x ^^. l = viewVar l x + infixl 8 ^^. {-# INLINE (^^.) #-} @@ -388,15 +438,15 @@ -- parallels to 'view' from lens. -- -- See documentation for '^^.' for more information, caveats, and warnings. -viewVar - :: forall b a s. (Backprop a, Backprop b, Reifies s W) - => Lens' b a - -> BVar s b - -> BVar s a +viewVar :: + forall b a s. + (Backprop a, Backprop b, Reifies s W) => + Lens' b a -> + BVar s b -> + BVar s a viewVar = E.viewVar E.addFunc E.zeroFunc {-# INLINE viewVar #-} - -- | An infix version of 'setVar', meant to evoke parallels to '.~' from -- lens. -- @@ -421,14 +471,14 @@ -- -- Note that this does not incurr the performance overhead issues of -- 'viewVar' and '^^.', and is fairly cheap. --- -(.~~) - :: (Backprop a, Backprop b, Reifies s W) - => Lens' b a - -> BVar s a - -> BVar s b - -> BVar s b +(.~~) :: + (Backprop a, Backprop b, Reifies s W) => + Lens' b a -> + BVar s a -> + BVar s b -> + BVar s b l .~~ x = setVar l x + infixl 8 .~~ {-# INLINE (.~~) #-} @@ -436,12 +486,12 @@ -- parallels to "set" from lens. -- -- See documentation for '.~~' for more information. -setVar - :: (Backprop a, Backprop b, Reifies s W) - => Lens' b a - -> BVar s a - -> BVar s b - -> BVar s b +setVar :: + (Backprop a, Backprop b, Reifies s W) => + Lens' b a -> + BVar s a -> + BVar s b -> + BVar s b setVar = E.setVar E.addFunc E.addFunc E.zeroFunc {-# INLINE setVar #-} @@ -470,14 +520,14 @@ -- a 'setVar'. -- -- @since 0.2.4.0 --- -(%~~) - :: (Backprop a, Backprop b, Reifies s W) - => Lens' b a - -> (BVar s a -> BVar s a) - -> BVar s b - -> BVar s b +(%~~) :: + (Backprop a, Backprop b, Reifies s W) => + Lens' b a -> + (BVar s a -> BVar s a) -> + BVar s b -> + BVar s b l %~~ f = overVar l f + infixr 4 %~~ {-# INLINE (%~~) #-} @@ -486,16 +536,15 @@ -- information. -- -- @since 0.2.4.0 -overVar - :: (Backprop a, Backprop b, Reifies s W) - => Lens' b a - -> (BVar s a -> BVar s a) - -> BVar s b - -> BVar s b +overVar :: + (Backprop a, Backprop b, Reifies s W) => + Lens' b a -> + (BVar s a -> BVar s a) -> + BVar s b -> + BVar s b overVar = E.overVar E.addFunc E.addFunc E.zeroFunc E.zeroFunc {-# INLINE overVar #-} - -- | An infix version of 'previewVar', meant to evoke parallels to '^?' -- from lens. -- @@ -528,12 +577,14 @@ -- -- __NOTE__: Has the same potential of performance overhead issues as -- '^^.'; see documentation of '^^.' for more details. -(^^?) - :: forall b a s. (Backprop b, Backprop a, Reifies s W) - => BVar s b - -> Traversal' b a - -> Maybe (BVar s a) +(^^?) :: + forall b a s. + (Backprop b, Backprop a, Reifies s W) => + BVar s b -> + Traversal' b a -> + Maybe (BVar s a) v ^^? t = previewVar t v + infixl 8 ^^? {-# INLINE (^^?) #-} @@ -545,14 +596,16 @@ -- Is essentially '^^?' with 'fromJust', or '^^..' with 'head'. -- -- @since 0.2.1.0 -(^^?!) - :: forall b a s. (Backprop b, Backprop a, Reifies s W) - => BVar s b - -> Traversal' b a - -> BVar s a +(^^?!) :: + forall b a s. + (Backprop b, Backprop a, Reifies s W) => + BVar s b -> + Traversal' b a -> + BVar s a v ^^?! t = fromMaybe (error e) (previewVar t v) where e = "Numeric.Backprop.^^?!: Empty traversal" + infixl 8 ^^?! {-# INLINE (^^?!) #-} @@ -562,11 +615,12 @@ -- to be used wth 'Prism''s, or up-to-one target traversals. -- -- See documentation for '^^?' for more information, warnings, and caveats. -previewVar - :: forall b a s. (Backprop b, Backprop a, Reifies s W) - => Traversal' b a - -> BVar s b - -> Maybe (BVar s a) +previewVar :: + forall b a s. + (Backprop b, Backprop a, Reifies s W) => + Traversal' b a -> + BVar s b -> + Maybe (BVar s a) previewVar = E.previewVar E.addFunc E.zeroFunc {-# INLINE previewVar #-} @@ -593,12 +647,12 @@ -- -- __NOTE__: Has all of the performance overhead issues of 'sequenceVar'; -- see documentation for 'sequenceVar' for more information. --- -(^^..) - :: forall b a s. (Backprop b, Backprop a, Reifies s W) - => BVar s b - -> Traversal' b a - -> [BVar s a] +(^^..) :: + forall b a s. + (Backprop b, Backprop a, Reifies s W) => + BVar s b -> + Traversal' b a -> + [BVar s a] v ^^.. t = toListOfVar t v {-# INLINE (^^..) #-} @@ -607,11 +661,12 @@ -- -- See documentation for '^^..' for more information, warnings, and -- caveats. -toListOfVar - :: forall b a s. (Backprop b, Backprop a, Reifies s W) - => Traversal' b a - -> BVar s b - -> [BVar s a] +toListOfVar :: + forall b a s. + (Backprop b, Backprop a, Reifies s W) => + Traversal' b a -> + BVar s b -> + [BVar s a] toListOfVar = E.toListOfVar E.addFunc E.zeroFunc {-# INLINE toListOfVar #-} @@ -631,10 +686,10 @@ -- with cheap 'add' (like 'Double'), but may be costly for things like -- large matrices. See for for details. -sequenceVar - :: (Traversable t, Backprop a, Reifies s W) - => BVar s (t a) - -> t (BVar s a) +sequenceVar :: + (Traversable t, Backprop a, Reifies s W) => + BVar s (t a) -> + t (BVar s a) sequenceVar = E.sequenceVar E.addFunc E.zeroFunc {-# INLINE sequenceVar #-} @@ -652,10 +707,10 @@ -- a very small constant factor that consistent for all types. This -- reveals a general property of reverse-mode automatic differentiation; -- "many to one" is cheap, but "one to many" is expensive. -collectVar - :: (Foldable t, Functor t, Backprop a, Reifies s W) - => t (BVar s a) - -> BVar s (t a) +collectVar :: + (Foldable t, Functor t, Backprop a, Reifies s W) => + t (BVar s a) -> + BVar s (t a) collectVar = E.collectVar E.addFunc E.zeroFunc {-# INLINE collectVar #-} @@ -668,11 +723,11 @@ -- See "Numeric.Backprop#liftops" and documentation for 'liftOp' for more -- information, and "Numeric.Backprop.Op#prod" for a mini-tutorial on using -- 'Rec'. -liftOp - :: (RPureConstrained Backprop as, Reifies s W) - => Op as b - -> Rec (BVar s) as - -> BVar s b +liftOp :: + (RPureConstrained Backprop as, Reifies s W) => + Op as b -> + Rec (BVar s) as -> + BVar s b liftOp = E.liftOp E.addFuncs {-# INLINE liftOp #-} @@ -683,11 +738,11 @@ -- -- See "Numeric.Backprop#liftops" and documentation for 'liftOp' for more -- information. -liftOp1 - :: (Backprop a, Reifies s W) - => Op '[a] b - -> BVar s a - -> BVar s b +liftOp1 :: + (Backprop a, Reifies s W) => + Op '[a] b -> + BVar s a -> + BVar s b liftOp1 = E.liftOp1 E.addFunc {-# INLINE liftOp1 #-} @@ -698,12 +753,12 @@ -- -- See "Numeric.Backprop#liftops" and documentation for 'liftOp' for more -- information. -liftOp2 - :: (Backprop a, Backprop b, Reifies s W) - => Op '[a,b] c - -> BVar s a - -> BVar s b - -> BVar s c +liftOp2 :: + (Backprop a, Backprop b, Reifies s W) => + Op '[a, b] c -> + BVar s a -> + BVar s b -> + BVar s c liftOp2 = E.liftOp2 E.addFunc E.addFunc {-# INLINE liftOp2 #-} @@ -714,13 +769,13 @@ -- -- See "Numeric.Backprop#liftops" and documentation for 'liftOp' for more -- information. -liftOp3 - :: (Backprop a, Backprop b, Backprop c, Reifies s W) - => Op '[a,b,c] d - -> BVar s a - -> BVar s b - -> BVar s c - -> BVar s d +liftOp3 :: + (Backprop a, Backprop b, Backprop c, Reifies s W) => + Op '[a, b, c] d -> + BVar s a -> + BVar s b -> + BVar s c -> + BVar s d liftOp3 = E.liftOp3 E.addFunc E.addFunc E.addFunc {-# INLINE liftOp3 #-} @@ -736,12 +791,12 @@ -- Basically, don't use this for any "numeric" isomorphisms. -- -- @since 0.1.4.0 -isoVar - :: (Backprop a, Reifies s W) - => (a -> b) - -> (b -> a) - -> BVar s a - -> BVar s b +isoVar :: + (Backprop a, Reifies s W) => + (a -> b) -> + (b -> a) -> + BVar s a -> + BVar s b isoVar = E.isoVar E.addFunc {-# INLINE isoVar #-} @@ -753,13 +808,13 @@ -- join together 'BVar's as their fields. -- -- @since 0.1.4.0 -isoVar2 - :: (Backprop a, Backprop b, Reifies s W) - => (a -> b -> c) - -> (c -> (a, b)) - -> BVar s a - -> BVar s b - -> BVar s c +isoVar2 :: + (Backprop a, Backprop b, Reifies s W) => + (a -> b -> c) -> + (c -> (a, b)) -> + BVar s a -> + BVar s b -> + BVar s c isoVar2 = E.isoVar2 E.addFunc E.addFunc {-# INLINE isoVar2 #-} @@ -767,14 +822,14 @@ -- Useful for things like constructors. See 'isoVar' for caveats. -- -- @since 0.1.4.0 -isoVar3 - :: (Backprop a, Backprop b, Backprop c, Reifies s W) - => (a -> b -> c -> d) - -> (d -> (a, b, c)) - -> BVar s a - -> BVar s b - -> BVar s c - -> BVar s d +isoVar3 :: + (Backprop a, Backprop b, Backprop c, Reifies s W) => + (a -> b -> c -> d) -> + (d -> (a, b, c)) -> + BVar s a -> + BVar s b -> + BVar s c -> + BVar s d isoVar3 = E.isoVar3 E.addFunc E.addFunc E.addFunc {-# INLINE isoVar3 #-} @@ -787,12 +842,12 @@ -- join together 'BVar's as their fields. -- -- @since 0.1.4.0 -isoVarN - :: (RPureConstrained Backprop as, Reifies s W) - => (Rec Identity as -> b) - -> (b -> Rec Identity as) - -> Rec (BVar s) as - -> BVar s b +isoVarN :: + (RPureConstrained Backprop as, Reifies s W) => + (Rec Identity as -> b) -> + (b -> Rec Identity as) -> + Rec (BVar s) as -> + BVar s b isoVarN = E.isoVarN E.addFuncs {-# INLINE isoVarN #-} @@ -800,11 +855,11 @@ -- two-tuples. -- -- @since 0.2.1.0 -pattern T2 - :: (Backprop a, Backprop b, Reifies s W) - => BVar s a - -> BVar s b - -> BVar s (a, b) +pattern T2 :: + (Backprop a, Backprop b, Reifies s W) => + BVar s a -> + BVar s b -> + BVar s (a, b) pattern T2 x y <- (\xy -> (xy ^^. _1, xy ^^. _2) -> (x, y)) where T2 = isoVar2 (,) id @@ -816,12 +871,12 @@ -- three-tuples. -- -- @since 0.2.1.0 -pattern T3 - :: (Backprop a, Backprop b, Backprop c, Reifies s W) - => BVar s a - -> BVar s b - -> BVar s c - -> BVar s (a, b, c) +pattern T3 :: + (Backprop a, Backprop b, Backprop c, Reifies s W) => + BVar s a -> + BVar s b -> + BVar s c -> + BVar s (a, b, c) pattern T3 x y z <- (\xyz -> (xyz ^^. _1, xyz ^^. _2, xyz ^^. _3) -> (x, y, z)) where T3 = isoVar3 (,,) id @@ -953,17 +1008,19 @@ -- interested in helping out! -- -- @since 0.2.2.0 -splitBV - :: ( Generic (z f) - , Generic (z (BVar s)) - , E.BVGroup s as (Rep (z f)) (Rep (z (BVar s))) - , Backprop (z f) - , Backprop (Rep (z f) ()) - , RPureConstrained Backprop as - , Reifies s W - ) - => BVar s (z f) -- ^ 'BVar' of value - -> z (BVar s) -- ^ 'BVar's of fields +splitBV :: + ( Generic (z f) + , Generic (z (BVar s)) + , E.BVGroup s as (Rep (z f)) (Rep (z (BVar s))) + , Backprop (z f) + , Backprop (Rep (z f) ()) + , RPureConstrained Backprop as + , Reifies s W + ) => + -- | 'BVar' of value + BVar s (z f) -> + -- | 'BVar's of fields + z (BVar s) splitBV = E.splitBV E.addFunc E.addFuncs E.zeroFunc E.zeroFuncs {-# INLINE splitBV #-} @@ -987,17 +1044,19 @@ -- differentiation: "many to one" is cheap, but "one to many" is expensive. -- -- @since 0.2.2.0 -joinBV - :: ( Generic (z f) - , Generic (z (BVar s)) - , E.BVGroup s as (Rep (z f)) (Rep (z (BVar s))) - , Backprop (z f) - , Backprop (Rep (z f) ()) - , RPureConstrained Backprop as - , Reifies s W - ) - => z (BVar s) -- ^ 'BVar's of fields - -> BVar s (z f) -- ^ 'BVar' of combined value +joinBV :: + ( Generic (z f) + , Generic (z (BVar s)) + , E.BVGroup s as (Rep (z f)) (Rep (z (BVar s))) + , Backprop (z f) + , Backprop (Rep (z f) ()) + , RPureConstrained Backprop as + , Reifies s W + ) => + -- | 'BVar's of fields + z (BVar s) -> + -- | 'BVar' of combined value + BVar s (z f) joinBV = E.joinBV E.addFunc E.addFuncs E.zeroFunc E.zeroFuncs {-# INLINE joinBV #-} @@ -1005,16 +1064,16 @@ -- is a pattern for a @'BVar' s (z f)@ containing a @z ('BVar' s)@ -- -- @since 0.2.3.0 -pattern BV - :: ( Generic (z f) - , Generic (z (BVar s)) - , E.BVGroup s as (Rep (z f)) (Rep (z (BVar s))) - , Backprop (Rep (z f) ()) - , Backprop (z f) - , RPureConstrained Backprop as - , RecApplicative as - , Reifies s W - ) +pattern BV :: + ( Generic (z f) + , Generic (z (BVar s)) + , E.BVGroup s as (Rep (z f)) (Rep (z (BVar s))) + , Backprop (Rep (z f) ()) + , Backprop (z f) + , RPureConstrained Backprop as + , RecApplicative as + , Reifies s W + ) #if MIN_VERSION_base(4,10,0) => z (BVar s) -- ^ 'BVar's of fields -> BVar s (z f) -- ^ 'BVar' of combined value @@ -1022,7 +1081,7 @@ => z (BVar s) -> BVar s (z f) #endif -pattern BV v <- (splitBV->v) +pattern BV v <- (splitBV -> v) where BV = joinBV #if MIN_VERSION_base(4,10,0) diff --git a/src/Numeric/Backprop/Class.hs b/src/Numeric/Backprop/Class.hs --- a/src/Numeric/Backprop/Class.hs +++ b/src/Numeric/Backprop/Class.hs @@ -1,1074 +1,1146 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE DeriveDataTypeable #-} -{-# LANGUAGE DeriveFoldable #-} -{-# LANGUAGE DeriveFunctor #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE DeriveTraversable #-} -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} - --- | --- Module : Numeric.Backprop.Class --- Copyright : (c) Justin Le 2023 --- License : BSD3 --- --- Maintainer : justin@jle.im --- Stability : experimental --- Portability : non-portable --- --- Provides the 'Backprop' typeclass, a class for values that can be used --- for backpropagation. --- --- This class replaces the old (version 0.1) API relying on 'Num'. --- --- @since 0.2.0.0 - -module Numeric.Backprop.Class ( - -- * Backpropagatable types - Backprop(..) - -- * Derived methods - , zeroNum, addNum, oneNum - , zeroVec, addVec, oneVec, zeroVecNum, oneVecNum - , zeroFunctor, addIsList, addAsList, oneFunctor - , genericZero, genericAdd, genericOne - -- * Newtype - , ABP(..), NumBP(..), NumVec(..) - -- * Generics - , GZero, GAdd, GOne - ) where - -import Control.Applicative -import Control.DeepSeq -import Control.Monad -import Data.Coerce -import Data.Complex -import Data.Data -import Data.Foldable hiding (toList) -import Data.Functor.Compose -import Data.Functor.Identity -import Data.List.NonEmpty (NonEmpty(..)) -import Data.Monoid -import Data.Ratio -import Data.Vinyl -import Data.Vinyl.ARec -import Data.Vinyl.TypeLevel -import Data.Void -import Data.Word -import GHC.Exts -import GHC.Generics -import Numeric.Natural -import qualified Control.Arrow as Arr -import qualified Data.Functor.Product as DFP -import qualified Data.IntMap as IM -import qualified Data.Map as M -import qualified Data.Semigroup as SG -import qualified Data.Sequence as Seq -import qualified Data.Vector as V -import qualified Data.Vector.Generic as VG -import qualified Data.Vector.Primitive as VP -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vinyl.Functor as V -import qualified Data.Vinyl.XRec as V - --- | Class of values that can be backpropagated in general. --- --- For instances of 'Num', these methods can be given by 'zeroNum', --- 'addNum', and 'oneNum'. There are also generic options given in --- "Numeric.Backprop.Class" for functors, 'IsList' instances, and 'Generic' --- instances. --- --- @ --- instance 'Backprop' 'Double' where --- 'zero' = 'zeroNum' --- 'add' = 'addNum' --- 'one' = 'oneNum' --- @ --- --- If you leave the body of an instance declaration blank, GHC Generics --- will be used to derive instances if the type has a single constructor --- and each field is an instance of 'Backprop'. --- --- To ensure that backpropagation works in a sound way, should obey the --- laws: --- --- [/identity/] --- --- * @'add' x ('zero' y) = x@ --- --- * @'add' ('zero' x) y = y@ --- --- Also implies preservation of information, making @'zipWith' ('+')@ an --- illegal implementation for lists and vectors. --- --- This is only expected to be true up to potential "extra zeroes" in @x@ --- and @y@ in the result. --- --- [/commutativity/] --- --- * @'add' x y = 'add' y x@ --- --- [/associativity/] --- --- * @'add' x ('add' y z) = 'add' ('add' x y) z@ --- --- [/idempotence/] --- --- * @'zero' '.' 'zero' = 'zero'@ --- --- * @'one' '.' 'one' = 'one'@ --- --- [/unital/] --- --- * @'one' = 'gradBP' 'id'@ --- --- Note that not all values in the backpropagation process needs all of --- these methods: Only the "final result" needs 'one', for example. These --- are all grouped under one typeclass for convenience in defining --- instances, and also to talk about sensible laws. For fine-grained --- control, use the "explicit" versions of library functions (for example, --- in "Numeric.Backprop.Explicit") instead of 'Backprop' based ones. --- --- This typeclass replaces the reliance on 'Num' of the previous API --- (v0.1). 'Num' is strictly more powerful than 'Backprop', and is --- a stronger constraint on types than is necessary for proper --- backpropagating. In particular, 'fromInteger' is a problem for many --- types, preventing useful backpropagation for lists, variable-length --- vectors (like "Data.Vector") and variable-size matrices from linear --- algebra libraries like /hmatrix/ and /accelerate/. --- --- @since 0.2.0.0 -class Backprop a where - -- | "Zero out" all components of a value. For scalar values, this - -- should just be @'const' 0@. For vectors and matrices, this should - -- set all components to zero, the additive identity. - -- - -- Should be idempotent: - -- - -- * @'zero' '.' 'zero' = 'zero'@ - -- - -- Should be as /lazy/ as possible. This behavior is observed for - -- all instances provided by this library. - -- - -- See 'zeroNum' for a pre-built definition for instances of 'Num' and - -- 'zeroFunctor' for a definition for instances of 'Functor'. If left - -- blank, will automatically be 'genericZero', a pre-built definition - -- for instances of 'GHC.Generic' whose fields are all themselves - -- instances of 'Backprop'. - zero :: a -> a - -- | Add together two values of a type. To combine contributions of - -- gradients, so should be information-preserving: - -- - -- * @'add' x ('zero' y) = x@ - -- - -- * @'add' ('zero' x) y = y@ - -- - -- Should be as /strict/ as possible. This behavior is observed for - -- all instances provided by this library. - -- - -- See 'addNum' for a pre-built definition for instances of 'Num' and - -- 'addIsList' for a definition for instances of 'IsList'. If left - -- blank, will automatically be 'genericAdd', a pre-built definition - -- for instances of 'GHC.Generic' with one constructor whose fields are - -- all themselves instances of 'Backprop'. - add :: a -> a -> a - -- | "One" all components of a value. For scalar values, this should - -- just be @'const' 1@. For vectors and matrices, this should set all - -- components to one, the multiplicative identity. - -- - -- As the library uses it, the most important law is: - -- - -- * @'one' = 'gradBP' 'id'@ - -- - -- That is, @'one' x@ is the gradient of the identity function with - -- respect to its input. - -- - -- Ideally should be idempotent: - -- - -- * @'one' '.' 'one' = 'one'@ - -- - -- Should be as /lazy/ as possible. This behavior is observed for - -- all instances provided by this library. - -- - -- See 'oneNum' for a pre-built definition for instances of 'Num' and - -- 'oneFunctor' for a definition for instances of 'Functor'. If left - -- blank, will automatically be 'genericOne', a pre-built definition - -- for instances of 'GHC.Generic' whose fields are all themselves - -- instances of 'Backprop'. - one :: a -> a - - default zero :: (Generic a, GZero (Rep a)) => a -> a - zero = genericZero - {-# INLINE zero #-} - default add :: (Generic a, GAdd (Rep a)) => a -> a -> a - add = genericAdd - {-# INLINE add #-} - default one :: (Generic a, GOne (Rep a)) => a -> a - one = genericOne - {-# INLINE one #-} - --- | 'zero' using GHC Generics; works if all fields are instances of --- 'Backprop'. -genericZero :: (Generic a, GZero (Rep a)) => a -> a -genericZero = to . gzero . from -{-# INLINE genericZero #-} - --- | 'add' using GHC Generics; works if all fields are instances of --- 'Backprop', but only for values with single constructors. -genericAdd :: (Generic a, GAdd (Rep a)) => a -> a -> a -genericAdd x y = to $ gadd (from x) (from y) -{-# INLINE genericAdd #-} - --- | 'one' using GHC Generics; works if all fields are instaces of --- 'Backprop'. -genericOne :: (Generic a, GOne (Rep a)) => a -> a -genericOne = to . gone . from -{-# INLINE genericOne #-} - --- | 'zero' for instances of 'Num'. --- --- Is lazy in its argument. -zeroNum :: Num a => a -> a -zeroNum _ = 0 -{-# INLINE zeroNum #-} - --- | 'add' for instances of 'Num'. -addNum :: Num a => a -> a -> a -addNum = (+) -{-# INLINE addNum #-} - --- | 'one' for instances of 'Num'. --- --- Is lazy in its argument. -oneNum :: Num a => a -> a -oneNum _ = 1 -{-# INLINE oneNum #-} - --- | 'zero' for instances of 'VG.Vector'. -zeroVec :: (VG.Vector v a, Backprop a) => v a -> v a -zeroVec = VG.map zero -{-# INLINE zeroVec #-} - --- | 'add' for instances of 'VG.Vector'. Automatically pads the end of the --- shorter vector with zeroes. -addVec :: (VG.Vector v a, Backprop a) => v a -> v a -> v a -addVec x y = case compare lX lY of - LT -> let (y1,y2) = VG.splitAt (lY - lX) y - in VG.zipWith add x y1 VG.++ y2 - EQ -> VG.zipWith add x y - GT -> let (x1,x2) = VG.splitAt (lX - lY) x - in VG.zipWith add x1 y VG.++ x2 - where - lX = VG.length x - lY = VG.length y - --- | 'one' for instances of 'VG.Vector'. -oneVec :: (VG.Vector v a, Backprop a) => v a -> v a -oneVec = VG.map one -{-# INLINE oneVec #-} - --- | 'zero' for instances of 'VG.Vector' when the contained type is an --- instance of 'Num'. Is potentially more performant than 'zeroVec' when --- the vectors are larger. --- --- See 'NumVec' for a 'Backprop' instance for 'VG.Vector' instances that --- uses this for 'zero'. --- --- @since 0.2.4.0 -zeroVecNum :: (VG.Vector v a, Num a) => v a -> v a -zeroVecNum = flip VG.replicate 0 . VG.length -{-# INLINE zeroVecNum #-} - --- | 'one' for instances of 'VG.Vector' when the contained type is an --- instance of 'Num'. Is potentially more performant than 'oneVec' when --- the vectors are larger. --- --- See 'NumVec' for a 'Backprop' instance for 'VG.Vector' instances that --- uses this for 'one'. --- --- @since 0.2.4.0 -oneVecNum :: (VG.Vector v a, Num a) => v a -> v a -oneVecNum = flip VG.replicate 1 . VG.length -{-# INLINE oneVecNum #-} - --- | 'zero' for 'Functor' instances. -zeroFunctor :: (Functor f, Backprop a) => f a -> f a -zeroFunctor = fmap zero -{-# INLINE zeroFunctor #-} - --- | 'add' for instances of 'IsList'. Automatically pads the end of the --- "shorter" value with zeroes. -addIsList :: (IsList a, Backprop (Item a)) => a -> a -> a -addIsList = addAsList toList fromList -{-# INLINE addIsList #-} - --- | 'add' for types that are isomorphic to a list. --- Automatically pads the end of the "shorter" value with zeroes. -addAsList - :: Backprop b - => (a -> [b]) -- ^ convert to list (should form isomorphism) - -> ([b] -> a) -- ^ convert from list (should form isomorphism) - -> a - -> a - -> a -addAsList f g x y = g $ go (f x) (f y) - where - go = \case - [] -> id - o@(x':xs) -> \case - [] -> o - y':ys -> add x' y' : go xs ys - --- | 'one' for instances of 'Functor'. -oneFunctor :: (Functor f, Backprop a) => f a -> f a -oneFunctor = fmap one -{-# INLINE oneFunctor #-} - --- | A newtype wrapper over an instance of 'Num' that gives a free --- 'Backprop' instance. --- --- Useful for things like /DerivingVia/, or for avoiding orphan instances. --- --- @since 0.2.1.0 -newtype NumBP a = NumBP { runNumBP :: a } - deriving (Show, Read, Eq, Ord, Typeable, Data, Generic, Functor, Foldable, Traversable, Num, Fractional, Floating) - -instance NFData a => NFData (NumBP a) - -instance Applicative NumBP where - pure = NumBP - {-# INLINE pure #-} - f <*> x = NumBP $ (runNumBP f) (runNumBP x) - {-# INLINE (<*>) #-} - -instance Monad NumBP where - return = NumBP - {-# INLINE return #-} - x >>= f = f (runNumBP x) - {-# INLINE (>>=) #-} - -instance Num a => Backprop (NumBP a) where - zero = coerce (zeroNum :: a -> a) - {-# INLINE zero #-} - add = coerce (addNum :: a -> a -> a) - {-# INLINE add #-} - one = coerce (oneNum :: a -> a) - {-# INLINE one #-} - --- | Newtype wrapper around a @v a@ for @'VG.Vector' v a@, that gives --- a more efficient 'Backprop' instance for /long/ vectors when @a@ is an --- instance of 'Num'. The normal 'Backprop' instance for vectors will map --- 'zero' or 'one' over all items; this instance will completely ignore the --- contents of the original vector and instead produce a new vector of the --- same length, with all @0@ or @1@ using the 'Num' instance of @a@ --- (essentially using 'zeroVecNum' and 'oneVecNum' instead of 'zeroVec' and --- 'oneVec'). --- --- 'add' is essentially the same as normal, but using '+' instead of the --- type's 'add'. --- --- @since 0.2.4.0 -newtype NumVec v a = NumVec { runNumVec :: v a } - deriving (Show, Read, Eq, Ord, Typeable, Data, Generic, Functor, Applicative, Monad, Alternative, MonadPlus, Foldable, Traversable) - -instance NFData (v a) => NFData (NumVec v a) - -instance (VG.Vector v a, Num a) => Backprop (NumVec v a) where - zero = coerce $ zeroVecNum @v @a - add (NumVec x) (NumVec y) = NumVec $ case compare lX lY of - LT -> let (y1,y2) = VG.splitAt (lY - lX) y - in VG.zipWith (+) x y1 VG.++ y2 - EQ -> VG.zipWith (+) x y - GT -> let (x1,x2) = VG.splitAt (lX - lY) x - in VG.zipWith (+) x1 y VG.++ x2 - where - lX = VG.length x - lY = VG.length y - one = coerce $ oneVecNum @v @a - --- | A newtype wrapper over an @f a@ for @'Applicative' f@ that gives --- a free 'Backprop' instance (as well as 'Num' etc. instances). --- --- Useful for performing backpropagation over functions that require some --- monadic context (like 'IO') to perform. --- --- @since 0.2.1.0 -newtype ABP f a = ABP { runABP :: f a } - deriving (Show, Read, Eq, Ord, Typeable, Data, Generic, Functor, Applicative, Monad, Alternative, MonadPlus, Foldable, Traversable) - -instance NFData (f a) => NFData (ABP f a) - -instance (Applicative f, Backprop a) => Backprop (ABP f a) where - zero = fmap zero - {-# INLINE zero #-} - add = liftA2 add - {-# INLINE add #-} - one = fmap one - {-# INLINE one #-} - -instance (Applicative f, Num a) => Num (ABP f a) where - (+) = liftA2 (+) - {-# INLINE (+) #-} - (-) = liftA2 (-) - {-# INLINE (-) #-} - (*) = liftA2 (*) - {-# INLINE (*) #-} - negate = fmap negate - {-# INLINE negate #-} - abs = fmap abs - {-# INLINE abs #-} - signum = fmap signum - {-# INLINE signum #-} - fromInteger = pure . fromInteger - {-# INLINE fromInteger #-} - -instance (Applicative f, Fractional a) => Fractional (ABP f a) where - (/) = liftA2 (/) - {-# INLINE (/) #-} - recip = fmap recip - {-# INLINE recip #-} - fromRational = pure . fromRational - {-# INLINE fromRational #-} - -instance (Applicative f, Floating a) => Floating (ABP f a) where - pi = pure pi - {-# INLINE pi #-} - exp = fmap exp - {-# INLINE exp #-} - log = fmap log - {-# INLINE log #-} - sqrt = fmap sqrt - {-# INLINE sqrt #-} - (**) = liftA2 (**) - {-# INLINE (**) #-} - logBase = liftA2 logBase - {-# INLINE logBase #-} - sin = fmap sin - {-# INLINE sin #-} - cos = fmap cos - {-# INLINE cos #-} - tan = fmap tan - {-# INLINE tan #-} - asin = fmap asin - {-# INLINE asin #-} - acos = fmap acos - {-# INLINE acos #-} - atan = fmap atan - {-# INLINE atan #-} - sinh = fmap sinh - {-# INLINE sinh #-} - cosh = fmap cosh - {-# INLINE cosh #-} - tanh = fmap tanh - {-# INLINE tanh #-} - asinh = fmap asinh - {-# INLINE asinh #-} - acosh = fmap acosh - {-# INLINE acosh #-} - atanh = fmap atanh - {-# INLINE atanh #-} - - --- | Helper class for automatically deriving 'zero' using GHC Generics. -class GZero f where - gzero :: f t -> f t - -instance Backprop a => GZero (K1 i a) where - gzero (K1 x) = K1 (zero x) - {-# INLINE gzero #-} - -instance (GZero f, GZero g) => GZero (f :*: g) where - gzero (x :*: y) = gzero x :*: gzero y - {-# INLINE gzero #-} - -instance (GZero f, GZero g) => GZero (f :+: g) where - gzero (L1 x) = L1 (gzero x) - gzero (R1 x) = R1 (gzero x) - {-# INLINE gzero #-} - -instance GZero V1 where - gzero = \case {} - {-# INLINE gzero #-} - -instance GZero U1 where - gzero _ = U1 - {-# INLINE gzero #-} - -instance GZero f => GZero (M1 i c f) where - gzero (M1 x) = M1 (gzero x) - {-# INLINE gzero #-} - -instance GZero f => GZero (f :.: g) where - gzero (Comp1 x) = Comp1 (gzero x) - {-# INLINE gzero #-} - - --- | Helper class for automatically deriving 'add' using GHC Generics. -class GAdd f where - gadd :: f t -> f t -> f t - -instance Backprop a => GAdd (K1 i a) where - gadd (K1 x) (K1 y) = K1 (add x y) - {-# INLINE gadd #-} - -instance (GAdd f, GAdd g) => GAdd (f :*: g) where - gadd (x1 :*: y1) (x2 :*: y2) = x3 :*: y3 - where - !x3 = gadd x1 x2 - !y3 = gadd y1 y2 - {-# INLINE gadd #-} - -instance GAdd V1 where - gadd = \case {} - {-# INLINE gadd #-} - -instance GAdd U1 where - gadd _ _ = U1 - {-# INLINE gadd #-} - -instance GAdd f => GAdd (M1 i c f) where - gadd (M1 x) (M1 y) = M1 (gadd x y) - {-# INLINE gadd #-} - -instance GAdd f => GAdd (f :.: g) where - gadd (Comp1 x) (Comp1 y) = Comp1 (gadd x y) - {-# INLINE gadd #-} - - --- | Helper class for automatically deriving 'one' using GHC Generics. -class GOne f where - gone :: f t -> f t - -instance Backprop a => GOne (K1 i a) where - gone (K1 x) = K1 (one x) - {-# INLINE gone #-} - -instance (GOne f, GOne g) => GOne (f :*: g) where - gone (x :*: y) = gone x :*: gone y - {-# INLINE gone #-} - -instance (GOne f, GOne g) => GOne (f :+: g) where - gone (L1 x) = L1 (gone x) - gone (R1 x) = R1 (gone x) - {-# INLINE gone #-} - -instance GOne V1 where - gone = \case {} - {-# INLINE gone #-} - -instance GOne U1 where - gone _ = U1 - {-# INLINE gone #-} - -instance GOne f => GOne (M1 i c f) where - gone (M1 x) = M1 (gone x) - {-# INLINE gone #-} - -instance GOne f => GOne (f :.: g) where - gone (Comp1 x) = Comp1 (gone x) - {-# INLINE gone #-} - -instance Backprop Int where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - -instance Backprop Integer where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - --- | @since 0.2.1.0 -instance Backprop Natural where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - --- | @since 0.2.2.0 -instance Backprop Word8 where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - --- | @since 0.2.2.0 -instance Backprop Word where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - --- | @since 0.2.2.0 -instance Backprop Word16 where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - --- | @since 0.2.2.0 -instance Backprop Word32 where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - --- | @since 0.2.2.0 -instance Backprop Word64 where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - -instance Integral a => Backprop (Ratio a) where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - -instance RealFloat a => Backprop (Complex a) where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - -instance Backprop Float where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - -instance Backprop Double where - zero = zeroNum - {-# INLINE zero #-} - add = addNum - {-# INLINE add #-} - one = oneNum - {-# INLINE one #-} - -instance Backprop a => Backprop (V.Vector a) where - zero = zeroVec - {-# INLINE zero #-} - add = addVec - {-# INLINE add #-} - one = oneVec - {-# INLINE one #-} - -instance (VU.Unbox a, Backprop a) => Backprop (VU.Vector a) where - zero = zeroVec - {-# INLINE zero #-} - add = addVec - {-# INLINE add #-} - one = oneVec - {-# INLINE one #-} - -instance (VS.Storable a, Backprop a) => Backprop (VS.Vector a) where - zero = zeroVec - {-# INLINE zero #-} - add = addVec - {-# INLINE add #-} - one = oneVec - {-# INLINE one #-} - -instance (VP.Prim a, Backprop a) => Backprop (VP.Vector a) where - zero = zeroVec - {-# INLINE zero #-} - add = addVec - {-# INLINE add #-} - one = oneVec - {-# INLINE one #-} - --- | 'add' assumes the shorter list has trailing zeroes, and the result has --- the length of the longest input. -instance Backprop a => Backprop [a] where - zero = zeroFunctor - {-# INLINE zero #-} - add = addIsList - {-# INLINE add #-} - one = oneFunctor - {-# INLINE one #-} - --- | 'add' assumes the shorter list has trailing zeroes, and the result has --- the length of the longest input. -instance Backprop a => Backprop (NonEmpty a) where - zero = zeroFunctor - {-# INLINE zero #-} - add = addIsList - {-# INLINE add #-} - one = oneFunctor - {-# INLINE one #-} - --- | 'add' assumes the shorter sequence has trailing zeroes, and the result --- has the length of the longest input. -instance Backprop a => Backprop (Seq.Seq a) where - zero = zeroFunctor - {-# INLINE zero #-} - add = addIsList - {-# INLINE add #-} - one = oneFunctor - {-# INLINE one #-} - --- | 'Nothing' is treated the same as @'Just' 0@. However, 'zero', 'add', --- and 'one' preserve 'Nothing' if all inputs are also 'Nothing'. -instance Backprop a => Backprop (Maybe a) where - zero = zeroFunctor - {-# INLINE zero #-} - add x y = asum [ add <$> x <*> y - , x - , y - ] - {-# INLINE add #-} - one = oneFunctor - {-# INLINE one #-} - --- | 'add' is strict, but 'zero' and 'one' are lazy in their arguments. -instance Backprop () where - zero _ = () - add () () = () - one _ = () - --- | 'add' is strict -instance (Backprop a, Backprop b) => Backprop (a, b) where - zero (x, y) = (zero x, zero y) - {-# INLINE zero #-} - add (x1, y1) (x2, y2) = (x3, y3) - where - !x3 = add x1 x2 - !y3 = add y1 y2 - {-# INLINE add #-} - one (x, y) = (one x, one y) - {-# INLINE one #-} - --- | 'add' is strict -instance (Backprop a, Backprop b, Backprop c) => Backprop (a, b, c) where - zero (x, y, z) = (zero x, zero y, zero z) - {-# INLINE zero #-} - add (x1, y1, z1) (x2, y2, z2) = (x3, y3, z3) - where - !x3 = add x1 x2 - !y3 = add y1 y2 - !z3 = add z1 z2 - {-# INLINE add #-} - one (x, y, z) = (one x, one y, one z) - {-# INLINE one #-} - --- | 'add' is strict -instance (Backprop a, Backprop b, Backprop c, Backprop d) => Backprop (a, b, c, d) where - zero (x, y, z, w) = (zero x, zero y, zero z, zero w) - {-# INLINE zero #-} - add (x1, y1, z1, w1) (x2, y2, z2, w2) = (x3, y3, z3, w3) - where - !x3 = add x1 x2 - !y3 = add y1 y2 - !z3 = add z1 z2 - !w3 = add w1 w2 - {-# INLINE add #-} - one (x, y, z, w) = (one x, one y, one z, one w) - {-# INLINE one #-} - --- | 'add' is strict -instance (Backprop a, Backprop b, Backprop c, Backprop d, Backprop e) => Backprop (a, b, c, d, e) where - zero (x, y, z, w, v) = (zero x, zero y, zero z, zero w, zero v) - {-# INLINE zero #-} - add (x1, y1, z1, w1, v1) (x2, y2, z2, w2, v2) = (x3, y3, z3, w3, v3) - where - !x3 = add x1 x2 - !y3 = add y1 y2 - !z3 = add z1 z2 - !w3 = add w1 w2 - !v3 = add v1 v2 - {-# INLINE add #-} - one (x, y, z, w, v) = (one x, one y, one z, one w, one v) - {-# INLINE one #-} - -instance Backprop a => Backprop (Identity a) where - zero = coerce (zero @a) - add = coerce (add @a) - one = coerce (one @a) - -instance Backprop (Proxy a) where - zero _ = Proxy - {-# INLINE zero #-} - add _ _ = Proxy - {-# INLINE add #-} - one _ = Proxy - {-# INLINE one #-} - --- | @since 0.2.2.0 -instance Backprop w => Backprop (Const w a) where - zero = coerce (zero @w) - add = coerce (add @w) - one = coerce (one @w) - -instance Backprop Void where - zero = \case {} - {-# INLINE zero #-} - add = \case {} - {-# INLINE add #-} - one = \case {} - {-# INLINE one #-} - --- | 'zero' and 'one' replace all current values, and 'add' merges keys --- from both maps, adding in the case of double-occurrences. -instance (Backprop a, Ord k) => Backprop (M.Map k a) where - zero = zeroFunctor - {-# INLINE zero #-} - add = M.unionWith add - {-# INLINE add #-} - one = oneFunctor - {-# INLINE one #-} - --- | 'zero' and 'one' replace all current values, and 'add' merges keys --- from both maps, adding in the case of double-occurrences. -instance (Backprop a) => Backprop (IM.IntMap a) where - zero = zeroFunctor - {-# INLINE zero #-} - add = IM.unionWith add - {-# INLINE add #-} - one = oneFunctor - {-# INLINE one #-} - --- | @since 0.2.2.0 -instance Backprop a => Backprop (K1 i a p) where - zero = coerce (zero @a) - add = coerce (add @a) - one = coerce (one @a) - --- | @since 0.2.2.0 -instance Backprop (f p) => Backprop (M1 i c f p) where - zero = coerce (zero @(f p)) - add = coerce (add @(f p)) - one = coerce (one @(f p)) - --- | @since 0.2.2.0 -instance (Backprop (f p), Backprop (g p)) => Backprop ((f :*: g) p) - --- | @since 0.2.6.3 -instance (Backprop (f (g a))) => Backprop ((f :.: g) a) where - zero = coerce (zero @(f (g a))) - add = coerce (add @(f (g a))) - one = coerce (one @(f (g a))) - --- | @since 0.2.2.0 -instance Backprop (V1 p) - --- | @since 0.2.2.0 -instance Backprop (U1 p) - --- | @since 0.2.2.0 -instance Backprop a => Backprop (Sum a) where - zero = coerce (zero @a) - add = coerce (add @a) - one = coerce (one @a) - --- | @since 0.2.2.0 -instance Backprop a => Backprop (Product a) where - zero = coerce (zero @a) - add = coerce (add @a) - one = coerce (one @a) - -#if !MIN_VERSION_base(4,16,0) --- | @since 0.2.2.0 -instance Backprop a => Backprop (SG.Option a) where - zero = coerce (zero @(Maybe a)) - add = coerce (add @(Maybe a)) - one = coerce (one @(Maybe a)) -#endif - --- | @since 0.2.2.0 -instance Backprop a => Backprop (SG.First a) where - zero = coerce (zero @a) - add = coerce (add @a) - one = coerce (one @a) - --- | @since 0.2.2.0 -instance Backprop a => Backprop (SG.Last a) where - zero = coerce (zero @a) - add = coerce (add @a) - one = coerce (one @a) - --- | @since 0.2.2.0 -instance Backprop a => Backprop (First a) where - zero = coerce (zero @(Maybe a)) - add = coerce (add @(Maybe a)) - one = coerce (one @(Maybe a)) - --- | @since 0.2.2.0 -instance Backprop a => Backprop (Data.Monoid.Last a) where - zero = coerce (zero @(Maybe a)) - add = coerce (add @(Maybe a)) - one = coerce (one @(Maybe a)) - --- | @since 0.2.2.0 -instance Backprop a => Backprop (Dual a) where - zero = coerce (zero @a) - add = coerce (add @a) - one = coerce (one @a) - --- | @since 0.2.2.0 -instance (Backprop a, Backprop b) => Backprop (SG.Arg a b) - --- | @since 0.2.2.0 -instance (Backprop (f a), Backprop (g a)) => Backprop (DFP.Product f g a) - --- | @since 0.2.2.0 -instance Backprop (f (g a)) => Backprop (Compose f g a) where - zero = coerce (zero @(f (g a))) - add = coerce (add @(f (g a))) - one = coerce (one @(f (g a))) - --- | 'add' adds together results; 'zero' and 'one' act on results. --- --- @since 0.2.2.0 -instance Backprop a => Backprop (r -> a) where - zero = zeroFunctor - {-# INLINE zero #-} - add = liftA2 add - {-# INLINE add #-} - one = oneFunctor - {-# INLINE one #-} - --- | @since 0.2.2.0 -instance (Backprop a, Applicative m) => Backprop (Arr.Kleisli m r a) where - zero (Arr.Kleisli f) = Arr.Kleisli ((fmap . fmap) zero f) - {-# INLINE zero #-} - add (Arr.Kleisli f) (Arr.Kleisli g) = Arr.Kleisli $ \x -> - add <$> f x <*> g x - {-# INLINE add #-} - one (Arr.Kleisli f) = Arr.Kleisli ((fmap . fmap) one f) - {-# INLINE one #-} - --- | @since 0.2.6.3 -instance (ReifyConstraint Backprop f rs, RMap rs, RApply rs) => Backprop (Rec f rs) where - zero = rmap (\case V.Compose (Dict x) -> zero x) - . reifyConstraint @Backprop - {-# INLINE zero #-} - add xs = rzipWith (\x -> \case V.Compose (Dict y) -> add x y) xs - . reifyConstraint @Backprop - {-# INLINE add #-} - one = rmap (\case V.Compose (Dict x) -> one x) - . reifyConstraint @Backprop - {-# INLINE one #-} - --- | @since 0.2.6.3 -#if MIN_VERSION_vinyl(0,14,2) -instance (ReifyConstraint Backprop f rs, RMap rs, RApply rs, RecApplicative rs, NatToInt (RLength rs), RPureConstrained (IndexableField rs) rs, ToARec rs) -#else -instance (ReifyConstraint Backprop f rs, RMap rs, RApply rs, RecApplicative rs, NatToInt (RLength rs), RPureConstrained (IndexableField rs) rs) -#endif - => Backprop (ARec f rs) where - zero = toARec . zero . fromARec - {-# INLINE zero #-} - add xs ys = toARec $ add (fromARec xs) (fromARec ys) - {-# INLINE add #-} - one = toARec . zero . fromARec - {-# INLINE one #-} - --- | @since 0.2.6.3 -instance (ReifyConstraint Backprop f rs, RMap rs, RApply rs, VS.Storable (Rec f rs)) - => Backprop (SRec f rs) where - zero = toSRec . zero . fromSRec - {-# INLINE zero #-} - add xs ys = toSRec $ add (fromSRec xs) (fromSRec ys) - {-# INLINE add #-} - one = toSRec . zero . fromSRec - {-# INLINE one #-} - --- | @since 0.2.6.3 -instance (ReifyConstraint Backprop f rs, RMap rs, RApply rs, IsoXRec f rs) - => Backprop (XRec f rs) where - zero = toXRec . zero . fromXRec - {-# INLINE zero #-} - add xs ys = toXRec $ add (fromXRec xs) (fromXRec ys) - {-# INLINE add #-} - one = toXRec . zero . fromXRec - {-# INLINE one #-} - --- | @since 0.2.6.3 -instance Backprop a => Backprop (V.Identity a) where - zero = coerce (zero @a) - add = coerce (add @a) - one = coerce (one @a) - --- | @since 0.2.6.3 -instance Backprop a => Backprop (V.Thunk a) where - zero (V.Thunk x) = V.Thunk (zero x) - add (V.Thunk x) (V.Thunk y) = V.Thunk (add x y) - one (V.Thunk x) = V.Thunk (one x) - --- | @since 0.2.6.3 -instance Backprop (op (f a) (g a)) => Backprop (V.Lift op f g a) where - zero = coerce (zero @(op (f a) (g a))) - add = coerce (add @(op (f a) (g a))) - one = coerce (one @(op (f a) (g a))) - --- | @since 0.2.6.3 -instance Backprop t => Backprop (V.ElField '(s, t)) where - zero (V.Field x) = V.Field (zero x) - add (V.Field x) (V.Field y) = V.Field (add x y) - one (V.Field x) = V.Field (one x) - --- | @since 0.2.6.3 -instance Backprop (f (g a)) => Backprop (V.Compose f g a) where - zero = coerce (zero @(f (g a))) - add = coerce (add @(f (g a))) - one = coerce (one @(f (g a))) - --- | @since 0.2.6.3 -instance Backprop w => Backprop (V.Const w a) where - zero = coerce (zero @w) - add = coerce (add @w) - one = coerce (one @w) - --- | @since 0.2.6.3 -instance Backprop (V.HKD t a) => Backprop (V.XData t a) where - zero = coerce (zero @(V.HKD t a)) - add = coerce (add @(V.HKD t a)) - one = coerce (one @(V.HKD t a)) - --- | @since 0.2.6.3 -instance Backprop (SField field) where - zero _ = SField - add _ _ = SField - one _ = SField - --- | @since 0.2.6.3 -instance Backprop (Label field) where - zero _ = Label - add _ _ = Label - one _ = Label +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE EmptyCase #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : Numeric.Backprop.Class +-- Copyright : (c) Justin Le 2023 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Provides the 'Backprop' typeclass, a class for values that can be used +-- for backpropagation. +-- +-- This class replaces the old (version 0.1) API relying on 'Num'. +-- +-- @since 0.2.0.0 +module Numeric.Backprop.Class ( + -- * Backpropagatable types + Backprop (..), + + -- * Derived methods + zeroNum, + addNum, + oneNum, + zeroVec, + addVec, + oneVec, + zeroVecNum, + oneVecNum, + zeroFunctor, + addIsList, + addAsList, + oneFunctor, + genericZero, + genericAdd, + genericOne, + + -- * Newtype + ABP (..), + NumBP (..), + NumVec (..), + + -- * Generics + GZero, + GAdd, + GOne, +) where + +import Control.Applicative +import qualified Control.Arrow as Arr +import Control.DeepSeq +import Control.Monad +import Data.Coerce +import Data.Complex +import Data.Data +import Data.Functor.Compose +import Data.Functor.Identity +import qualified Data.Functor.Product as DFP +import qualified Data.IntMap as IM +import Data.List.NonEmpty (NonEmpty (..)) +import qualified Data.Map as M +import Data.Monoid +import Data.Ratio +import qualified Data.Semigroup as SG +import qualified Data.Sequence as Seq +import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Primitive as VP +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as VU +import Data.Vinyl +import Data.Vinyl.ARec +import qualified Data.Vinyl.Functor as V +import Data.Vinyl.TypeLevel +import qualified Data.Vinyl.XRec as V +import Data.Void +import Data.Word +import GHC.Exts +import GHC.Generics +import Numeric.Natural + +-- | Class of values that can be backpropagated in general. +-- +-- For instances of 'Num', these methods can be given by 'zeroNum', +-- 'addNum', and 'oneNum'. There are also generic options given in +-- "Numeric.Backprop.Class" for functors, 'IsList' instances, and 'Generic' +-- instances. +-- +-- @ +-- instance 'Backprop' 'Double' where +-- 'zero' = 'zeroNum' +-- 'add' = 'addNum' +-- 'one' = 'oneNum' +-- @ +-- +-- If you leave the body of an instance declaration blank, GHC Generics +-- will be used to derive instances if the type has a single constructor +-- and each field is an instance of 'Backprop'. +-- +-- To ensure that backpropagation works in a sound way, should obey the +-- laws: +-- +-- [/identity/] +-- +-- * @'add' x ('zero' y) = x@ +-- +-- * @'add' ('zero' x) y = y@ +-- +-- Also implies preservation of information, making @'zipWith' ('+')@ an +-- illegal implementation for lists and vectors. +-- +-- This is only expected to be true up to potential "extra zeroes" in @x@ +-- and @y@ in the result. +-- +-- [/commutativity/] +-- +-- * @'add' x y = 'add' y x@ +-- +-- [/associativity/] +-- +-- * @'add' x ('add' y z) = 'add' ('add' x y) z@ +-- +-- [/idempotence/] +-- +-- * @'zero' '.' 'zero' = 'zero'@ +-- +-- * @'one' '.' 'one' = 'one'@ +-- +-- [/unital/] +-- +-- * @'one' = 'gradBP' 'id'@ +-- +-- Note that not all values in the backpropagation process needs all of +-- these methods: Only the "final result" needs 'one', for example. These +-- are all grouped under one typeclass for convenience in defining +-- instances, and also to talk about sensible laws. For fine-grained +-- control, use the "explicit" versions of library functions (for example, +-- in "Numeric.Backprop.Explicit") instead of 'Backprop' based ones. +-- +-- This typeclass replaces the reliance on 'Num' of the previous API +-- (v0.1). 'Num' is strictly more powerful than 'Backprop', and is +-- a stronger constraint on types than is necessary for proper +-- backpropagating. In particular, 'fromInteger' is a problem for many +-- types, preventing useful backpropagation for lists, variable-length +-- vectors (like "Data.Vector") and variable-size matrices from linear +-- algebra libraries like /hmatrix/ and /accelerate/. +-- +-- @since 0.2.0.0 +class Backprop a where + -- | "Zero out" all components of a value. For scalar values, this + -- should just be @'const' 0@. For vectors and matrices, this should + -- set all components to zero, the additive identity. + -- + -- Should be idempotent: + -- + -- * @'zero' '.' 'zero' = 'zero'@ + -- + -- Should be as /lazy/ as possible. This behavior is observed for + -- all instances provided by this library. + -- + -- See 'zeroNum' for a pre-built definition for instances of 'Num' and + -- 'zeroFunctor' for a definition for instances of 'Functor'. If left + -- blank, will automatically be 'genericZero', a pre-built definition + -- for instances of 'GHC.Generic' whose fields are all themselves + -- instances of 'Backprop'. + zero :: a -> a + + -- | Add together two values of a type. To combine contributions of + -- gradients, so should be information-preserving: + -- + -- * @'add' x ('zero' y) = x@ + -- + -- * @'add' ('zero' x) y = y@ + -- + -- Should be as /strict/ as possible. This behavior is observed for + -- all instances provided by this library. + -- + -- See 'addNum' for a pre-built definition for instances of 'Num' and + -- 'addIsList' for a definition for instances of 'IsList'. If left + -- blank, will automatically be 'genericAdd', a pre-built definition + -- for instances of 'GHC.Generic' with one constructor whose fields are + -- all themselves instances of 'Backprop'. + add :: a -> a -> a + + -- | "One" all components of a value. For scalar values, this should + -- just be @'const' 1@. For vectors and matrices, this should set all + -- components to one, the multiplicative identity. + -- + -- As the library uses it, the most important law is: + -- + -- * @'one' = 'gradBP' 'id'@ + -- + -- That is, @'one' x@ is the gradient of the identity function with + -- respect to its input. + -- + -- Ideally should be idempotent: + -- + -- * @'one' '.' 'one' = 'one'@ + -- + -- Should be as /lazy/ as possible. This behavior is observed for + -- all instances provided by this library. + -- + -- See 'oneNum' for a pre-built definition for instances of 'Num' and + -- 'oneFunctor' for a definition for instances of 'Functor'. If left + -- blank, will automatically be 'genericOne', a pre-built definition + -- for instances of 'GHC.Generic' whose fields are all themselves + -- instances of 'Backprop'. + one :: a -> a + + default zero :: (Generic a, GZero (Rep a)) => a -> a + zero = genericZero + {-# INLINE zero #-} + default add :: (Generic a, GAdd (Rep a)) => a -> a -> a + add = genericAdd + {-# INLINE add #-} + default one :: (Generic a, GOne (Rep a)) => a -> a + one = genericOne + {-# INLINE one #-} + +-- | 'zero' using GHC Generics; works if all fields are instances of +-- 'Backprop'. +genericZero :: (Generic a, GZero (Rep a)) => a -> a +genericZero = to . gzero . from +{-# INLINE genericZero #-} + +-- | 'add' using GHC Generics; works if all fields are instances of +-- 'Backprop', but only for values with single constructors. +genericAdd :: (Generic a, GAdd (Rep a)) => a -> a -> a +genericAdd x y = to $ gadd (from x) (from y) +{-# INLINE genericAdd #-} + +-- | 'one' using GHC Generics; works if all fields are instaces of +-- 'Backprop'. +genericOne :: (Generic a, GOne (Rep a)) => a -> a +genericOne = to . gone . from +{-# INLINE genericOne #-} + +-- | 'zero' for instances of 'Num'. +-- +-- Is lazy in its argument. +zeroNum :: Num a => a -> a +zeroNum _ = 0 +{-# INLINE zeroNum #-} + +-- | 'add' for instances of 'Num'. +addNum :: Num a => a -> a -> a +addNum = (+) +{-# INLINE addNum #-} + +-- | 'one' for instances of 'Num'. +-- +-- Is lazy in its argument. +oneNum :: Num a => a -> a +oneNum _ = 1 +{-# INLINE oneNum #-} + +-- | 'zero' for instances of 'VG.Vector'. +zeroVec :: (VG.Vector v a, Backprop a) => v a -> v a +zeroVec = VG.map zero +{-# INLINE zeroVec #-} + +-- | 'add' for instances of 'VG.Vector'. Automatically pads the end of the +-- shorter vector with zeroes. +addVec :: (VG.Vector v a, Backprop a) => v a -> v a -> v a +addVec x y = case compare lX lY of + LT -> + let (y1, y2) = VG.splitAt (lY - lX) y + in VG.zipWith add x y1 VG.++ y2 + EQ -> VG.zipWith add x y + GT -> + let (x1, x2) = VG.splitAt (lX - lY) x + in VG.zipWith add x1 y VG.++ x2 + where + lX = VG.length x + lY = VG.length y + +-- | 'one' for instances of 'VG.Vector'. +oneVec :: (VG.Vector v a, Backprop a) => v a -> v a +oneVec = VG.map one +{-# INLINE oneVec #-} + +-- | 'zero' for instances of 'VG.Vector' when the contained type is an +-- instance of 'Num'. Is potentially more performant than 'zeroVec' when +-- the vectors are larger. +-- +-- See 'NumVec' for a 'Backprop' instance for 'VG.Vector' instances that +-- uses this for 'zero'. +-- +-- @since 0.2.4.0 +zeroVecNum :: (VG.Vector v a, Num a) => v a -> v a +zeroVecNum = flip VG.replicate 0 . VG.length +{-# INLINE zeroVecNum #-} + +-- | 'one' for instances of 'VG.Vector' when the contained type is an +-- instance of 'Num'. Is potentially more performant than 'oneVec' when +-- the vectors are larger. +-- +-- See 'NumVec' for a 'Backprop' instance for 'VG.Vector' instances that +-- uses this for 'one'. +-- +-- @since 0.2.4.0 +oneVecNum :: (VG.Vector v a, Num a) => v a -> v a +oneVecNum = flip VG.replicate 1 . VG.length +{-# INLINE oneVecNum #-} + +-- | 'zero' for 'Functor' instances. +zeroFunctor :: (Functor f, Backprop a) => f a -> f a +zeroFunctor = fmap zero +{-# INLINE zeroFunctor #-} + +-- | 'add' for instances of 'IsList'. Automatically pads the end of the +-- "shorter" value with zeroes. +addIsList :: (IsList a, Backprop (Item a)) => a -> a -> a +addIsList = addAsList toList fromList +{-# INLINE addIsList #-} + +-- | 'add' for types that are isomorphic to a list. +-- Automatically pads the end of the "shorter" value with zeroes. +addAsList :: + Backprop b => + -- | convert to list (should form isomorphism) + (a -> [b]) -> + -- | convert from list (should form isomorphism) + ([b] -> a) -> + a -> + a -> + a +addAsList f g x y = g $ go (f x) (f y) + where + go = \case + [] -> id + o@(x' : xs) -> \case + [] -> o + y' : ys -> add x' y' : go xs ys + +-- | 'one' for instances of 'Functor'. +oneFunctor :: (Functor f, Backprop a) => f a -> f a +oneFunctor = fmap one +{-# INLINE oneFunctor #-} + +-- | A newtype wrapper over an instance of 'Num' that gives a free +-- 'Backprop' instance. +-- +-- Useful for things like /DerivingVia/, or for avoiding orphan instances. +-- +-- @since 0.2.1.0 +newtype NumBP a = NumBP {runNumBP :: a} + deriving + ( Show + , Read + , Eq + , Ord + , Typeable + , Data + , Generic + , Functor + , Foldable + , Traversable + , Num + , Fractional + , Floating + ) + +instance NFData a => NFData (NumBP a) + +instance Applicative NumBP where + pure = NumBP + {-# INLINE pure #-} + f <*> x = NumBP $ runNumBP f (runNumBP x) + {-# INLINE (<*>) #-} + +instance Monad NumBP where + x >>= f = f (runNumBP x) + {-# INLINE (>>=) #-} + +instance Num a => Backprop (NumBP a) where + zero = coerce (zeroNum :: a -> a) + {-# INLINE zero #-} + add = coerce (addNum :: a -> a -> a) + {-# INLINE add #-} + one = coerce (oneNum :: a -> a) + {-# INLINE one #-} + +-- | Newtype wrapper around a @v a@ for @'VG.Vector' v a@, that gives +-- a more efficient 'Backprop' instance for /long/ vectors when @a@ is an +-- instance of 'Num'. The normal 'Backprop' instance for vectors will map +-- 'zero' or 'one' over all items; this instance will completely ignore the +-- contents of the original vector and instead produce a new vector of the +-- same length, with all @0@ or @1@ using the 'Num' instance of @a@ +-- (essentially using 'zeroVecNum' and 'oneVecNum' instead of 'zeroVec' and +-- 'oneVec'). +-- +-- 'add' is essentially the same as normal, but using '+' instead of the +-- type's 'add'. +-- +-- @since 0.2.4.0 +newtype NumVec v a = NumVec {runNumVec :: v a} + deriving + ( Show + , Read + , Eq + , Ord + , Typeable + , Data + , Generic + , Functor + , Applicative + , Monad + , Alternative + , MonadPlus + , Foldable + , Traversable + ) + +instance NFData (v a) => NFData (NumVec v a) + +instance (VG.Vector v a, Num a) => Backprop (NumVec v a) where + zero = coerce $ zeroVecNum @v @a + add (NumVec x) (NumVec y) = NumVec $ case compare lX lY of + LT -> + let (y1, y2) = VG.splitAt (lY - lX) y + in VG.zipWith (+) x y1 VG.++ y2 + EQ -> VG.zipWith (+) x y + GT -> + let (x1, x2) = VG.splitAt (lX - lY) x + in VG.zipWith (+) x1 y VG.++ x2 + where + lX = VG.length x + lY = VG.length y + one = coerce $ oneVecNum @v @a + +-- | A newtype wrapper over an @f a@ for @'Applicative' f@ that gives +-- a free 'Backprop' instance (as well as 'Num' etc. instances). +-- +-- Useful for performing backpropagation over functions that require some +-- monadic context (like 'IO') to perform. +-- +-- @since 0.2.1.0 +newtype ABP f a = ABP {runABP :: f a} + deriving + ( Show + , Read + , Eq + , Ord + , Typeable + , Data + , Generic + , Functor + , Applicative + , Monad + , Alternative + , MonadPlus + , Foldable + , Traversable + ) + +instance NFData (f a) => NFData (ABP f a) + +instance (Applicative f, Backprop a) => Backprop (ABP f a) where + zero = fmap zero + {-# INLINE zero #-} + add = liftA2 add + {-# INLINE add #-} + one = fmap one + {-# INLINE one #-} + +instance (Applicative f, Num a) => Num (ABP f a) where + (+) = liftA2 (+) + {-# INLINE (+) #-} + (-) = liftA2 (-) + {-# INLINE (-) #-} + (*) = liftA2 (*) + {-# INLINE (*) #-} + negate = fmap negate + {-# INLINE negate #-} + abs = fmap abs + {-# INLINE abs #-} + signum = fmap signum + {-# INLINE signum #-} + fromInteger = pure . fromInteger + {-# INLINE fromInteger #-} + +instance (Applicative f, Fractional a) => Fractional (ABP f a) where + (/) = liftA2 (/) + {-# INLINE (/) #-} + recip = fmap recip + {-# INLINE recip #-} + fromRational = pure . fromRational + {-# INLINE fromRational #-} + +instance (Applicative f, Floating a) => Floating (ABP f a) where + pi = pure pi + {-# INLINE pi #-} + exp = fmap exp + {-# INLINE exp #-} + log = fmap log + {-# INLINE log #-} + sqrt = fmap sqrt + {-# INLINE sqrt #-} + (**) = liftA2 (**) + {-# INLINE (**) #-} + logBase = liftA2 logBase + {-# INLINE logBase #-} + sin = fmap sin + {-# INLINE sin #-} + cos = fmap cos + {-# INLINE cos #-} + tan = fmap tan + {-# INLINE tan #-} + asin = fmap asin + {-# INLINE asin #-} + acos = fmap acos + {-# INLINE acos #-} + atan = fmap atan + {-# INLINE atan #-} + sinh = fmap sinh + {-# INLINE sinh #-} + cosh = fmap cosh + {-# INLINE cosh #-} + tanh = fmap tanh + {-# INLINE tanh #-} + asinh = fmap asinh + {-# INLINE asinh #-} + acosh = fmap acosh + {-# INLINE acosh #-} + atanh = fmap atanh + {-# INLINE atanh #-} + +-- | Helper class for automatically deriving 'zero' using GHC Generics. +class GZero f where + gzero :: f t -> f t + +instance Backprop a => GZero (K1 i a) where + gzero (K1 x) = K1 (zero x) + {-# INLINE gzero #-} + +instance (GZero f, GZero g) => GZero (f :*: g) where + gzero (x :*: y) = gzero x :*: gzero y + {-# INLINE gzero #-} + +instance (GZero f, GZero g) => GZero (f :+: g) where + gzero (L1 x) = L1 (gzero x) + gzero (R1 x) = R1 (gzero x) + {-# INLINE gzero #-} + +instance GZero V1 where + gzero = \case {} + {-# INLINE gzero #-} + +instance GZero U1 where + gzero _ = U1 + {-# INLINE gzero #-} + +instance GZero f => GZero (M1 i c f) where + gzero (M1 x) = M1 (gzero x) + {-# INLINE gzero #-} + +instance GZero f => GZero (f :.: g) where + gzero (Comp1 x) = Comp1 (gzero x) + {-# INLINE gzero #-} + +-- | Helper class for automatically deriving 'add' using GHC Generics. +class GAdd f where + gadd :: f t -> f t -> f t + +instance Backprop a => GAdd (K1 i a) where + gadd (K1 x) (K1 y) = K1 (add x y) + {-# INLINE gadd #-} + +instance (GAdd f, GAdd g) => GAdd (f :*: g) where + gadd (x1 :*: y1) (x2 :*: y2) = x3 :*: y3 + where + !x3 = gadd x1 x2 + !y3 = gadd y1 y2 + {-# INLINE gadd #-} + +instance GAdd V1 where + gadd = \case {} + {-# INLINE gadd #-} + +instance GAdd U1 where + gadd _ _ = U1 + {-# INLINE gadd #-} + +instance GAdd f => GAdd (M1 i c f) where + gadd (M1 x) (M1 y) = M1 (gadd x y) + {-# INLINE gadd #-} + +instance GAdd f => GAdd (f :.: g) where + gadd (Comp1 x) (Comp1 y) = Comp1 (gadd x y) + {-# INLINE gadd #-} + +-- | Helper class for automatically deriving 'one' using GHC Generics. +class GOne f where + gone :: f t -> f t + +instance Backprop a => GOne (K1 i a) where + gone (K1 x) = K1 (one x) + {-# INLINE gone #-} + +instance (GOne f, GOne g) => GOne (f :*: g) where + gone (x :*: y) = gone x :*: gone y + {-# INLINE gone #-} + +instance (GOne f, GOne g) => GOne (f :+: g) where + gone (L1 x) = L1 (gone x) + gone (R1 x) = R1 (gone x) + {-# INLINE gone #-} + +instance GOne V1 where + gone = \case {} + {-# INLINE gone #-} + +instance GOne U1 where + gone _ = U1 + {-# INLINE gone #-} + +instance GOne f => GOne (M1 i c f) where + gone (M1 x) = M1 (gone x) + {-# INLINE gone #-} + +instance GOne f => GOne (f :.: g) where + gone (Comp1 x) = Comp1 (gone x) + {-# INLINE gone #-} + +instance Backprop Int where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +instance Backprop Integer where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +-- | @since 0.2.1.0 +instance Backprop Natural where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +-- | @since 0.2.2.0 +instance Backprop Word8 where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +-- | @since 0.2.2.0 +instance Backprop Word where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +-- | @since 0.2.2.0 +instance Backprop Word16 where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +-- | @since 0.2.2.0 +instance Backprop Word32 where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +-- | @since 0.2.2.0 +instance Backprop Word64 where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +instance Integral a => Backprop (Ratio a) where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +instance RealFloat a => Backprop (Complex a) where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +instance Backprop Float where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +instance Backprop Double where + zero = zeroNum + {-# INLINE zero #-} + add = addNum + {-# INLINE add #-} + one = oneNum + {-# INLINE one #-} + +instance Backprop a => Backprop (V.Vector a) where + zero = zeroVec + {-# INLINE zero #-} + add = addVec + {-# INLINE add #-} + one = oneVec + {-# INLINE one #-} + +instance (VU.Unbox a, Backprop a) => Backprop (VU.Vector a) where + zero = zeroVec + {-# INLINE zero #-} + add = addVec + {-# INLINE add #-} + one = oneVec + {-# INLINE one #-} + +instance (VS.Storable a, Backprop a) => Backprop (VS.Vector a) where + zero = zeroVec + {-# INLINE zero #-} + add = addVec + {-# INLINE add #-} + one = oneVec + {-# INLINE one #-} + +instance (VP.Prim a, Backprop a) => Backprop (VP.Vector a) where + zero = zeroVec + {-# INLINE zero #-} + add = addVec + {-# INLINE add #-} + one = oneVec + {-# INLINE one #-} + +-- | 'add' assumes the shorter list has trailing zeroes, and the result has +-- the length of the longest input. +instance Backprop a => Backprop [a] where + zero = zeroFunctor + {-# INLINE zero #-} + add = addIsList + {-# INLINE add #-} + one = oneFunctor + {-# INLINE one #-} + +-- | 'add' assumes the shorter list has trailing zeroes, and the result has +-- the length of the longest input. +instance Backprop a => Backprop (NonEmpty a) where + zero = zeroFunctor + {-# INLINE zero #-} + add = addIsList + {-# INLINE add #-} + one = oneFunctor + {-# INLINE one #-} + +-- | 'add' assumes the shorter sequence has trailing zeroes, and the result +-- has the length of the longest input. +instance Backprop a => Backprop (Seq.Seq a) where + zero = zeroFunctor + {-# INLINE zero #-} + add = addIsList + {-# INLINE add #-} + one = oneFunctor + {-# INLINE one #-} + +-- | 'Nothing' is treated the same as @'Just' 0@. However, 'zero', 'add', +-- and 'one' preserve 'Nothing' if all inputs are also 'Nothing'. +instance Backprop a => Backprop (Maybe a) where + zero = zeroFunctor + {-# INLINE zero #-} + add x y = (add <$> x <*> y) <|> x <|> y + {-# INLINE add #-} + one = oneFunctor + {-# INLINE one #-} + +-- | 'add' is strict, but 'zero' and 'one' are lazy in their arguments. +instance Backprop () where + zero _ = () + add () () = () + one _ = () + +-- | 'add' is strict +instance (Backprop a, Backprop b) => Backprop (a, b) where + zero (x, y) = (zero x, zero y) + {-# INLINE zero #-} + add (x1, y1) (x2, y2) = (x3, y3) + where + !x3 = add x1 x2 + !y3 = add y1 y2 + {-# INLINE add #-} + one (x, y) = (one x, one y) + {-# INLINE one #-} + +-- | 'add' is strict +instance (Backprop a, Backprop b, Backprop c) => Backprop (a, b, c) where + zero (x, y, z) = (zero x, zero y, zero z) + {-# INLINE zero #-} + add (x1, y1, z1) (x2, y2, z2) = (x3, y3, z3) + where + !x3 = add x1 x2 + !y3 = add y1 y2 + !z3 = add z1 z2 + {-# INLINE add #-} + one (x, y, z) = (one x, one y, one z) + {-# INLINE one #-} + +-- | 'add' is strict +instance (Backprop a, Backprop b, Backprop c, Backprop d) => Backprop (a, b, c, d) where + zero (x, y, z, w) = (zero x, zero y, zero z, zero w) + {-# INLINE zero #-} + add (x1, y1, z1, w1) (x2, y2, z2, w2) = (x3, y3, z3, w3) + where + !x3 = add x1 x2 + !y3 = add y1 y2 + !z3 = add z1 z2 + !w3 = add w1 w2 + {-# INLINE add #-} + one (x, y, z, w) = (one x, one y, one z, one w) + {-# INLINE one #-} + +-- | 'add' is strict +instance (Backprop a, Backprop b, Backprop c, Backprop d, Backprop e) => Backprop (a, b, c, d, e) where + zero (x, y, z, w, v) = (zero x, zero y, zero z, zero w, zero v) + {-# INLINE zero #-} + add (x1, y1, z1, w1, v1) (x2, y2, z2, w2, v2) = (x3, y3, z3, w3, v3) + where + !x3 = add x1 x2 + !y3 = add y1 y2 + !z3 = add z1 z2 + !w3 = add w1 w2 + !v3 = add v1 v2 + {-# INLINE add #-} + one (x, y, z, w, v) = (one x, one y, one z, one w, one v) + {-# INLINE one #-} + +instance Backprop a => Backprop (Identity a) where + zero = coerce (zero @a) + add = coerce (add @a) + one = coerce (one @a) + +instance Backprop (Proxy a) where + zero _ = Proxy + {-# INLINE zero #-} + add _ _ = Proxy + {-# INLINE add #-} + one _ = Proxy + {-# INLINE one #-} + +-- | @since 0.2.2.0 +instance Backprop w => Backprop (Const w a) where + zero = coerce (zero @w) + add = coerce (add @w) + one = coerce (one @w) + +instance Backprop Void where + zero = \case {} + {-# INLINE zero #-} + add = \case {} + {-# INLINE add #-} + one = \case {} + {-# INLINE one #-} + +-- | 'zero' and 'one' replace all current values, and 'add' merges keys +-- from both maps, adding in the case of double-occurrences. +instance (Backprop a, Ord k) => Backprop (M.Map k a) where + zero = zeroFunctor + {-# INLINE zero #-} + add = M.unionWith add + {-# INLINE add #-} + one = oneFunctor + {-# INLINE one #-} + +-- | 'zero' and 'one' replace all current values, and 'add' merges keys +-- from both maps, adding in the case of double-occurrences. +instance Backprop a => Backprop (IM.IntMap a) where + zero = zeroFunctor + {-# INLINE zero #-} + add = IM.unionWith add + {-# INLINE add #-} + one = oneFunctor + {-# INLINE one #-} + +-- | @since 0.2.2.0 +instance Backprop a => Backprop (K1 i a p) where + zero = coerce (zero @a) + add = coerce (add @a) + one = coerce (one @a) + +-- | @since 0.2.2.0 +instance Backprop (f p) => Backprop (M1 i c f p) where + zero = coerce (zero @(f p)) + add = coerce (add @(f p)) + one = coerce (one @(f p)) + +-- | @since 0.2.2.0 +instance (Backprop (f p), Backprop (g p)) => Backprop ((f :*: g) p) + +-- | @since 0.2.6.3 +instance Backprop (f (g a)) => Backprop ((f :.: g) a) where + zero = coerce (zero @(f (g a))) + add = coerce (add @(f (g a))) + one = coerce (one @(f (g a))) + +-- | @since 0.2.2.0 +instance Backprop (V1 p) + +-- | @since 0.2.2.0 +instance Backprop (U1 p) + +-- | @since 0.2.2.0 +instance Backprop a => Backprop (Sum a) where + zero = coerce (zero @a) + add = coerce (add @a) + one = coerce (one @a) + +-- | @since 0.2.2.0 +instance Backprop a => Backprop (Product a) where + zero = coerce (zero @a) + add = coerce (add @a) + one = coerce (one @a) + +#if !MIN_VERSION_base(4,16,0) +-- | @since 0.2.2.0 +instance Backprop a => Backprop (SG.Option a) where + zero = coerce (zero @(Maybe a)) + add = coerce (add @(Maybe a)) + one = coerce (one @(Maybe a)) +#endif + +-- | @since 0.2.2.0 +instance Backprop a => Backprop (SG.First a) where + zero = coerce (zero @a) + add = coerce (add @a) + one = coerce (one @a) + +-- | @since 0.2.2.0 +instance Backprop a => Backprop (SG.Last a) where + zero = coerce (zero @a) + add = coerce (add @a) + one = coerce (one @a) + +-- | @since 0.2.2.0 +instance Backprop a => Backprop (First a) where + zero = coerce (zero @(Maybe a)) + add = coerce (add @(Maybe a)) + one = coerce (one @(Maybe a)) + +-- | @since 0.2.2.0 +instance Backprop a => Backprop (Data.Monoid.Last a) where + zero = coerce (zero @(Maybe a)) + add = coerce (add @(Maybe a)) + one = coerce (one @(Maybe a)) + +-- | @since 0.2.2.0 +instance Backprop a => Backprop (Dual a) where + zero = coerce (zero @a) + add = coerce (add @a) + one = coerce (one @a) + +-- | @since 0.2.2.0 +instance (Backprop a, Backprop b) => Backprop (SG.Arg a b) + +-- | @since 0.2.2.0 +instance (Backprop (f a), Backprop (g a)) => Backprop (DFP.Product f g a) + +-- | @since 0.2.2.0 +instance Backprop (f (g a)) => Backprop (Compose f g a) where + zero = coerce (zero @(f (g a))) + add = coerce (add @(f (g a))) + one = coerce (one @(f (g a))) + +-- | 'add' adds together results; 'zero' and 'one' act on results. +-- +-- @since 0.2.2.0 +instance Backprop a => Backprop (r -> a) where + zero = zeroFunctor + {-# INLINE zero #-} + add = liftA2 add + {-# INLINE add #-} + one = oneFunctor + {-# INLINE one #-} + +-- | @since 0.2.2.0 +instance (Backprop a, Applicative m) => Backprop (Arr.Kleisli m r a) where + zero (Arr.Kleisli f) = Arr.Kleisli ((fmap . fmap) zero f) + {-# INLINE zero #-} + add (Arr.Kleisli f) (Arr.Kleisli g) = Arr.Kleisli $ \x -> + add <$> f x <*> g x + {-# INLINE add #-} + one (Arr.Kleisli f) = Arr.Kleisli ((fmap . fmap) one f) + {-# INLINE one #-} + +-- | @since 0.2.6.3 +instance (ReifyConstraint Backprop f rs, RMap rs, RApply rs) => Backprop (Rec f rs) where + zero = + rmap (\case V.Compose (Dict x) -> zero x) + . reifyConstraint @Backprop + {-# INLINE zero #-} + add xs = + rzipWith (\x -> \case V.Compose (Dict y) -> add x y) xs + . reifyConstraint @Backprop + {-# INLINE add #-} + one = + rmap (\case V.Compose (Dict x) -> one x) + . reifyConstraint @Backprop + {-# INLINE one #-} + +-- | @since 0.2.6.3 +#if MIN_VERSION_vinyl(0,14,2) +instance (ReifyConstraint Backprop f rs, RMap rs, RApply rs, RecApplicative rs, NatToInt (RLength rs), RPureConstrained (IndexableField rs) rs, ToARec rs) + => Backprop (ARec f rs) where + zero = toARec . zero . fromARec + {-# INLINE zero #-} + add xs ys = toARec $ add (fromARec xs) (fromARec ys) + {-# INLINE add #-} + one = toARec . zero . fromARec + {-# INLINE one #-} +#else +instance (ReifyConstraint Backprop f rs, RMap rs, RApply rs, RecApplicative rs, NatToInt (RLength rs), RPureConstrained (IndexableField rs) rs) + => Backprop (ARec f rs) where + zero = toARec . zero . fromARec + {-# INLINE zero #-} + add xs ys = toARec $ add (fromARec xs) (fromARec ys) + {-# INLINE add #-} + one = toARec . zero . fromARec + {-# INLINE one #-} +#endif + +-- | @since 0.2.6.3 +instance + (ReifyConstraint Backprop f rs, RMap rs, RApply rs, VS.Storable (Rec f rs)) => + Backprop (SRec f rs) + where + zero = toSRec . zero . fromSRec + {-# INLINE zero #-} + add xs ys = toSRec $ add (fromSRec xs) (fromSRec ys) + {-# INLINE add #-} + one = toSRec . zero . fromSRec + {-# INLINE one #-} + +-- | @since 0.2.6.3 +instance + (ReifyConstraint Backprop f rs, RMap rs, RApply rs, IsoXRec f rs) => + Backprop (XRec f rs) + where + zero = toXRec . zero . fromXRec + {-# INLINE zero #-} + add xs ys = toXRec $ add (fromXRec xs) (fromXRec ys) + {-# INLINE add #-} + one = toXRec . zero . fromXRec + {-# INLINE one #-} + +-- | @since 0.2.6.3 +instance Backprop a => Backprop (V.Identity a) where + zero = coerce (zero @a) + add = coerce (add @a) + one = coerce (one @a) + +-- | @since 0.2.6.3 +instance Backprop a => Backprop (V.Thunk a) where + zero (V.Thunk x) = V.Thunk (zero x) + add (V.Thunk x) (V.Thunk y) = V.Thunk (add x y) + one (V.Thunk x) = V.Thunk (one x) + +-- | @since 0.2.6.3 +instance Backprop (op (f a) (g a)) => Backprop (V.Lift op f g a) where + zero = coerce (zero @(op (f a) (g a))) + add = coerce (add @(op (f a) (g a))) + one = coerce (one @(op (f a) (g a))) + +-- | @since 0.2.6.3 +instance Backprop t => Backprop (V.ElField '(s, t)) where + zero (V.Field x) = V.Field (zero x) + add (V.Field x) (V.Field y) = V.Field (add x y) + one (V.Field x) = V.Field (one x) + +-- | @since 0.2.6.3 +instance Backprop (f (g a)) => Backprop (V.Compose f g a) where + zero = coerce (zero @(f (g a))) + add = coerce (add @(f (g a))) + one = coerce (one @(f (g a))) + +-- | @since 0.2.6.3 +instance Backprop w => Backprop (V.Const w a) where + zero = coerce (zero @w) + add = coerce (add @w) + one = coerce (one @w) + +-- | @since 0.2.6.3 +instance Backprop (V.HKD t a) => Backprop (V.XData t a) where + zero = coerce (zero @(V.HKD t a)) + add = coerce (add @(V.HKD t a)) + one = coerce (one @(V.HKD t a)) + +-- | @since 0.2.6.3 +instance Backprop (SField field) where + zero _ = SField + add _ _ = SField + one _ = SField + +-- | @since 0.2.6.3 +instance Backprop (Label field) where + zero _ = Label + add _ _ = Label + one _ = Label diff --git a/src/Numeric/Backprop/Explicit.hs b/src/Numeric/Backprop/Explicit.hs --- a/src/Numeric/Backprop/Explicit.hs +++ b/src/Numeric/Backprop/Explicit.hs @@ -1,18 +1,16 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE EmptyCase #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE FunctionalDependencies #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} -{-# OPTIONS_HADDOCK not-home #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_HADDOCK not-home #-} -- | -- Module : Numeric.Backprop.Explicit @@ -41,62 +39,122 @@ -- reflect changes in underlying implementation. -- -- @since 0.2.0.0 - module Numeric.Backprop.Explicit ( - -- * Types - BVar, W, Backprop(..), ABP(..), NumBP(..) - -- * Explicit 'zero', 'add', and 'one' - , ZeroFunc(..), zfNum, zfNums, zeroFunc, zeroFuncs, zfFunctor - , AddFunc(..), afNum, afNums, addFunc, addFuncs - , OneFunc(..), ofNum, ofNums, oneFunc, oneFuncs, ofFunctor - -- * Running - , backprop, evalBP, gradBP, backpropWith - -- ** Multiple inputs - , evalBP0 - , backprop2, evalBP2, gradBP2, backpropWith2 - , backpropN, evalBPN, gradBPN, backpropWithN, RPureConstrained - -- * Manipulating 'BVar' - , constVar, auto, coerceVar - , viewVar, setVar, overVar - , sequenceVar, collectVar - , previewVar, toListOfVar - -- ** With Isomorphisms - , isoVar, isoVar2, isoVar3, isoVarN - -- ** With 'Op's - , liftOp - , liftOp1, liftOp2, liftOp3 - -- ** Generics - , splitBV - , joinBV - , BVGroup - -- * 'Op' - , Op(..) - -- ** Creation - , op0, opConst, idOp - , bpOp - -- *** Giving gradients directly - , op1, op2, op3 - -- *** From Isomorphisms - , opCoerce, opTup, opIso, opIsoN, opLens - -- *** No gradients - , noGrad1, noGrad - -- * Utility - , Reifies - ) where + -- * Types + BVar, + W, + Backprop (..), + ABP (..), + NumBP (..), -import Data.Bifunctor -import Data.Functor.Identity -import Data.Reflection -import Data.Type.Util -import Data.Vinyl.Core -import Data.Vinyl.TypeLevel -import GHC.Generics as G -import Lens.Micro -import Numeric.Backprop.Class -import Numeric.Backprop.Internal -import Numeric.Backprop.Op -import Unsafe.Coerce + -- * Explicit 'zero', 'add', and 'one' + ZeroFunc (..), + zfNum, + zfNums, + zeroFunc, + zeroFuncs, + zfFunctor, + AddFunc (..), + afNum, + afNums, + addFunc, + addFuncs, + OneFunc (..), + ofNum, + ofNums, + oneFunc, + oneFuncs, + ofFunctor, + -- * Running + backprop, + evalBP, + gradBP, + backpropWith, + + -- ** Multiple inputs + evalBP0, + backprop2, + evalBP2, + gradBP2, + backpropWith2, + backpropN, + evalBPN, + gradBPN, + backpropWithN, + RPureConstrained, + + -- * Manipulating 'BVar' + constVar, + auto, + coerceVar, + viewVar, + setVar, + overVar, + sequenceVar, + collectVar, + previewVar, + toListOfVar, + + -- ** With Isomorphisms + isoVar, + isoVar2, + isoVar3, + isoVarN, + + -- ** With 'Op's + liftOp, + liftOp1, + liftOp2, + liftOp3, + + -- ** Generics + splitBV, + joinBV, + BVGroup, + + -- * 'Op' + Op (..), + + -- ** Creation + op0, + opConst, + idOp, + bpOp, + + -- *** Giving gradients directly + op1, + op2, + op3, + + -- *** From Isomorphisms + opCoerce, + opTup, + opIso, + opIsoN, + opLens, + + -- *** No gradients + noGrad1, + noGrad, + + -- * Utility + Reifies, +) where + +import Data.Bifunctor +import Data.Functor.Identity +import Data.Reflection +import Data.Type.Util +import Data.Vinyl.Core +import Data.Vinyl.TypeLevel +import GHC.Generics as G +import Lens.Micro +import Numeric.Backprop.Class +import Numeric.Backprop.Internal +import Numeric.Backprop.Op +import Unsafe.Coerce + -- | 'ZeroFunc's for every item in a type level list based on their -- 'Num' instances -- @@ -161,42 +219,45 @@ {-# INLINE auto #-} -- | 'Numeric.Backprop.backpropN', but with explicit 'zero' and 'one'. -backpropN - :: forall as b. () - => Rec ZeroFunc as - -> OneFunc b - -> (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> (b, Rec Identity as) +backpropN :: + forall as b. + () => + Rec ZeroFunc as -> + OneFunc b -> + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + (b, Rec Identity as) backpropN zfs ob f xs = case backpropWithN zfs f xs of - (y, g) -> (y, g (runOF ob y)) + (y, g) -> (y, g (runOF ob y)) {-# INLINE backpropN #-} -- | 'Numeric.Backprop.backprop', but with explicit 'zero' and 'one'. -backprop - :: ZeroFunc a - -> OneFunc b - -> (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> (b, a) -backprop zfa ofb f = second (\case Identity x :& RNil -> x) - . backpropN (zfa :& RNil) ofb (f . (\case x :& RNil -> x)) - . (:& RNil) - . Identity +backprop :: + ZeroFunc a -> + OneFunc b -> + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + (b, a) +backprop zfa ofb f = + second (\case Identity x :& RNil -> x) + . backpropN (zfa :& RNil) ofb (f . (\case x :& RNil -> x)) + . (:& RNil) + . Identity {-# INLINE backprop #-} -- | 'Numeric.Backprop.backpropWith', but with explicit 'zero'. -- -- Note that argument order changed in v0.2.4. -backpropWith - :: ZeroFunc a - -> (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> (b, b -> a) -backpropWith zfa f = second ((\case Identity x :& RNil -> x) .) - . backpropWithN (zfa :& RNil) (f . (\case x :& RNil -> x)) - . (:& RNil) - . Identity +backpropWith :: + ZeroFunc a -> + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + (b, b -> a) +backpropWith zfa f = + second ((\case Identity x :& RNil -> x) .) + . backpropWithN (zfa :& RNil) (f . (\case x :& RNil -> x)) + . (:& RNil) + . Identity {-# INLINE backpropWith #-} -- | 'evalBP' but with no arguments. Useful when everything is just given @@ -214,42 +275,45 @@ -- -- See documentation of 'Numeric.Backprop.backprop' for more information. evalBP :: (forall s. Reifies s W => BVar s a -> BVar s b) -> a -> b -evalBP f = evalBPN (f . (\case x :& RNil -> x)) . (:& RNil) . Identity +evalBP f = evalBPN (f . (\case x :& RNil -> x)) . (:& RNil) . Identity {-# INLINE evalBP #-} -- | 'Numeric.Backprop.gradBP', but with explicit 'zero' and 'one'. -gradBP - :: ZeroFunc a - -> OneFunc b - -> (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> a +gradBP :: + ZeroFunc a -> + OneFunc b -> + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + a gradBP zfa ofb f = snd . backprop zfa ofb f {-# INLINE gradBP #-} -- | 'Numeric.Backprop.gradBP', Nbut with explicit 'zero' and 'one'. -gradBPN - :: Rec ZeroFunc as - -> OneFunc b - -> (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> Rec Identity as +gradBPN :: + Rec ZeroFunc as -> + OneFunc b -> + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + Rec Identity as gradBPN zfas ofb f = snd . backpropN zfas ofb f {-# INLINE gradBPN #-} -- | 'Numeric.Backprop.backprop2', but with explicit 'zero' and 'one'. -backprop2 - :: ZeroFunc a - -> ZeroFunc b - -> OneFunc c - -> (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (c, (a, b)) -backprop2 zfa zfb ofc f x y = second (\(Identity dx :& Identity dy :& RNil) -> (dx, dy)) $ - backpropN (zfa :& zfb :& RNil) ofc - (\(x' :& y' :& RNil) -> f x' y') - (Identity x :& Identity y :& RNil) +backprop2 :: + ZeroFunc a -> + ZeroFunc b -> + OneFunc c -> + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + (c, (a, b)) +backprop2 zfa zfb ofc f x y = + second (\(Identity dx :& Identity dy :& RNil) -> (dx, dy)) $ + backpropN + (zfa :& zfb :& RNil) + ofc + (\(x' :& y' :& RNil) -> f x' y') + (Identity x :& Identity y :& RNil) {-# INLINE backprop2 #-} -- | 'Numeric.Backprop.backpropWith2', but with explicit 'zero'. @@ -257,114 +321,118 @@ -- Note that argument order changed in v0.2.4. -- -- @since 0.2.0.0 -backpropWith2 - :: ZeroFunc a - -> ZeroFunc b - -> (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (c, c -> (a, b)) -backpropWith2 zfa zfb f x y = second ((\(Identity dx :& Identity dy :& RNil) -> (dx, dy)) .) $ - backpropWithN (zfa :& zfb :& RNil) - (\(x' :& y' :& RNil) -> f x' y') - (Identity x :& Identity y :& RNil) +backpropWith2 :: + ZeroFunc a -> + ZeroFunc b -> + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + (c, c -> (a, b)) +backpropWith2 zfa zfb f x y = + second ((\(Identity dx :& Identity dy :& RNil) -> (dx, dy)) .) $ + backpropWithN + (zfa :& zfb :& RNil) + (\(x' :& y' :& RNil) -> f x' y') + (Identity x :& Identity y :& RNil) {-# INLINE backpropWith2 #-} -- | 'evalBP' for a two-argument function. See -- 'Numeric.Backprop.backprop2' for notes. -evalBP2 - :: (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> c -evalBP2 f x y = evalBPN (\(x' :& y' :& RNil) -> f x' y') $ Identity x - :& Identity y - :& RNil +evalBP2 :: + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + c +evalBP2 f x y = + evalBPN (\(x' :& y' :& RNil) -> f x' y') $ + Identity x + :& Identity y + :& RNil {-# INLINE evalBP2 #-} -- | 'Numeric.Backprop.gradBP2' with explicit 'zero' and 'one'. -gradBP2 - :: ZeroFunc a - -> ZeroFunc b - -> OneFunc c - -> (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (a, b) +gradBP2 :: + ZeroFunc a -> + ZeroFunc b -> + OneFunc c -> + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + (a, b) gradBP2 zfa zfb ofc f x = snd . backprop2 zfa zfb ofc f x {-# INLINE gradBP2 #-} -- | 'Numeric.Backprop.bpOp' with explicit 'zero'. -bpOp - :: Rec ZeroFunc as - -> (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Op as b +bpOp :: + Rec ZeroFunc as -> + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Op as b bpOp zfs f = Op (backpropWithN zfs f) {-# INLINE bpOp #-} -- | 'Numeric.Backprop.overVar' with explicit 'add' and 'zero'. -- -- @since 0.2.4.0 -overVar - :: Reifies s W - => AddFunc a - -> AddFunc b - -> ZeroFunc a - -> ZeroFunc b - -> Lens' b a - -> (BVar s a -> BVar s a) - -> BVar s b - -> BVar s b +overVar :: + Reifies s W => + AddFunc a -> + AddFunc b -> + ZeroFunc a -> + ZeroFunc b -> + Lens' b a -> + (BVar s a -> BVar s a) -> + BVar s b -> + BVar s b overVar afa afb zfa zfb l f x = setVar afa afb zfa l (f (viewVar afa zfb l x)) x {-# INLINE overVar #-} -- | 'Numeric.Backprop.isoVar' with explicit 'add' and 'zero'. -isoVar - :: Reifies s W - => AddFunc a - -> (a -> b) - -> (b -> a) - -> BVar s a - -> BVar s b +isoVar :: + Reifies s W => + AddFunc a -> + (a -> b) -> + (b -> a) -> + BVar s a -> + BVar s b isoVar af f g = liftOp1 af (opIso f g) {-# INLINE isoVar #-} -- | 'Numeric.Backprop.isoVar2' with explicit 'add' and 'zero'. -isoVar2 - :: Reifies s W - => AddFunc a - -> AddFunc b - -> (a -> b -> c) - -> (c -> (a, b)) - -> BVar s a - -> BVar s b - -> BVar s c +isoVar2 :: + Reifies s W => + AddFunc a -> + AddFunc b -> + (a -> b -> c) -> + (c -> (a, b)) -> + BVar s a -> + BVar s b -> + BVar s c isoVar2 afa afb f g = liftOp2 afa afb (opIso2 f g) {-# INLINE isoVar2 #-} -- | 'Numeric.Backprop.isoVar3' with explicit 'add' and 'zero'. -isoVar3 - :: Reifies s W - => AddFunc a - -> AddFunc b - -> AddFunc c - -> (a -> b -> c -> d) - -> (d -> (a, b, c)) - -> BVar s a - -> BVar s b - -> BVar s c - -> BVar s d +isoVar3 :: + Reifies s W => + AddFunc a -> + AddFunc b -> + AddFunc c -> + (a -> b -> c -> d) -> + (d -> (a, b, c)) -> + BVar s a -> + BVar s b -> + BVar s c -> + BVar s d isoVar3 afa afb afc f g = liftOp3 afa afb afc (opIso3 f g) {-# INLINE isoVar3 #-} -- | 'Numeric.Backprop.isoVarN' with explicit 'add' and 'zero'. -isoVarN - :: Reifies s W - => Rec AddFunc as - -> (Rec Identity as -> b) - -> (b -> Rec Identity as) - -> Rec (BVar s) as - -> BVar s b +isoVarN :: + Reifies s W => + Rec AddFunc as -> + (Rec Identity as -> b) -> + (b -> Rec Identity as) -> + Rec (BVar s) as -> + BVar s b isoVarN afs f g = liftOp afs (opIsoN f g) {-# INLINE isoVarN #-} @@ -379,132 +447,155 @@ -- -- @since 0.2.2.0 class BVGroup s as i o | o -> i, i -> as where - -- | Helper method for generically "splitting" 'BVar's out of - -- constructors inside a 'BVar'. See 'splitBV'. - gsplitBV :: Rec AddFunc as -> Rec ZeroFunc as -> BVar s (i ()) -> o () - -- | Helper method for generically "joining" 'BVar's inside - -- a constructor into a 'BVar'. See 'joinBV'. - gjoinBV :: Rec AddFunc as -> Rec ZeroFunc as -> o () -> BVar s (i ()) + -- | Helper method for generically "splitting" 'BVar's out of + -- constructors inside a 'BVar'. See 'splitBV'. + gsplitBV :: Rec AddFunc as -> Rec ZeroFunc as -> BVar s (i ()) -> o () + -- | Helper method for generically "joining" 'BVar's inside + -- a constructor into a 'BVar'. See 'joinBV'. + gjoinBV :: Rec AddFunc as -> Rec ZeroFunc as -> o () -> BVar s (i ()) + instance BVGroup s '[] (K1 i a) (K1 i (BVar s a)) where - gsplitBV _ _ = K1 . coerceVar - {-# INLINE gsplitBV #-} - gjoinBV _ _ = coerceVar . unK1 - {-# INLINE gjoinBV #-} + gsplitBV _ _ = K1 . coerceVar + {-# INLINE gsplitBV #-} + gjoinBV _ _ = coerceVar . unK1 + {-# INLINE gjoinBV #-} -instance BVGroup s as i o - => BVGroup s as (M1 p c i) (M1 p c o) where - gsplitBV afs zfs = M1 . gsplitBV afs zfs . coerceVar @_ @(i ()) - {-# INLINE gsplitBV #-} - gjoinBV afs zfs = coerceVar @(i ()) . gjoinBV afs zfs . unM1 - {-# INLINE gjoinBV #-} +instance + BVGroup s as i o => + BVGroup s as (M1 p c i) (M1 p c o) + where + gsplitBV afs zfs = M1 . gsplitBV afs zfs . coerceVar @_ @(i ()) + {-# INLINE gsplitBV #-} + gjoinBV afs zfs = coerceVar @(i ()) . gjoinBV afs zfs . unM1 + {-# INLINE gjoinBV #-} instance BVGroup s '[] V1 V1 where - gsplitBV _ _ = unsafeCoerce - {-# INLINE gsplitBV #-} - gjoinBV _ _ = \case - {-# INLINE gjoinBV #-} + gsplitBV _ _ = unsafeCoerce + {-# INLINE gsplitBV #-} + gjoinBV _ _ = \case {} + {-# INLINE gjoinBV #-} instance BVGroup s '[] U1 U1 where - gsplitBV _ _ _ = U1 - {-# INLINE gsplitBV #-} - gjoinBV _ _ _ = constVar U1 - {-# INLINE gjoinBV #-} + gsplitBV _ _ _ = U1 + {-# INLINE gsplitBV #-} + gjoinBV _ _ _ = constVar U1 + {-# INLINE gjoinBV #-} -instance ( Reifies s W - , BVGroup s as i1 o1 - , BVGroup s bs i2 o2 - , cs ~ (as ++ bs) - , RecApplicative as - ) => BVGroup s (i1 () ': i2 () ': cs) (i1 :*: i2) (o1 :*: o2) where - gsplitBV (afa :& afb :& afs) (zfa :& zfb :& zfs) xy = x :*: y - where - (afas, afbs) = splitRec afs - (zfas, zfbs) = splitRec zfs - zfab = ZF $ \(xx :*: yy) -> runZF zfa xx :*: runZF zfb yy - x = gsplitBV afas zfas . viewVar afa zfab p1 $ xy - y = gsplitBV afbs zfbs . viewVar afb zfab p2 $ xy - {-# INLINE gsplitBV #-} - gjoinBV (afa :& afb :& afs) (_ :& _ :& zfs) (x :*: y) - = isoVar2 afa afb (:*:) unP - (gjoinBV afas zfas x) - (gjoinBV afbs zfbs y) - where - (afas, afbs) = splitRec afs - (zfas, zfbs) = splitRec zfs - unP (xx :*: yy) = (xx, yy) - {-# INLINE gjoinBV #-} +instance + ( Reifies s W + , BVGroup s as i1 o1 + , BVGroup s bs i2 o2 + , cs ~ (as ++ bs) + , RecApplicative as + ) => + BVGroup s (i1 () ': i2 () ': cs) (i1 :*: i2) (o1 :*: o2) + where + gsplitBV (afa :& afb :& afs) (zfa :& zfb :& zfs) xy = x :*: y + where + (afas, afbs) = splitRec afs + (zfas, zfbs) = splitRec zfs + zfab = ZF $ \(xx :*: yy) -> runZF zfa xx :*: runZF zfb yy + x = gsplitBV afas zfas . viewVar afa zfab p1 $ xy + y = gsplitBV afbs zfbs . viewVar afb zfab p2 $ xy + {-# INLINE gsplitBV #-} + gjoinBV (afa :& afb :& afs) (_ :& _ :& zfs) (x :*: y) = + isoVar2 + afa + afb + (:*:) + unP + (gjoinBV afas zfas x) + (gjoinBV afbs zfbs y) + where + (afas, afbs) = splitRec afs + (zfas, zfbs) = splitRec zfs + unP (xx :*: yy) = (xx, yy) + {-# INLINE gjoinBV #-} -- | This instance is possible but it is not clear when it would be useful -instance ( Reifies s W - , BVGroup s as i1 o1 - , BVGroup s bs i2 o2 - , cs ~ (as ++ bs) - , RecApplicative as - ) => BVGroup s (i1 () ': i2 () ': cs) (i1 :+: i2) (o1 :+: o2) where - gsplitBV (afa :& afb :& afs) (zfa :& zfb :& zfs) xy = - case previewVar afa zf s1 xy of - Just x -> L1 $ gsplitBV afas zfas x - Nothing -> case previewVar afb zf s2 xy of - Just y -> R1 $ gsplitBV afbs zfbs y - Nothing -> error "Numeric.Backprop.gsplitBV: Internal error occurred" - where - zf = ZF $ \case - L1 xx -> L1 $ runZF zfa xx - R1 yy -> R1 $ runZF zfb yy - (afas, afbs) = splitRec afs - (zfas, zfbs) = splitRec zfs - {-# INLINE gsplitBV #-} - gjoinBV (afa :& afb :& afs) (zfa :& zfb :& zfs) = \case - L1 x -> liftOp1 afa (op1 (\xx -> (L1 xx, \case L1 d -> d; R1 _ -> runZF zfa xx))) - (gjoinBV afas zfas x) - R1 y -> liftOp1 afb (op1 (\yy -> (R1 yy, \case L1 _ -> runZF zfb yy; R1 d -> d))) - (gjoinBV afbs zfbs y) - where - (afas, afbs) = splitRec afs - (zfas, zfbs) = splitRec zfs - {-# INLINE gjoinBV #-} +instance + ( Reifies s W + , BVGroup s as i1 o1 + , BVGroup s bs i2 o2 + , cs ~ (as ++ bs) + , RecApplicative as + ) => + BVGroup s (i1 () ': i2 () ': cs) (i1 :+: i2) (o1 :+: o2) + where + gsplitBV (afa :& afb :& afs) (zfa :& zfb :& zfs) xy = + case previewVar afa zf s1 xy of + Just x -> L1 $ gsplitBV afas zfas x + Nothing -> case previewVar afb zf s2 xy of + Just y -> R1 $ gsplitBV afbs zfbs y + Nothing -> error "Numeric.Backprop.gsplitBV: Internal error occurred" + where + zf = ZF $ \case + L1 xx -> L1 $ runZF zfa xx + R1 yy -> R1 $ runZF zfb yy + (afas, afbs) = splitRec afs + (zfas, zfbs) = splitRec zfs + {-# INLINE gsplitBV #-} + gjoinBV (afa :& afb :& afs) (zfa :& zfb :& zfs) = \case + L1 x -> + liftOp1 + afa + (op1 (\xx -> (L1 xx, \case L1 d -> d; R1 _ -> runZF zfa xx))) + (gjoinBV afas zfas x) + R1 y -> + liftOp1 + afb + (op1 (\yy -> (R1 yy, \case L1 _ -> runZF zfb yy; R1 d -> d))) + (gjoinBV afbs zfbs y) + where + (afas, afbs) = splitRec afs + (zfas, zfbs) = splitRec zfs + {-# INLINE gjoinBV #-} -- | 'Numeric.Backprop.splitBV' with explicit 'add' and 'zero'. -- -- @since 0.2.2.0 -splitBV - :: forall z f s as. - ( Generic (z f) - , Generic (z (BVar s)) - , BVGroup s as (Rep (z f)) (Rep (z (BVar s))) - , Reifies s W - ) - => AddFunc (Rep (z f) ()) - -> Rec AddFunc as - -> ZeroFunc (z f) - -> Rec ZeroFunc as - -> BVar s (z f) -- ^ 'BVar' of value - -> z (BVar s) -- ^ 'BVar's of fields +splitBV :: + forall z f s as. + ( Generic (z f) + , Generic (z (BVar s)) + , BVGroup s as (Rep (z f)) (Rep (z (BVar s))) + , Reifies s W + ) => + AddFunc (Rep (z f) ()) -> + Rec AddFunc as -> + ZeroFunc (z f) -> + Rec ZeroFunc as -> + -- | 'BVar' of value + BVar s (z f) -> + -- | 'BVar's of fields + z (BVar s) splitBV af afs zf zfs = - G.to - . gsplitBV afs zfs - . viewVar af zf (lens (from @(z f) @()) (const G.to)) + G.to + . gsplitBV afs zfs + . viewVar af zf (lens (from @(z f) @()) (const G.to)) {-# INLINE splitBV #-} -- | 'Numeric.Backprop.joinBV' with explicit 'add' and 'zero'. -- -- @since 0.2.2.0 -joinBV - :: forall z f s as. - ( Generic (z f) - , Generic (z (BVar s)) - , BVGroup s as (Rep (z f)) (Rep (z (BVar s))) - , Reifies s W - ) - => AddFunc (z f) - -> Rec AddFunc as - -> ZeroFunc (Rep (z f) ()) - -> Rec ZeroFunc as - -> z (BVar s) -- ^ 'BVar's of fields - -> BVar s (z f) -- ^ 'BVar' of combined value +joinBV :: + forall z f s as. + ( Generic (z f) + , Generic (z (BVar s)) + , BVGroup s as (Rep (z f)) (Rep (z (BVar s))) + , Reifies s W + ) => + AddFunc (z f) -> + Rec AddFunc as -> + ZeroFunc (Rep (z f) ()) -> + Rec ZeroFunc as -> + -- | 'BVar's of fields + z (BVar s) -> + -- | 'BVar' of combined value + BVar s (z f) joinBV af afs zf zfs = - viewVar af zf (lens G.to (const from)) - . gjoinBV afs zfs - . from @(z (BVar s)) @() + viewVar af zf (lens G.to (const from)) + . gjoinBV afs zfs + . from @(z (BVar s)) @() {-# INLINE joinBV #-} diff --git a/src/Numeric/Backprop/Internal.hs b/src/Numeric/Backprop/Internal.hs --- a/src/Numeric/Backprop/Internal.hs +++ b/src/Numeric/Backprop/Internal.hs @@ -1,873 +1,974 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DeriveDataTypeable #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeInType #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_HADDOCK not-home #-} - --- | --- Module : Numeric.Backprop.Internal --- Copyright : (c) Justin Le 2023 --- 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 ( - BVar - , W - , backpropWithN, evalBPN - , constVar - , liftOp, liftOp1, liftOp2, liftOp3 - , viewVar, setVar, sequenceVar, collectVar, previewVar, toListOfVar - , coerceVar - -- * Func wrappers - , ZeroFunc(..), zfNum, zeroFunc - , AddFunc(..), afNum, addFunc - , OneFunc(..), ofNum, oneFunc - -- * Debug - , debugSTN - , debugIR - ) where - -import Control.DeepSeq -import Control.Exception -import Control.Monad -import Control.Monad.ST -import Control.Monad.Trans.State -import Data.Bifunctor -import Data.Coerce -import Data.Foldable -import Data.Function -import Data.Functor.Identity -import Data.IORef -import Data.Kind -import Data.Maybe -import Data.Monoid hiding (Any(..)) -import Data.Proxy -import Data.Reflection -import Data.Type.Util -import Data.Typeable -import Data.Vinyl.Core -import GHC.Exts (Any) -import GHC.Generics as G -import Lens.Micro -import Lens.Micro.Extras -import Numeric.Backprop.Class -import Numeric.Backprop.Op -import System.IO.Unsafe -import Unsafe.Coerce -import qualified Data.Vector as V -import qualified Data.Vector.Mutable as MV -import qualified Data.Vinyl.Recursive as VR -import qualified Data.Vinyl.XRec as X - --- | "Zero out" all components of a value. For scalar values, this should --- just be @'const' 0@. For vectors and matrices, this should set all --- components to zero, the additive identity. --- --- Should be idempotent: Applying the function twice is the same as --- applying it just once. --- --- Each type should ideally only have one 'ZeroFunc'. This coherence --- constraint is given by the typeclass 'Backprop'. --- --- @since 0.2.0.0 -newtype ZeroFunc a = ZF { runZF :: a -> a } - --- | Add together two values of a type. To combine contributions of --- gradients, so should ideally be information-preserving. --- --- See laws for 'Backprop' for the laws this should be expected to --- preserve. Namely, it should be commutative and associative, with an --- identity for a valid 'ZeroFunc'. --- --- Each type should ideally only have one 'AddFunc'. This coherence --- constraint is given by the typeclass 'Backprop'. --- --- @since 0.2.0.0 -newtype AddFunc a = AF { runAF :: a -> a -> a } - --- | "One" all components of a value. For scalar values, this should --- just be @'const' 1@. For vectors and matrices, this should set all --- components to one, the multiplicative identity. --- --- Should be idempotent: Applying the function twice is the same as --- applying it just once. --- --- Each type should ideally only have one 'OneFunc'. This coherence --- constraint is given by the typeclass 'Backprop'. --- --- @since 0.2.0.0 -newtype OneFunc a = OF { runOF :: a -> a } - --- | If a type has a 'Num' instance, this is the canonical 'ZeroFunc'. --- --- @since 0.2.0.0 -zfNum :: Num a => ZeroFunc a -zfNum = ZF (const 0) -{-# INLINE zfNum #-} - --- | If a type has a 'Num' instance, this is the canonical 'AddFunc'. --- --- @since 0.2.0.0 -afNum :: Num a => AddFunc a -afNum = AF (+) -{-# INLINE afNum #-} - --- | If a type has a 'Num' instance, this is the canonical 'OneFunc'. --- --- @since 0.2.0.0 -ofNum :: Num a => OneFunc a -ofNum = OF (const 1) -{-# INLINE ofNum #-} - --- | A @'BVar' s a@ is a value of type @a@ that can be "backpropagated". --- --- Functions referring to 'BVar's are tracked by the library and can be --- automatically differentiated to get their gradients and results. --- --- For simple numeric values, you can use its 'Num', 'Fractional', and --- 'Floating' instances to manipulate them as if they were the numbers they --- represent. --- --- If @a@ contains items, the items can be accessed and extracted using --- lenses. A @'Lens'' b a@ can be used to access an @a@ inside a @b@, using --- '^^.' ('Numeric.Backprop.viewVar'): --- --- @ --- ('^.') :: a -> 'Lens'' a b -> b --- ('^^.') :: 'BVar' s a -> 'Lens'' a b -> 'BVar' s b --- @ --- --- There is also '^^?' ('Numeric.Backprop.previewVar'), to use a 'Prism'' --- or 'Traversal'' to extract a target that may or may not be present --- (which can implement pattern matching), '^^..' --- ('Numeric.Backprop.toListOfVar') to use a 'Traversal'' to extract /all/ --- targets inside a 'BVar', and '.~~' ('setVar') to set and update values --- inside a 'BVar'. --- --- If you have control over your data type definitions, you can also use --- 'Numeric.Backprop.splitBV' and 'Numeric.Backprop.joinBV' to manipulate --- data types by easily extracting fields out of a 'BVar' of data types and --- creating 'BVar's of data types out of 'BVar's of their fields. See --- "Numeric.Backprop#hkd" for a tutorial on this use pattern. --- --- For more complex operations, libraries can provide functions on 'BVar's --- using 'Numeric.Backprop.liftOp' and related functions. This is how you --- can create primitive functions that users can use to manipulate your --- library's values. See --- for a detailed --- guide. --- --- For example, the /hmatrix/ library has a matrix-vector multiplication --- function, @#> :: L m n -> R n -> L m@. --- --- A library could instead provide a function @#> :: 'BVar' (L m n) -> BVar --- (R n) -> BVar (R m)@, which the user can then use to manipulate their --- 'BVar's of @L m n@s and @R n@s, etc. --- --- See "Numeric.Backprop#liftops" and documentation for --- 'Numeric.Backprop.liftOp' for more information. --- -data BVar s a = BV { _bvRef :: !(BRef s) - , _bvVal :: !a - } - --- | @since 0.1.5.1 -deriving instance Typeable (BVar s a) - --- | @since 0.2.6.3 -instance X.IsoHKD (BVar s) a - -data BRef (s :: Type) = BRInp !Int - | BRIx !Int - | BRC - deriving (Generic, Show) - -instance NFData (BRef s) - --- | This will force the value inside, as well. -instance NFData a => NFData (BVar s a) where - rnf (BV r v) = force r `seq` force v `seq` () - --- | Project out a constant value if the 'BVar' refers to one. -bvConst :: BVar s a -> Maybe a -bvConst (BV BRC !x) = Just x -bvConst _ = Nothing -{-# INLINE bvConst #-} - -forceBVar :: BVar s a -> () -forceBVar (BV r !_) = force r `seq` () -{-# INLINE forceBVar #-} - -data InpRef :: Type -> Type where - IR :: { _irIx :: !(BVar s b) - , _irAdd :: !(a -> b -> b) - , _irEmbed :: !(a -> b) - } - -> InpRef a - -forceInpRef :: InpRef a -> () -forceInpRef (IR v !_ !_) = forceBVar v `seq` () -{-# INLINE forceInpRef #-} - --- | Debugging string for an 'InpRef'. -debugIR :: InpRef a -> String -debugIR IR{..} = show (_bvRef _irIx) - -data TapeNode :: Type -> Type where - TN :: { _tnInputs :: !(Rec InpRef as) - , _tnGrad :: !(a -> Rec Identity as) - } - -> TapeNode a - -forceTapeNode :: TapeNode a -> () -forceTapeNode (TN inps !_) = VR.rfoldMap forceInpRef inps `seq` () -{-# INLINE forceTapeNode #-} - -data SomeTapeNode :: Type where - STN :: { _stnNode :: !(TapeNode a) - } - -> SomeTapeNode - -forceSomeTapeNode :: SomeTapeNode -> () -forceSomeTapeNode (STN n) = forceTapeNode n - --- | Debugging string for a 'SomeTapeMode'. -debugSTN :: SomeTapeNode -> String -debugSTN (STN TN{..}) = show . VR.rfoldMap ((:[]) . debugIR) $ _tnInputs - --- | An ephemeral Wengert Tape in the environment. Used internally to --- track of the computational graph of variables. --- --- For the end user, one can just imagine @'Reifies' s 'W'@ as a required --- constraint on @s@ that allows backpropagation to work. -newtype W = W { wRef :: IORef (Int, [SomeTapeNode]) } - -initWengert :: IO W -initWengert = W <$> newIORef (0,[]) -{-# INLINE initWengert #-} - -insertNode - :: TapeNode a - -> a -- ^ val - -> W - -> IO (BVar s a) -insertNode tn !x !w = fmap ((`BV` x) . BRIx) . atomicModifyIORef' (wRef w) $ \(!n,!t) -> - let n' = n + 1 - t' = STN tn : t - in forceTapeNode tn `seq` n' `seq` t' `seq` ((n', t'), n) -{-# INLINE insertNode #-} - --- | Lift a value into a 'BVar' representing a constant value. --- --- This value will not be considered an input, and its gradients will not --- be backpropagated. -constVar :: a -> BVar s a -constVar = BV BRC -{-# INLINE constVar #-} - -liftOp_ - :: forall s as b. Reifies s W - => Rec AddFunc as - -> Op as b - -> Rec (BVar s) as - -> IO (BVar s b) -liftOp_ afs o !vs = case rtraverse (fmap Identity . bvConst) vs of - Just xs -> return $ constVar (evalOp o xs) - Nothing -> insertNode tn y (reflect (Proxy @s)) - where - (y,g) = runOpWith o (VR.rmap (Identity . _bvVal) vs) - tn = TN { _tnInputs = VR.rzipWith go afs vs - , _tnGrad = g - } - go :: forall a. AddFunc a -> BVar s a -> InpRef a - go af !v = forceBVar v `seq` IR v (runAF af) id - {-# INLINE go #-} -{-# INLINE liftOp_ #-} - --- | 'Numeric.Backprop.liftOp', but with explicit 'add' and 'zero'. -liftOp - :: forall as b s. Reifies s W - => Rec AddFunc as - -> Op as b - -> Rec (BVar s) as - -> BVar s b -liftOp afs o !vs = unsafePerformIO $ liftOp_ afs o vs -{-# INLINE liftOp #-} - -liftOp1_ - :: forall a b s. Reifies s W - => AddFunc a - -> Op '[a] b - -> BVar s a - -> IO (BVar s b) -liftOp1_ _ o (bvConst->Just x) = return . constVar . evalOp o $ (Identity x :& RNil) -liftOp1_ af o v = forceBVar v `seq` insertNode tn y (reflect (Proxy @s)) - where - (y,g) = runOpWith o (Identity (_bvVal v) :& RNil) - tn = TN { _tnInputs = IR v (runAF af) id :& RNil - , _tnGrad = g - } -{-# INLINE liftOp1_ #-} - --- | 'Numeric.Backprop.liftOp1', but with explicit 'add' and 'zero'. -liftOp1 - :: forall a b s. Reifies s W - => AddFunc a - -> Op '[a] b - -> BVar s a - -> BVar s b -liftOp1 af o !v = unsafePerformIO $ liftOp1_ af o v -{-# INLINE liftOp1 #-} - -liftOp2_ - :: forall a b c s. Reifies s W - => AddFunc a - -> AddFunc b - -> Op '[a,b] c - -> BVar s a - -> BVar s b - -> IO (BVar s c) -liftOp2_ _ _ o (bvConst->Just x) (bvConst->Just y) - = return . constVar . evalOp o $ Identity x :& Identity y :& RNil -liftOp2_ afa afb o v u = forceBVar v - `seq` forceBVar u - `seq` insertNode tn y (reflect (Proxy @s)) - where - (y,g) = runOpWith o $ Identity (_bvVal v) - :& Identity (_bvVal u) - :& RNil - tn = TN { _tnInputs = IR v (runAF afa) id :& IR u (runAF afb) id :& RNil - , _tnGrad = g - } -{-# INLINE liftOp2_ #-} - --- | 'Numeric.Backprop.liftOp2', but with explicit 'add' and 'zero'. -liftOp2 - :: forall a b c s. Reifies s W - => AddFunc a - -> AddFunc b - -> Op '[a,b] c - -> BVar s a - -> BVar s b - -> BVar s c -liftOp2 afa afb o !v !u = unsafePerformIO $ liftOp2_ afa afb o v u -{-# INLINE liftOp2 #-} - -liftOp3_ - :: forall a b c d s. Reifies s W - => AddFunc a - -> AddFunc b - -> AddFunc c - -> Op '[a,b,c] d - -> BVar s a - -> BVar s b - -> BVar s c - -> IO (BVar s d) -liftOp3_ _ _ _ o (bvConst->Just x) (bvConst->Just y) (bvConst->Just z) - = return . constVar . evalOp o $ Identity x - :& Identity y - :& Identity z - :& RNil -liftOp3_ afa afb afc o v u w = forceBVar v - `seq` forceBVar u - `seq` forceBVar w - `seq` insertNode tn y (reflect (Proxy @s)) - where - (y, g) = runOpWith o $ Identity (_bvVal v) - :& Identity (_bvVal u) - :& Identity (_bvVal w) - :& RNil - tn = TN { _tnInputs = IR v (runAF afa) id - :& IR u (runAF afb) id - :& IR w (runAF afc) id - :& RNil - , _tnGrad = g - } -{-# INLINE liftOp3_ #-} - --- | 'Numeric.Backprop.liftOp3', but with explicit 'add' and 'zero'. -liftOp3 - :: forall a b c d s. Reifies s W - => AddFunc a - -> AddFunc b - -> AddFunc c - -> Op '[a,b,c] d - -> BVar s a - -> BVar s b - -> BVar s c - -> BVar s d -liftOp3 afa afb afc o !v !u !w = unsafePerformIO $ liftOp3_ afa afb afc o v u w -{-# INLINE liftOp3 #-} - --- TODO: can we get the zero and add func from the bvar? -viewVar_ - :: forall a b s. Reifies s W - => AddFunc a - -> ZeroFunc b - -> Lens' b a - -> BVar s b - -> IO (BVar s a) -viewVar_ af z l v = forceBVar v `seq` insertNode tn y (reflect (Proxy @s)) - where - x = _bvVal v - y = x ^. l - tn = TN { _tnInputs = IR v (over l . runAF af) (\g -> set l g (runZF z x)) - :& RNil - , _tnGrad = (:& RNil) . Identity - } -{-# INLINE viewVar_ #-} - --- | 'Numeric.Backprop.viewVar', but with explicit 'add' and 'zero'. -viewVar - :: forall a b s. Reifies s W - => AddFunc a - -> ZeroFunc b - -> Lens' b a - -> BVar s b - -> BVar s a -viewVar af z l !v = unsafePerformIO $ viewVar_ af z l v -{-# INLINE viewVar #-} - --- TODO: can zero and add func be gotten from the input bvars? -setVar_ - :: forall a b s. Reifies s W - => AddFunc a - -> AddFunc b - -> ZeroFunc a - -> Lens' b a - -> BVar s a - -> BVar s b - -> IO (BVar s b) -setVar_ afa afb za l w v = forceBVar v - `seq` forceBVar w - `seq` insertNode tn y (reflect (Proxy @s)) - where - y = _bvVal v & l .~ _bvVal w - tn = TN { _tnInputs = IR w (runAF afa) id - :& IR v (runAF afb) id - :& RNil - , _tnGrad = \d -> let (dw,dv) = l (\x -> (x, runZF za x)) d - in Identity dw :& Identity dv :& RNil - } -{-# INLINE setVar_ #-} - --- | 'Numeric.Backprop.setVar', but with explicit 'add' and 'zero'. -setVar - :: forall a b s. Reifies s W - => AddFunc a - -> AddFunc b - -> ZeroFunc a - -> Lens' b a - -> BVar s a - -> BVar s b - -> BVar s b -setVar afa afb za l !w !v = unsafePerformIO $ setVar_ afa afb za l w v -{-# INLINE setVar #-} - --- | 'Numeric.Backprop.sequenceVar', but with explicit 'add' and 'zero'. -sequenceVar - :: forall t a s. (Reifies s W, Traversable t) - => AddFunc a - -> ZeroFunc a - -> BVar s (t a) - -> t (BVar s a) -sequenceVar af z !v = unsafePerformIO $ - traverseVar' af (ZF (fmap (runZF z))) id traverse v -{-# INLINE sequenceVar #-} - --- TODO: can add funcs and zeros be had from bvars and Functor instance? -collectVar_ - :: forall t a s. (Reifies s W, Foldable t, Functor t) - => AddFunc a - -> ZeroFunc a - -> t (BVar s a) - -> IO (BVar s (t a)) -collectVar_ af z !vs = withVec (toList vs) $ \(vVec :: VecT n (BVar s) a) -> do - let tn :: TapeNode (t a) - tn = TN - { _tnInputs = vecToRec (vmap (\v -> IR v (runAF af) id) vVec) - , _tnGrad = vecToRec - . zipVecList (\v -> Identity . fromMaybe (runZF z (_bvVal v))) vVec - . toList - } - traverse_ (evaluate . forceBVar) vs - insertNode tn (_bvVal <$> vs) (reflect (Proxy @s)) -{-# INLINE collectVar_ #-} - --- | 'Numeric.Backprop.collectVar', but with explicit 'add' and 'zero'. -collectVar - :: forall t a s. (Reifies s W, Foldable t, Functor t) - => AddFunc a - -> ZeroFunc a - -> t (BVar s a) - -> BVar s (t a) -collectVar af z !vs = unsafePerformIO $ collectVar_ af z vs -{-# INLINE collectVar #-} - -traverseVar' - :: forall b a f s. (Reifies s W, Traversable f) - => AddFunc a - -> ZeroFunc b - -> (b -> f a) - -> Traversal' b a - -> BVar s b - -> IO (f (BVar s a)) -traverseVar' af z f t v = forceBVar v - `seq` itraverse go (f x) - where - x = _bvVal v - go :: Int -> a -> IO (BVar s a) - go i y = insertNode tn y (reflect (Proxy @s)) - where - tn = TN { _tnInputs = IR v (over (ixt t i) . runAF af) - (\g -> set (ixt t i) g (runZF z x)) - :& RNil - , _tnGrad = (:& RNil) . Identity - } - {-# INLINE go #-} -{-# INLINE traverseVar' #-} - --- | 'Numeric.Backprop.previewVar', but with explicit 'add' and 'zero'. -previewVar - :: forall b a s. Reifies s W - => AddFunc a - -> ZeroFunc b - -> Traversal' b a - -> BVar s b - -> Maybe (BVar s a) -previewVar af z t !v = unsafePerformIO $ - traverseVar' af z (preview t) t v -{-# INLINE previewVar #-} - --- | 'Numeric.Backprop.toListOfVar', but with explicit 'add' and 'zero'. -toListOfVar - :: forall b a s. Reifies s W - => AddFunc a - -> ZeroFunc b - -> Traversal' b a - -> BVar s b - -> [BVar s a] -toListOfVar af z t !v = unsafePerformIO $ - traverseVar' af z (toListOf t) t v -{-# INLINE toListOfVar #-} - --- | Coerce a 'BVar' contents. Useful for things like newtype wrappers. --- --- @since 0.1.5.2 -coerceVar - :: Coercible a b - => BVar s a - -> BVar s b -coerceVar v@(BV r x) = forceBVar v `seq` BV r (coerce x) - -data Runner s = R { _rDelta :: !(MV.MVector s (Maybe Any)) - , _rInputs :: !(MV.MVector s (Maybe Any)) - } - -initRunner - :: (Int, [SomeTapeNode]) - -> (Int, [Maybe Any]) - -> ST s (Runner s) -initRunner (n, stns) (nx,xs) = do - delts <- MV.new n - for_ (zip [n-1,n-2..] stns) $ \(i, STN (TN{} :: TapeNode c)) -> - MV.write delts i $ unsafeCoerce (Nothing @c) - inps <- MV.new nx - for_ (zip [0..] xs) . uncurry $ \i z -> - MV.write inps i z - return $ R delts inps -{-# INLINE initRunner #-} - -gradRunner - :: forall b s. () - => b -- ^ one - -> Runner s - -> (Int, [SomeTapeNode]) - -> ST s () -gradRunner o R{..} (n,stns) = do - when (n > 0) $ - MV.write _rDelta (n - 1) (unsafeCoerce (Just o)) - zipWithM_ go [n-1,n-2..] stns - where - go :: Int -> SomeTapeNode -> ST s () - go i (STN (TN{..} :: TapeNode c)) = do - delt <- MV.read _rDelta i - forM_ delt $ \d -> do - let gs = _tnGrad (unsafeCoerce d) - rzipWithM_ propagate _tnInputs gs - {-# INLINE go #-} - propagate :: forall x. InpRef x -> Identity x -> ST s () - propagate (IR v (+*) e) (Identity d) = case _bvRef v of - BRInp i -> flip (MV.modify _rInputs) i $ - unsafeCoerce . bumpMaybe d (+*) e . unsafeCoerce - BRIx i -> flip (MV.modify _rDelta) i $ - unsafeCoerce . bumpMaybe d (+*) e . unsafeCoerce - BRC -> return () - {-# INLINE propagate #-} -{-# INLINE gradRunner #-} - -bumpMaybe - :: a -- ^ val - -> (a -> b -> b) -- ^ add - -> (a -> b) -- ^ embed - -> Maybe b - -> Maybe b -bumpMaybe x (+*) e = \case - Nothing -> Just (e x) - Just y -> Just (x +* y) -{-# INLINE bumpMaybe #-} - -seqEither :: Either a (b, [SomeTapeNode]) -> Either a (b, [SomeTapeNode]) -seqEither e@(Left !_) = e -seqEither e@(Right (!_,foldMap forceSomeTapeNode->(!_))) = e -{-# INLINE seqEither #-} - --- | 'Numeric.Backprop.backpropWithN', but with explicit 'zero' and 'one'. --- --- Note that argument order changed in v0.2.4. --- --- @since 0.2.0.0 -backpropWithN - :: forall as b. () - => Rec ZeroFunc as - -> (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> (b, b -> Rec Identity as) -backpropWithN zfs f !xs = (y, g') - where - !(seqEither->(!tp0),!y) = unsafePerformIO $ fillWengert f xs - g' :: b -> Rec Identity as - g' = case tp0 of - Left i -> setInput i - Right tp -> g tp - {-# INLINE g' #-} - g :: (Int, [SomeTapeNode]) -> b -> Rec Identity as - g tp o = runST $ do - r <- initRunner tp . bimap getSum (`appEndo` []) - . VR.rfoldMap go -- TODO: use strict tuple? - $ xs - gradRunner o r tp - delts <- toList <$> V.freeze (_rInputs r) - return . fromMaybe (internalError "backpropN") $ - fillRec (\z -> maybe z (Identity . unsafeCoerce)) - (VR.rzipWith (fmap . runZF) zfs xs) - delts - where - go :: forall a. Identity a -> (Sum Int, Endo [Maybe Any]) - go _ = (1, Endo (unsafeCoerce (Nothing @a) :)) - {-# INLINE go #-} - setInput :: Int -> b -> Rec Identity as - setInput !i !x = go zfs xs 0 - where - go :: Rec ZeroFunc bs -> Rec Identity bs -> Int -> Rec Identity bs - go = \case - RNil -> \_ _ -> RNil - z :& zs -> \case - q :& qs -> \(!j) -> - if j == i - then Identity (unsafeCoerce x) :& VR.rzipWith coerce zs qs - else coerce z q :& go zs qs (j + 1) - {-# INLINE setInput #-} -{-# INLINE backpropWithN #-} - --- | 'evalBP' generalized to multiple inputs of different types. See --- documentation for 'Numeric.Backprop.backpropN' for more details. -evalBPN - :: forall as b. () - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> b -evalBPN f = snd . unsafePerformIO . fillWengert f -{-# INLINE evalBPN #-} - -fillWengert - :: forall as b. () - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> IO (Either Int (Int, [SomeTapeNode]), b) -fillWengert f xs = do - w <- initWengert - (i, o) <- reify w $ \(Proxy :: Proxy s) -> do - let oVar = f (inpRec @s) - evaluate (forceBVar oVar) - let isInput = case _bvRef oVar of - BRInp i -> Just i - _ -> Nothing - pure (isInput, _bvVal oVar) - t <- case i of - Nothing -> Right <$> readIORef (wRef w) - Just i' -> pure $ Left i' - pure (t, o) - where - inpRec :: forall s. Rec (BVar s) as - inpRec = evalState (rtraverse (state . go . runIdentity) xs) 0 - where - go :: a -> Int -> (BVar s a, Int) - go x i = (BV (BRInp i) x, i + 1) - {-# INLINE go #-} - {-# INLINE inpRec #-} -{-# INLINE fillWengert #-} - - -instance (Num a, Reifies s W) => Num (BVar s a) where - (+) = liftOp2 afNum afNum (+.) - {-# INLINE (+) #-} - (-) = liftOp2 afNum afNum (-.) - {-# INLINE (-) #-} - (*) = liftOp2 afNum afNum (*.) - {-# INLINE (*) #-} - negate = liftOp1 afNum negateOp - {-# INLINE negate #-} - signum = liftOp1 afNum signumOp - {-# INLINE signum #-} - abs = liftOp1 afNum absOp - {-# INLINE abs #-} - fromInteger = constVar . fromInteger - {-# INLINE fromInteger #-} - -instance (Fractional a, Reifies s W) => Fractional (BVar s a) where - (/) = liftOp2 afNum afNum (/.) - {-# INLINE (/) #-} - recip = liftOp1 afNum recipOp - {-# INLINE recip #-} - fromRational = constVar . fromRational - {-# INLINE fromRational #-} - -instance (Floating a, Reifies s W) => Floating (BVar s a) where - pi = constVar pi - {-# INLINE pi #-} - exp = liftOp1 afNum expOp - {-# INLINE exp #-} - log = liftOp1 afNum logOp - {-# INLINE log #-} - sqrt = liftOp1 afNum sqrtOp - {-# INLINE sqrt #-} - (**) = liftOp2 afNum afNum (**.) - {-# INLINE (**) #-} - logBase = liftOp2 afNum afNum logBaseOp - {-# INLINE logBase #-} - sin = liftOp1 afNum sinOp - {-# INLINE sin #-} - cos = liftOp1 afNum cosOp - {-# INLINE cos #-} - tan = liftOp1 afNum tanOp - {-# INLINE tan #-} - asin = liftOp1 afNum asinOp - {-# INLINE asin #-} - acos = liftOp1 afNum acosOp - {-# INLINE acos #-} - atan = liftOp1 afNum atanOp - {-# INLINE atan #-} - sinh = liftOp1 afNum sinhOp - {-# INLINE sinh #-} - cosh = liftOp1 afNum coshOp - {-# INLINE cosh #-} - tanh = liftOp1 afNum tanhOp - {-# INLINE tanh #-} - asinh = liftOp1 afNum asinhOp - {-# INLINE asinh #-} - acosh = liftOp1 afNum acoshOp - {-# INLINE acosh #-} - atanh = liftOp1 afNum atanhOp - {-# INLINE atanh #-} - --- | Compares the values inside the 'BVar'. --- --- @since 0.1.5.0 -instance Eq a => Eq (BVar s a) where - (==) = (==) `on` _bvVal - (/=) = (/=) `on` _bvVal - --- | Compares the values inside the 'BVar'. --- --- @since 0.1.5.0 -instance Ord a => Ord (BVar s a) where - compare = compare `on` _bvVal - (<) = (<) `on` _bvVal - (<=) = (<=) `on` _bvVal - (>) = (>) `on` _bvVal - (>=) = (>=) `on` _bvVal - --- Some utility functions to get around a lens dependency -itraverse - :: forall t a b f. (Traversable t, Monad f) - => (Int -> a -> f b) -> t a -> f (t b) -itraverse f xs = evalStateT (traverse (StateT . go) xs) 0 - where - go :: a -> Int -> f (b, Int) - go x i = (,i+1) <$> f i x -{-# INLINE itraverse #-} - -ixi :: Int -> Lens' [a] a -ixi _ _ [] = internalError "ixi" -ixi 0 f (x:xs) = (:xs) <$> f x -ixi n f (x:xs) = (x:) <$> ixi (n - 1) f xs -{-# INLINE ixi #-} - -ixt :: forall b a. Traversal' b a -> Int -> Lens' b a -ixt t i f xs = stuff <$> ixi i f contents - where - contents = xs ^.. t - stuff = evalState (traverseOf t (state . const go) xs) - where - go :: [a] -> (a, [a]) - go [] = internalError "ixt" - go (y:ys) = (y, ys) -{-# INLINE ixt #-} - --- | @since 0.2.2.0 -instance (Backprop a, Reifies s W) => Backprop (BVar s a) where - zero = liftOp1 addFunc . op1 $ \x -> (zero x, zero) - {-# INLINE zero #-} - add = liftOp2 addFunc addFunc . op2 $ \x y -> - ( add x y - , \d -> (d, d) - ) - {-# INLINE add #-} - one = liftOp1 addFunc . op1 $ \x -> (one x, zero) - {-# INLINE one #-} - --- | The canonical 'ZeroFunc' for instances of 'Backprop'. --- --- @since 0.2.0.0 -zeroFunc :: Backprop a => ZeroFunc a -zeroFunc = ZF zero -{-# INLINE zeroFunc #-} - --- | The canonical 'AddFunc' for instances of 'Backprop'. --- --- @since 0.2.0.0 -addFunc :: Backprop a => AddFunc a -addFunc = AF add -{-# INLINE addFunc #-} - --- | The canonical 'OneFunc' for instances of 'Backprop'. --- --- @since 0.2.0.0 -oneFunc :: Backprop a => OneFunc a -oneFunc = OF one -{-# INLINE oneFunc #-} - -internalError :: String -> a -internalError m = errorWithoutStackTrace $ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- Module : Numeric.Backprop.Internal +-- Copyright : (c) Justin Le 2023 +-- 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. +-- +-- Do not import this unless you want to mess with the internals and know what +-- you're doing. Otherwise, use "Numeric.Backprop", "Numeric.Backprop.Op", +-- etc. +-- +-- @since 0.2.7.0 +module Numeric.Backprop.Internal ( + BVar, + W, + backpropWithN, + evalBPN, + constVar, + liftOp, + liftOp1, + liftOp2, + liftOp3, + viewVar, + setVar, + sequenceVar, + collectVar, + previewVar, + toListOfVar, + coerceVar, + + -- * Func wrappers + ZeroFunc (..), + zfNum, + zeroFunc, + AddFunc (..), + afNum, + addFunc, + OneFunc (..), + ofNum, + oneFunc, + + -- * Debug + debugSTN, + debugIR, +) where + +import Control.DeepSeq +import Control.Exception +import Control.Monad +import Control.Monad.ST +import Control.Monad.Trans.State +import Data.Bifunctor +import Data.Coerce +import Data.Foldable +import Data.Function +import Data.Functor.Identity +import Data.IORef +import Data.Kind +import Data.Maybe +import Data.Monoid hiding (Any (..)) +import Data.Proxy +import Data.Reflection +import Data.Type.Util +import Data.Typeable +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as MV +import Data.Vinyl.Core +import qualified Data.Vinyl.Recursive as VR +import qualified Data.Vinyl.XRec as X +import GHC.Exts (Any) +import GHC.Generics as G +import Lens.Micro +import Lens.Micro.Extras +import Numeric.Backprop.Class +import Numeric.Backprop.Op +import System.IO.Unsafe +import Unsafe.Coerce + +-- | "Zero out" all components of a value. For scalar values, this should +-- just be @'const' 0@. For vectors and matrices, this should set all +-- components to zero, the additive identity. +-- +-- Should be idempotent: Applying the function twice is the same as +-- applying it just once. +-- +-- Each type should ideally only have one 'ZeroFunc'. This coherence +-- constraint is given by the typeclass 'Backprop'. +-- +-- @since 0.2.0.0 +newtype ZeroFunc a = ZF {runZF :: a -> a} + +-- | Add together two values of a type. To combine contributions of +-- gradients, so should ideally be information-preserving. +-- +-- See laws for 'Backprop' for the laws this should be expected to +-- preserve. Namely, it should be commutative and associative, with an +-- identity for a valid 'ZeroFunc'. +-- +-- Each type should ideally only have one 'AddFunc'. This coherence +-- constraint is given by the typeclass 'Backprop'. +-- +-- @since 0.2.0.0 +newtype AddFunc a = AF {runAF :: a -> a -> a} + +-- | "One" all components of a value. For scalar values, this should +-- just be @'const' 1@. For vectors and matrices, this should set all +-- components to one, the multiplicative identity. +-- +-- Should be idempotent: Applying the function twice is the same as +-- applying it just once. +-- +-- Each type should ideally only have one 'OneFunc'. This coherence +-- constraint is given by the typeclass 'Backprop'. +-- +-- @since 0.2.0.0 +newtype OneFunc a = OF {runOF :: a -> a} + +-- | If a type has a 'Num' instance, this is the canonical 'ZeroFunc'. +-- +-- @since 0.2.0.0 +zfNum :: Num a => ZeroFunc a +zfNum = ZF (const 0) +{-# INLINE zfNum #-} + +-- | If a type has a 'Num' instance, this is the canonical 'AddFunc'. +-- +-- @since 0.2.0.0 +afNum :: Num a => AddFunc a +afNum = AF (+) +{-# INLINE afNum #-} + +-- | If a type has a 'Num' instance, this is the canonical 'OneFunc'. +-- +-- @since 0.2.0.0 +ofNum :: Num a => OneFunc a +ofNum = OF (const 1) +{-# INLINE ofNum #-} + +-- | A @'BVar' s a@ is a value of type @a@ that can be "backpropagated". +-- +-- Functions referring to 'BVar's are tracked by the library and can be +-- automatically differentiated to get their gradients and results. +-- +-- For simple numeric values, you can use its 'Num', 'Fractional', and +-- 'Floating' instances to manipulate them as if they were the numbers they +-- represent. +-- +-- If @a@ contains items, the items can be accessed and extracted using +-- lenses. A @'Lens'' b a@ can be used to access an @a@ inside a @b@, using +-- '^^.' ('Numeric.Backprop.viewVar'): +-- +-- @ +-- ('^.') :: a -> 'Lens'' a b -> b +-- ('^^.') :: 'BVar' s a -> 'Lens'' a b -> 'BVar' s b +-- @ +-- +-- There is also '^^?' ('Numeric.Backprop.previewVar'), to use a 'Prism'' +-- or 'Traversal'' to extract a target that may or may not be present +-- (which can implement pattern matching), '^^..' +-- ('Numeric.Backprop.toListOfVar') to use a 'Traversal'' to extract /all/ +-- targets inside a 'BVar', and '.~~' ('setVar') to set and update values +-- inside a 'BVar'. +-- +-- If you have control over your data type definitions, you can also use +-- 'Numeric.Backprop.splitBV' and 'Numeric.Backprop.joinBV' to manipulate +-- data types by easily extracting fields out of a 'BVar' of data types and +-- creating 'BVar's of data types out of 'BVar's of their fields. See +-- "Numeric.Backprop#hkd" for a tutorial on this use pattern. +-- +-- For more complex operations, libraries can provide functions on 'BVar's +-- using 'Numeric.Backprop.liftOp' and related functions. This is how you +-- can create primitive functions that users can use to manipulate your +-- library's values. See +-- for a detailed +-- guide. +-- +-- For example, the /hmatrix/ library has a matrix-vector multiplication +-- function, @#> :: L m n -> R n -> L m@. +-- +-- A library could instead provide a function @#> :: 'BVar' (L m n) -> BVar +-- (R n) -> BVar (R m)@, which the user can then use to manipulate their +-- 'BVar's of @L m n@s and @R n@s, etc. +-- +-- See "Numeric.Backprop#liftops" and documentation for +-- 'Numeric.Backprop.liftOp' for more information. +data BVar s a = BV + { _bvRef :: !(BRef s) + , _bvVal :: !a + } + +-- | @since 0.1.5.1 +deriving instance Typeable (BVar s a) + +-- | @since 0.2.6.3 +instance X.IsoHKD (BVar s) a + +data BRef (s :: Type) + = BRInp !Int + | BRIx !Int + | BRC + deriving (Generic, Show) + +instance NFData (BRef s) + +-- | This will force the value inside, as well. +instance NFData a => NFData (BVar s a) where + rnf (BV r v) = force r `seq` force v `seq` () + +-- | Project out a constant value if the 'BVar' refers to one. +bvConst :: BVar s a -> Maybe a +bvConst (BV BRC !x) = Just x +bvConst _ = Nothing +{-# INLINE bvConst #-} + +forceBVar :: BVar s a -> () +forceBVar (BV r !_) = force r `seq` () +{-# INLINE forceBVar #-} + +data InpRef :: Type -> Type where + IR :: + { _irIx :: !(BVar s b) + , _irAdd :: !(a -> b -> b) + , _irEmbed :: !(a -> b) + } -> + InpRef a + +forceInpRef :: InpRef a -> () +forceInpRef (IR v !_ !_) = forceBVar v `seq` () +{-# INLINE forceInpRef #-} + +-- | Debugging string for an 'InpRef'. +debugIR :: InpRef a -> String +debugIR IR{..} = show (_bvRef _irIx) + +data TapeNode :: Type -> Type where + TN :: + { _tnInputs :: !(Rec InpRef as) + , _tnGrad :: !(a -> Rec Identity as) + } -> + TapeNode a + +forceTapeNode :: TapeNode a -> () +forceTapeNode (TN inps !_) = VR.rfoldMap forceInpRef inps `seq` () +{-# INLINE forceTapeNode #-} + +data SomeTapeNode :: Type where + STN :: + { _stnNode :: !(TapeNode a) + } -> + SomeTapeNode + +forceSomeTapeNode :: SomeTapeNode -> () +forceSomeTapeNode (STN n) = forceTapeNode n + +-- | Debugging string for a 'SomeTapeMode'. +debugSTN :: SomeTapeNode -> String +debugSTN (STN TN{..}) = show . VR.rfoldMap ((: []) . debugIR) $ _tnInputs + +-- | An ephemeral Wengert Tape in the environment. Used internally to +-- track of the computational graph of variables. +-- +-- For the end user, one can just imagine @'Reifies' s 'W'@ as a required +-- constraint on @s@ that allows backpropagation to work. +newtype W = W {wRef :: IORef (Int, [SomeTapeNode])} + +initWengert :: IO W +initWengert = W <$> newIORef (0, []) +{-# INLINE initWengert #-} + +insertNode :: + TapeNode a -> + -- | val + a -> + W -> + IO (BVar s a) +insertNode tn !x !w = fmap ((`BV` x) . BRIx) . atomicModifyIORef' (wRef w) $ \(!n, !t) -> + let n' = n + 1 + t' = STN tn : t + in forceTapeNode tn `seq` n' `seq` t' `seq` ((n', t'), n) +{-# INLINE insertNode #-} + +-- | Lift a value into a 'BVar' representing a constant value. +-- +-- This value will not be considered an input, and its gradients will not +-- be backpropagated. +constVar :: a -> BVar s a +constVar = BV BRC +{-# INLINE constVar #-} + +liftOp_ :: + forall s as b. + Reifies s W => + Rec AddFunc as -> + Op as b -> + Rec (BVar s) as -> + IO (BVar s b) +liftOp_ afs o !vs = case rtraverse (fmap Identity . bvConst) vs of + Just xs -> return $ constVar (evalOp o xs) + Nothing -> insertNode tn y (reflect (Proxy @s)) + where + (y, g) = runOpWith o (VR.rmap (Identity . _bvVal) vs) + tn = + TN + { _tnInputs = VR.rzipWith go afs vs + , _tnGrad = g + } + go :: forall a. AddFunc a -> BVar s a -> InpRef a + go af !v = forceBVar v `seq` IR v (runAF af) id + {-# INLINE go #-} +{-# INLINE liftOp_ #-} + +-- | 'Numeric.Backprop.liftOp', but with explicit 'add' and 'zero'. +liftOp :: + forall as b s. + Reifies s W => + Rec AddFunc as -> + Op as b -> + Rec (BVar s) as -> + BVar s b +liftOp afs o !vs = unsafePerformIO $ liftOp_ afs o vs +{-# INLINE liftOp #-} + +liftOp1_ :: + forall a b s. + Reifies s W => + AddFunc a -> + Op '[a] b -> + BVar s a -> + IO (BVar s b) +liftOp1_ _ o (bvConst -> Just x) = return . constVar . evalOp o $ (Identity x :& RNil) +liftOp1_ af o v = forceBVar v `seq` insertNode tn y (reflect (Proxy @s)) + where + (y, g) = runOpWith o (Identity (_bvVal v) :& RNil) + tn = + TN + { _tnInputs = IR v (runAF af) id :& RNil + , _tnGrad = g + } +{-# INLINE liftOp1_ #-} + +-- | 'Numeric.Backprop.liftOp1', but with explicit 'add' and 'zero'. +liftOp1 :: + forall a b s. + Reifies s W => + AddFunc a -> + Op '[a] b -> + BVar s a -> + BVar s b +liftOp1 af o !v = unsafePerformIO $ liftOp1_ af o v +{-# INLINE liftOp1 #-} + +liftOp2_ :: + forall a b c s. + Reifies s W => + AddFunc a -> + AddFunc b -> + Op '[a, b] c -> + BVar s a -> + BVar s b -> + IO (BVar s c) +liftOp2_ _ _ o (bvConst -> Just x) (bvConst -> Just y) = + return . constVar . evalOp o $ Identity x :& Identity y :& RNil +liftOp2_ afa afb o v u = + forceBVar v `seq` + forceBVar u `seq` + insertNode tn y (reflect (Proxy @s)) + where + (y, g) = + runOpWith o $ + Identity (_bvVal v) + :& Identity (_bvVal u) + :& RNil + tn = + TN + { _tnInputs = IR v (runAF afa) id :& IR u (runAF afb) id :& RNil + , _tnGrad = g + } +{-# INLINE liftOp2_ #-} + +-- | 'Numeric.Backprop.liftOp2', but with explicit 'add' and 'zero'. +liftOp2 :: + forall a b c s. + Reifies s W => + AddFunc a -> + AddFunc b -> + Op '[a, b] c -> + BVar s a -> + BVar s b -> + BVar s c +liftOp2 afa afb o !v !u = unsafePerformIO $ liftOp2_ afa afb o v u +{-# INLINE liftOp2 #-} + +liftOp3_ :: + forall a b c d s. + Reifies s W => + AddFunc a -> + AddFunc b -> + AddFunc c -> + Op '[a, b, c] d -> + BVar s a -> + BVar s b -> + BVar s c -> + IO (BVar s d) +liftOp3_ _ _ _ o (bvConst -> Just x) (bvConst -> Just y) (bvConst -> Just z) = + return . constVar . evalOp o $ + Identity x + :& Identity y + :& Identity z + :& RNil +liftOp3_ afa afb afc o v u w = + forceBVar v `seq` + forceBVar u `seq` + forceBVar w `seq` + insertNode tn y (reflect (Proxy @s)) + where + (y, g) = + runOpWith o $ + Identity (_bvVal v) + :& Identity (_bvVal u) + :& Identity (_bvVal w) + :& RNil + tn = + TN + { _tnInputs = + IR v (runAF afa) id + :& IR u (runAF afb) id + :& IR w (runAF afc) id + :& RNil + , _tnGrad = g + } +{-# INLINE liftOp3_ #-} + +-- | 'Numeric.Backprop.liftOp3', but with explicit 'add' and 'zero'. +liftOp3 :: + forall a b c d s. + Reifies s W => + AddFunc a -> + AddFunc b -> + AddFunc c -> + Op '[a, b, c] d -> + BVar s a -> + BVar s b -> + BVar s c -> + BVar s d +liftOp3 afa afb afc o !v !u !w = unsafePerformIO $ liftOp3_ afa afb afc o v u w +{-# INLINE liftOp3 #-} + +-- TODO: can we get the zero and add func from the bvar? +viewVar_ :: + forall a b s. + Reifies s W => + AddFunc a -> + ZeroFunc b -> + Lens' b a -> + BVar s b -> + IO (BVar s a) +viewVar_ af z l v = forceBVar v `seq` insertNode tn y (reflect (Proxy @s)) + where + x = _bvVal v + y = x ^. l + tn = + TN + { _tnInputs = + IR v (over l . runAF af) (\g -> set l g (runZF z x)) + :& RNil + , _tnGrad = (:& RNil) . Identity + } +{-# INLINE viewVar_ #-} + +-- | 'Numeric.Backprop.viewVar', but with explicit 'add' and 'zero'. +viewVar :: + forall a b s. + Reifies s W => + AddFunc a -> + ZeroFunc b -> + Lens' b a -> + BVar s b -> + BVar s a +viewVar af z l !v = unsafePerformIO $ viewVar_ af z l v +{-# INLINE viewVar #-} + +-- TODO: can zero and add func be gotten from the input bvars? +setVar_ :: + forall a b s. + Reifies s W => + AddFunc a -> + AddFunc b -> + ZeroFunc a -> + Lens' b a -> + BVar s a -> + BVar s b -> + IO (BVar s b) +setVar_ afa afb za l w v = + forceBVar v `seq` + forceBVar w `seq` + insertNode tn y (reflect (Proxy @s)) + where + y = _bvVal v & l .~ _bvVal w + tn = + TN + { _tnInputs = + IR w (runAF afa) id + :& IR v (runAF afb) id + :& RNil + , _tnGrad = \d -> + let (dw, dv) = l (\x -> (x, runZF za x)) d + in Identity dw :& Identity dv :& RNil + } +{-# INLINE setVar_ #-} + +-- | 'Numeric.Backprop.setVar', but with explicit 'add' and 'zero'. +setVar :: + forall a b s. + Reifies s W => + AddFunc a -> + AddFunc b -> + ZeroFunc a -> + Lens' b a -> + BVar s a -> + BVar s b -> + BVar s b +setVar afa afb za l !w !v = unsafePerformIO $ setVar_ afa afb za l w v +{-# INLINE setVar #-} + +-- | 'Numeric.Backprop.sequenceVar', but with explicit 'add' and 'zero'. +sequenceVar :: + forall t a s. + (Reifies s W, Traversable t) => + AddFunc a -> + ZeroFunc a -> + BVar s (t a) -> + t (BVar s a) +sequenceVar af z !v = + unsafePerformIO $ + traverseVar' af (ZF (fmap (runZF z))) id traverse v +{-# INLINE sequenceVar #-} + +-- TODO: can add funcs and zeros be had from bvars and Functor instance? +collectVar_ :: + forall t a s. + (Reifies s W, Foldable t, Functor t) => + AddFunc a -> + ZeroFunc a -> + t (BVar s a) -> + IO (BVar s (t a)) +collectVar_ af z !vs = withVec (toList vs) $ \(vVec :: VecT n (BVar s) a) -> do + let tn :: TapeNode (t a) + tn = + TN + { _tnInputs = vecToRec (vmap (\v -> IR v (runAF af) id) vVec) + , _tnGrad = + vecToRec + . zipVecList (\v -> Identity . fromMaybe (runZF z (_bvVal v))) vVec + . toList + } + traverse_ (evaluate . forceBVar) vs + insertNode tn (_bvVal <$> vs) (reflect (Proxy @s)) +{-# INLINE collectVar_ #-} + +-- | 'Numeric.Backprop.collectVar', but with explicit 'add' and 'zero'. +collectVar :: + forall t a s. + (Reifies s W, Foldable t, Functor t) => + AddFunc a -> + ZeroFunc a -> + t (BVar s a) -> + BVar s (t a) +collectVar af z !vs = unsafePerformIO $ collectVar_ af z vs +{-# INLINE collectVar #-} + +traverseVar' :: + forall b a f s. + (Reifies s W, Traversable f) => + AddFunc a -> + ZeroFunc b -> + (b -> f a) -> + Traversal' b a -> + BVar s b -> + IO (f (BVar s a)) +traverseVar' af z f t v = + forceBVar v `seq` + itraverse go (f x) + where + x = _bvVal v + go :: Int -> a -> IO (BVar s a) + go i y = insertNode tn y (reflect (Proxy @s)) + where + tn = + TN + { _tnInputs = + IR + v + (over (ixt t i) . runAF af) + (\g -> set (ixt t i) g (runZF z x)) + :& RNil + , _tnGrad = (:& RNil) . Identity + } + {-# INLINE go #-} +{-# INLINE traverseVar' #-} + +-- | 'Numeric.Backprop.previewVar', but with explicit 'add' and 'zero'. +previewVar :: + forall b a s. + Reifies s W => + AddFunc a -> + ZeroFunc b -> + Traversal' b a -> + BVar s b -> + Maybe (BVar s a) +previewVar af z t !v = + unsafePerformIO $ + traverseVar' af z (preview t) t v +{-# INLINE previewVar #-} + +-- | 'Numeric.Backprop.toListOfVar', but with explicit 'add' and 'zero'. +toListOfVar :: + forall b a s. + Reifies s W => + AddFunc a -> + ZeroFunc b -> + Traversal' b a -> + BVar s b -> + [BVar s a] +toListOfVar af z t !v = + unsafePerformIO $ + traverseVar' af z (toListOf t) t v +{-# INLINE toListOfVar #-} + +-- | Coerce a 'BVar' contents. Useful for things like newtype wrappers. +-- +-- @since 0.1.5.2 +coerceVar :: + Coercible a b => + BVar s a -> + BVar s b +coerceVar v@(BV r x) = forceBVar v `seq` BV r (coerce x) + +data Runner s = R + { _rDelta :: !(MV.MVector s (Maybe Any)) + , _rInputs :: !(MV.MVector s (Maybe Any)) + } + +initRunner :: + (Int, [SomeTapeNode]) -> + (Int, [Maybe Any]) -> + ST s (Runner s) +initRunner (n, stns) (nx, xs) = do + delts <- MV.new n + for_ (zip [n - 1, n - 2 ..] stns) $ \(i, STN (TN{} :: TapeNode c)) -> + MV.write delts i $ unsafeCoerce (Nothing @c) + inps <- MV.new nx + itraverse_ (MV.write inps) xs + return $ R delts inps +{-# INLINE initRunner #-} + +gradRunner :: + forall b s. + () => + -- | one + b -> + Runner s -> + (Int, [SomeTapeNode]) -> + ST s () +gradRunner o R{..} (n, stns) = do + when (n > 0) $ + MV.write _rDelta (n - 1) (unsafeCoerce (Just o)) + zipWithM_ go [n - 1, n - 2 ..] stns + where + go :: Int -> SomeTapeNode -> ST s () + go i (STN (TN{..} :: TapeNode c)) = do + delt <- MV.read _rDelta i + forM_ delt $ \d -> do + let gs = _tnGrad (unsafeCoerce d) + rzipWithM_ propagate _tnInputs gs + {-# INLINE go #-} + propagate :: forall x. InpRef x -> Identity x -> ST s () + propagate (IR v (+*) e) (Identity d) = case _bvRef v of + BRInp i -> + flip (MV.modify _rInputs) i $ + unsafeCoerce . bumpMaybe d (+*) e . unsafeCoerce + BRIx i -> + flip (MV.modify _rDelta) i $ + unsafeCoerce . bumpMaybe d (+*) e . unsafeCoerce + BRC -> return () + {-# INLINE propagate #-} +{-# INLINE gradRunner #-} + +bumpMaybe :: + -- | val + a -> + -- | add + (a -> b -> b) -> + -- | embed + (a -> b) -> + Maybe b -> + Maybe b +bumpMaybe x (+*) e = \case + Nothing -> Just (e x) + Just y -> Just (x +* y) +{-# INLINE bumpMaybe #-} + +seqEither :: Either a (b, [SomeTapeNode]) -> Either a (b, [SomeTapeNode]) +seqEither e@(Left !_) = e +seqEither e@(Right (!_, foldMap forceSomeTapeNode -> (!_))) = e +{-# INLINE seqEither #-} + +-- | 'Numeric.Backprop.backpropWithN', but with explicit 'zero' and 'one'. +-- +-- Note that argument order changed in v0.2.4. +-- +-- @since 0.2.0.0 +backpropWithN :: + forall as b. + () => + Rec ZeroFunc as -> + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + (b, b -> Rec Identity as) +backpropWithN zfs f !xs = (y, g') + where + !(seqEither -> (!tp0), !y) = unsafePerformIO $ fillWengert f xs + g' :: b -> Rec Identity as + g' = case tp0 of + Left i -> setInput i + Right tp -> g tp + {-# INLINE g' #-} + g :: (Int, [SomeTapeNode]) -> b -> Rec Identity as + g tp o = runST $ do + r <- + initRunner tp + . bimap getSum (`appEndo` []) + . VR.rfoldMap go -- TODO: use strict tuple? + $ xs + gradRunner o r tp + delts <- toList <$> V.freeze (_rInputs r) + return . fromMaybe (internalError "backpropN") $ + fillRec + (\z -> maybe z (Identity . unsafeCoerce)) + (VR.rzipWith (fmap . runZF) zfs xs) + delts + where + go :: forall a. Identity a -> (Sum Int, Endo [Maybe Any]) + go _ = (1, Endo (unsafeCoerce (Nothing @a) :)) + {-# INLINE go #-} + setInput :: Int -> b -> Rec Identity as + setInput !i !x = go zfs xs 0 + where + go :: Rec ZeroFunc bs -> Rec Identity bs -> Int -> Rec Identity bs + go = \case + RNil -> \_ _ -> RNil + z :& zs -> \case + q :& qs -> \(!j) -> + if j == i + then Identity (unsafeCoerce x) :& VR.rzipWith coerce zs qs + else coerce z q :& go zs qs (j + 1) + {-# INLINE setInput #-} +{-# INLINE backpropWithN #-} + +-- | 'evalBP' generalized to multiple inputs of different types. See +-- documentation for 'Numeric.Backprop.backpropN' for more details. +evalBPN :: + forall as b. + () => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + b +evalBPN f = snd . unsafePerformIO . fillWengert f +{-# INLINE evalBPN #-} + +fillWengert :: + forall as b. + () => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + IO (Either Int (Int, [SomeTapeNode]), b) +fillWengert f xs = do + w <- initWengert + (i, o) <- reify w $ \(Proxy :: Proxy s) -> do + let oVar = f (inpRec @s) + evaluate (forceBVar oVar) + let isInput = case _bvRef oVar of + BRInp i -> Just i + _ -> Nothing + pure (isInput, _bvVal oVar) + t <- case i of + Nothing -> Right <$> readIORef (wRef w) + Just i' -> pure $ Left i' + pure (t, o) + where + inpRec :: forall s. Rec (BVar s) as + inpRec = evalState (rtraverse (state . go . runIdentity) xs) 0 + where + go :: a -> Int -> (BVar s a, Int) + go x i = (BV (BRInp i) x, i + 1) + {-# INLINE go #-} + {-# INLINE inpRec #-} +{-# INLINE fillWengert #-} + +instance (Num a, Reifies s W) => Num (BVar s a) where + (+) = liftOp2 afNum afNum (+.) + {-# INLINE (+) #-} + (-) = liftOp2 afNum afNum (-.) + {-# INLINE (-) #-} + (*) = liftOp2 afNum afNum (*.) + {-# INLINE (*) #-} + negate = liftOp1 afNum negateOp + {-# INLINE negate #-} + signum = liftOp1 afNum signumOp + {-# INLINE signum #-} + abs = liftOp1 afNum absOp + {-# INLINE abs #-} + fromInteger = constVar . fromInteger + {-# INLINE fromInteger #-} + +instance (Fractional a, Reifies s W) => Fractional (BVar s a) where + (/) = liftOp2 afNum afNum (/.) + {-# INLINE (/) #-} + recip = liftOp1 afNum recipOp + {-# INLINE recip #-} + fromRational = constVar . fromRational + {-# INLINE fromRational #-} + +instance (Floating a, Reifies s W) => Floating (BVar s a) where + pi = constVar pi + {-# INLINE pi #-} + exp = liftOp1 afNum expOp + {-# INLINE exp #-} + log = liftOp1 afNum logOp + {-# INLINE log #-} + sqrt = liftOp1 afNum sqrtOp + {-# INLINE sqrt #-} + (**) = liftOp2 afNum afNum (**.) + {-# INLINE (**) #-} + logBase = liftOp2 afNum afNum logBaseOp + {-# INLINE logBase #-} + sin = liftOp1 afNum sinOp + {-# INLINE sin #-} + cos = liftOp1 afNum cosOp + {-# INLINE cos #-} + tan = liftOp1 afNum tanOp + {-# INLINE tan #-} + asin = liftOp1 afNum asinOp + {-# INLINE asin #-} + acos = liftOp1 afNum acosOp + {-# INLINE acos #-} + atan = liftOp1 afNum atanOp + {-# INLINE atan #-} + sinh = liftOp1 afNum sinhOp + {-# INLINE sinh #-} + cosh = liftOp1 afNum coshOp + {-# INLINE cosh #-} + tanh = liftOp1 afNum tanhOp + {-# INLINE tanh #-} + asinh = liftOp1 afNum asinhOp + {-# INLINE asinh #-} + acosh = liftOp1 afNum acoshOp + {-# INLINE acosh #-} + atanh = liftOp1 afNum atanhOp + {-# INLINE atanh #-} + +-- | Compares the values inside the 'BVar'. +-- +-- @since 0.1.5.0 +instance Eq a => Eq (BVar s a) where + (==) = (==) `on` _bvVal + (/=) = (/=) `on` _bvVal + +-- | Compares the values inside the 'BVar'. +-- +-- @since 0.1.5.0 +instance Ord a => Ord (BVar s a) where + compare = compare `on` _bvVal + (<) = (<) `on` _bvVal + (<=) = (<=) `on` _bvVal + (>) = (>) `on` _bvVal + (>=) = (>=) `on` _bvVal + +-- Some utility functions to get around a lens dependency +itraverse :: + forall t a b f. + (Traversable t, Monad f) => + (Int -> a -> f b) -> t a -> f (t b) +itraverse f xs = evalStateT (traverse (StateT . go) xs) 0 + where + go :: a -> Int -> f (b, Int) + go x i = (,i + 1) <$> f i x +{-# INLINE itraverse #-} + +-- Some utility functions to get around a lens dependency +itraverse_ :: + forall t a b f. + (Foldable t, Monad f) => + (Int -> a -> f b) -> t a -> f () +itraverse_ f xs = traverse_ (uncurry f) (zip [0 ..] (toList xs)) +{-# INLINE itraverse_ #-} + +ixi :: Int -> Lens' [a] a +ixi _ _ [] = internalError "ixi" +ixi 0 f (x : xs) = (: xs) <$> f x +ixi n f (x : xs) = (x :) <$> ixi (n - 1) f xs +{-# INLINE ixi #-} + +ixt :: forall b a. Traversal' b a -> Int -> Lens' b a +ixt t i f xs = stuff <$> ixi i f contents + where + contents = xs ^.. t + stuff = evalState (traverseOf t (state . const go) xs) + where + go :: [a] -> (a, [a]) + go [] = internalError "ixt" + go (y : ys) = (y, ys) +{-# INLINE ixt #-} + +-- | @since 0.2.2.0 +instance (Backprop a, Reifies s W) => Backprop (BVar s a) where + zero = liftOp1 addFunc . op1 $ \x -> (zero x, zero) + {-# INLINE zero #-} + add = liftOp2 addFunc addFunc . op2 $ \x y -> + ( add x y + , \d -> (d, d) + ) + {-# INLINE add #-} + one = liftOp1 addFunc . op1 $ \x -> (one x, zero) + {-# INLINE one #-} + +-- | The canonical 'ZeroFunc' for instances of 'Backprop'. +-- +-- @since 0.2.0.0 +zeroFunc :: Backprop a => ZeroFunc a +zeroFunc = ZF zero +{-# INLINE zeroFunc #-} + +-- | The canonical 'AddFunc' for instances of 'Backprop'. +-- +-- @since 0.2.0.0 +addFunc :: Backprop a => AddFunc a +addFunc = AF add +{-# INLINE addFunc #-} + +-- | The canonical 'OneFunc' for instances of 'Backprop'. +-- +-- @since 0.2.0.0 +oneFunc :: Backprop a => OneFunc a +oneFunc = OF one +{-# INLINE oneFunc #-} + +internalError :: String -> a +internalError m = + errorWithoutStackTrace $ "Numeric.Backprop.Internal." ++ m ++ ": unexpected shape involved in gradient computation" diff --git a/src/Numeric/Backprop/Num.hs b/src/Numeric/Backprop/Num.hs --- a/src/Numeric/Backprop/Num.hs +++ b/src/Numeric/Backprop/Num.hs @@ -1,9 +1,8 @@ -{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RankNTypes #-} -{-# OPTIONS_HADDOCK not-home #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# OPTIONS_HADDOCK not-home #-} -- | -- Module : Numeric.Backprop.Num @@ -43,50 +42,95 @@ -- functions. -- -- @since 0.2.0.0 - module Numeric.Backprop.Num ( - -- * Types - BVar, W - -- * Running - , backprop, E.evalBP, gradBP, backpropWith - -- ** Multiple inputs - , E.evalBP0 - , backprop2, E.evalBP2, gradBP2, backpropWith2 - , backpropN, E.evalBPN, gradBPN, backpropWithN - -- * Manipulating 'BVar' - , E.constVar, E.auto, E.coerceVar - , (^^.), (.~~), (%~~), (^^?), (^^..), (^^?!) - , viewVar, setVar, overVar - , sequenceVar, collectVar - , previewVar, toListOfVar - -- ** With Isomorphisms - , isoVar, isoVar2, isoVar3, isoVarN - -- ** With 'Op's - , liftOp - , liftOp1, liftOp2, liftOp3 - -- * 'Op' - , Op(..) - -- ** Creation - , op0, opConst, idOp - , bpOp - -- *** Giving gradients directly - , op1, op2, op3 - -- *** From Isomorphisms - , opCoerce, opTup, opIso, opIsoN, opLens - -- *** No gradients - , noGrad1, noGrad - -- * Utility - , Reifies - ) where + -- * Types + BVar, + W, -import Data.Functor.Identity -import Data.Maybe -import Data.Reflection -import Data.Vinyl -import Lens.Micro -import Numeric.Backprop.Explicit (BVar, W) -import Numeric.Backprop.Op + -- * Running + backprop, + E.evalBP, + gradBP, + backpropWith, + + -- ** Multiple inputs + E.evalBP0, + backprop2, + E.evalBP2, + gradBP2, + backpropWith2, + backpropN, + E.evalBPN, + gradBPN, + backpropWithN, + + -- * Manipulating 'BVar' + E.constVar, + E.auto, + E.coerceVar, + (^^.), + (.~~), + (%~~), + (^^?), + (^^..), + (^^?!), + viewVar, + setVar, + overVar, + sequenceVar, + collectVar, + previewVar, + toListOfVar, + + -- ** With Isomorphisms + isoVar, + isoVar2, + isoVar3, + isoVarN, + + -- ** With 'Op's + liftOp, + liftOp1, + liftOp2, + liftOp3, + + -- * 'Op' + Op (..), + + -- ** Creation + op0, + opConst, + idOp, + bpOp, + + -- *** Giving gradients directly + op1, + op2, + op3, + + -- *** From Isomorphisms + opCoerce, + opTup, + opIso, + opIsoN, + opLens, + + -- *** No gradients + noGrad1, + noGrad, + + -- * Utility + Reifies, +) where + +import Data.Functor.Identity +import Data.Maybe +import Data.Reflection +import Data.Vinyl +import Lens.Micro +import Numeric.Backprop.Explicit (BVar, W) import qualified Numeric.Backprop.Explicit as E +import Numeric.Backprop.Op -- | 'Numeric.Backprop.backpropN', but with 'Num' constraints instead of -- 'Backprop' constraints. @@ -99,12 +143,11 @@ -- If you stick to /concerete/, monomorphic usage of this (with specific -- types, typed into source code, known at compile-time), then -- @'AllPureConstrained' 'Num' as@ should be fulfilled automatically. --- -backpropN - :: (RPureConstrained Num as, Num b) - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> (b, Rec Identity as) +backpropN :: + (RPureConstrained Num as, Num b) => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + (b, Rec Identity as) backpropN = E.backpropN E.zfNums E.ofNum {-# INLINE backpropN #-} @@ -116,11 +159,11 @@ -- Note that argument order changed in v0.2.4. -- -- @since 0.2.0.0 -backpropWithN - :: RPureConstrained Num as - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> (b, b -> Rec Identity as) +backpropWithN :: + RPureConstrained Num as => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + (b, b -> Rec Identity as) backpropWithN = E.backpropWithN E.zfNums {-# INLINE backpropWithN #-} @@ -129,11 +172,11 @@ -- -- See module documentation for "Numeric.Backprop.Num" for information on -- using this with tuples. -backprop - :: (Num a, Num b) - => (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> (b, a) +backprop :: + (Num a, Num b) => + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + (b, a) backprop = E.backprop E.zfNum E.ofNum {-# INLINE backprop #-} @@ -146,42 +189,42 @@ -- Note that argument order changed in v0.2.4. -- -- @since 0.2.0.0 -backpropWith - :: Num a - => (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> (b, b -> a) +backpropWith :: + Num a => + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + (b, b -> a) backpropWith = E.backpropWith E.zfNum {-# INLINE backpropWith #-} -- | 'Numeric.Backprop.gradBP', but with 'Num' constraints instead of -- 'Backprop' constraints. -gradBP - :: (Num a, Num b) - => (forall s. Reifies s W => BVar s a -> BVar s b) - -> a - -> a +gradBP :: + (Num a, Num b) => + (forall s. Reifies s W => BVar s a -> BVar s b) -> + a -> + a gradBP = E.gradBP E.zfNum E.ofNum {-# INLINE gradBP #-} -- | 'Numeric.Backprop.gradBPN', but with 'Num' constraints instead of -- 'Backprop' constraints. -gradBPN - :: (RPureConstrained Num as, Num b) - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Rec Identity as - -> Rec Identity as +gradBPN :: + (RPureConstrained Num as, Num b) => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Rec Identity as -> + Rec Identity as gradBPN = E.gradBPN E.zfNums E.ofNum {-# INLINE gradBPN #-} -- | 'Numeric.Backprop.backprop2', but with 'Num' constraints instead of -- 'Backprop' constraints. -backprop2 - :: (Num a, Num b, Num c) - => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (c, (a, b)) +backprop2 :: + (Num a, Num b, Num c) => + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + (c, (a, b)) backprop2 = E.backprop2 E.zfNum E.zfNum E.ofNum {-# INLINE backprop2 #-} @@ -191,77 +234,82 @@ -- Note that argument order changed in v0.2.4. -- -- @since 0.2.0.0 -backpropWith2 - :: (Num a, Num b) - => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (c, c -> (a, b)) -- ^ Takes function giving gradient of final result given the output of function +backpropWith2 :: + (Num a, Num b) => + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + -- | Takes function giving gradient of final result given the output of function + (c, c -> (a, b)) backpropWith2 = E.backpropWith2 E.zfNum E.zfNum {-# INLINE backpropWith2 #-} -- | 'Numeric.Backprop.gradBP2', but with 'Num' constraints instead of -- 'Backprop' constraints. -gradBP2 - :: (Num a, Num b, Num c) - => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) - -> a - -> b - -> (a, b) +gradBP2 :: + (Num a, Num b, Num c) => + (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) -> + a -> + b -> + (a, b) gradBP2 = E.gradBP2 E.zfNum E.zfNum E.ofNum {-# INLINE gradBP2 #-} -- | 'Numeric.Backprop.bpOp', but with 'Num' constraints instead of -- 'Backprop' constraints. -bpOp - :: RPureConstrained Num as - => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) - -> Op as b +bpOp :: + RPureConstrained Num as => + (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) -> + Op as b bpOp = E.bpOp E.zfNums {-# INLINE bpOp #-} -- | 'Numeric.Backprop.^^.', but with 'Num' constraints instead of -- 'Backprop' constraints. -(^^.) - :: forall b a s. (Num a, Num b, Reifies s W) - => BVar s b - -> Lens' b a - -> BVar s a +(^^.) :: + forall b a s. + (Num a, Num b, Reifies s W) => + BVar s b -> + Lens' b a -> + BVar s a x ^^. l = viewVar l x + infixl 8 ^^. {-# INLINE (^^.) #-} -- | 'Numeric.Backprop.viewVar', but with 'Num' constraints instead of -- 'Backprop' constraints. -viewVar - :: forall b a s. (Num a, Num b, Reifies s W) - => Lens' b a - -> BVar s b - -> BVar s a +viewVar :: + forall b a s. + (Num a, Num b, Reifies s W) => + Lens' b a -> + BVar s b -> + BVar s a viewVar = E.viewVar E.afNum E.zfNum {-# INLINE viewVar #-} - -- | 'Numeric.Backprop..~~', but with 'Num' constraints instead of -- 'Backprop' constraints. -(.~~) - :: (Num a, Num b, Reifies s W) - => Lens' b a - -> BVar s a - -> BVar s b - -> BVar s b +(.~~) :: + (Num a, Num b, Reifies s W) => + Lens' b a -> + BVar s a -> + BVar s b -> + BVar s b l .~~ x = setVar l x + infixl 8 .~~ {-# INLINE (.~~) #-} -- | 'Numeric.Backprop.setVar', but with 'Num' constraints instead of -- 'Backprop' constraints. -setVar - :: forall a b s. (Num a, Num b, Reifies s W) - => Lens' b a - -> BVar s a - -> BVar s b - -> BVar s b +setVar :: + forall a b s. + (Num a, Num b, Reifies s W) => + Lens' b a -> + BVar s a -> + BVar s b -> + BVar s b setVar = E.setVar E.afNum E.afNum E.zfNum {-# INLINE setVar #-} @@ -269,14 +317,14 @@ -- 'Backprop' constraints. -- -- @since 0.2.4.0 --- -(%~~) - :: (Num a, Num b, Reifies s W) - => Lens' b a - -> (BVar s a -> BVar s a) - -> BVar s b - -> BVar s b +(%~~) :: + (Num a, Num b, Reifies s W) => + Lens' b a -> + (BVar s a -> BVar s a) -> + BVar s b -> + BVar s b l %~~ f = overVar l f + infixr 4 %~~ {-# INLINE (%~~) #-} @@ -284,12 +332,12 @@ -- 'Backprop' constraints. -- -- @since 0.2.4.0 -overVar - :: (Num a, Num b, Reifies s W) - => Lens' b a - -> (BVar s a -> BVar s a) - -> BVar s b - -> BVar s b +overVar :: + (Num a, Num b, Reifies s W) => + Lens' b a -> + (BVar s a -> BVar s a) -> + BVar s b -> + BVar s b overVar = E.overVar E.afNum E.afNum E.zfNum E.zfNum {-# INLINE overVar #-} @@ -311,12 +359,14 @@ -- myPrism :: 'Prism'' c (a, b) -- myPrism . 'iso' 'tupT2' 't2Tup' :: 'Prism'' c ('T2' a b) -- @ -(^^?) - :: forall b a s. (Num b, Num a, Reifies s W) - => BVar s b - -> Traversal' b a - -> Maybe (BVar s a) +(^^?) :: + forall b a s. + (Num b, Num a, Reifies s W) => + BVar s b -> + Traversal' b a -> + Maybe (BVar s a) v ^^? t = previewVar t v + infixl 8 ^^? {-# INLINE (^^?) #-} @@ -326,14 +376,16 @@ -- Like 'Numeric.Backprop.^^?!', is *UNSAFE*. -- -- @since 0.2.1.0 -(^^?!) - :: forall b a s. (Num b, Num a, Reifies s W) - => BVar s b - -> Traversal' b a - -> BVar s a +(^^?!) :: + forall b a s. + (Num b, Num a, Reifies s W) => + BVar s b -> + Traversal' b a -> + BVar s a v ^^?! t = fromMaybe (error e) (previewVar t v) where e = "Numeric.Backprop.Num.^^?!: Empty traversal" + infixl 8 ^^?! {-# INLINE (^^?!) #-} @@ -341,31 +393,34 @@ -- 'Backprop' constraints. -- -- See documentation for '^^?' for more information and important notes. -previewVar - :: forall b a s. (Num b, Num a, Reifies s W) - => Traversal' b a - -> BVar s b - -> Maybe (BVar s a) +previewVar :: + forall b a s. + (Num b, Num a, Reifies s W) => + Traversal' b a -> + BVar s b -> + Maybe (BVar s a) previewVar = E.previewVar E.afNum E.zfNum {-# INLINE previewVar #-} -- | 'Numeric.Backprop.^^..', but with 'Num' constraints instead of -- 'Backprop' constraints. -(^^..) - :: forall b a s. (Num b, Num a, Reifies s W) - => BVar s b - -> Traversal' b a - -> [BVar s a] +(^^..) :: + forall b a s. + (Num b, Num a, Reifies s W) => + BVar s b -> + Traversal' b a -> + [BVar s a] v ^^.. t = toListOfVar t v {-# INLINE (^^..) #-} -- | 'Numeric.Backprop.toListOfVar', but with 'Num' constraints instead of -- 'Backprop' constraints. -toListOfVar - :: forall b a s. (Num b, Num a, Reifies s W) - => Traversal' b a - -> BVar s b - -> [BVar s a] +toListOfVar :: + forall b a s. + (Num b, Num a, Reifies s W) => + Traversal' b a -> + BVar s b -> + [BVar s a] toListOfVar = E.toListOfVar E.afNum E.zfNum {-# INLINE toListOfVar #-} @@ -373,10 +428,10 @@ -- 'Backprop' constraints. -- -- Since v0.2.4, requires a 'Num' constraint on @t a@. -sequenceVar - :: (Traversable t, Num a, Reifies s W) - => BVar s (t a) - -> t (BVar s a) +sequenceVar :: + (Traversable t, Num a, Reifies s W) => + BVar s (t a) -> + t (BVar s a) sequenceVar = E.sequenceVar E.afNum E.zfNum {-# INLINE sequenceVar #-} @@ -384,99 +439,99 @@ -- 'Backprop' constraints. -- -- Prior to v0.2.3, required a 'Num' constraint on @t a@. -collectVar - :: (Foldable t, Functor t, Num a, Reifies s W) - => t (BVar s a) - -> BVar s (t a) +collectVar :: + (Foldable t, Functor t, Num a, Reifies s W) => + t (BVar s a) -> + BVar s (t a) collectVar = E.collectVar E.afNum E.zfNum {-# INLINE collectVar #-} -- | 'Numeric.Backprop.liftOp', but with 'Num' constraints instead of -- 'Backprop' constraints. -liftOp - :: (RPureConstrained Num as, Reifies s W) - => Op as b - -> Rec (BVar s) as - -> BVar s b +liftOp :: + (RPureConstrained Num as, Reifies s W) => + Op as b -> + Rec (BVar s) as -> + BVar s b liftOp = E.liftOp E.afNums {-# INLINE liftOp #-} -- | 'Numeric.Backprop.liftOp1', but with 'Num' constraints instead of -- 'Backprop' constraints. -liftOp1 - :: (Num a, Reifies s W) - => Op '[a] b - -> BVar s a - -> BVar s b +liftOp1 :: + (Num a, Reifies s W) => + Op '[a] b -> + BVar s a -> + BVar s b liftOp1 = E.liftOp1 E.afNum {-# INLINE liftOp1 #-} -- | 'Numeric.Backprop.liftOp2', but with 'Num' constraints instead of -- 'Backprop' constraints. -liftOp2 - :: (Num a, Num b, Reifies s W) - => Op '[a,b] c - -> BVar s a - -> BVar s b - -> BVar s c +liftOp2 :: + (Num a, Num b, Reifies s W) => + Op '[a, b] c -> + BVar s a -> + BVar s b -> + BVar s c liftOp2 = E.liftOp2 E.afNum E.afNum {-# INLINE liftOp2 #-} -- | 'Numeric.Backprop.liftOp3', but with 'Num' constraints instead of -- 'Backprop' constraints. -liftOp3 - :: (Num a, Num b, Num c, Reifies s W) - => Op '[a,b,c] d - -> BVar s a - -> BVar s b - -> BVar s c - -> BVar s d +liftOp3 :: + (Num a, Num b, Num c, Reifies s W) => + Op '[a, b, c] d -> + BVar s a -> + BVar s b -> + BVar s c -> + BVar s d liftOp3 = E.liftOp3 E.afNum E.afNum E.afNum {-# INLINE liftOp3 #-} -- | 'Numeric.Backprop.isoVar', but with 'Num' constraints instead of -- 'Backprop' constraints. -isoVar - :: (Num a, Reifies s W) - => (a -> b) - -> (b -> a) - -> BVar s a - -> BVar s b +isoVar :: + (Num a, Reifies s W) => + (a -> b) -> + (b -> a) -> + BVar s a -> + BVar s b isoVar = E.isoVar E.afNum {-# INLINE isoVar #-} -- | 'Numeric.Backprop.isoVar', but with 'Num' constraints instead of -- 'Backprop' constraints. -isoVar2 - :: (Num a, Num b, Reifies s W) - => (a -> b -> c) - -> (c -> (a, b)) - -> BVar s a - -> BVar s b - -> BVar s c +isoVar2 :: + (Num a, Num b, Reifies s W) => + (a -> b -> c) -> + (c -> (a, b)) -> + BVar s a -> + BVar s b -> + BVar s c isoVar2 = E.isoVar2 E.afNum E.afNum {-# INLINE isoVar2 #-} -- | 'Numeric.Backprop.isoVar3', but with 'Num' constraints instead of -- 'Backprop' constraints. -isoVar3 - :: (Num a, Num b, Num c, Reifies s W) - => (a -> b -> c -> d) - -> (d -> (a, b, c)) - -> BVar s a - -> BVar s b - -> BVar s c - -> BVar s d +isoVar3 :: + (Num a, Num b, Num c, Reifies s W) => + (a -> b -> c -> d) -> + (d -> (a, b, c)) -> + BVar s a -> + BVar s b -> + BVar s c -> + BVar s d isoVar3 = E.isoVar3 E.afNum E.afNum E.afNum {-# INLINE isoVar3 #-} -- | 'Numeric.Backprop.isoVarN', but with 'Num' constraints instead of -- 'Backprop' constraints. -isoVarN - :: (RPureConstrained Num as, Reifies s W) - => (Rec Identity as -> b) - -> (b -> Rec Identity as) - -> Rec (BVar s) as - -> BVar s b +isoVarN :: + (RPureConstrained Num as, Reifies s W) => + (Rec Identity as -> b) -> + (b -> Rec Identity as) -> + Rec (BVar s) as -> + BVar s b isoVarN = E.isoVarN E.afNums {-# INLINE isoVarN #-} diff --git a/src/Numeric/Backprop/Op.hs b/src/Numeric/Backprop/Op.hs --- a/src/Numeric/Backprop/Op.hs +++ b/src/Numeric/Backprop/Op.hs @@ -1,12 +1,12 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE UndecidableInstances #-} -- | @@ -42,51 +42,92 @@ -- -- See also -- for writing Ops manually on your own numerical functions. --- - module Numeric.Backprop.Op ( -- * Implementation -- $opdoc - -- * Types - -- ** Op and Synonyms - Op(..) + Op (..), + -- ** Tuple Types#prod# -- $prod - , Rec(..) + Rec (..), + -- * Running + -- ** Pure - , runOp, evalOp, gradOp, gradOpWith + runOp, + evalOp, + gradOp, + gradOpWith, + -- * Creation - , op0, opConst, idOp - , opLens + op0, + opConst, + idOp, + opLens, + -- ** Giving gradients directly - , op1, op2, op3 + op1, + op2, + op3, + -- ** From Isomorphisms - , opCoerce, opTup, opIso, opIso2, opIso3, opIsoN + opCoerce, + opTup, + opIso, + opIso2, + opIso3, + opIsoN, + -- ** No gradient - , noGrad1, noGrad + noGrad1, + noGrad, + -- * Manipulation - , composeOp, composeOp1, (~.) + composeOp, + composeOp1, + (~.), + -- * Utility + -- ** Numeric Ops#numops# -- $numops - , (+.), (-.), (*.), negateOp, absOp, signumOp - , (/.), recipOp - , expOp, logOp, sqrtOp, (**.), logBaseOp - , sinOp, cosOp, tanOp, asinOp, acosOp, atanOp - , sinhOp, coshOp, tanhOp, asinhOp, acoshOp, atanhOp - ) where + (+.), + (-.), + (*.), + negateOp, + absOp, + signumOp, + (/.), + recipOp, + expOp, + logOp, + sqrtOp, + (**.), + logBaseOp, + sinOp, + cosOp, + tanOp, + asinOp, + acosOp, + atanOp, + sinhOp, + coshOp, + tanhOp, + asinhOp, + acoshOp, + atanhOp, +) where -import Control.Applicative -import Data.Bifunctor -import Data.Coerce -import Data.Functor.Identity -import Data.List -import Data.Type.Util -import Data.Vinyl.Core -import Lens.Micro -import Lens.Micro.Extras -import qualified Data.Vinyl.Recursive as VR +import Control.Applicative +import Data.Bifunctor +import Data.Coerce +import Data.Functor.Identity +import Data.List (foldl') +import Data.Type.Util +import Data.Vinyl.Core +import qualified Data.Vinyl.Recursive as VR +import Lens.Micro +import Lens.Micro.Extras -- $opdoc -- 'Op's contain information on a function as well as its gradient, but @@ -181,61 +222,71 @@ -- -- To /use/ an 'Op' with the backprop library, see 'liftOp', 'liftOp1', -- 'liftOp2', and 'liftOp3'. -newtype Op as a = - -- | Construct an 'Op' by giving a function creating the +newtype Op as 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 'Op' expect. - Op { -- | Run the function that the 'Op' encodes, returning - -- a continuation to compute the gradient, given the total - -- derivative of @a@. See documentation for "Numeric.Backprop.Op" - -- for more information. - runOpWith :: Rec Identity as -> (a, a -> Rec Identity as) - } + Op + { runOpWith :: Rec Identity as -> (a, a -> Rec Identity as) + -- ^ Run the function that the 'Op' encodes, returning + -- a continuation to compute the gradient, given the total + -- derivative of @a@. See documentation for "Numeric.Backprop.Op" + -- for more information. + } -- | Helper wrapper used for the implementation of 'composeOp'. -newtype OpCont as a = OC { runOpCont :: a -> Rec Identity as } +newtype OpCont as a = OC {runOpCont :: a -> Rec Identity as} -- | Compose 'Op's together, like 'sequence' for functions, or @liftAN@. -- -- That is, given an @'Op' as b1@, an @'Op' as b2@, and an @'Op' as b3@, it -- can compose them with an @'Op' '[b1,b2,b3] c@ to create an @'Op' as -- c@. -composeOp - :: forall as bs c. (RPureConstrained Num as) - => Rec (Op as) bs -- ^ 'Rec' of 'Op's taking @as@ and returning - -- different @b@ in @bs@ - -> Op bs c -- ^ 'OpM' taking eac of the @bs@ from the - -- input 'Rec'. - -> Op as c -- ^ Composed 'Op' +composeOp :: + forall as bs c. + RPureConstrained Num as => + -- | 'Rec' of 'Op's taking @as@ and returning + -- different @b@ in @bs@ + Rec (Op as) bs -> + -- | 'OpM' taking eac of the @bs@ from the + -- input 'Rec'. + Op bs c -> + -- | Composed 'Op' + Op as c composeOp os o = Op $ \xs -> - let (ys, conts) = runzipWith (bimap Identity OC . flip runOpWith xs) os - (z, gFz) = runOpWith o ys - gFunc g0 = - let g1 = gFz g0 - g2s :: Rec (Const (Rec Identity as)) bs - g2s = VR.rzipWith (\oc (Identity g) -> Const $ runOpCont oc g) - conts g1 - in VR.rmap (\(Dict x) -> Identity x) - . foldl' (VR.rzipWith (\(Dict !x) (Identity y) -> - let q = x + y in q `seq` Dict q - ) - ) - (rpureConstrained @Num (Dict @Num 0)) - . VR.rfoldMap ((:[]) . getConst) - $ g2s - in (z, gFunc) + let (ys, conts) = runzipWith (bimap Identity OC . flip runOpWith xs) os + (z, gFz) = runOpWith o ys + gFunc g0 = + let g1 = gFz g0 + g2s :: Rec (Const (Rec Identity as)) bs + g2s = + VR.rzipWith + (\oc (Identity g) -> Const $ runOpCont oc g) + conts + g1 + in VR.rmap (\(Dict x) -> Identity x) + . foldl' + ( VR.rzipWith + ( \(Dict !x) (Identity y) -> + let q = x + y in q `seq` Dict q + ) + ) + (rpureConstrained @Num (Dict @Num 0)) + . VR.rfoldMap ((: []) . getConst) + $ g2s + in (z, gFunc) -- | Convenient wrapper over 'composeOp' for the case where the second -- function only takes one input, so the two 'Op's can be directly piped -- together, like for '.'. -composeOp1 - :: RPureConstrained Num as - => Op as b - -> Op '[b] c - -> Op as c +composeOp1 :: + RPureConstrained Num as => + Op as b -> + Op '[b] c -> + Op as c composeOp1 = composeOp . (:& RNil) -- | Convenient infix synonym for (flipped) 'composeOp1'. Meant to be used @@ -248,15 +299,15 @@ -- f '~.' g :: Op '[a, a] c -- @ infixr 9 ~. -(~.) - :: (RPureConstrained Num as) - => Op '[b] c - -> Op as b - -> Op as c + +(~.) :: + RPureConstrained Num as => + Op '[b] c -> + Op as b -> + Op as c (~.) = flip composeOp1 {-# INLINE (~.) #-} - -- | Run the function that an 'Op' encodes, to get the result. -- -- >>> runOp (op2 (*)) (3 :& 5 :& RNil) @@ -279,11 +330,15 @@ -- -- See the module documentaiton for "Numeric.Backprop.Op" for more -- information. -gradOpWith - :: Op as a -- ^ 'Op' to run - -> Rec Identity as -- ^ Inputs to run it with - -> a -- ^ The total derivative of the result. - -> Rec Identity as -- ^ The gradient +gradOpWith :: + -- | 'Op' to run + Op as a -> + -- | Inputs to run it with + Rec Identity as -> + -- | The total derivative of the result. + a -> + -- | The gradient + Rec Identity as gradOpWith o = snd . runOpWith o {-# INLINE gradOpWith #-} @@ -297,7 +352,6 @@ -- @ -- 'gradOp' o xs = 'gradOpWith' o xs 1 -- @ --- gradOp :: Num a => Op as a -> Rec Identity as -> Rec Identity as gradOp o i = gradOpWith o i 1 {-# INLINE gradOp #-} @@ -330,9 +384,9 @@ -- @since 0.1.3.0 noGrad1 :: (a -> b) -> Op '[a] b noGrad1 f = op1 $ \x -> - ( f x - , \_ -> errorWithoutStackTrace "Numeric.Backprop.Op.noGrad1: no gradient defined" - ) + ( f x + , \_ -> errorWithoutStackTrace "Numeric.Backprop.Op.noGrad1: no gradient defined" + ) {-# INLINE noGrad1 #-} -- | Create an 'Op' with no gradient. Can be evaluated with 'evalOp', but @@ -350,9 +404,9 @@ -- @since 0.1.3.0 noGrad :: (Rec Identity as -> b) -> Op as b noGrad f = Op $ \xs -> - ( f xs - , \_ -> errorWithoutStackTrace "Numeric.Backprop.Op.noGrad: no gradient defined" - ) + ( f xs + , \_ -> errorWithoutStackTrace "Numeric.Backprop.Op.noGrad: no gradient defined" + ) {-# INLINE noGrad #-} -- | An 'Op' that just returns whatever it receives. The identity @@ -362,7 +416,7 @@ -- 'idOp' = 'opIso' 'id' 'id' -- @ idOp :: Op '[a] a -idOp = op1 $ \x -> (x, id) +idOp = op1 (,id) {-# INLINE idOp #-} -- | An 'Op' that takes @as@ and returns exactly the input tuple. @@ -370,7 +424,7 @@ -- >>> gradOp' opTup (1 :& 2 :& 3 :& RNil) -- (1 :& 2 :& 3 :& RNil, 1 :& 1 :& 1 :& RNil) opTup :: Op as (Rec Identity as) -opTup = Op $ \xs -> (xs, id) +opTup = Op (,id) {-# INLINE opTup #-} -- | An 'Op' that runs the input value through an isomorphism. @@ -378,7 +432,7 @@ -- Warning: This is unsafe! It assumes that the isomorphisms themselves -- have derivative 1, so will break for things like 'exp' & 'log'. -- Basically, don't use this for any "numeric" isomorphisms. -opIso :: (a -> b) -> (b -> a) -> Op '[ a ] b +opIso :: (a -> b) -> (b -> a) -> Op '[a] b opIso to' from' = op1 $ \x -> (to' x, from') {-# INLINE opIso #-} @@ -414,7 +468,7 @@ -- Warning: This is unsafe! It assumes that it extracts a specific value -- unchanged, with derivative 1, so will break for things that numerically -- manipulate things before returning them. -opLens :: Num a => Lens' a b -> Op '[ a ] b +opLens :: Num a => Lens' a b -> Op '[a] b opLens l = op1 $ \x -> (view l x, \d -> set l d 0) {-# INLINE opLens #-} @@ -423,12 +477,15 @@ -- -- >>> gradOp' (opConst 10) (1 :& 2 :& 3 :& RNil) -- (10, 0 :& 0 :& 0 :& RNil) -opConst - :: forall as a. RPureConstrained Num as - => a - -> Op as a -opConst x = Op $ const - (x, const $ rpureConstrained @Num 0) +opConst :: + forall as a. + RPureConstrained Num as => + a -> + Op as a +opConst x = + Op $ + const + (x, const $ rpureConstrained @Num 0) {-# INLINE opConst #-} -- | Create an 'Op' that takes no inputs and always returns the given @@ -444,7 +501,7 @@ -- 'opConst''. op0 :: a -> Op '[] a op0 x = Op $ \case - RNil -> (x, const RNil) + RNil -> (x, const RNil) {-# INLINE op0 #-} -- | Create an 'Op' of a function taking one input, by giving its explicit @@ -482,13 +539,13 @@ -- -- Remember that, generally, end users shouldn't directly construct 'Op's; -- they should be provided by libraries or generated automatically. -op1 - :: (a -> (b, b -> a)) - -> Op '[a] b +op1 :: + (a -> (b, b -> a)) -> + Op '[a] b op1 f = Op $ \case - Identity x :& RNil -> - let (y, dx) = f x - in (y, \(!d) -> (:& RNil) . Identity . dx $ d) + Identity x :& RNil -> + let (y, dx) = f x + in (y, \(!d) -> (:& RNil) . Identity . dx $ d) {-# INLINE op1 #-} -- | Create an 'Op' of a function taking two inputs, by giving its explicit @@ -528,86 +585,86 @@ -- -- Remember that, generally, end users shouldn't directly construct 'Op's; -- they should be provided by libraries or generated automatically. -op2 - :: (a -> b -> (c, c -> (a, b))) - -> Op '[a,b] c +op2 :: + (a -> b -> (c, c -> (a, b))) -> + Op '[a, b] c op2 f = Op $ \case - Identity x :& Identity y :& RNil -> - let (z, dxdy) = f x y - in (z, (\(!dx,!dy) -> Identity dx :& Identity dy :& RNil) . dxdy) + Identity x :& Identity y :& RNil -> + let (z, dxdy) = f x y + in (z, (\(!dx, !dy) -> Identity dx :& Identity dy :& RNil) . dxdy) {-# INLINE op2 #-} -- | 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, d -> (a, b, c))) - -> Op '[a,b,c] d +op3 :: + (a -> b -> c -> (d, d -> (a, b, c))) -> + Op '[a, b, c] d op3 f = Op $ \case - Identity x :& Identity y :& Identity z :& RNil -> - let (q, dxdydz) = f x y z - in (q, (\(!dx, !dy, !dz) -> Identity dx :& Identity dy :& Identity dz :& RNil) . dxdydz) + Identity x :& Identity y :& Identity z :& RNil -> + let (q, dxdydz) = f x y z + in (q, (\(!dx, !dy, !dz) -> Identity dx :& Identity dy :& Identity dz :& RNil) . dxdydz) {-# INLINE op3 #-} instance (RPureConstrained Num as, Num a) => Num (Op as a) where - o1 + o2 = composeOp (o1 :& o2 :& RNil) (+.) - {-# INLINE (+) #-} - o1 - o2 = composeOp (o1 :& o2 :& RNil) (-.) - {-# INLINE (-) #-} - o1 * o2 = composeOp (o1 :& o2 :& RNil) (*.) - {-# INLINE (*) #-} - negate o = composeOp (o :& RNil) negateOp - {-# INLINE negate #-} - signum o = composeOp (o :& RNil) signumOp - {-# INLINE signum #-} - abs o = composeOp (o :& RNil) absOp - {-# INLINE abs #-} - fromInteger x = opConst (fromInteger x) - {-# INLINE fromInteger #-} + o1 + o2 = composeOp (o1 :& o2 :& RNil) (+.) + {-# INLINE (+) #-} + o1 - o2 = composeOp (o1 :& o2 :& RNil) (-.) + {-# INLINE (-) #-} + o1 * o2 = composeOp (o1 :& o2 :& RNil) (*.) + {-# INLINE (*) #-} + negate o = composeOp (o :& RNil) negateOp + {-# INLINE negate #-} + signum o = composeOp (o :& RNil) signumOp + {-# INLINE signum #-} + abs o = composeOp (o :& RNil) absOp + {-# INLINE abs #-} + fromInteger x = opConst (fromInteger x) + {-# INLINE fromInteger #-} instance (RPureConstrained Num as, Fractional a) => Fractional (Op as a) where - o1 / o2 = composeOp (o1 :& o2 :& RNil) (/.) - recip o = composeOp (o :& RNil) recipOp - {-# INLINE recip #-} - fromRational x = opConst (fromRational x) - {-# INLINE fromRational #-} + o1 / o2 = composeOp (o1 :& o2 :& RNil) (/.) + recip o = composeOp (o :& RNil) recipOp + {-# INLINE recip #-} + fromRational x = opConst (fromRational x) + {-# INLINE fromRational #-} instance (RPureConstrained Num as, Floating a) => Floating (Op as a) where - pi = opConst pi - {-# INLINE pi #-} - exp o = composeOp (o :& RNil) expOp - {-# INLINE exp #-} - log o = composeOp (o :& RNil) logOp - {-# INLINE log #-} - sqrt o = composeOp (o :& RNil) sqrtOp - {-# INLINE sqrt #-} - o1 ** o2 = composeOp (o1 :& o2 :& RNil) (**.) - {-# INLINE (**) #-} - logBase o1 o2 = composeOp (o1 :& o2 :& RNil) logBaseOp - {-# INLINE logBase #-} - sin o = composeOp (o :& RNil) sinOp - {-# INLINE sin #-} - cos o = composeOp (o :& RNil) cosOp - {-# INLINE cos #-} - tan o = composeOp (o :& RNil) tanOp - {-# INLINE tan #-} - asin o = composeOp (o :& RNil) asinOp - {-# INLINE asin #-} - acos o = composeOp (o :& RNil) acosOp - {-# INLINE acos #-} - atan o = composeOp (o :& RNil) atanOp - {-# INLINE atan #-} - sinh o = composeOp (o :& RNil) sinhOp - {-# INLINE sinh #-} - cosh o = composeOp (o :& RNil) coshOp - {-# INLINE cosh #-} - tanh o = composeOp (o :& RNil) tanhOp - {-# INLINE tanh #-} - asinh o = composeOp (o :& RNil) asinhOp - {-# INLINE asinh #-} - acosh o = composeOp (o :& RNil) acoshOp - {-# INLINE acosh #-} - atanh o = composeOp (o :& RNil) atanhOp - {-# INLINE atanh #-} + pi = opConst pi + {-# INLINE pi #-} + exp o = composeOp (o :& RNil) expOp + {-# INLINE exp #-} + log o = composeOp (o :& RNil) logOp + {-# INLINE log #-} + sqrt o = composeOp (o :& RNil) sqrtOp + {-# INLINE sqrt #-} + o1 ** o2 = composeOp (o1 :& o2 :& RNil) (**.) + {-# INLINE (**) #-} + logBase o1 o2 = composeOp (o1 :& o2 :& RNil) logBaseOp + {-# INLINE logBase #-} + sin o = composeOp (o :& RNil) sinOp + {-# INLINE sin #-} + cos o = composeOp (o :& RNil) cosOp + {-# INLINE cos #-} + tan o = composeOp (o :& RNil) tanOp + {-# INLINE tan #-} + asin o = composeOp (o :& RNil) asinOp + {-# INLINE asin #-} + acos o = composeOp (o :& RNil) acosOp + {-# INLINE acos #-} + atan o = composeOp (o :& RNil) atanOp + {-# INLINE atan #-} + sinh o = composeOp (o :& RNil) sinhOp + {-# INLINE sinh #-} + cosh o = composeOp (o :& RNil) coshOp + {-# INLINE cosh #-} + tanh o = composeOp (o :& RNil) tanhOp + {-# INLINE tanh #-} + asinh o = composeOp (o :& RNil) asinhOp + {-# INLINE asinh #-} + acosh o = composeOp (o :& RNil) acoshOp + {-# INLINE acosh #-} + atanh o = composeOp (o :& RNil) atanhOp + {-# INLINE atanh #-} -- $numops -- @@ -632,32 +689,33 @@ -- | 'Op' for multiplication (*.) :: Num a => Op '[a, a] a -(*.) = op2 $ \x y -> (x * y, \g -> (y*g, x*g)) +(*.) = op2 $ \x y -> (x * y, \g -> (y * g, x * g)) {-# INLINE (*.) #-} -- | 'Op' for division (/.) :: Fractional a => Op '[a, a] a -(/.) = op2 $ \x y -> (x / y, \g -> (g/y, -g*x/(y*y))) +(/.) = op2 $ \x y -> (x / y, \g -> (g / y, -(g * x / (y * y)))) {-# INLINE (/.) #-} -- | 'Op' for exponentiation (**.) :: Floating a => Op '[a, a] a -(**.) = op2 $ \x y -> ( x ** y - , let dx = y*x**(y-1) - dy = x**y*log x - in \g -> (g*dx, g*dy) - ) +(**.) = op2 $ \x y -> + ( x ** y + , let dx = y * x ** (y - 1) + dy = x ** y * log x + in \g -> (g * dx, g * dy) + ) {-# INLINE (**.) #-} -- | 'Op' for negation negateOp :: Num a => Op '[a] a negateOp = op1 $ \x -> (negate x, negate) -{-# INLINE negateOp #-} +{-# INLINE negateOp #-} -- | 'Op' for 'signum' signumOp :: Num a => Op '[a] a signumOp = op1 $ \x -> (signum x, const 0) -{-# INLINE signumOp #-} +{-# INLINE signumOp #-} -- | 'Op' for absolute value absOp :: Num a => Op '[a] a @@ -666,7 +724,7 @@ -- | 'Op' for multiplicative inverse recipOp :: Fractional a => Op '[a] a -recipOp = op1 $ \x -> (recip x, (/(x*x)) . negate) +recipOp = op1 $ \x -> (recip x, (/ (x * x)) . negate) {-# INLINE recipOp #-} -- | 'Op' for 'exp' @@ -676,7 +734,7 @@ -- | 'Op' for the natural logarithm logOp :: Floating a => Op '[a] a -logOp = op1 $ \x -> (log x, (/x)) +logOp = op1 $ \x -> (log x, (/ x)) {-# INLINE logOp #-} -- | 'Op' for square root @@ -686,10 +744,11 @@ -- | 'Op' for 'logBase' logBaseOp :: Floating a => Op '[a, a] a -logBaseOp = op2 $ \x y -> ( logBase x y - , let dx = - logBase x y / (log x * x) - in \g -> (g*dx, g/(y * log x)) - ) +logBaseOp = op2 $ \x y -> + ( logBase x y + , let dx = -(logBase x y / (log x * x)) + in \g -> (g * dx, g / (y * log x)) + ) {-# INLINE logBaseOp #-} -- | 'Op' for sine @@ -704,22 +763,22 @@ -- | 'Op' for tangent tanOp :: Floating a => Op '[a] a -tanOp = op1 $ \x -> (tan x, (/ cos x^(2::Int))) +tanOp = op1 $ \x -> (tan x, (/ cos x ^ (2 :: Int))) {-# INLINE tanOp #-} -- | 'Op' for arcsine asinOp :: Floating a => Op '[a] a -asinOp = op1 $ \x -> (asin x, (/ sqrt(1 - x*x))) +asinOp = op1 $ \x -> (asin x, (/ sqrt (1 - x * x))) {-# INLINE asinOp #-} -- | 'Op' for arccosine acosOp :: Floating a => Op '[a] a -acosOp = op1 $ \x -> (acos x, (/ sqrt (1 - x*x)) . negate) +acosOp = op1 $ \x -> (acos x, (/ sqrt (1 - x * x)) . negate) {-# INLINE acosOp #-} -- | 'Op' for arctangent atanOp :: Floating a => Op '[a] a -atanOp = op1 $ \x -> (atan x, (/ (x*x + 1))) +atanOp = op1 $ \x -> (atan x, (/ (x * x + 1))) {-# INLINE atanOp #-} -- | 'Op' for hyperbolic sine @@ -734,22 +793,22 @@ -- | 'Op' for hyperbolic tangent tanhOp :: Floating a => Op '[a] a -tanhOp = op1 $ \x -> (tanh x, (/ cosh x^(2::Int))) +tanhOp = op1 $ \x -> (tanh x, (/ cosh x ^ (2 :: Int))) {-# INLINE tanhOp #-} -- | 'Op' for hyperbolic arcsine asinhOp :: Floating a => Op '[a] a -asinhOp = op1 $ \x -> (asinh x, (/ sqrt (x*x + 1))) +asinhOp = op1 $ \x -> (asinh x, (/ sqrt (x * x + 1))) {-# INLINE asinhOp #-} -- | 'Op' for hyperbolic arccosine acoshOp :: Floating a => Op '[a] a -acoshOp = op1 $ \x -> (acosh x, (/ sqrt (x*x - 1))) +acoshOp = op1 $ \x -> (acosh x, (/ sqrt (x * x - 1))) {-# INLINE acoshOp #-} -- | 'Op' for hyperbolic arctangent atanhOp :: Floating a => Op '[a] a -atanhOp = op1 $ \x -> (atanh x, (/ (1 - x*x))) +atanhOp = op1 $ \x -> (atanh x, (/ (1 - x * x))) {-# INLINE atanhOp #-} -- $prod @@ -776,4 +835,3 @@ -- z :: f c -- x :& y :& z :& RNil :: Rec f '[a, b, c] -- @ --- diff --git a/src/Prelude/Backprop.hs b/src/Prelude/Backprop.hs --- a/src/Prelude/Backprop.hs +++ b/src/Prelude/Backprop.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleContexts #-} -- | -- Module : Prelude.Backprop @@ -20,87 +20,101 @@ -- allowing you to provide 'zero', 'add', and 'one' explicitly. -- -- @since 0.1.3.0 --- - module Prelude.Backprop ( -- * Foldable and Traversable - sum - , product - , length - , minimum - , maximum - , traverse - , toList - , mapAccumL - , mapAccumR - , foldr, foldl' + sum, + product, + length, + minimum, + maximum, + traverse, + toList, + mapAccumL, + mapAccumR, + foldr, + foldl', + -- * Functor and Applicative - , fmap, fmapConst - , (<$>), (<$), ($>) - , pure - , liftA2 - , liftA3 + fmap, + fmapConst, + (<$>), + (<$), + ($>), + pure, + liftA2, + liftA3, + -- * Numeric - , fromIntegral - , realToFrac - , round - , fromIntegral' + fromIntegral, + realToFrac, + round, + fromIntegral', + -- * Misc - , E.coerce - ) where + E.coerce, +) where -import Numeric.Backprop -import Prelude (Num(..), Fractional(..), Ord(..), Functor, Foldable, Traversable, Applicative) +import Numeric.Backprop import qualified Numeric.Backprop.Explicit as E -import qualified Prelude as P import qualified Prelude.Backprop.Explicit as E +import Prelude ( + Applicative, + Foldable, + Fractional (..), + Functor, + Num (..), + Ord (..), + Traversable, + ) +import qualified Prelude as P -- | Lifted 'P.sum'. More efficient than going through 'toList'. -sum :: (Foldable t, Functor t, Backprop (t a), Num a, Reifies s W) - => BVar s (t a) - -> BVar s a +sum :: + (Foldable t, Functor t, Backprop (t a), Num a, Reifies s W) => + BVar s (t a) -> + BVar s a sum = E.sum E.addFunc {-# INLINE sum #-} -- | Lifted 'P.pure'. -pure - :: (Foldable t, Applicative t, Backprop a, Reifies s W) - => BVar s a - -> BVar s (t a) +pure :: + (Foldable t, Applicative t, Backprop a, Reifies s W) => + BVar s a -> + BVar s (t a) pure = E.pure E.addFunc E.zeroFunc {-# INLINE pure #-} -- | Lifted 'P.product'. More efficient than going through 'toList'. -product - :: (Foldable t, Functor t, Backprop (t a), Fractional a, Reifies s W) - => BVar s (t a) - -> BVar s a +product :: + (Foldable t, Functor t, Backprop (t a), Fractional a, Reifies s W) => + BVar s (t a) -> + BVar s a product = E.product E.addFunc {-# INLINE product #-} -- | Lifted 'P.length'. More efficient than going through 'toList'. -length - :: (Foldable t, Backprop (t a), Num b, Reifies s W) - => BVar s (t a) - -> BVar s b +length :: + (Foldable t, Backprop (t a), Num b, Reifies s W) => + BVar s (t a) -> + BVar s b length = E.length E.addFunc E.zeroFunc {-# INLINE length #-} -- | Lifted 'P.minimum'. Undefined for situations where 'P.minimum' would -- be undefined. More efficient than going through 'toList'. -minimum - :: (Foldable t, Functor t, Backprop a, Ord a, Backprop (t a), Reifies s W) - => BVar s (t a) - -> BVar s a +minimum :: + (Foldable t, Functor t, Backprop a, Ord a, Backprop (t a), Reifies s W) => + BVar s (t a) -> + BVar s a minimum = E.minimum E.addFunc E.zeroFunc {-# INLINE minimum #-} -- | Lifted 'P.maximum'. Undefined for situations where 'P.maximum' would -- be undefined. More efficient than going through 'toList'. -maximum - :: (Foldable t, Functor t, Backprop a, Ord a, Backprop (t a), Reifies s W) - => BVar s (t a) - -> BVar s a +maximum :: + (Foldable t, Functor t, Backprop a, Ord a, Backprop (t a), Reifies s W) => + BVar s (t a) -> + BVar s a maximum = E.maximum E.addFunc E.zeroFunc {-# INLINE maximum #-} @@ -108,12 +122,12 @@ -- list 'P.foldr', and is only here for convenience. -- -- @since 0.2.3.0 -foldr - :: (Traversable t, Backprop a, Reifies s W) - => (BVar s a -> BVar s b -> BVar s b) - -> BVar s b - -> BVar s (t a) - -> BVar s b +foldr :: + (Traversable t, Backprop a, Reifies s W) => + (BVar s a -> BVar s b -> BVar s b) -> + BVar s b -> + BVar s (t a) -> + BVar s b foldr = E.foldr E.addFunc E.zeroFunc {-# INLINE foldr #-} @@ -121,22 +135,22 @@ -- list 'P.foldl'', and is only here for convenience. -- -- @since 0.2.3.0 -foldl' - :: (Traversable t, Backprop a, Reifies s W) - => (BVar s b -> BVar s a -> BVar s b) - -> BVar s b - -> BVar s (t a) - -> BVar s b +foldl' :: + (Traversable t, Backprop a, Reifies s W) => + (BVar s b -> BVar s a -> BVar s b) -> + BVar s b -> + BVar s (t a) -> + BVar s b foldl' = E.foldl' E.addFunc E.zeroFunc {-# INLINE foldl' #-} -- | Lifted 'P.fmap'. Lifts backpropagatable functions to be -- backpropagatable functions on 'Traversable' 'Functor's. -fmap - :: (Traversable f, Backprop a, Backprop b, Reifies s W) - => (BVar s a -> BVar s b) - -> BVar s (f a) - -> BVar s (f b) +fmap :: + (Traversable f, Backprop a, Backprop b, Reifies s W) => + (BVar s a -> BVar s b) -> + BVar s (f a) -> + BVar s (f b) fmap = E.fmap E.addFunc E.addFunc E.zeroFunc E.zeroFunc {-# INLINE fmap #-} @@ -150,108 +164,136 @@ -- but much more efficient. -- -- @since 0.2.4.0 -fmapConst - :: (Functor f, Foldable f, Backprop b, Backprop (f a), Reifies s W) - => BVar s b - -> BVar s (f a) - -> BVar s (f b) +fmapConst :: + (Functor f, Foldable f, Backprop b, Backprop (f a), Reifies s W) => + BVar s b -> + BVar s (f a) -> + BVar s (f b) fmapConst = E.fmapConst E.addFunc E.addFunc E.zeroFunc E.zeroFunc {-# INLINE fmapConst #-} -- | Alias for 'fmap'. -(<$>) - :: (Traversable f, Backprop a, Backprop b, Reifies s W) - => (BVar s a -> BVar s b) - -> BVar s (f a) - -> BVar s (f b) +(<$>) :: + (Traversable f, Backprop a, Backprop b, Reifies s W) => + (BVar s a -> BVar s b) -> + BVar s (f a) -> + BVar s (f b) (<$>) = fmap + infixl 4 <$> {-# INLINE (<$>) #-} -- | Alias for 'fmapConst'. -- -- @since 0.2.4.0 -(<$) - :: (Traversable f, Backprop b, Backprop (f a), Reifies s W) - => BVar s b - -> BVar s (f a) - -> BVar s (f b) +(<$) :: + (Traversable f, Backprop b, Backprop (f a), Reifies s W) => + BVar s b -> + BVar s (f a) -> + BVar s (f b) (<$) = fmapConst + infixl 4 <$ {-# INLINE (<$) #-} -- | Alias for @'flip' 'fmapConst'@. -- -- @since 0.2.4.0 -($>) - :: (Traversable f, Backprop b, Backprop (f a), Reifies s W) - => BVar s (f a) - -> BVar s b - -> BVar s (f b) +($>) :: + (Traversable f, Backprop b, Backprop (f a), Reifies s W) => + BVar s (f a) -> + BVar s b -> + BVar s (f b) xs $> x = x <$ xs + infixl 4 $> {-# INLINE ($>) #-} -- | Lifted 'P.traverse'. Lifts backpropagatable functions to be -- backpropagatable functions on 'Traversable' 'Functor's. -traverse - :: (Traversable t, Applicative f, Foldable f, Backprop a, Backprop b, Backprop (t b), Reifies s W) - => (BVar s a -> f (BVar s b)) - -> BVar s (t a) - -> BVar s (f (t b)) -traverse = E.traverse E.addFunc E.addFunc E.addFunc - E.zeroFunc E.zeroFunc +traverse :: + (Traversable t, Applicative f, Foldable f, Backprop a, Backprop b, Backprop (t b), Reifies s W) => + (BVar s a -> f (BVar s b)) -> + BVar s (t a) -> + BVar s (f (t b)) +traverse = + E.traverse + E.addFunc + E.addFunc + E.addFunc + E.zeroFunc + E.zeroFunc {-# INLINE traverse #-} -- | Lifted 'P.liftA2'. Lifts backpropagatable functions to be -- backpropagatable functions on 'Traversable' 'Applicative's. -liftA2 - :: ( Traversable f, Applicative f - , Backprop a, Backprop b, Backprop c - , Reifies s W - ) - => (BVar s a -> BVar s b -> BVar s c) - -> BVar s (f a) - -> BVar s (f b) - -> BVar s (f c) -liftA2 = E.liftA2 E.addFunc E.addFunc E.addFunc - E.zeroFunc E.zeroFunc E.zeroFunc +liftA2 :: + ( Traversable f + , Applicative f + , Backprop a + , Backprop b + , Backprop c + , Reifies s W + ) => + (BVar s a -> BVar s b -> BVar s c) -> + BVar s (f a) -> + BVar s (f b) -> + BVar s (f c) +liftA2 = + E.liftA2 + E.addFunc + E.addFunc + E.addFunc + E.zeroFunc + E.zeroFunc + E.zeroFunc {-# INLINE liftA2 #-} -- | Lifted 'P.liftA3'. Lifts backpropagatable functions to be -- backpropagatable functions on 'Traversable' 'Applicative's. -liftA3 - :: ( Traversable f - , Applicative f - , Backprop a, Backprop b, Backprop c, Backprop d - , Reifies s W - ) - => (BVar s a -> BVar s b -> BVar s c -> BVar s d) - -> BVar s (f a) - -> BVar s (f b) - -> BVar s (f c) - -> BVar s (f d) -liftA3 = E.liftA3 E.addFunc E.addFunc E.addFunc E.addFunc - E.zeroFunc E.zeroFunc E.zeroFunc E.zeroFunc +liftA3 :: + ( Traversable f + , Applicative f + , Backprop a + , Backprop b + , Backprop c + , Backprop d + , Reifies s W + ) => + (BVar s a -> BVar s b -> BVar s c -> BVar s d) -> + BVar s (f a) -> + BVar s (f b) -> + BVar s (f c) -> + BVar s (f d) +liftA3 = + E.liftA3 + E.addFunc + E.addFunc + E.addFunc + E.addFunc + E.zeroFunc + E.zeroFunc + E.zeroFunc + E.zeroFunc {-# INLINE liftA3 #-} -- | Lifted conversion between two 'P.Integral' instances. -- -- @since 0.2.1.0 -fromIntegral - :: (Backprop a, P.Integral a, P.Integral b, Reifies s W) - => BVar s a - -> BVar s b +fromIntegral :: + (Backprop a, P.Integral a, P.Integral b, Reifies s W) => + BVar s a -> + BVar s b fromIntegral = E.fromIntegral E.addFunc {-# INLINE fromIntegral #-} -- | Lifted conversion between two 'Fractional' and 'P.Real' instances. -- -- @since 0.2.1.0 -realToFrac - :: (Backprop a, Fractional a, P.Real a, Fractional b, P.Real b, Reifies s W) - => BVar s a - -> BVar s b +realToFrac :: + (Backprop a, Fractional a, P.Real a, Fractional b, P.Real b, Reifies s W) => + BVar s a -> + BVar s b realToFrac = E.realToFrac E.addFunc {-# INLINE realToFrac #-} @@ -261,10 +303,10 @@ -- but does not do this for convenience reasons. -- -- @since 0.2.3.0 -round - :: (P.RealFrac a, P.Integral b, Reifies s W) - => BVar s a - -> BVar s b +round :: + (P.RealFrac a, P.Integral b, Reifies s W) => + BVar s a -> + BVar s b round = E.round E.afNum {-# INLINE round #-} @@ -277,10 +319,10 @@ -- reasons. -- -- @since 0.2.3.0 -fromIntegral' - :: (P.Integral a, P.RealFrac b, Reifies s W) - => BVar s a - -> BVar s b +fromIntegral' :: + (P.Integral a, P.RealFrac b, Reifies s W) => + BVar s a -> + BVar s b fromIntegral' = E.fromIntegral' E.afNum {-# INLINE fromIntegral' #-} @@ -293,10 +335,10 @@ -- @'P.minimum' . 'toList'.@ -- -- @since 0.2.2.0 -toList - :: (Traversable t, Backprop a, Reifies s W) - => BVar s (t a) - -> [BVar s a] +toList :: + (Traversable t, Backprop a, Reifies s W) => + BVar s (t a) -> + [BVar s a] toList = E.toList E.addFunc E.zeroFunc {-# INLINE toList #-} @@ -305,12 +347,12 @@ -- Prior to v0.2.3, required a 'Backprop' constraint on @t b@. -- -- @since 0.2.2.0 -mapAccumL - :: (Traversable t, Backprop b, Backprop c, Reifies s W) - => (BVar s a -> BVar s b -> (BVar s a, BVar s c)) - -> BVar s a - -> BVar s (t b) - -> (BVar s a, BVar s (t c)) +mapAccumL :: + (Traversable t, Backprop b, Backprop c, Reifies s W) => + (BVar s a -> BVar s b -> (BVar s a, BVar s c)) -> + BVar s a -> + BVar s (t b) -> + (BVar s a, BVar s (t c)) mapAccumL = E.mapAccumL E.addFunc E.addFunc E.zeroFunc E.zeroFunc {-# INLINE mapAccumL #-} @@ -319,12 +361,11 @@ -- Prior to v0.2.3, required a 'Backprop' constraint on @t b@. -- -- @since 0.2.2.0 -mapAccumR - :: (Traversable t, Backprop b, Backprop c, Reifies s W) - => (BVar s a -> BVar s b -> (BVar s a, BVar s c)) - -> BVar s a - -> BVar s (t b) - -> (BVar s a, BVar s (t c)) +mapAccumR :: + (Traversable t, Backprop b, Backprop c, Reifies s W) => + (BVar s a -> BVar s b -> (BVar s a, BVar s c)) -> + BVar s a -> + BVar s (t b) -> + (BVar s a, BVar s (t c)) mapAccumR = E.mapAccumR E.addFunc E.addFunc E.zeroFunc E.zeroFunc {-# INLINE mapAccumR #-} - diff --git a/src/Prelude/Backprop/Explicit.hs b/src/Prelude/Backprop/Explicit.hs --- a/src/Prelude/Backprop/Explicit.hs +++ b/src/Prelude/Backprop/Explicit.hs @@ -1,5 +1,5 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# OPTIONS_HADDOCK not-home #-} +{-# LANGUAGE FlexibleContexts #-} +{-# OPTIONS_HADDOCK not-home #-} -- | -- Module : Prelude.Backprop.Explicit @@ -20,252 +20,271 @@ -- reflect changes in underlying implementation. -- -- @since 0.2.0.0 - module Prelude.Backprop.Explicit ( -- * Foldable and Traversable - sum - , product - , length - , minimum - , maximum - , traverse - , toList - , mapAccumL - , mapAccumR - , foldr, foldl' + sum, + product, + length, + minimum, + maximum, + traverse, + toList, + mapAccumL, + mapAccumR, + foldr, + foldl', + -- * Functor and Applicative - , fmap, fmapConst - , pure - , liftA2 - , liftA3 + fmap, + fmapConst, + pure, + liftA2, + liftA3, + -- * Numeric - , fromIntegral - , realToFrac - , round - , fromIntegral' + fromIntegral, + realToFrac, + round, + fromIntegral', + -- * Misc - , coerce - ) where + coerce, +) where -import Data.Bifunctor -import Numeric.Backprop.Explicit -import Prelude (Num(..), Fractional(..), Eq(..), Ord(..), Functor, Foldable, Traversable, Applicative, (.), ($)) -import qualified Control.Applicative as P -import qualified Data.Coerce as C -import qualified Data.Foldable as P -import qualified Data.Traversable as P -import qualified Prelude as P +import qualified Control.Applicative as P +import Data.Bifunctor +import qualified Data.Coerce as C +import qualified Data.Foldable as P +import qualified Data.Traversable as P +import Numeric.Backprop.Explicit +import Prelude ( + Applicative, + Eq (..), + Foldable, + Fractional (..), + Functor, + Num (..), + Ord (..), + Traversable, + ($), + (.), + ) +import qualified Prelude as P -- | 'Prelude.Backprop.sum', but taking explicit 'add' and 'zero'. -sum :: (Foldable t, Functor t, Num a, Reifies s W) - => AddFunc (t a) - -> BVar s (t a) - -> BVar s a +sum :: + (Foldable t, Functor t, Num a, Reifies s W) => + AddFunc (t a) -> + BVar s (t a) -> + BVar s a sum af = liftOp1 af . op1 $ \xs -> - ( P.sum xs - , (P.<$ xs) - ) + ( P.sum xs + , (P.<$ xs) + ) {-# INLINE sum #-} -- | 'Prelude.Backprop.pure', but taking explicit 'add' and 'zero'. -pure - :: (Foldable t, Applicative t, Reifies s W) - => AddFunc a - -> ZeroFunc a - -> BVar s a - -> BVar s (t a) +pure :: + (Foldable t, Applicative t, Reifies s W) => + AddFunc a -> + ZeroFunc a -> + BVar s a -> + BVar s (t a) pure af zfa = liftOp1 af . op1 $ \x -> - ( P.pure x - , \d -> case P.toList d of - [] -> runZF zfa x - e:es -> P.foldl' (runAF af) e es - ) + ( P.pure x + , \d -> case P.toList d of + [] -> runZF zfa x + e : es -> P.foldl' (runAF af) e es + ) {-# INLINE pure #-} -- | 'Prelude.Backprop.product', but taking explicit 'add' and 'zero'. -product - :: (Foldable t, Functor t, Fractional a, Reifies s W) - => AddFunc (t a) - -> BVar s (t a) - -> BVar s a +product :: + (Foldable t, Functor t, Fractional a, Reifies s W) => + AddFunc (t a) -> + BVar s (t a) -> + BVar s a product af = liftOp1 af . op1 $ \xs -> - let p = P.product xs - in ( p - , \d -> (\x -> p * d / x) P.<$> xs - ) + let p = P.product xs + in ( p + , \d -> (\x -> p * d / x) P.<$> xs + ) {-# INLINE product #-} -- | 'Prelude.Backprop.length', but taking explicit 'add' and 'zero'. -length - :: (Foldable t, Num b, Reifies s W) - => AddFunc (t a) - -> ZeroFunc (t a) - -> BVar s (t a) - -> BVar s b +length :: + (Foldable t, Num b, Reifies s W) => + AddFunc (t a) -> + ZeroFunc (t a) -> + BVar s (t a) -> + BVar s b length af zfa = liftOp1 af . op1 $ \xs -> - ( P.fromIntegral (P.length xs) - , P.const (runZF zfa xs) - ) + ( P.fromIntegral (P.length xs) + , P.const (runZF zfa xs) + ) {-# INLINE length #-} -- | 'Prelude.Backprop.minimum', but taking explicit 'add' and 'zero'. -minimum - :: (Foldable t, Functor t, Ord a, Reifies s W) - => AddFunc (t a) - -> ZeroFunc a - -> BVar s (t a) - -> BVar s a +minimum :: + (Foldable t, Functor t, Ord a, Reifies s W) => + AddFunc (t a) -> + ZeroFunc a -> + BVar s (t a) -> + BVar s a minimum af zf = liftOp1 af . op1 $ \xs -> - let m = P.minimum xs - in ( m - , \d -> (\x -> if x == m then d else runZF zf x) P.<$> xs - ) + let m = P.minimum xs + in ( m + , \d -> (\x -> if x == m then d else runZF zf x) P.<$> xs + ) {-# INLINE minimum #-} -- | 'Prelude.Backprop.maximum', but taking explicit 'add' and 'zero'. -maximum - :: (Foldable t, Functor t, Ord a, Reifies s W) - => AddFunc (t a) - -> ZeroFunc a - -> BVar s (t a) - -> BVar s a +maximum :: + (Foldable t, Functor t, Ord a, Reifies s W) => + AddFunc (t a) -> + ZeroFunc a -> + BVar s (t a) -> + BVar s a maximum af zf = liftOp1 af . op1 $ \xs -> - let m = P.maximum xs - in ( m - , \d -> (\x -> if x == m then d else runZF zf x) P.<$> xs - ) + let m = P.maximum xs + in ( m + , \d -> (\x -> if x == m then d else runZF zf x) P.<$> xs + ) {-# INLINE maximum #-} -- | 'Prelude.Backprop.foldr', but taking explicit 'add' and 'zero'. -- -- @since 0.2.3.0 -foldr - :: (Traversable t, Reifies s W) - => AddFunc a - -> ZeroFunc a - -> (BVar s a -> BVar s b -> BVar s b) - -> BVar s b - -> BVar s (t a) - -> BVar s b +foldr :: + (Traversable t, Reifies s W) => + AddFunc a -> + ZeroFunc a -> + (BVar s a -> BVar s b -> BVar s b) -> + BVar s b -> + BVar s (t a) -> + BVar s b foldr af z f x = P.foldr f x . toList af z {-# INLINE foldr #-} -- | 'Prelude.Backprop.foldl'', but taking explicit 'add' and 'zero'. -- -- @since 0.2.3.0 -foldl' - :: (Traversable t, Reifies s W) - => AddFunc a - -> ZeroFunc a - -> (BVar s b -> BVar s a -> BVar s b) - -> BVar s b - -> BVar s (t a) - -> BVar s b +foldl' :: + (Traversable t, Reifies s W) => + AddFunc a -> + ZeroFunc a -> + (BVar s b -> BVar s a -> BVar s b) -> + BVar s b -> + BVar s (t a) -> + BVar s b foldl' af z f x = P.foldl' f x . toList af z {-# INLINE foldl' #-} -- | 'Prelude.Backprop.fmap', but taking explicit 'add' and 'zero'. -fmap - :: (Traversable f, Reifies s W) - => AddFunc a - -> AddFunc b - -> ZeroFunc a - -> ZeroFunc b - -> (BVar s a -> BVar s b) - -> BVar s (f a) - -> BVar s (f b) +fmap :: + (Traversable f, Reifies s W) => + AddFunc a -> + AddFunc b -> + ZeroFunc a -> + ZeroFunc b -> + (BVar s a -> BVar s b) -> + BVar s (f a) -> + BVar s (f b) fmap afa afb zfa zfb f = collectVar afb zfb . P.fmap f . sequenceVar afa zfa {-# INLINE fmap #-} -- | 'Prelude.Backprop.fmapConst', but taking explicit 'add' and 'zero'. -- -- @since 0.2.4.0 -fmapConst - :: (Functor f, Foldable f, Reifies s W) - => AddFunc (f a) - -> AddFunc b - -> ZeroFunc (f a) - -> ZeroFunc b - -> BVar s b - -> BVar s (f a) - -> BVar s (f b) +fmapConst :: + (Functor f, Foldable f, Reifies s W) => + AddFunc (f a) -> + AddFunc b -> + ZeroFunc (f a) -> + ZeroFunc b -> + BVar s b -> + BVar s (f a) -> + BVar s (f b) fmapConst afa afb zfa zfb = liftOp2 afb afa . op2 $ \x xs -> - ( x P.<$ xs - , \d -> ( case P.toList d of - [] -> runZF zfb x - e:es -> P.foldl' (runAF afb) e es - , runZF zfa xs - ) - ) + ( x P.<$ xs + , \d -> + ( case P.toList d of + [] -> runZF zfb x + e : es -> P.foldl' (runAF afb) e es + , runZF zfa xs + ) + ) {-# INLINE fmapConst #-} -- | 'Prelude.Backprop.traverse', but taking explicit 'add' and 'zero'. -traverse - :: (Traversable t, Applicative f, Foldable f, Reifies s W) - => AddFunc a - -> AddFunc b - -> AddFunc (t b) - -> ZeroFunc a - -> ZeroFunc b - -> (BVar s a -> f (BVar s b)) - -> BVar s (t a) - -> BVar s (f (t b)) -traverse afa afb aftb zfa zfb f - = collectVar aftb zftb - . P.fmap (collectVar afb zfb) - . P.traverse f - . sequenceVar afa zfa +traverse :: + (Traversable t, Applicative f, Foldable f, Reifies s W) => + AddFunc a -> + AddFunc b -> + AddFunc (t b) -> + ZeroFunc a -> + ZeroFunc b -> + (BVar s a -> f (BVar s b)) -> + BVar s (t a) -> + BVar s (f (t b)) +traverse afa afb aftb zfa zfb f = + collectVar aftb zftb + . P.fmap (collectVar afb zfb) + . P.traverse f + . sequenceVar afa zfa where zftb = ZF $ P.fmap (runZF zfb) {-# INLINE zftb #-} {-# INLINE traverse #-} -- | 'Prelude.Backprop.liftA2', but taking explicit 'add' and 'zero'. -liftA2 - :: ( Traversable f - , Applicative f - , Reifies s W - ) - => AddFunc a - -> AddFunc b - -> AddFunc c - -> ZeroFunc a - -> ZeroFunc b - -> ZeroFunc c - -> (BVar s a -> BVar s b -> BVar s c) - -> BVar s (f a) - -> BVar s (f b) - -> BVar s (f c) -liftA2 afa afb afc zfa zfb zfc f x y - = collectVar afc zfc - $ f P.<$> sequenceVar afa zfa x - P.<*> sequenceVar afb zfb y +liftA2 :: + ( Traversable f + , Applicative f + , Reifies s W + ) => + AddFunc a -> + AddFunc b -> + AddFunc c -> + ZeroFunc a -> + ZeroFunc b -> + ZeroFunc c -> + (BVar s a -> BVar s b -> BVar s c) -> + BVar s (f a) -> + BVar s (f b) -> + BVar s (f c) +liftA2 afa afb afc zfa zfb zfc f x y = + collectVar afc zfc $ + f + P.<$> sequenceVar afa zfa x + P.<*> sequenceVar afb zfb y {-# INLINE liftA2 #-} -- | 'Prelude.Backprop.liftA3', but taking explicit 'add' and 'zero'. -liftA3 - :: ( Traversable f - , Applicative f - , Reifies s W - ) - => AddFunc a - -> AddFunc b - -> AddFunc c - -> AddFunc d - -> ZeroFunc a - -> ZeroFunc b - -> ZeroFunc c - -> ZeroFunc d - -> (BVar s a -> BVar s b -> BVar s c -> BVar s d) - -> BVar s (f a) - -> BVar s (f b) - -> BVar s (f c) - -> BVar s (f d) -liftA3 afa afb afc afd zfa zfb zfc zfd f x y z - = collectVar afd zfd - $ f P.<$> sequenceVar afa zfa x - P.<*> sequenceVar afb zfb y - P.<*> sequenceVar afc zfc z +liftA3 :: + ( Traversable f + , Applicative f + , Reifies s W + ) => + AddFunc a -> + AddFunc b -> + AddFunc c -> + AddFunc d -> + ZeroFunc a -> + ZeroFunc b -> + ZeroFunc c -> + ZeroFunc d -> + (BVar s a -> BVar s b -> BVar s c -> BVar s d) -> + BVar s (f a) -> + BVar s (f b) -> + BVar s (f c) -> + BVar s (f d) +liftA3 afa afb afc afd zfa zfb zfc zfd f x y z = + collectVar afd zfd $ + f + P.<$> sequenceVar afa zfa x + P.<*> sequenceVar afb zfb y + P.<*> sequenceVar afc zfc z {-# INLINE liftA3 #-} -- | Coerce items inside a 'BVar'. @@ -276,33 +295,33 @@ -- | 'Prelude.Backprop.fromIntegral', but taking explicit 'add' and 'zero'. -- -- @since 0.2.1.0 -fromIntegral - :: (P.Integral a, P.Integral b, Reifies s W) - => AddFunc a - -> BVar s a - -> BVar s b +fromIntegral :: + (P.Integral a, P.Integral b, Reifies s W) => + AddFunc a -> + BVar s a -> + BVar s b fromIntegral af = isoVar af P.fromIntegral P.fromIntegral {-# INLINE fromIntegral #-} -- | 'Prelude.Backprop.realToFrac', but taking explicit 'add' and 'zero'. -- -- @since 0.2.1.0 -realToFrac - :: (Fractional a, P.Real a, Fractional b, P.Real b, Reifies s W) - => AddFunc a - -> BVar s a - -> BVar s b +realToFrac :: + (Fractional a, P.Real a, Fractional b, P.Real b, Reifies s W) => + AddFunc a -> + BVar s a -> + BVar s b realToFrac af = isoVar af P.realToFrac P.realToFrac {-# INLINE realToFrac #-} -- | 'Prelude.Backprop.round', but taking explicit 'add' and 'zero'. -- -- @since 0.2.3.0 -round - :: (P.RealFrac a, P.Integral b, Reifies s W) - => AddFunc a - -> BVar s a - -> BVar s b +round :: + (P.RealFrac a, P.Integral b, Reifies s W) => + AddFunc a -> + BVar s a -> + BVar s b round af = isoVar af P.round P.fromIntegral {-# INLINE round #-} @@ -310,60 +329,60 @@ -- 'zero'. -- -- @since 0.2.3.0 -fromIntegral' - :: (P.Integral a, P.RealFrac b, Reifies s W) - => AddFunc a - -> BVar s a - -> BVar s b +fromIntegral' :: + (P.Integral a, P.RealFrac b, Reifies s W) => + AddFunc a -> + BVar s a -> + BVar s b fromIntegral' af = isoVar af P.fromIntegral P.round {-# INLINE fromIntegral' #-} -- | 'Prelude.Backprop.length', but taking explicit 'add' and 'zero'. -- -- @since 0.2.2.0 -toList - :: (Traversable t, Reifies s W) - => AddFunc a - -> ZeroFunc a - -> BVar s (t a) - -> [BVar s a] +toList :: + (Traversable t, Reifies s W) => + AddFunc a -> + ZeroFunc a -> + BVar s (t a) -> + [BVar s a] toList af z = toListOfVar af (ZF (P.fmap (runZF z))) P.traverse {-# INLINE toList #-} -- | 'Prelude.Backprop.mapAccumL', but taking explicit 'add' and 'zero'. -- -- @since 0.2.2.0 -mapAccumL - :: (Traversable t, Reifies s W) - => AddFunc b - -> AddFunc c - -> ZeroFunc b - -> ZeroFunc c - -> (BVar s a -> BVar s b -> (BVar s a, BVar s c)) - -> BVar s a - -> BVar s (t b) - -> (BVar s a, BVar s (t c)) +mapAccumL :: + (Traversable t, Reifies s W) => + AddFunc b -> + AddFunc c -> + ZeroFunc b -> + ZeroFunc c -> + (BVar s a -> BVar s b -> (BVar s a, BVar s c)) -> + BVar s a -> + BVar s (t b) -> + (BVar s a, BVar s (t c)) mapAccumL afb afc zfb zfc f s = - second (collectVar afc zfc) - . P.mapAccumL f s - . sequenceVar afb zfb + second (collectVar afc zfc) + . P.mapAccumL f s + . sequenceVar afb zfb {-# INLINE mapAccumL #-} -- | 'Prelude.Backprop.mapAccumR', but taking explicit 'add' and 'zero'. -- -- @since 0.2.2.0 -mapAccumR - :: (Traversable t, Reifies s W) - => AddFunc b - -> AddFunc c - -> ZeroFunc b - -> ZeroFunc c - -> (BVar s a -> BVar s b -> (BVar s a, BVar s c)) - -> BVar s a - -> BVar s (t b) - -> (BVar s a, BVar s (t c)) +mapAccumR :: + (Traversable t, Reifies s W) => + AddFunc b -> + AddFunc c -> + ZeroFunc b -> + ZeroFunc c -> + (BVar s a -> BVar s b -> (BVar s a, BVar s c)) -> + BVar s a -> + BVar s (t b) -> + (BVar s a, BVar s (t c)) mapAccumR afb afc zfb zfc f s = - second (collectVar afc zfc) - . P.mapAccumR f s - . sequenceVar afb zfb + second (collectVar afc zfc) + . P.mapAccumR f s + . sequenceVar afb zfb {-# INLINE mapAccumR #-} diff --git a/src/Prelude/Backprop/Num.hs b/src/Prelude/Backprop/Num.hs --- a/src/Prelude/Backprop/Num.hs +++ b/src/Prelude/Backprop/Num.hs @@ -1,5 +1,5 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# OPTIONS_HADDOCK not-home #-} +{-# LANGUAGE FlexibleContexts #-} +{-# OPTIONS_HADDOCK not-home #-} -- | -- Module : Prelude.Backprop.Num @@ -14,90 +14,105 @@ -- 'Num' instances for all types involved instead of 'Backprop' instances. -- -- @since 0.2.0.0 - module Prelude.Backprop.Num ( -- * Foldable and Traversable - sum - , product - , length - , minimum - , maximum - , traverse - , toList - , mapAccumL - , mapAccumR - , foldr, foldl' + sum, + product, + length, + minimum, + maximum, + traverse, + toList, + mapAccumL, + mapAccumR, + foldr, + foldl', + -- * Functor and Applicative - , fmap, fmapConst - , (<$>), (<$), ($>) - , pure - , liftA2 - , liftA3 + fmap, + fmapConst, + (<$>), + (<$), + ($>), + pure, + liftA2, + liftA3, + -- * Numeric - , fromIntegral - , realToFrac - , round - , fromIntegral' + fromIntegral, + realToFrac, + round, + fromIntegral', + -- * Misc - , E.coerce - ) where + E.coerce, +) where -import Numeric.Backprop.Num -import Prelude (Num(..), Fractional(..), Ord(..), Functor, Foldable, Traversable, Applicative) import qualified Numeric.Backprop.Explicit as E -import qualified Prelude as P +import Numeric.Backprop.Num import qualified Prelude.Backprop.Explicit as E +import Prelude ( + Applicative, + Foldable, + Fractional (..), + Functor, + Num (..), + Ord (..), + Traversable, + ) +import qualified Prelude as P -- | 'Prelude.Backprop.sum', but with 'Num' constraints instead of -- 'Backprop' constraints. -sum :: (Foldable t, Functor t, Num (t a), Num a, Reifies s W) - => BVar s (t a) - -> BVar s a +sum :: + (Foldable t, Functor t, Num (t a), Num a, Reifies s W) => + BVar s (t a) -> + BVar s a sum = E.sum E.afNum {-# INLINE sum #-} -- | 'Prelude.Backprop.pure', but with 'Num' constraints instead of -- 'Backprop' constraints. -pure - :: (Foldable t, Applicative t, Num a, Reifies s W) - => BVar s a - -> BVar s (t a) +pure :: + (Foldable t, Applicative t, Num a, Reifies s W) => + BVar s a -> + BVar s (t a) pure = E.pure E.afNum E.zfNum {-# INLINE pure #-} -- | 'Prelude.Backprop.product', but with 'Num' constraints instead of -- 'Backprop' constraints. -product - :: (Foldable t, Functor t, Num (t a), Fractional a, Reifies s W) - => BVar s (t a) - -> BVar s a +product :: + (Foldable t, Functor t, Num (t a), Fractional a, Reifies s W) => + BVar s (t a) -> + BVar s a product = E.product E.afNum {-# INLINE product #-} -- | 'Prelude.Backprop.length', but with 'Num' constraints instead of -- 'Backprop' constraints. -length - :: (Foldable t, Num (t a), Num b, Reifies s W) - => BVar s (t a) - -> BVar s b +length :: + (Foldable t, Num (t a), Num b, Reifies s W) => + BVar s (t a) -> + BVar s b length = E.length E.afNum E.zfNum {-# INLINE length #-} -- | 'Prelude.Backprop.minimum', but with 'Num' constraints instead of -- 'Backprop' constraints. -minimum - :: (Foldable t, Functor t, Num a, Ord a, Num (t a), Reifies s W) - => BVar s (t a) - -> BVar s a +minimum :: + (Foldable t, Functor t, Num a, Ord a, Num (t a), Reifies s W) => + BVar s (t a) -> + BVar s a minimum = E.minimum E.afNum E.zfNum {-# INLINE minimum #-} -- | 'Prelude.Backprop.maximum', but with 'Num' constraints instead of -- 'Backprop' constraints. -maximum - :: (Foldable t, Functor t, Num a, Ord a, Num (t a), Reifies s W) - => BVar s (t a) - -> BVar s a +maximum :: + (Foldable t, Functor t, Num a, Ord a, Num (t a), Reifies s W) => + BVar s (t a) -> + BVar s a maximum = E.maximum E.afNum E.zfNum {-# INLINE maximum #-} @@ -105,12 +120,12 @@ -- 'Backprop' constraints. -- -- @since 0.2.3.0 -foldr - :: (Traversable t, Num a, Reifies s W) - => (BVar s a -> BVar s b -> BVar s b) - -> BVar s b - -> BVar s (t a) - -> BVar s b +foldr :: + (Traversable t, Num a, Reifies s W) => + (BVar s a -> BVar s b -> BVar s b) -> + BVar s b -> + BVar s (t a) -> + BVar s b foldr = E.foldr E.afNum E.zfNum {-# INLINE foldr #-} @@ -118,22 +133,22 @@ -- 'Backprop' constraints. -- -- @since 0.2.3.0 -foldl' - :: (Traversable t, Num a, Reifies s W) - => (BVar s b -> BVar s a -> BVar s b) - -> BVar s b - -> BVar s (t a) - -> BVar s b +foldl' :: + (Traversable t, Num a, Reifies s W) => + (BVar s b -> BVar s a -> BVar s b) -> + BVar s b -> + BVar s (t a) -> + BVar s b foldl' = E.foldl' E.afNum E.zfNum {-# INLINE foldl' #-} -- | 'Prelude.Backprop.fmap', but with 'Num' constraints instead of -- 'Backprop' constraints. -fmap - :: (Traversable f, Num a, Num b, Reifies s W) - => (BVar s a -> BVar s b) - -> BVar s (f a) - -> BVar s (f b) +fmap :: + (Traversable f, Num a, Num b, Reifies s W) => + (BVar s a -> BVar s b) -> + BVar s (f a) -> + BVar s (f b) fmap = E.fmap E.afNum E.afNum E.zfNum E.zfNum {-# INLINE fmap #-} @@ -141,45 +156,48 @@ -- 'Backprop' constraints. -- -- @since 0.2.4.0 -fmapConst - :: (Functor f, Foldable f, Num b, Num (f a), Reifies s W) - => BVar s b - -> BVar s (f a) - -> BVar s (f b) +fmapConst :: + (Functor f, Foldable f, Num b, Num (f a), Reifies s W) => + BVar s b -> + BVar s (f a) -> + BVar s (f b) fmapConst = E.fmapConst E.afNum E.afNum E.zfNum E.zfNum {-# INLINE fmapConst #-} -- | Alias for 'fmap'. -(<$>) - :: (Traversable f, Num a, Num b, Reifies s W) - => (BVar s a -> BVar s b) - -> BVar s (f a) - -> BVar s (f b) +(<$>) :: + (Traversable f, Num a, Num b, Reifies s W) => + (BVar s a -> BVar s b) -> + BVar s (f a) -> + BVar s (f b) (<$>) = fmap + infixl 4 <$> {-# INLINE (<$>) #-} -- | Alias for 'fmapConst'. -- -- @since 0.2.4.0 -(<$) - :: (Functor f, Foldable f, Num b, Num (f a), Reifies s W) - => BVar s b - -> BVar s (f a) - -> BVar s (f b) +(<$) :: + (Functor f, Foldable f, Num b, Num (f a), Reifies s W) => + BVar s b -> + BVar s (f a) -> + BVar s (f b) (<$) = fmapConst + infixl 4 <$ {-# INLINE (<$) #-} -- | Alias for @'flip' 'fmapConst'@. -- -- @since 0.2.4.0 -($>) - :: (Functor f, Foldable f, Num b, Num (f a), Reifies s W) - => BVar s (f a) - -> BVar s b - -> BVar s (f b) +($>) :: + (Functor f, Foldable f, Num b, Num (f a), Reifies s W) => + BVar s (f a) -> + BVar s b -> + BVar s (f b) xs $> x = x <$ xs + infixl 4 $> {-# INLINE ($>) #-} @@ -188,54 +206,67 @@ -- -- See for -- a fixed-length vector type with a very appropriate 'Num' instance! -traverse - :: (Traversable t, Applicative f, Foldable f, Num a, Num b, Num (t b), Reifies s W) - => (BVar s a -> f (BVar s b)) - -> BVar s (t a) - -> BVar s (f (t b)) +traverse :: + (Traversable t, Applicative f, Foldable f, Num a, Num b, Num (t b), Reifies s W) => + (BVar s a -> f (BVar s b)) -> + BVar s (t a) -> + BVar s (f (t b)) traverse = E.traverse E.afNum E.afNum E.afNum E.zfNum E.zfNum {-# INLINE traverse #-} -- | 'Prelude.Backprop.liftA2', but with 'Num' constraints instead of -- 'Backprop' constraints. -liftA2 - :: ( Traversable f - , Applicative f - , Num a, Num b, Num c - , Reifies s W - ) - => (BVar s a -> BVar s b -> BVar s c) - -> BVar s (f a) - -> BVar s (f b) - -> BVar s (f c) +liftA2 :: + ( Traversable f + , Applicative f + , Num a + , Num b + , Num c + , Reifies s W + ) => + (BVar s a -> BVar s b -> BVar s c) -> + BVar s (f a) -> + BVar s (f b) -> + BVar s (f c) liftA2 = E.liftA2 E.afNum E.afNum E.afNum E.zfNum E.zfNum E.zfNum {-# INLINE liftA2 #-} -- | 'Prelude.Backprop.liftA3', but with 'Num' constraints instead of -- 'Backprop' constraints. -liftA3 - :: ( Traversable f - , Applicative f - , Num a, Num b, Num c, Num d - , Reifies s W - ) - => (BVar s a -> BVar s b -> BVar s c -> BVar s d) - -> BVar s (f a) - -> BVar s (f b) - -> BVar s (f c) - -> BVar s (f d) -liftA3 = E.liftA3 E.afNum E.afNum E.afNum E.afNum - E.zfNum E.zfNum E.zfNum E.zfNum +liftA3 :: + ( Traversable f + , Applicative f + , Num a + , Num b + , Num c + , Num d + , Reifies s W + ) => + (BVar s a -> BVar s b -> BVar s c -> BVar s d) -> + BVar s (f a) -> + BVar s (f b) -> + BVar s (f c) -> + BVar s (f d) +liftA3 = + E.liftA3 + E.afNum + E.afNum + E.afNum + E.afNum + E.zfNum + E.zfNum + E.zfNum + E.zfNum {-# INLINE liftA3 #-} -- | 'Prelude.Backprop.fromIntegral', but with 'Num' constraints instead of -- 'Backprop' constraints. -- -- @since 0.2.1.0 -fromIntegral - :: (P.Integral a, P.Integral b, Reifies s W) - => BVar s a - -> BVar s b +fromIntegral :: + (P.Integral a, P.Integral b, Reifies s W) => + BVar s a -> + BVar s b fromIntegral = E.fromIntegral E.afNum {-# INLINE fromIntegral #-} @@ -243,10 +274,10 @@ -- 'Backprop' constraints. -- -- @since 0.2.1.0 -realToFrac - :: (Fractional a, P.Real a, Fractional b, P.Real b, Reifies s W) - => BVar s a - -> BVar s b +realToFrac :: + (Fractional a, P.Real a, Fractional b, P.Real b, Reifies s W) => + BVar s a -> + BVar s b realToFrac = E.realToFrac E.afNum {-# INLINE realToFrac #-} @@ -254,10 +285,10 @@ -- 'Backprop' constraints. -- -- @since 0.2.3.0 -round - :: (P.RealFrac a, P.Integral b, Reifies s W) - => BVar s a - -> BVar s b +round :: + (P.RealFrac a, P.Integral b, Reifies s W) => + BVar s a -> + BVar s b round = E.round E.afNum {-# INLINE round #-} @@ -265,10 +296,10 @@ -- of 'Backprop' constraints. -- -- @since 0.2.3.0 -fromIntegral' - :: (P.Integral a, P.RealFrac b, Reifies s W) - => BVar s a - -> BVar s b +fromIntegral' :: + (P.Integral a, P.RealFrac b, Reifies s W) => + BVar s a -> + BVar s b fromIntegral' = E.fromIntegral' E.afNum {-# INLINE fromIntegral' #-} @@ -276,10 +307,10 @@ -- 'Backprop' constraints. -- -- @since 0.2.2.0 -toList - :: (Traversable t, Num a, Reifies s W) - => BVar s (t a) - -> [BVar s a] +toList :: + (Traversable t, Num a, Reifies s W) => + BVar s (t a) -> + [BVar s a] toList = E.toList E.afNum E.zfNum {-# INLINE toList #-} @@ -289,12 +320,12 @@ -- Prior to v0.2.3, required a 'Num' constraint on @t b@. -- -- @since 0.2.2.0 -mapAccumL - :: (Traversable t, Num b, Num c, Reifies s W) - => (BVar s a -> BVar s b -> (BVar s a, BVar s c)) - -> BVar s a - -> BVar s (t b) - -> (BVar s a, BVar s (t c)) +mapAccumL :: + (Traversable t, Num b, Num c, Reifies s W) => + (BVar s a -> BVar s b -> (BVar s a, BVar s c)) -> + BVar s a -> + BVar s (t b) -> + (BVar s a, BVar s (t c)) mapAccumL = E.mapAccumL E.afNum E.afNum E.zfNum E.zfNum {-# INLINE mapAccumL #-} @@ -304,12 +335,11 @@ -- Prior to v0.2.3, required a 'Num' constraint on @t b@. -- -- @since 0.2.2.0 -mapAccumR - :: (Traversable t, Num b, Num c, Reifies s W) - => (BVar s a -> BVar s b -> (BVar s a, BVar s c)) - -> BVar s a - -> BVar s (t b) - -> (BVar s a, BVar s (t c)) +mapAccumR :: + (Traversable t, Num b, Num c, Reifies s W) => + (BVar s a -> BVar s b -> (BVar s a, BVar s c)) -> + BVar s a -> + BVar s (t b) -> + (BVar s a, BVar s (t c)) mapAccumR = E.mapAccumR E.afNum E.afNum E.zfNum E.zfNum {-# INLINE mapAccumR #-} -