simplifyFix

1 files changed, 23 insertions, 13 deletions

diff --git a/test/Main.hs b/test/Main.hs
index a3a614a..e7dda69 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -50,21 +50,26 @@ primalEnv :: SList STy env' -> SList STy (D1E env')
 primalEnv SNil = SNil
 primalEnv (t `SCons` env) = d1 t `SCons` primalEnv env
 
-diffCHAD :: Int -> SList STy env -> Ex env (TScal TF64)
+data SimplIters = SimplIters Int | SimplFix
+  deriving (Show)
+
+diffCHAD :: SimplIters -> SList STy env -> Ex env (TScal TF64)
          -> Ex (D1E env) (TPair (TScal TF64) (Tup (D2E env)))
 diffCHAD = \simplIters env term ->
-  case (mapMergeNoAccum env, mapMergeOnlyMerge env) of
-    (Refl, Refl) ->
+  case (mapMergeNoAccum env, mapMergeOnlyMerge env, envKnown (primalEnv env)) of
+    (Refl, Refl, Dict) ->
       let descr = makeMergeDescr env
-      in case envKnown (primalEnv env) of
-           Dict -> simplifyN simplIters $ freezeRet descr (drev descr term) (EConst ext STF64 1.0)
+          simpl = case simplIters of
+                    SimplIters n -> simplifyN n
+                    SimplFix -> simplifyFix
+      in simpl $ freezeRet descr (drev descr term) (EConst ext STF64 1.0)
   where
     makeMergeDescr :: SList STy env' -> Descr env' (MapMerge env')
     makeMergeDescr SNil = DTop
     makeMergeDescr (t `SCons` env) = makeMergeDescr env `DPush` (t, SMerge)
 
 -- In addition to the gradient, also returns the pretty-printed differentiated term.
-gradientByCHAD :: forall env. Int -> SList STy env -> Ex env (TScal TF64) -> SList Value env -> (String, (Double, SList Value (D2E env)))
+gradientByCHAD :: forall env. SimplIters -> SList STy env -> Ex env (TScal TF64) -> SList Value env -> (String, (Double, SList Value (D2E env)))
 gradientByCHAD = \simplIters env term input ->
   case (mapMergeNoAccum env, mapMergeOnlyMerge env) of
     (Refl, Refl) ->
@@ -88,7 +93,7 @@ gradientByCHAD = \simplIters env term input ->
       STAccum{} -> error "Accumulators not allowed in input program"
 
 -- In addition to the gradient, also returns the pretty-printed differentiated term.
-gradientByCHAD' :: Int -> SList STy env -> Ex env (TScal TF64) -> SList Value env -> (String, (Double, SList Value (TanE env)))
+gradientByCHAD' :: SimplIters -> SList STy env -> Ex env (TScal TF64) -> SList Value env -> (String, (Double, SList Value (TanE env)))
 gradientByCHAD' = \simplIters env term input -> second (second (toTanE env input)) $ gradientByCHAD simplIters env term input
   where
     toTanE :: SList STy env -> SList Value env -> SList Value (D2E env) -> SList Value (TanE env)
@@ -212,8 +217,9 @@ adTestGen expr envGenerator = property $ do
   input <- forAllWith (showEnv env) envGenerator
   let outPrimal = interpretOpen False input expr
       gradFwd = gradientByForward knownEnv expr input
-      (ppdterm, (outChad, gradCHAD)) = gradientByCHAD' 0 knownEnv expr input
-      (ppdterm_S, (outChad_S, gradCHAD_S)) = gradientByCHAD' 20 knownEnv expr input
+      (ppdterm, (outChad, gradCHAD)) = gradientByCHAD' (SimplIters 0) knownEnv expr input
+      (ppdterm_S, (outChad_S, gradCHAD_S)) = gradientByCHAD' SimplFix knownEnv expr input
+      (ppdterm_S20, _) = gradientByCHAD' (SimplIters 20) knownEnv expr input
       scFwd = envScalars env gradFwd
       scCHAD = envScalars env gradCHAD
       scCHAD_S = envScalars env gradCHAD_S
@@ -221,6 +227,7 @@ adTestGen expr envGenerator = property $ do
   annotate (ppExpr knownEnv expr)
   annotate ppdterm
   annotate ppdterm_S
+  diff ppdterm_S20 (==) ppdterm_S
   diff outChad closeIsh outChad_S
   diff outPrimal closeIsh outChad_S
   diff scCHAD (\x y -> and (zipWith closeIsh x y)) scCHAD_S
@@ -235,6 +242,12 @@ term_build1_sum = fromNamed $ lambda #x $ body $
   idx0 $ sum1i $
     build (SS SZ) (shape #x) $ #idx :-> #x ! #idx
 
+term_pairs :: Ex [TScal TF64, TScal TF64] (TScal TF64)
+term_pairs = fromNamed $ lambda #x $ lambda #y $ body $
+  let_ #p (pair #x #y) $
+  let_ #q (pair (snd_ #p * fst_ #p + #y) #x) $
+    fst_ #q * #x + snd_ #q * fst_ #p
+
 tests :: IO Bool
 tests = checkSequential $ Group "AD"
   [("id", adTest $ fromNamed $ lambda #x $ body $ #x)
@@ -246,10 +259,7 @@ tests = checkSequential $ Group "AD"
   ,("sum-replicate", adTest $ fromNamed $ lambda #x $ body $
       idx0 $ sum1i $ replicate1i 10 #x)
 
-  ,("pairs", adTest $ fromNamed $ lambda #x $ lambda #y $ body $
-      let_ #p (pair #x #y) $
-      let_ #q (pair (snd_ #p * fst_ #p + #y) #x) $
-        fst_ #q * #x + snd_ #q * fst_ #p)
+  ,("pairs", adTest term_pairs)
 
   ,("build0 const", adTest $ fromNamed $ lambda @(TScal TF64) #x $ body $
       idx0 $ build SZ nil $ #idx :-> const_ 0.0)