summaryrefslogtreecommitdiff
path: root/src/CHAD.hs
diff options
context:
space:
mode:
Diffstat (limited to 'src/CHAD.hs')
-rw-r--r--src/CHAD.hs141
1 files changed, 70 insertions, 71 deletions
diff --git a/src/CHAD.hs b/src/CHAD.hs
index 241825e..7cd4c26 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -1094,42 +1094,39 @@ drev des accumMap sd = \case
case lemAppendNil @e_binds of { Refl ->
let tapety = tapeTy (subList (bindingsBinds e0) subtapeE) in
let collectexpr = bindingsCollectTape (bindingsBinds e0) subtapeE in
- Ret (BTop `BPush` (shty, drevPrimal des she)
- `BPush` (STArr ndim (STPair (d1 eltty) tapety)
- ,EBuild ext ndim
- (EVar ext shty IZ)
- (letBinds (fst (weakenBindings weakenExpr (autoWeak (#ix (shty `SCons` SNil)
- &. #sh (shty `SCons` SNil)
- &. #d1env (desD1E des)
- &. #d1env' (desD1E usedDes))
- (#ix :++: LPreW #d1env' #d1env (wUndoSubenv subD1eUsed))
- (#ix :++: #sh :++: #d1env))
- e0)) $
- let w = autoWeak (#ix (shty `SCons` SNil)
- &. #sh (shty `SCons` SNil)
- &. #e0 (bindingsBinds e0)
- &. #d1env (desD1E des)
- &. #d1env' (desD1E usedDes))
- (#e0 :++: #ix :++: LPreW #d1env' #d1env (wUndoSubenv subD1eUsed))
- (#e0 :++: #ix :++: #sh :++: #d1env)
- w' = w .> wCopies (bindingsBinds e0) (WClosed @(shty : D1E env'))
- in EPair ext (weakenExpr w e1) (collectexpr w')))
- `BPush` (STArr ndim tapety, emap (ESnd ext (EVar ext (STPair (d1 eltty) tapety) IZ))
- (EVar ext (STArr ndim (STPair (d1 eltty) tapety)) IZ)))
- (SEYesR (SENo (SEYesR SETop)))
- (emap (EFst ext (EVar ext (STPair (d1 eltty) tapety) IZ))
- (EVar ext (STArr ndim (STPair (d1 eltty) tapety)) (IS IZ)))
+ let mergePrimalSub = subenvD1E (selectSub SMerge des `subenvCompose` subMergeUsed `subenvCompose` proSub) in
+ let mergePrimalBindings = collectBindings (d1e (descrList des)) mergePrimalSub in
+ Ret (mergePrimalBindings
+ `BPush` (shty, weakenExpr (wSinks (d1e envPro)) (drevPrimal des she))
+ `BPush` (STArr ndim (STPair (d1 eltty) tapety)
+ ,EBuild ext ndim
+ (EVar ext shty IZ)
+ (letBinds (fst (weakenBindings weakenExpr (autoWeak (#ix (shty `SCons` SNil)
+ &. #sh (shty `SCons` SNil)
+ &. #propr (d1e envPro)
+ &. #d1env (desD1E des)
+ &. #d1env' (desD1E usedDes))
+ (#ix :++: LPreW #d1env' #d1env (wUndoSubenv subD1eUsed))
+ (#ix :++: #sh :++: #propr :++: #d1env))
+ e0)) $
+ let w = autoWeak (#ix (shty `SCons` SNil)
+ &. #sh (shty `SCons` SNil)
+ &. #e0 (bindingsBinds e0)
+ &. #propr (d1e envPro)
+ &. #d1env (desD1E des)
+ &. #d1env' (desD1E usedDes))
+ (#e0 :++: #ix :++: LPreW #d1env' #d1env (wUndoSubenv subD1eUsed))
+ (#e0 :++: #ix :++: #sh :++: #propr :++: #d1env)
+ w' = w .> wCopies (bindingsBinds e0) (WClosed @(shty : D1E env'))
+ in EPair ext (weakenExpr w e1) (collectexpr w')))
+ `BPush` (STArr ndim tapety, emap (ESnd ext (evar IZ)) (EVar ext (STArr ndim (STPair (d1 eltty) tapety)) IZ)))
+ (SEYesR (SENo (SEYesR (subenvAll (d1e envPro)))))
+ (emap (EFst ext (evar IZ)) (EVar ext (STArr ndim (STPair (d1 eltty) tapety)) (IS IZ)))
(subenvMap (\t Refl -> spDense t) (d2eM (select SMerge des)) (subenvD2E (subenvCompose subMergeUsed proSub)))
- (let sinkOverEnvPro = wSinks @(sd : TArr ndim (Tape e_tape) : Tup (Replicate ndim TIx) : D2AcE (Select env sto "accum")) (d2ace envPro) in
+ (let sinkOverEnvPro = wSinks @(sd : TArr ndim (Tape e_tape) : Tup (Replicate ndim TIx) : Append (D1E envPro) (D2AcE (Select env sto "accum"))) (d2ace envPro) in
ESnd ext $
uninvertTup (d2e envPro) (STArr ndim STNil) $
- -- TODO: what's happening here is that because of the sparsity
- -- rewrite, makeAccumulators needs primals where it previously
- -- didn't. The build derivative is currently not saving those
- -- primals, so the hole below cannot currently be filled. The
- -- appropriate primals (waves hands) need to be stored, so that a
- -- weakening can be provided here.
- makeAccumulators @_ @_ @(TArr ndim TNil) (_ (subenvCompose subMergeUsed proSub)) envPro $
+ makeAccumulators @_ @_ @(TArr ndim TNil) (WSink .> WSink .> WSink .> wRaiseAbove (d1e envPro) (d2ace (select SAccum des))) envPro $
EBuild ext ndim (EVar ext shty (sinkOverEnvPro @> IS (IS IZ))) $
-- the cotangent for this element
ELet ext (EIdx ext (EVar ext (STArr ndim (applySparse sdElt (d2 eltty))) (WSink .> sinkOverEnvPro @> IZ))
@@ -1148,10 +1145,11 @@ drev des accumMap sd = \case
&. #darr (auto1 @(TArr ndim sdElt))
&. #tapearr (auto1 @(TArr ndim (Tape e_tape)))
&. #sh (auto1 @shty)
+ &. #propr (d1e envPro)
&. #d2acUsed (d2ace (select SAccum usedDes))
&. #d2acEnv (d2ace (select SAccum des)))
(#pro :++: #d :++: #etape :++: LPreW #d2acUsed #d2acEnv (wUndoSubenv subAccumUsed))
- ((#etape :++: #prerebinds) :++: #tape :++: #d :++: #ix :++: #pro :++: #darr :++: #tapearr :++: #sh :++: #d2acEnv)
+ ((#etape :++: #prerebinds) :++: #tape :++: #d :++: #ix :++: #pro :++: #darr :++: #tapearr :++: #sh :++: #propr :++: #d2acEnv)
.> wPro (subList (bindingsBinds e0) subtapeE))
e2)
}}}
@@ -1167,32 +1165,34 @@ drev des accumMap sd = \case
weakenExpr (WCopy WSink) e2)
EReplicate1Inner _ en e
- -- We're allowed to ignore en2 here because the output of 'ei' is discrete.
- | Rets binds subtape (RetPair en1 _ _ `SCons` RetPair e1 sub e2 `SCons` SNil)
- <- retConcat des $ drev des accumMap en `SCons` drev des accumMap e `SCons` SNil
+ -- We're allowed to differentiate 'en' as primal-only here because its output is discrete.
+ | SpArr sdElt <- sd
, let STArr ndim eltty = typeOf e ->
- Ret binds
- subtape
- (EReplicate1Inner ext en1 e1)
- sub
- (EMaybe ext
- (zeroTup (subList (select SMerge des) sub))
- (ELet ext (EJust ext (EFold1Inner ext Commut
- (EPlus ext (d2M eltty) (EVar ext (d2 eltty) (IS IZ)) (EVar ext (d2 eltty) IZ))
- (ezeroD2 eltty)
- (EVar ext (STArr (SS ndim) (d2 eltty)) IZ))) $
- weakenExpr (WCopy (WSink .> WSink)) e2)
- (EVar ext (d2 (STArr (SS ndim) eltty)) IZ))
+ -- This pessimistic sparsity union is because the array might have been empty, in which case we need to generate a zero.
+ sparsePlusS ST ST (d2M eltty) sdElt SpAbsent $ \sdElt' (Inj inj1) (Inj inj2) _ ->
+ case drev des accumMap (SpArr sdElt') e of { Ret binds subtape e1 sub e2 ->
+ Ret binds
+ subtape
+ (EReplicate1Inner ext (weakenExpr (wSinks (bindingsBinds binds)) (drevPrimal des en)) e1)
+ sub
+ (ELet ext (EFold1Inner ext Commut
+ (sparsePlus (d2M eltty) sdElt'
+ (EVar ext (applySparse sdElt' (d2 eltty)) (IS IZ))
+ (EVar ext (applySparse sdElt' (d2 eltty)) IZ))
+ (inj2 (ENil ext))
+ (emap (inj1 (evar IZ)) $ EVar ext (STArr (SS ndim) (applySparse sdElt (d2 eltty))) IZ)) $
+ weakenExpr (WCopy WSink) e2)
+ }
EIdx0 _ e
- | Ret e0 subtape e1 sub e2 <- drev des accumMap e
+ | Ret e0 subtape e1 sub e2 <- drev des accumMap (SpArr sd) e
, STArr _ t <- typeOf e ->
Ret e0
subtape
(EIdx0 ext e1)
sub
- (ELet ext (EJust ext (EUnit ext (EVar ext (d2 t) IZ))) $
- weakenExpr (WCopy WSink) e2)
+ (ELet ext (EUnit ext (EVar ext (applySparse sd (d2 t)) IZ)) $
+ weakenExpr (WCopy WSink) e2)
EIdx1{} -> error "CHAD of EIdx1: Please use EIdx instead"
{-
@@ -1214,26 +1214,25 @@ drev des accumMap sd = \case
-}
EIdx _ e ei
- -- We're allowed to ignore ei2 here because the output of 'ei' is discrete.
- | Rets binds subtape (RetPair e1 sub e2 `SCons` RetPair ei1 _ _ `SCons` SNil)
- <- retConcat des $ drev des accumMap e `SCons` drev des accumMap ei `SCons` SNil
- , STArr n eltty <- typeOf e
+ -- We're allowed to differentiate ei as primal because its output is discrete.
+ | STArr n eltty <- typeOf e
, Refl <- indexTupD1Id n
- , Refl <- lemZeroInfoD2 eltty
- , let tIxN = tTup (sreplicate n tIx) ->
- Ret (binds `BPush` (STArr n (d1 eltty), e1)
- `BPush` (tIxN, EShape ext (EVar ext (typeOf e1) IZ))
- `BPush` (tIxN, weakenExpr (WSink .> WSink) ei1))
- (SEYesR (SEYesR (SENo subtape)))
- (EIdx ext (EVar ext (STArr n (d1 eltty)) (IS (IS IZ)))
- (EVar ext (tTup (sreplicate n tIx)) IZ))
- sub
- (ELet ext (EOneHot ext (d2M (STArr n eltty)) (SAPJust (SAPArrIdx SAPHere))
- (EPair ext (EPair ext (EVar ext tIxN (IS IZ))
- (EBuild ext n (EVar ext tIxN (IS (IS IZ))) (ENil ext)))
- (ENil ext))
- (EVar ext (d2 eltty) IZ)) $
- weakenExpr (WCopy (WSink .> WSink .> WSink)) e2)
+ , let tIxN = tTup (sreplicate n tIx) ->
+ sparsePlusS ST ST (d2M eltty) sd SpAbsent $ \sd' (Inj inj1) (Inj inj2) _ ->
+ case drev des accumMap (SpArr sd') e of { Ret binds subtape e1 sub e2 ->
+ Ret (binds `BPush` (STArr n (d1 eltty), e1)
+ `BPush` (tIxN, EShape ext (EVar ext (typeOf e1) IZ))
+ `BPush` (tIxN, weakenExpr (WSink .> WSink .> wSinks (bindingsBinds binds)) (drevPrimal des ei)))
+ (SEYesR (SEYesR (SENo subtape)))
+ (EIdx ext (EVar ext (STArr n (d1 eltty)) (IS (IS IZ)))
+ (EVar ext (tTup (sreplicate n tIx)) IZ))
+ sub
+ (ELet ext (EOneHot ext (SMTArr n (applySparse sd' (d2M eltty))) (SAPArrIdx SAPHere)
+ (EPair ext (EPair ext (EVar ext tIxN (IS IZ))
+ (EBuild ext n (EVar ext tIxN (IS (IS IZ))) (ENil ext))) (ENil ext))
+ (inj1 $ EVar ext (applySparse sd (d2 eltty)) IZ)) $
+ weakenExpr (WCopy (WSink .> WSink .> WSink)) e2)
+ }
EShape _ e
-- Allowed to ignore e2 here because the output of EShape is discrete,
generated by cgit v1.2.3-70-g09d2 (git 2.48.1) at 2025-06-13 15:29:10 +0200