diff options
-rw-r--r-- | test/Main.hs | 28 |
1 files changed, 14 insertions, 14 deletions
diff --git a/test/Main.hs b/test/Main.hs index b95f765..9dad35e 100644 --- a/test/Main.hs +++ b/test/Main.hs @@ -24,7 +24,7 @@ import qualified Hedgehog.Range as Range import Test.Framework import Array -import AST +import AST hiding ((.>)) import AST.Pretty import AST.UnMonoid import CHAD.Top @@ -388,7 +388,7 @@ gen_neural = do lay3 <- genArray tR (ShNil `ShCons` n2) return (input `SCons` lay3 `SCons` lay2 `SCons` lay1 `SCons` SNil) -term_build1_sum :: Ex '[TArr N1 R] R +term_build1_sum :: Ex '[TVec R] R term_build1_sum = fromNamed $ lambda #x $ body $ idx0 $ sum1i $ build (SS SZ) (shape #x) $ #idx :-> #x ! #idx @@ -399,7 +399,7 @@ term_pairs = fromNamed $ lambda #x $ lambda #y $ body $ let_ #q (pair (snd_ #p * fst_ #p + #y) #x) $ fst_ #q * #x + snd_ #q * fst_ #p -term_sparse :: Ex '[TArr N1 R] R +term_sparse :: Ex '[TVec R] R term_sparse = fromNamed $ lambda #inp $ body $ let_ #n (snd_ (shape #inp)) $ let_ #arr (build1 #n (#i :-> #inp ! pair nil #i)) $ @@ -408,7 +408,7 @@ term_sparse = fromNamed $ lambda #inp $ body $ let_ #c (build1 #n (#i :-> #arr ! pair nil 4)) $ idx0 (sum1i #a) + idx0 (sum1i #b) + idx0 (sum1i #c) -term_regression_simpl1 :: Ex '[TArr N1 R] R +term_regression_simpl1 :: Ex '[TVec R] R term_regression_simpl1 = fromNamed $ lambda #q $ body $ idx0 $ sum1i $ build (SS SZ) (shape #q) $ #idx :-> let_ #j (snd_ #idx) $ @@ -416,8 +416,8 @@ term_regression_simpl1 = fromNamed $ lambda #q $ body $ (#q ! pair nil 0) (if_ (#j .== #j) 1.0 2.0) -term_mulmatvec :: Ex [TArr N1 R, TArr N2 R] R -term_mulmatvec = fromNamed $ lambda @(TArr N2 _) #mat $ lambda @(TArr N1 _) #vec $ body $ +term_mulmatvec :: Ex [TVec R, TMat R] R +term_mulmatvec = fromNamed $ lambda #mat $ lambda #vec $ body $ idx0 $ sum1i $ let_ #hei (snd_ (fst_ (shape #mat))) $ let_ #wid (snd_ (shape #mat)) $ @@ -441,7 +441,7 @@ tests_Compile = testGroup "Compile" ,compileTestTp "accum Arr 1 f64" (() :& C "" 3) $ fromNamed $ lambda #b $ lambda #x $ body $ let_ #len (snd_ (shape #x)) $ - with @(TArr N1 R) nothing $ #ac :-> + with @(TVec R) nothing $ #ac :-> let_ #_ (if_ #b (accum (SAPArrIdx SAPHere (SS SZ)) (pair (pair (pair nil 2) (pair nil #len)) nil) 6.0 #ac) nil) $ let_ #_ (accum SAPHere nil (just #x) #ac) $ @@ -487,19 +487,19 @@ tests_AD = testGroup "AD" ,adTest "build1-sum" term_build1_sum - ,adTest "build2-sum" $ fromNamed $ lambda @(TArr N2 _) #x $ body $ + ,adTest "build2-sum" $ fromNamed $ lambda @(TMat _) #x $ body $ idx0 $ sum1i . sum1i $ build (SS (SS SZ)) (shape #x) $ #idx :-> #x ! #idx ,adTestCon "maximum" (\(Value a `SCons` _) -> let _ `ShCons` n = arrayShape a in n > 0) $ - fromNamed $ lambda @(TArr N2 R) #x $ body $ + fromNamed $ lambda @(TMat R) #x $ body $ idx0 $ sum1i $ maximum1i #x ,adTestCon "minimum" (\(Value a `SCons` _) -> let _ `ShCons` n = arrayShape a in n > 0) $ - fromNamed $ lambda @(TArr N2 R) #x $ body $ + fromNamed $ lambda @(TMat R) #x $ body $ idx0 $ sum1i $ minimum1i #x - ,adTest "unused" $ fromNamed $ lambda @(TArr N1 R) #x $ body $ + ,adTest "unused" $ fromNamed $ lambda @(TVec R) #x $ body $ let_ #a (build1 (snd_ (shape #x)) (#i :-> #x ! pair nil #i)) $ 42 @@ -510,8 +510,8 @@ tests_AD = testGroup "AD" -- Regression test for refcounts when indexing in nested arrays ,adTestTp "regression-idx1" (C "" 1 :$ C "" 1) $ - fromNamed $ lambda @(TArr N2 R) #L $ body $ - if_ (const_ @TI64 1 Language..> 0) + fromNamed $ lambda @(TMat R) #L $ body $ + if_ (const_ @TI64 1 .> 0) (idx0 $ sum1i (build1 1 $ #_ :-> idx0 (sum1i (build1 1 $ #_ :-> #L ! pair (pair nil 0) 0 * #L ! pair (pair nil 0) 0)))) @@ -522,7 +522,7 @@ tests_AD = testGroup "AD" ,adTestGen "neural-unMonoid" (unMonoid (simplifyFix Example.neural)) gen_neural ,adTestTp "logsumexp" (C "" 1) $ - fromNamed $ lambda @(TArr N1 _) #vec $ body $ + fromNamed $ lambda @(TVec _) #vec $ body $ let_ #m (maximum1i #vec) $ log (idx0 (sum1i (map_ (#x :-> exp (#x - idx0 #m)) #vec))) + idx0 #m |