aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--src/CHAD/AST.hs6
1 files changed, 4 insertions, 2 deletions
diff --git a/src/CHAD/AST.hs b/src/CHAD/AST.hs
index ce9eb20..b795070 100644
--- a/src/CHAD/AST.hs
+++ b/src/CHAD/AST.hs
@@ -583,8 +583,10 @@ eshapeProd (SS n) e =
eflatten :: Ex env (TArr n t) -> Ex env (TArr N1 t)
eflatten e =
let STArr n _ = typeOf e
- in elet e $
- EReshape ext (SS SZ) (EPair ext (ENil ext) (eshapeProd n (EShape ext (evar IZ)))) (evar IZ)
+ in case n of
+ SS SZ -> e
+ _ -> elet e $
+ EReshape ext (SS SZ) (EPair ext (ENil ext) (eshapeProd n (EShape ext (evar IZ)))) (evar IZ)
-- ezeroD2 :: STy t -> Ex env (ZeroInfo (D2 t)) -> Ex env (D2 t)
-- ezeroD2 t ezi = EZero ext (d2M t) ezi