diff options
Diffstat (limited to 'src/CHAD.hs')
| -rw-r--r-- | src/CHAD.hs | 15 |
1 files changed, 14 insertions, 1 deletions
diff --git a/src/CHAD.hs b/src/CHAD.hs index 8080ec0..2f05807 100644 --- a/src/CHAD.hs +++ b/src/CHAD.hs @@ -874,6 +874,20 @@ drev des = \case (EVar ext (d2 (opt2 op)) IZ)) (weakenExpr (WCopy (wSinks' @[_,_])) e2)) + ECustom _ _ _ storety _ pr du a b + -- allowed to ignore a2 because 'a' is the part of the input that is inactive + | Rets binds subtape (RetPair a1 _ _ `SCons` RetPair b1 bsub b2 `SCons` SNil) + <- retConcat des $ drev des a `SCons` drev des b `SCons` SNil -> + Ret (binds `BPush` (typeOf a1, a1) + `BPush` (typeOf b1, weakenExpr WSink b1) + `BPush` (typeOf pr, weakenExpr (WCopy (WCopy WClosed)) pr) + `BPush` (storety, ESnd ext (EVar ext (typeOf pr) IZ))) + (SEYes (SENo (SENo (SENo subtape)))) + (EFst ext (EVar ext (typeOf pr) (IS IZ))) + bsub + (ELet ext (weakenExpr (WCopy (WCopy WClosed)) du) $ + weakenExpr (WCopy (WSink .> WSink)) b2) + EError t s -> Ret BTop SETop @@ -888,7 +902,6 @@ drev des = \case (subenvNone (select SMerge des)) (ENil ext) - -- TODO: merge the e0 and e1 builds in a single build just like they are merged into a single case in D[case]0, then it can really store only the parts that need to be preserved until D[build]2 EBuild _ (ndim :: SNat ndim) she (orige :: Ex _ eltty) | Ret (she0 :: Bindings _ _ she_binds) _ she1 _ _ <- drev des she -- allowed to ignore she2 here because she has a discrete result , let eltty = typeOf orige |