Add AD for snest and sunNest

Mikolaj · Nov 20, 2024 · 95d1420 · 95d1420
1 parent ad2ee2f
commit 95d1420
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/src/HordeAd/Core/Delta.hs b/src/HordeAd/Core/Delta.hs
@@ -570,6 +570,12 @@ data Delta :: Target -> TensorKindType -> Type where
     -- TODO: this is a haddock for Gather1; fix.
   CastS :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2, KnownShS sh)
         => Delta target (TKS sh r1) -> Delta target (TKS sh r2)
+  NestS :: (GoodScalar r, KnownShS sh1, KnownShS sh2, KnownShS (sh1 ++ sh2))
+        => Delta target (TKS (sh1 ++ sh2) r)
+        -> Delta target (TKS2 sh1 (TKS sh2 r))
+  UnNestS :: (GoodScalar r, KnownShS sh1, KnownShS sh2, KnownShS (sh1 ++ sh2))
+          => Delta target (TKS2 sh1 (TKS sh2 r))
+          -> Delta target (TKS (sh1 ++ sh2) r)
   SFromR :: forall sh r target. (GoodScalar r, KnownShS sh, KnownNat (Rank sh))
          => Delta target (TKR (Rank sh) r)
          -> Delta target (TKS sh r)
@@ -694,6 +700,8 @@ shapeDeltaFull = \case
   ReshapeS{} -> FTKS knownShS FTKScalar
   GatherS{} -> FTKS knownShS FTKScalar
   CastS{} -> FTKS knownShS FTKScalar
+  NestS{} -> FTKS knownShS (FTKS knownShS FTKScalar)
+  UnNestS{} -> FTKS knownShS FTKScalar
   SFromR{} -> FTKS knownShS FTKScalar
   SFromH{} -> FTKS knownShS FTKScalar
 
@@ -1207,6 +1215,10 @@ evalSame !s !c = \case
   CastS @r1 @_ @sh d ->
     evalSRuntimeSpecialized s (toADTensorKindShared (stensorKind @(TKS sh r1))
                                $ scast c) d
+  NestS d ->
+    evalSame s (sunNest c) d
+  UnNestS d ->
+    evalSame s (snest knownShS c) d
   SFromR @sh (RFromS @sh2 d) ->
     case sameShape @sh @sh2 of
       Just Refl -> evalSame s c d
@@ -1554,6 +1566,8 @@ fwdSame params s = \case
       case sameTensorKind @(TKS sh r1) @(ADTensorKind (TKS sh r1)) of
         Just Refl -> second scast $ fwdSame params s d
         _ -> (s, repConstant 0 $ aDTensorKind $ shapeDeltaFull d0)
+  NestS d -> second (snest knownShS) $ fwdSame params s d
+  UnNestS d -> second sunNest $ fwdSame params s d
   SFromR @sh (RFromS @sh2 d) ->
     case sameShape @sh @sh2 of
       Just Refl -> fwdSame params s d

diff --git a/src/HordeAd/Core/OpsADVal.hs b/src/HordeAd/Core/OpsADVal.hs
@@ -408,6 +408,9 @@ instance (ADReadyNoLet target, ShareTensor target, ShareTensor (PrimalOf target)
     let v = sfromIntegral u
     in fromPrimalADVal v
   sconcrete t = fromPrimalADVal (sconcrete t)
+  snest sh (D u u') | Dict <- Nested.Internal.Shape.shsKnownShS sh =
+    dD (snest sh u) (NestS u')
+  sunNest (D u u') = dD (sunNest u) (UnNestS u')
   sfromR :: forall r sh. (GoodScalar r, KnownShS sh, KnownNat (Rank sh))
          => ADVal target (TKR (Rank sh) r) -> ADVal target (TKS sh r)
   sfromR (D u u') = dDnotShared (sfromR u) (dSFromR u')