From 79e072eddf0ec2a97ca455c27cb5ff6f2132bbab Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Tue, 22 Oct 2024 22:02:06 +0200
Subject: Differentiate Replicate

---
 src/AST.hs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

(limited to 'src/AST.hs')

diff --git a/src/AST.hs b/src/AST.hs
index 94c8537..e2702ab 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -83,7 +83,7 @@ data Expr x env t where
   EConstArr :: Show (ScalRep t) => x (TArr n (TScal t)) -> SNat n -> SScalTy t -> Array n (ScalRep t) -> Expr x env (TArr n (TScal t))
   EBuild1 :: x (TArr (S Z) t) -> Expr x env TIx -> Expr x (TIx : env) t -> Expr x env (TArr (S Z) t)
   EBuild :: x (TArr n t) -> SNat n -> Expr x env (Tup (Replicate n TIx)) -> Expr x (Tup (Replicate n TIx) : env) t -> Expr x env (TArr n t)
-  EFold1Inner :: x (TArr n t) -> Expr x (t : t : env) t -> Expr x env (TArr (S n) t) -> Expr x env (TArr n t)
+  EFold1Inner :: x (TArr n t) -> Expr x (t : t : env) t -> Expr x env t -> Expr x env (TArr (S n) t) -> Expr x env (TArr n t)
   ESum1Inner :: ScalIsNumeric t ~ True => x (TArr n (TScal t)) -> Expr x env (TArr (S n) (TScal t)) -> Expr x env (TArr n (TScal t))
   EUnit :: x (TArr Z t) -> Expr x env t -> Expr x env (TArr Z t)
   EReplicate1Inner :: x (TArr (S n) t) -> Expr x env TIx -> Expr x env (TArr n t) -> Expr x env (TArr (S n) t)
@@ -186,7 +186,7 @@ typeOf = \case
   EConstArr _ n t _ -> STArr n (STScal t)
   EBuild1 _ _ e -> STArr (SS SZ) (typeOf e)
   EBuild _ n _ e -> STArr n (typeOf e)
-  EFold1Inner _ _ e | STArr (SS n) t <- typeOf e -> STArr n t
+  EFold1Inner _ _ _ e | STArr (SS n) t <- typeOf e -> STArr n t
   ESum1Inner _ e | STArr (SS n) t <- typeOf e -> STArr n t
   EUnit _ e -> STArr SZ (typeOf e)
   EReplicate1Inner _ _ e | STArr n t <- typeOf e -> STArr (SS n) t
@@ -260,7 +260,7 @@ subst' f w = \case
   EConstArr x n t a -> EConstArr x n t a
   EBuild1 x a b -> EBuild1 x (subst' f w a) (subst' (sinkF f) (WCopy w) b)
   EBuild x n a b -> EBuild x n (subst' f w a) (subst' (sinkF f) (WCopy w) b)
-  EFold1Inner x a b -> EFold1Inner x (subst' (sinkF (sinkF f)) (WCopy (WCopy w)) a) (subst' f w b)
+  EFold1Inner x a b c -> EFold1Inner x (subst' (sinkF (sinkF f)) (WCopy (WCopy w)) a) (subst' f w b) (subst' f w c)
   ESum1Inner x e -> ESum1Inner x (subst' f w e)
   EUnit x e -> EUnit x (subst' f w e)
   EReplicate1Inner x a b -> EReplicate1Inner x (subst' f w a) (subst' f w b)
-- 
cgit v1.2.3-70-g09d2