Differentiate Replicate

1 files changed, 2 insertions, 2 deletions

diff --git a/src/Language/AST.hs b/src/Language/AST.hs
index af5a5a2..0945dd9 100644
--- a/src/Language/AST.hs
+++ b/src/Language/AST.hs
@@ -43,7 +43,7 @@ data NExpr env t where
   NEConstArr :: Show (ScalRep t) => SNat n -> SScalTy t -> Array n (ScalRep t) -> NExpr env (TArr n (TScal t))
   NEBuild1 :: NExpr env TIx -> Var name TIx -> NExpr ('(name, TIx) : env) t -> NExpr env (TArr (S Z) t)
   NEBuild :: SNat n -> NExpr env (Tup (Replicate n TIx)) -> Var name (Tup (Replicate n TIx)) -> NExpr ('(name, Tup (Replicate n TIx)) : env) t -> NExpr env (TArr n t)
-  NEFold1Inner :: Var name1 t -> Var name2 t -> NExpr ('(name2, t) : '(name1, t) : env) t -> NExpr env (TArr (S n) t) -> NExpr env (TArr n t)
+  NEFold1Inner :: Var name1 t -> Var name2 t -> NExpr ('(name2, t) : '(name1, t) : env) t -> NExpr env t -> NExpr env (TArr (S n) t) -> NExpr env (TArr n t)
   NESum1Inner :: ScalIsNumeric t ~ True => NExpr env (TArr (S n) (TScal t)) -> NExpr env (TArr n (TScal t))
   NEUnit :: NExpr env t -> NExpr env (TArr Z t)
   NEReplicate1Inner :: NExpr env TIx -> NExpr env (TArr n t) -> NExpr env (TArr (S n) t)
@@ -123,7 +123,7 @@ fromNamedExpr val = \case
   NEConstArr n t x -> EConstArr ext n t x
   NEBuild1 a n b -> EBuild1 ext (go a) (lambda val n b)
   NEBuild k a n b -> EBuild ext k (go a) (lambda val n b)
-  NEFold1Inner n1 n2 a b -> EFold1Inner ext (lambda2 val n1 n2 a) (go b)
+  NEFold1Inner n1 n2 a b c -> EFold1Inner ext (lambda2 val n1 n2 a) (go b) (go c)
   NESum1Inner e -> ESum1Inner ext (go e)
   NEUnit e -> EUnit ext (go e)
   NEReplicate1Inner a b -> EReplicate1Inner ext (go a) (go b)