WIP fold: Implement D[fold1i]

Still need to handle the new primitives in the rest of the library

1 files changed, 14 insertions, 0 deletions

diff --git a/src/AST.hs b/src/AST.hs
index a10f1ae..275abcd 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -62,6 +62,7 @@ data Expr x env t where
   -- array operations
   EConstArr :: Show (ScalRep t) => x (TArr n (TScal t)) -> SNat n -> SScalTy t -> Array n (ScalRep t) -> Expr x env (TArr n (TScal 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)
+  -- bottommost t in 't : t : env' is the rightmost argument (environments grow to the right)
   EFold1Inner :: x (TArr n t) -> Commutative -> 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)
@@ -69,6 +70,19 @@ data Expr x env t where
   EMaximum1Inner :: ScalIsNumeric t ~ True => x (TArr n (TScal t)) -> Expr x env (TArr (S n) (TScal t)) -> Expr x env (TArr n (TScal t))
   EMinimum1Inner :: ScalIsNumeric t ~ True => x (TArr n (TScal t)) -> Expr x env (TArr (S n) (TScal t)) -> Expr x env (TArr n (TScal t))
 
+  EFold1InnerD1 :: x (TPair (TArr n t1) (TArr (S n) t1)) -> Commutative -> Expr x (t1 : t1 : env) (TPair t1 tape) -> Expr x env t1 -> Expr x env (TArr (S n) t1)
+                -> Expr x env (TPair (TArr n t1)  -- normal primal fold output
+                                     (TArr (S n) (TPair t1 tape)))  -- bag-of-goodies: zip (prescanl) (the tape stores)
+  -- TODO: as-is, the primal input array is mostly unused; it is used only if the combination function returns sparse cotangents that need to be expanded, and nowhere else. That's wasteful storage.
+  EFold1InnerD2 :: x (TArr (S n) t2) -> Commutative
+                -> SMTy t2  -- t2 must be a monoid in order to be able to add all inner-vector contributions to the single x0
+                   -- TODO: `fold1i (*)` should have zero tape stores since primals are directly made available here, but this is not yet true
+                -> Expr x (t2 : t1 : t1 : tape : env) (TPair t2 t2)  -- reverse derivative of function (should contribute to free variables via accumulation)
+                -> Expr x env (TArr (S n) t1)  -- primal input array
+                -> Expr x env (TArr (S n) (TPair t1 tape))  -- bag-of-goodies from EFold1InnerD1
+                -> Expr x env (TArr n t2)  -- incoming cotangent
+                -> Expr x env (TPair t2 (TArr (S n) t2))  -- outgoing cotangents to x0 and input array
+
   -- expression operations
   EConst :: Show (ScalRep t) => x (TScal t) -> SScalTy t -> ScalRep t -> Expr x env (TScal t)
   EIdx0 :: x t -> Expr x env (TArr Z t) -> Expr x env t