From a76ec3bcbdea7beaf9066e4ce0b8c5868f571cdb Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Tue, 5 Nov 2024 22:45:04 +0100
Subject: Generate EOneHot in D[EIdx]
This generates a one-hot for the zero-dimensional inner array because
indexing one level further to the actual element is too difficult. But
this should simplify away fine.
---
src/CHAD.hs | 10 ++++++----
1 file changed, 6 insertions(+), 4 deletions(-)
(limited to 'src/CHAD.hs')
diff --git a/src/CHAD.hs b/src/CHAD.hs
index b3e2358..6b0627d 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -1156,10 +1156,12 @@ drev des = \case
(EIdx ext (EVar ext (STArr n (d1 eltty)) (IS (IS IZ)))
(EVar ext (tTup (sreplicate n tIx)) IZ))
sub
- (ELet ext (EBuild ext n (EVar ext (tTup (sreplicate n tIx)) (IS (IS IZ))) $
- ECase ext (EOp ext OIf (eidxEq n (EVar ext tIxN IZ) (EVar ext tIxN (IS (IS IZ)))))
- (EVar ext (d2 eltty) (IS (IS IZ)))
- (EZero eltty)) $
+ (ELet ext (EOneHot (STArr n eltty) n
+ (arrIdxToAcIdx (d2 eltty) n $ EVar ext tIxN (IS IZ))
+ (case n of SZ -> EUnit ext (EVar ext (d2 eltty) IZ)
+ SS{} | Refl <- lemAcValArrN (d2 eltty) n ->
+ EPair ext (EVar ext tIxN (IS (IS IZ)))
+ (EUnit ext (EVar ext (d2 eltty) IZ)))) $
weakenExpr (WCopy (WSink .> WSink .> WSink)) e2)
EShape _ e
--
cgit v1.2.3-70-g09d2