summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorTom Smeding <tom@tomsmeding.com>2024-08-29 14:31:03 +0200
committerTom Smeding <tom@tomsmeding.com>2024-08-29 14:31:03 +0200
commitde804fd0eedc13fb5079270a2a21279518d4bb0b (patch)
tree3160418964ddc07b44da545e898501b38fef4e75
parenteabee174af9393350c09965d6b383744710154ed (diff)
Generalise freezeRet a little bit
-rw-r--r--src/CHAD.hs4
1 files changed, 2 insertions, 2 deletions
diff --git a/src/CHAD.hs b/src/CHAD.hs
index 5ae1a07..2513f84 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -613,14 +613,14 @@ sD1eEnv (DPush d _) = SCons (Const ()) (sD1eEnv d)
freezeRet :: Descr env sto
-> Ret env sto t
- -> (forall env'. Ex env' (D2 t)) -- the incoming cotangent value
+ -> Ex (D1E env) (D2 t) -- the incoming cotangent value
-> Ex (D1E env) (TPair (D1 t) (TEVM (D2E (Select env sto "accum")) (Tup (D2E (Select env sto "merge")))))
freezeRet descr (Ret e0 e1 sub e2) d =
let e2' = weakenExpr (WCopy (wRaiseAbove (bindingsBinds e0) (sD1eEnv descr))) e2
in letBinds e0 $
EPair ext
e1
- (ELet ext d
+ (ELet ext (weakenExpr (sinkWithBindings e0) d)
(EMBind e2'
(EMReturn (d2e (select SAccum descr))
(expandSubenvZeros (select SMerge descr) sub (EVar ext (tTup (d2e (subList (select SMerge descr) sub))) IZ)))))