summaryrefslogtreecommitdiff
path: root/test
diff options
context:
space:
mode:
Diffstat (limited to 'test')
-rw-r--r--test/Main.hs126
1 files changed, 89 insertions, 37 deletions
diff --git a/test/Main.hs b/test/Main.hs
index 1b83a2e..0a57cbf 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -305,7 +305,9 @@ adTestGen name expr envGenerator =
testGroupCollapse name
[adTestGenPrimal env envGenerator expr exprS primalfun primalSfun
,adTestGenFwd env envGenerator exprS
- ,adTestGenChad env envGenerator expr exprS primalSfun]
+ ,testGroup "chad"
+ [adTestGenChad "default" defaultConfig env envGenerator expr exprS primalSfun
+ ,adTestGenChad "accum" (chcSetAccum defaultConfig) env envGenerator expr exprS primalSfun]]
adTestGenPrimal :: SList STy env -> Gen (SList Value env)
-> Ex env R -> Ex env R
@@ -336,24 +338,28 @@ adTestGenFwd env envGenerator exprS =
diff outDNI1 (closeIsh' 1e-8) outDNC1
diff outDNI2 (closeIsh' 1e-8) outDNC2
-adTestGenChad :: forall env. SList STy env -> Gen (SList Value env)
+adTestGenChad :: forall env. String -> CHADConfig -> SList STy env -> Gen (SList Value env)
-> Ex env R -> Ex env R
-> (SList Value env -> IO Double)
-> TestTree
-adTestGenChad env envGenerator expr exprS primalSfun | Dict <- envKnown env =
- let dtermChad0 = ELet ext (EConst ext STF64 1.0) $ chad' defaultConfig env expr
+adTestGenChad testname config env envGenerator expr exprS primalSfun | Dict <- envKnown env =
+ let dtermChad0 = ELet ext (EConst ext STF64 1.0) $ chad' config env expr
dtermChadS = simplifyFix dtermChad0
- dtermSChad0 = ELet ext (EConst ext STF64 1.0) $ chad' defaultConfig env exprS
+ dtermChadSUS = simplifyFix $ unMonoid dtermChadS
+ dtermSChad0 = ELet ext (EConst ext STF64 1.0) $ chad' config env exprS
dtermSChadS = simplifyFix dtermSChad0
+ dtermSChadSUS = simplifyFix $ unMonoid dtermSChadS
in
withResource (makeFwdADArtifactCompile env exprS) (\_ -> pure ()) $ \fwdartifactC ->
- withCompiled env (simplifyFix (unMonoid dtermSChadS)) $ \dcompSChadS ->
- testProperty "chad" $ property $ do
+ withCompiled env dtermSChadSUS $ \dcompSChadSUS ->
+ testProperty testname $ property $ do
annotate (concat (unSList (\t -> ppSTy 0 t ++ " -> ") env) ++ ppSTy 0 (typeOf expr))
- -- pack Text for less GC pressure (these values are retained for some reason)
+ -- check simplifier convergence; pack Text for less GC pressure (these values are retained for some reason)
diff (T.pack (ppExpr env dtermChadS)) (==) (T.pack (ppExpr env (simplifyN 20 dtermChad0)))
+ diff (T.pack (ppExpr env dtermChadSUS)) (==) (T.pack (ppExpr env (simplifyN 20 $ unMonoid $ simplifyN 20 dtermChad0)))
diff (T.pack (ppExpr env dtermSChadS)) (==) (T.pack (ppExpr env (simplifyN 20 dtermSChad0)))
+ diff (T.pack (ppExpr env dtermSChadSUS)) (==) (T.pack (ppExpr env (simplifyN 20 $ unMonoid $ simplifyN 20 dtermSChad0)))
input <- forAllWith (showEnv env) envGenerator
outPrimal <- evalIO $ primalSfun input
@@ -363,17 +369,21 @@ adTestGenChad env envGenerator expr exprS primalSfun | Dict <- envKnown env =
let tansFwd = TypedEnv (tanenv env) $ gradientByForward fwdartifactC input
- let (outChad0 , gradChad0) = second unpackGrad $ interpretOpen False env input dtermChad0
- (outChadS , gradChadS) = second unpackGrad $ interpretOpen False env input dtermChadS
- (outSChad0, gradSChad0) = second unpackGrad $ interpretOpen False env input dtermSChad0
- (outSChadS, gradSChadS) = second unpackGrad $ interpretOpen False env input dtermSChadS
- tansChad = TypedEnv (tanenv env) $ toTanE env input gradChad0
- tansChadS = TypedEnv (tanenv env) $ toTanE env input gradChadS
- tansSChad = TypedEnv (tanenv env) $ toTanE env input gradSChad0
- tansSChadS = TypedEnv (tanenv env) $ toTanE env input gradSChadS
+ let (outChad0 , gradChad0) = second unpackGrad $ interpretOpen False env input dtermChad0
+ (outChadS , gradChadS) = second unpackGrad $ interpretOpen False env input dtermChadS
+ (outChadSUS , gradChadSUS) = second unpackGrad $ interpretOpen False env input dtermChadSUS
+ (outSChad0 , gradSChad0) = second unpackGrad $ interpretOpen False env input dtermSChad0
+ (outSChadS , gradSChadS) = second unpackGrad $ interpretOpen False env input dtermSChadS
+ (outSChadSUS, gradSChadSUS) = second unpackGrad $ interpretOpen False env input dtermSChadSUS
+ tansChad = TypedEnv (tanenv env) $ toTanE env input gradChad0
+ tansChadS = TypedEnv (tanenv env) $ toTanE env input gradChadS
+ tansChadSUS = TypedEnv (tanenv env) $ toTanE env input gradChadSUS
+ tansSChad = TypedEnv (tanenv env) $ toTanE env input gradSChad0
+ tansSChadS = TypedEnv (tanenv env) $ toTanE env input gradSChadS
+ tansSChadSUS = TypedEnv (tanenv env) $ toTanE env input gradSChadSUS
- (outCompSChadS, gradCompSChadS) <- second unpackGrad <$> evalIO (dcompSChadS input)
- let tansCompSChadS = TypedEnv (tanenv env) $ toTanE env input gradCompSChadS
+ (outCompSChadSUS, gradCompSChadSUS) <- second unpackGrad <$> evalIO (dcompSChadSUS input)
+ let tansCompSChadSUS = TypedEnv (tanenv env) $ toTanE env input gradCompSChadSUS
-- annotate (showEnv (d2e env) gradChad0)
-- annotate (showEnv (d2e env) gradChadS)
@@ -381,17 +391,21 @@ adTestGenChad env envGenerator expr exprS primalSfun | Dict <- envKnown env =
-- annotate (ppExpr env dtermChad0)
-- annotate (ppExpr env dtermChadS)
annotate (ppExpr env (simplifyFix (unMonoid dtermSChadS)))
- diff outChad0 closeIsh outPrimal
- diff outChadS closeIsh outPrimal
- diff outSChad0 closeIsh outPrimal
- diff outSChadS closeIsh outPrimal
- diff outCompSChadS closeIsh outPrimal
+ diff outChad0 closeIsh outPrimal
+ diff outChadS closeIsh outPrimal
+ diff outChadSUS closeIsh outPrimal
+ diff outSChad0 closeIsh outPrimal
+ diff outSChadS closeIsh outPrimal
+ diff outSChadSUS closeIsh outPrimal
+ diff outCompSChadSUS closeIsh outPrimal
let closeIshE' e1 e2 = closeIshE (tanenv env) (unTypedEnv e1) (unTypedEnv e2)
- diff tansChad closeIshE' tansFwd
- diff tansChadS closeIshE' tansFwd
- diff tansSChad closeIshE' tansFwd
- diff tansSChadS closeIshE' tansFwd
- diff tansCompSChadS closeIshE' tansFwd
+ diff tansChad closeIshE' tansFwd
+ diff tansChadS closeIshE' tansFwd
+ diff tansChadSUS closeIshE' tansFwd
+ diff tansSChad closeIshE' tansFwd
+ diff tansSChadS closeIshE' tansFwd
+ diff tansSChadSUS closeIshE' tansFwd
+ diff tansCompSChadSUS closeIshE' tansFwd
withCompiled :: SList STy env -> Ex env t -> ((SList Value env -> IO (Rep t)) -> TestTree) -> TestTree
withCompiled env expr = withResource (compile env expr) (\_ -> pure ())
@@ -435,11 +449,30 @@ gen_neural = do
lay3 <- genArray tR (ShNil `ShCons` n2)
return (input `SCons` lay3 `SCons` lay2 `SCons` lay1 `SCons` SNil)
+term_build0 :: Ex '[TArr N0 R] R
+term_build0 = fromNamed $ lambda @(TArr N0 _) #x $ body $
+ idx0 $
+ build SZ (shape #x) $ #idx :-> #x ! #idx
+
term_build1_sum :: Ex '[TVec R] R
term_build1_sum = fromNamed $ lambda #x $ body $
idx0 $ sum1i $
build (SS SZ) (shape #x) $ #idx :-> #x ! #idx
+term_build1_idx :: Ex '[TVec R] R
+term_build1_idx = fromNamed $ lambda @(TVec _) #x $ body $
+ let_ #n (snd_ (shape #x)) $
+ idx0 $ sum1i $
+ build1 (#n `idiv` 2) (#i :-> #x ! pair nil (2 * #i))
+
+term_idx_coprod :: Ex '[TVec (TEither R R)] R
+term_idx_coprod = fromNamed $ lambda @(TVec (TEither R R)) #x $ body $
+ let_ #n (snd_ (shape #x)) $
+ idx0 $ sum1i $ build1 #n $ #i :->
+ case_ (#x ! pair nil #i)
+ (#a :-> #a * 2)
+ (#b :-> #b * 3)
+
term_pairs :: Ex [R, R] R
term_pairs = fromNamed $ lambda #x $ lambda #y $ body $
let_ #p (pair #x #y) $
@@ -502,22 +535,22 @@ tests_Compile = testGroup "Compile"
,compileTest "accum (f64,f64)" $ fromNamed $ lambda #b $ lambda #x $ body $
with @(TPair R R) (pair 0.0 0.0) $ #ac :->
- let_ #_ (if_ #b (accum (SAPFst SAPHere) (pair nil nil) 3.0 #ac) nil) $
+ let_ #_ (if_ #b (accum (SAPFst SAPHere) nil 3.0 #ac) nil) $
let_ #_ (accum SAPHere nil #x #ac) $
- let_ #_ (accum (SAPSnd SAPHere) (pair nil nil) 4.0 #ac) $
+ let_ #_ (accum (SAPSnd SAPHere) nil 4.0 #ac) $
nil
,compileTest "accum (Maybe (f64,f64))" $ fromNamed $ lambda #b $ lambda #x $ body $
- with @(TMaybe (TPair R R)) nothing $ #ac :->
- let_ #_ (if_ #b (accum (SAPJust (SAPFst SAPHere)) (pair nil nil) 3.0 #ac) nil) $
+ with @(TMaybe (TPair R R)) (just (pair 0 0)) $ #ac :->
+ let_ #_ (if_ #b (accum (SAPJust (SAPFst SAPHere)) nil 3.0 #ac) nil) $
let_ #_ (accum SAPHere nil #x #ac) $
- let_ #_ (accum (SAPJust (SAPSnd SAPHere)) (pair nil nil) 4.0 #ac) $
+ let_ #_ (accum (SAPJust (SAPSnd SAPHere)) nil 4.0 #ac) $
nil
,compileTestTp "accum Arr 1 f64" (() :& C "" 3) $ fromNamed $ lambda #b $ lambda @(TVec R) #x $ body $
let_ #len (snd_ (shape #x)) $
with @(TVec R) (build1 #len (#_ :-> 0)) $ #ac :->
- let_ #_ (if_ #b (accum (SAPArrIdx SAPHere) (pair (pair (pair nil 2) (build1 #len (#_ :-> nil))) nil) 6.0 #ac)
+ let_ #_ (if_ #b (accum (SAPArrIdx SAPHere) (pair (pair nil 2) nil) 6.0 #ac)
nil) $
let_ #_ (accum SAPHere nil #x #ac) $
nil
@@ -556,9 +589,7 @@ tests_AD = testGroup "AD"
,adTest "build0 const" $ fromNamed $ lambda @R #x $ body $
idx0 $ build SZ nil $ #idx :-> const_ 0.0
- ,adTest "build0" $ fromNamed $ lambda @(TArr N0 _) #x $ body $
- idx0 $
- build SZ (shape #x) $ #idx :-> #x ! #idx
+ ,adTest "build0" term_build0
,adTest "build1-sum" term_build1_sum
@@ -566,6 +597,27 @@ tests_AD = testGroup "AD"
idx0 $ sum1i . sum1i $
build (SS (SS SZ)) (shape #x) $ #idx :-> #x ! #idx
+ ,adTest "build1-idx" term_build1_idx
+
+ ,adTest "idx-pair" $ fromNamed $ lambda @(TVec (TPair R R)) #x $ body $
+ let_ #n (snd_ (shape #x)) $
+ idx0 $ sum1i $ build1 #n $ #i :->
+ let_ #p (#x ! pair nil #i) $
+ 3 * fst_ #p + 2 * snd_ #p
+
+ ,adTest "idx-coprod" $ term_idx_coprod
+
+ ,adTest "idx-coprod-pair" $ fromNamed $ lambda @(TVec R) #arr $ body $
+ let_ #n (snd_ (shape #arr)) $
+ let_ #b (build1 #n (#i :-> let_ #x (#arr ! pair nil #i) $
+ if_ (#x .>= 1) (pair (inl (pair #x (7 * #x))) (2 * #x))
+ (pair (inr (3 * #x)) (exp #x)))) $
+ idx0 $ sum1i $ build1 #n $ #i :->
+ let_ #p (#b ! pair nil #i) $
+ case_ (fst_ #p)
+ (#a :-> fst_ #a * 2 + snd_ #a * snd_ #p)
+ (#b :-> #b * 4)
+
,adTestCon "maximum" (\(Value a `SCons` _) -> let _ `ShCons` n = arrayShape a in n > 0) $
fromNamed $ lambda @(TMat R) #x $ body $
idx0 $ sum1i $ maximum1i #x
generated by cgit v1.2.3-70-g09d2 (git 2.51.2) at 2025-12-15 01:38:41 +0100