Differentiate Replicate

1 files changed, 4 insertions, 3 deletions

diff --git a/src/Interpreter.hs b/src/Interpreter.hs
index 8ce1b0e..b818eb0 100644
--- a/src/Interpreter.hs
+++ b/src/Interpreter.hs
@@ -75,11 +75,12 @@ interpret' env = \case
   EBuild _ dim a b -> do
     sh <- unTupRepIdx ShNil ShCons dim <$> interpret' env a
     arrayGenerateM sh (\idx -> interpret' (Value (tupRepIdx ixUncons dim idx) `SCons` env) b)
-  EFold1Inner _ a b -> do
+  EFold1Inner _ a b c -> do
     let f = \x y -> interpret' (Value y `SCons` Value x `SCons` env) a
-    arr <- interpret' env b
+    x0 <- interpret' env b
+    arr <- interpret' env c
     let sh `ShCons` n = arrayShape arr
-    arrayGenerateM sh $ \idx -> foldl1M f [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n - 1]]
+    arrayGenerateM sh $ \idx -> foldM f x0 [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n - 1]]
   ESum1Inner _ e -> do
     arr <- interpret' env e
     let STArr _ (STScal t) = typeOf e