Try debugging crash

5 files changed, 113 insertions, 60 deletions

diff --git a/ad-dual.cabal b/ad-dual.cabal
index 5d3ca39..fe14d31 100644
--- a/ad-dual.cabal
+++ b/ad-dual.cabal
@@ -5,41 +5,56 @@ license:         BSD-3-Clause
 author:          Tom Smeding
 build-type:      Simple
 
+common common
+  build-depends: base >= 4.14.3
+  default-language: Haskell2010
+  ghc-options: -Wall
+
 library
+  import: common
   exposed-modules:
     Numeric.ADDual
     Numeric.ADDual.Internal
   other-modules:
   build-depends:
-    base >= 4.14.3,
     transformers,
     vector
   hs-source-dirs: src
-  default-language: Haskell2010
-  ghc-options: -Wall
+
+library ad-dual-examples
+  import: common
+  exposed-modules:
+    Numeric.ADDual.Examples
+  build-depends:
+    deepseq,
+    hedgehog,
+    vector
+  hs-source-dirs: examples
 
 test-suite test
+  import: common
   type: exitcode-stdio-1.0
   main-is: Main.hs
   build-depends:
-    base,
     ad-dual,
+    ad-dual-examples,
+    ad,
+    hedgehog,
     tasty,
-    tasty-hunit
+    tasty-hedgehog,
+    tasty-hunit,
+    vector
   hs-source-dirs: test
-  default-language: Haskell2010
-  ghc-options: -Wall
 
 benchmark bench
+  import: common
   type: exitcode-stdio-1.0
   main-is: Main.hs
   build-depends:
-    base,
     ad-dual,
+    ad-dual-examples,
     ad,
     criterion,
     deepseq,
     vector
   hs-source-dirs: bench
-  default-language: Haskell2010
-  ghc-options: -Wall
diff --git a/bench/Main.hs b/bench/Main.hs
index cb5e829..a11f4e8 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -1,56 +1,13 @@
-{-# LANGUAGE DeriveTraversable #-}
 {-# LANGUAGE TypeApplications #-}
-{-# LANGUAGE DeriveGeneric #-}
 module Main where
 
-import Control.DeepSeq
 import Criterion
 import Criterion.Main
-import qualified Data.Vector as V
-import GHC.Generics (Generic)
 
 import qualified Numeric.ADDual as ADD
+import Numeric.ADDual.Examples
 
 
-type Matrix s = V.Vector s
-
-data FNeural a = FNeural [(Matrix a, V.Vector a)] (V.Vector a)
-  deriving (Show, Functor, Foldable, Traversable, Generic)
-
-instance NFData a => NFData (FNeural a)
-
-fneural :: (Floating a, Ord a) => FNeural a -> a
-fneural (FNeural layers input) =
-  let dotp v1 v2 = V.sum (V.zipWith (*) v1 v2)
-
-      mat @. vec =
-        let n = V.length vec
-            m = V.length mat `div` n
-        in V.fromListN m $ map (\i -> dotp (V.slice (n*i) n mat) vec) [0 .. m-1]
-      (+.) = V.zipWith (+)
-
-      relu x = if x >= 0.0 then x else 0.0
-      safeSoftmax vec = let m = V.maximum vec
-                            factor = V.sum (V.map (\z -> exp (z - m)) vec)
-                        in V.map (\z -> exp (z - m) / factor) vec
-      forward [] x = safeSoftmax x
-      forward ((weights, bias) : lys) x =
-        let x' = V.map relu ((weights @. x) +. bias)
-        in forward lys x'
-  in V.sum $ forward layers input
-
-makeNeuralInput :: FNeural Double
-makeNeuralInput =
-  let genMatrix nin nout =
-        V.fromListN (nin*nout) [sin (fromIntegral @Int (i+j))
-                               | i <- [0..nout-1], j <- [0..nin-1]]
-      genVector nout = V.fromListN nout [sin (0.41 * fromIntegral @Int i) | i <- [0..nout-1]]
-      -- 50 inputs; 2 hidden layers (100; 50); final softmax, then sum the outputs.
-      nIn = 50; n1 = 100; n2 = 50
-  in FNeural [(genMatrix nIn n1, genVector n1)
-             ,(genMatrix n1 n2, genVector n2)]
-             (genVector nIn)
-
 main :: IO ()
 main = defaultMain
   [env (pure makeNeuralInput) $ \input ->
diff --git a/examples/Numeric/ADDual/Examples.hs b/examples/Numeric/ADDual/Examples.hs
new file mode 100644
index 0000000..d6aa6d2
--- /dev/null
+++ b/examples/Numeric/ADDual/Examples.hs
@@ -0,0 +1,65 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DeriveTraversable #-}
+{-# LANGUAGE TypeApplications #-}
+module Numeric.ADDual.Examples where
+
+import Control.DeepSeq
+import Control.Monad (replicateM)
+import qualified Data.Vector as V
+import GHC.Generics (Generic)
+import Hedgehog (Gen)
+import qualified Hedgehog.Gen as Gen
+import qualified Hedgehog.Range as Range
+
+
+type Matrix s = V.Vector s
+
+data FNeural a = FNeural [(Matrix a, V.Vector a)] (V.Vector a)
+  deriving (Show, Functor, Foldable, Traversable, Generic)
+
+instance NFData a => NFData (FNeural a)
+
+fneural :: (Floating a, Ord a) => FNeural a -> a
+fneural (FNeural layers input) =
+  let dotp v1 v2 = V.sum (V.zipWith (*) v1 v2)
+
+      mat @. vec =
+        let n = V.length vec
+            m = V.length mat `div` n
+        in V.fromListN m $ map (\i -> dotp (V.slice (n*i) n mat) vec) [0 .. m-1]
+      (+.) = V.zipWith (+)
+
+      relu x = if x >= 0.0 then x else 0.0
+      safeSoftmax vec = let m = V.maximum vec
+                            factor = V.sum (V.map (\z -> exp (z - m)) vec)
+                        in V.map (\z -> exp (z - m) / factor) vec
+      forward [] x = safeSoftmax x
+      forward ((weights, bias) : lys) x =
+        let x' = V.map relu ((weights @. x) +. bias)
+        in forward lys x'
+  in V.sum $ forward layers input
+
+makeNeuralInput :: FNeural Double
+makeNeuralInput =
+  let genMatrix nin nout =
+        V.fromListN (nin*nout) [sin (fromIntegral @Int (i+j))
+                               | i <- [0..nout-1], j <- [0..nin-1]]
+      genVector nout = V.fromListN nout [sin (0.41 * fromIntegral @Int i) | i <- [0..nout-1]]
+      -- 50 inputs; 2 hidden layers (100; 50); final softmax, then sum the outputs.
+      nIn = 50; n1 = 100; n2 = 50
+  in FNeural [(genMatrix nIn n1, genVector n1)
+             ,(genMatrix n1 n2, genVector n2)]
+             (genVector nIn)
+
+genNeuralInput :: Gen (FNeural Double)
+genNeuralInput = do
+  let genScalar = Gen.double (Range.linearFracFrom 0 (-1) 1)
+      genMatrix nin nout = V.fromListN (nin*nout) <$> replicateM (nin*nout) genScalar
+      genVector nout = V.fromListN nout <$> replicateM nout genScalar
+  nIn <- Gen.integral (Range.linear 1 80)
+  n1 <- Gen.integral (Range.linear 1 100)
+  n2 <- Gen.integral (Range.linear 1 80)
+  m1 <- genMatrix nIn n1; v1 <- genVector n1
+  m2 <- genMatrix n1 n2; v2 <- genVector n2
+  inp <- genVector nIn
+  pure $ FNeural [(m1, v1), (m2, v2)] inp
diff --git a/src/Numeric/ADDual/Internal.hs b/src/Numeric/ADDual/Internal.hs
index 8228694..d8e6e0a 100644
--- a/src/Numeric/ADDual/Internal.hs
+++ b/src/Numeric/ADDual/Internal.hs
@@ -15,18 +15,18 @@ import qualified Data.Vector.Storable as VS
 import qualified Data.Vector.Storable.Mutable as VSM
 import Foreign.Ptr
 import Foreign.Storable
+import GHC.Stack
 import GHC.Exts (withDict)
 import System.IO.Unsafe
 
--- import System.IO (hPutStrLn, stderr)
-
--- import Numeric.ADDual.IDGen
+import System.IO (hPutStrLn, stderr)
 
 
 -- TODO: full vjp (just some more Traversable mess)
 {-# NOINLINE gradient' #-}
 gradient' :: forall a f. (Traversable f, Num a, Storable a)
-          -- => Show a  -- TODO: remove
+          => HasCallStack
+          => Show a  -- TODO: remove
           => (forall s. Taping s a => f (Dual s a) -> Dual s a)
           -> f a -> a -> (a, f a)
 gradient' f inp topctg = unsafePerformIO $ do
@@ -50,6 +50,7 @@ gradient' f inp topctg = unsafePerformIO $ do
 
   let backpropagate i chunk@(Chunk ci0 vec) tape
         | i >= ci0 = do
+            -- hPutStrLn stderr $ "read at ci0=" ++ show ci0 ++ " i=" ++ show i
             ctg <- VSM.read accums i
             Contrib i1 dx i2 dy <- VSM.read vec (i - ci0)
             when (i1 /= -1) $ VSM.modify accums (+ ctg*dx) i1
@@ -174,7 +175,9 @@ data WriteTapeAction a = WTANewvec (VSM.IOVector (Contrib a))
 writeTape :: forall a s proxy. (Num a, Storable a, Taping s a) => proxy s -> Int -> a -> Int -> a -> Int
 writeTape _ i1 dx i2 dy = unsafePerformIO $ writeTapeIO (Proxy @s) i1 dx i2 dy
 
-writeTapeIO :: forall a s proxy. (Num a, Storable a, Taping s a) => proxy s -> Int -> a -> Int -> a -> IO Int
+writeTapeIO :: forall a s proxy. (Num a, Storable a, Taping s a)
+            => HasCallStack
+            => proxy s -> Int -> a -> Int -> a -> IO Int
 writeTapeIO _ i1 dx i2 dy = do
   MLog idref (Chunk ci0 vec) _ <- readIORef (getTape @s)
   let n = VSM.length vec
diff --git a/test/Main.hs b/test/Main.hs
index aade6a5..86b6011 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -1,13 +1,26 @@
 {-# LANGUAGE TypeApplications #-}
 module Main where
 
+import qualified Data.Vector as V
+import Hedgehog
 import Test.Tasty
+import Test.Tasty.Hedgehog
 import Test.Tasty.HUnit
 
 import Numeric.ADDual
+import Numeric.ADDual.Examples
 
 
 main :: IO ()
 main = defaultMain $ testGroup "Tests"
   [testCase "product [1..5]" $
-      gradient' @Double product [1..5] 1 @?= (120, [120, 60, 40, 30, 24])]
+      gradient' @Double product [1..5] 1 @?= (120, [120, 60, 40, 30, 24])
+  ,testCase "neural one" $
+      let problem = FNeural
+            [(V.replicate 6 0.0, V.replicate 6 0.0), (V.replicate 24 0.0, V.replicate 4 0.0)]
+            (V.replicate 1 0.0)
+      in fst (gradient' @Double fneural problem 1) @?= 0.0
+  ,testProperty "neural run" $ property $ do
+      input <- forAll genNeuralInput
+      let (res, _grad) = gradient' fneural input 1
+      res === fneural input]