diff options
Diffstat (limited to 'src/CHAD.hs')
| -rw-r--r-- | src/CHAD.hs | 12 |
1 files changed, 5 insertions, 7 deletions
diff --git a/src/CHAD.hs b/src/CHAD.hs index 6b0627d..ffbdcac 100644 --- a/src/CHAD.hs +++ b/src/CHAD.hs @@ -689,22 +689,20 @@ retConcat descr (SCons (Ret (b :: Bindings _ _ shbinds1) (subtape :: Subenv _ ta freezeRet :: Descr env sto -> Ret env sto t - -> Ex (D1E env) (D2 t) -- the incoming cotangent value - -> Ex (Append (D2AcE (Select env sto "accum")) (D1E env)) (TPair (D1 t) (Tup (D2E (Select env sto "merge")))) -freezeRet descr (Ret e0 subtape e1 sub e2 :: Ret _ _ t) d = - let (e0', wInsertD2Ac) = weakenBindings weakenExpr (wSinks (d2ace (select SAccum descr))) e0 + -> Ex (D2 t : Append (D2AcE (Select env sto "accum")) (D1E env)) (TPair (D1 t) (Tup (D2E (Select env sto "merge")))) +freezeRet descr (Ret e0 subtape e1 sub e2 :: Ret _ _ t) = + let (e0', wInsertD2Ac) = weakenBindings weakenExpr (WSink .> wSinks (d2ace (select SAccum descr))) e0 e2' = weakenExpr (WCopy (wCopies (subList (bindingsBinds e0) subtape) (wRaiseAbove (d2ace (select SAccum descr)) (sD1eEnv descr)))) e2 in letBinds e0' $ EPair ext (weakenExpr wInsertD2Ac e1) - (ELet ext (weakenExpr (sinkWithBindings e0 .> wSinks (d2ace (select SAccum descr))) d) $ - ELet ext (weakenExpr (autoWeak (#d (auto1 @(D2 t)) + (ELet ext (weakenExpr (autoWeak (#d (auto1 @(D2 t)) &. #tape (subList (bindingsBinds e0) subtape) &. #shbinds (bindingsBinds e0) &. #d2ace (d2ace (select SAccum descr)) &. #tl (sD1eEnv descr)) (#d :++: LPreW #tape #shbinds (wUndoSubenv subtape) :++: #d2ace :++: #tl) - (#d :++: #shbinds :++: #d2ace :++: #tl)) + (#shbinds :++: #d :++: #d2ace :++: #tl)) e2') $ expandSubenvZeros (select SMerge descr) sub (EVar ext (tTup (d2e (subList (select SMerge descr) sub))) IZ)) |