Add some simple fold tests

1 files changed, 13 insertions, 0 deletions

diff --git a/test/Main.hs b/test/Main.hs
index 9c992ec..cb10ed6 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -669,6 +669,19 @@ tests_AD = testGroup "AD"
   ,adTestGen "gmm" (Example.gmmObjective False) gen_gmm
 
   ,adTestTp "uniform-free" (C "" 0 :& ()) Example.exUniformFree
+
+  ,adTest "fold-sum" $ fromNamed $ lambda @(TArr N1 R) #a $ body $
+    idx0 $ fold1i (#x :-> #y :-> #x + #y) 0 #a
+
+  ,adTest "fold-prod" $ fromNamed $ lambda @(TArr N1 R) #a $ body $
+    idx0 $ fold1i (#x :-> #y :-> #x * #y) 1 #a
+
+  ,adTest "fold-freevar" $ fromNamed $ lambda @(TArr N1 R) #a $ body $
+    let_ #v 2 $
+    idx0 $ fold1i (#x :-> #y :-> #x * #y + #v) 1 #a
+
+  ,adTestTp "fold-freearr" (C "" 1) $ fromNamed $ lambda @(TArr N1 R) #a $ body $
+    idx0 $ fold1i (#x :-> #y :-> #x * #y + #a ! pair nil 0) 1 #a
   ]
 
 main :: IO ()