From a9e6c72eff3bee8d45e0d906e8cd027066e04793 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Mon, 26 Jan 2026 23:36:56 +0100
Subject: Optimise eflatten

---
 src/CHAD/AST.hs | 6 ++++--
 1 file 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
-- 
cgit v1.2.3-70-g09d2