diff options
author | Tom Smeding <tom@tomsmeding.com> | 2024-08-29 14:31:03 +0200 |
---|---|---|
committer | Tom Smeding <tom@tomsmeding.com> | 2024-08-29 14:31:03 +0200 |
commit | de804fd0eedc13fb5079270a2a21279518d4bb0b (patch) | |
tree | 3160418964ddc07b44da545e898501b38fef4e75 | |
parent | eabee174af9393350c09965d6b383744710154ed (diff) |
Generalise freezeRet a little bit
-rw-r--r-- | src/CHAD.hs | 4 |
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))))) |