horde-ad-0.3.0.0: test/tool/CrossTesting.hs
-- | Testing harness that differentiates a single objective function using
-- over a twenty different pipeline variants and cross-checks the results.
module CrossTesting
( rev', assertEqualUpToEpsilon'
, t16, t16b, t48, t128, t128b, t128c
, rrev1, rfwd1, srev1, sfwd1
) where
import Prelude
import Data.Proxy (Proxy (Proxy))
import GHC.Exts (IsList (..))
import GHC.TypeLits (KnownNat)
import System.IO.Unsafe (unsafePerformIO)
import Test.Tasty.HUnit hiding (assert)
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped.Shape
import HordeAd.ADEngine
( IncomingCotangentHandling (..)
, cjvp
, jvp
, revInterpretArtifact
, revProduceArtifactWithoutInterpretation
)
import HordeAd.Core.Adaptor
import HordeAd.Core.Ast
import HordeAd.Core.AstEngine
import HordeAd.Core.AstEnv
import HordeAd.Core.AstFreshId
import HordeAd.Core.AstInterpret
import HordeAd.Core.AstTools
import HordeAd.Core.CarriersADVal
import HordeAd.Core.CarriersAst
import HordeAd.Core.CarriersConcrete
import HordeAd.Core.ConvertTensor (rfromK)
import HordeAd.Core.Ops
import HordeAd.Core.OpsADVal
import HordeAd.Core.TensorKind
import HordeAd.Core.Types
import HordeAd.Core.Unwind
import HordeAd.OpsTensor
import EqEpsilon
crevMaybeBoth
:: forall r m f src tgt.
( GoodScalar r, f ~ Concrete, X src ~ X (DValue src), KnownSTK (X src)
, AdaptableTarget (ADVal Concrete) src
, AdaptableTarget Concrete (DValue src)
, tgt ~ ADVal f (TKR m r) )
=> (src -> tgt)
-> DValue src
-> (f (TKR m r), f (ADTensorKind (X src)))
{-# INLINE crevMaybeBoth #-}
crevMaybeBoth f vals | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =
let g :: ADVal Concrete (X src) -> ADVal Concrete (TKR m r)
g = toTarget . f . fromTarget
valsH = toTarget vals
in crevOnParams Nothing g (tftk knownSTK valsH) valsH
rev' :: forall r m n v a w.
( KnownNat n, NumScalar r, NumScalar (ADTensorScalar r)
, v ~ Concrete (TKR m r)
, w ~ Concrete (ADTensorKind (TKR m r))
, a ~ Concrete (ADTensorKind (TKR n r)) )
=> (forall f. ADReady f => f (TKR n r) -> f (TKR m r))
-> Concrete (TKR n r)
-> ( ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a
, AstTensor AstMethodLet FullSpan (TKR m r), AstTensor AstMethodLet FullSpan (TKR m r)
, v, v, v, v, v, v, v, v, v, v, v, v, v, v
, a, a, a, a, a, a, a, a, a, a, a, a, a, a
, Concrete (TKR n r), w, w, w
, a, a, a, a, a )
, ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a
, AstTensor AstMethodLet FullSpan (TKR m r), AstTensor AstMethodLet FullSpan (TKR m r)
, v, v, v, v, v, v, v, v, v, v, v, v, v, v
, a, a, a, a, a, a, a, a, a, a, a, a, a, a
, Concrete (TKR n r), w, w, w
, a, a, a, a, a )
)
{-# NOINLINE rev' #-}
rev' f vals = unsafePerformIO $ do
setTotalSharing False
!resNormalSharing <- rev1 f vals
setTotalSharing True
!resTotalSharing <- rev1 f vals
setTotalSharing False
return (resNormalSharing, resTotalSharing)
rev1 :: forall r m n v a w.
( KnownNat n, NumScalar r, NumScalar (ADTensorScalar r)
, v ~ Concrete (TKR m r)
, w ~ Concrete (ADTensorKind (TKR m r))
, a ~ Concrete (ADTensorKind (TKR n r)) )
=> (forall f. ADReady f => f (TKR n r) -> f (TKR m r))
-> Concrete (TKR n r)
-> IO ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a
, AstTensor AstMethodLet FullSpan (TKR m r), AstTensor AstMethodLet FullSpan (TKR m r)
, v, v, v, v, v, v, v, v, v, v, v, v, v, v
, a, a, a, a, a, a, a, a, a, a, a, a, a, a
, Concrete (TKR n r), w, w, w
, a, a, a, a, a )
{-# NOINLINE rev1 #-}
rev1 f !vals = do
let !value0 = f vals
ftk = tftk knownSTK vals
ftkz = adFTK ftk
g :: ADVal Concrete (TKR n r)
-> ADVal Concrete (TKR m r)
g inputs = f $ fromTarget inputs
!(!value1, !gradient1) = crevMaybeBoth g vals
gradientRrev1 = rrev1 @Concrete @r @n @m f vals
secondRrev1 = rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))
ftkz (rrev1 @_ @r @n @m f) vals
g9 :: ADVal (AstRaw FullSpan) (TKR n r)
-> ADVal (AstRaw FullSpan) (TKR m r)
g9 inputs = f @(ADVal (AstRaw FullSpan))
$ fromTarget inputs
-- fromTarget is fine, because primal of inputs is a variable,
-- hence it's duplicable
artifactsGradAst9 =
fst $ revProduceArtifactWithoutInterpretation
IgnoreIncomingCotangent g9 ftk
!(!value9, !gradient9) = revInterpretArtifact7 artifactsGradAst9
revInterpretArtifact7
:: AstArtifactRev (TKR n r) (TKR m r)
-> (Concrete (TKR m r), Concrete (ADTensorKind (TKR n r)))
revInterpretArtifact7 a1 = revInterpretArtifact a1 vals Nothing
hGeneral
:: (ADReady fgen, ADReady f1)
=> (f1 (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))
-> (AstTensor AstMethodLet FullSpan (TKR n r) -> f1 (TKR n r))
-> (AstTensor AstMethodLet FullSpan (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))
-> fgen (TKR n r)
-> fgen (TKR m r)
hGeneral fx1 fx2 gx inputs =
let (var, ast) = funToAst (FTKR (rshape vals) FTKScalar) (fx1 . f . fx2)
env = extendEnv var inputs emptyEnv
in interpretAstFull env (gx ast)
h :: ADReady f1
=> (f1 (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))
-> (AstTensor AstMethodLet FullSpan (TKR n r) -> f1 (TKR n r))
-> (AstTensor AstMethodLet FullSpan (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))
-> ADVal Concrete (TKR n r)
-> ADVal Concrete (TKR m r)
h fx1 fx2 gx inputs =
hGeneral @(ADVal Concrete) fx1 fx2 gx
(fromTarget inputs)
!(!value2, !gradient2) =
crevMaybeBoth (h id id id) vals
!(!value3, !gradient3) =
crevMaybeBoth (h id id simplifyInlineContract) vals
!(!value2UnSimp, !gradient2UnSimp) =
crevMaybeBoth (h unAstNoSimplify AstNoSimplify id) vals
gradientRrev2UnSimp =
rrev1 @Concrete @r @n @m @r
(hGeneral unAstNoSimplify AstNoSimplify id) vals
secondRrev2UnSimp =
rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))
ftkz (rrev1 @_ @r @n @m @r
(hGeneral unAstNoSimplify AstNoSimplify id)) vals
!(!value3UnSimp, !gradient3UnSimp) =
crevMaybeBoth (h (simplifyUserCode . unAstNoSimplify) AstNoSimplify simplifyInlineContract)
vals
gradientRrev3UnSimp =
rrev1 @Concrete @r @n @m @r
(hGeneral unAstNoSimplify AstNoSimplify simplifyInlineContract) vals
secondRrev3UnSimp =
rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))
ftkz (rrev1 @_ @r @n @m @r
(hGeneral unAstNoSimplify AstNoSimplify simplifyInlineContract)) vals
!(!value4, !gradient4) =
crevMaybeBoth (h unAstNoVectorize AstNoVectorize id)
vals
-- use the AstNoVectorize instance that does no vectorization
-- and then interpret the results as the Ast instance
gradientRrev4 =
rrev1 @Concrete @r @n @m @r
(hGeneral unAstNoVectorize AstNoVectorize id) vals
secondRrev4 =
rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))
ftkz (rrev1 @_ @r @n @m @r
(hGeneral unAstNoVectorize AstNoVectorize id)) vals
!(!value5, !gradient5) =
crevMaybeBoth (h unAstNoVectorize AstNoVectorize simplifyInlineContract)
vals
gradientRrev5 =
rrev1 @Concrete @r @n @m @r
(hGeneral unAstNoVectorize AstNoVectorize simplifyInlineContract) vals
secondRrev5 =
rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))
ftkz (rrev1 @_ @r @n @m @r
(hGeneral unAstNoVectorize AstNoVectorize simplifyInlineContract)) vals
astVectSimp = simplifyInlineContract $ snd $ funToAst (FTKR (rshape vals) FTKScalar) f
astSimp =
simplifyInlineContract $ simplifyInlineContract $ snd -- builds simplify with difficulty
$ funToAst (FTKR (rshape vals) FTKScalar) (unAstNoVectorize . f . AstNoVectorize)
-- Here comes the part with Ast gradients.
hAst :: ADReady f1
=> (f1 (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))
-> (AstTensor AstMethodLet FullSpan (TKR n r) -> f1 (TKR n r))
-> (AstTensor AstMethodLet FullSpan (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))
-> ADVal (AstRaw FullSpan) (TKR n r)
-> ADVal (AstRaw FullSpan) (TKR m r)
hAst fx1 fx2 gx inputs
= hGeneral @(ADVal (AstRaw FullSpan))
fx1 fx2 gx (fromTarget inputs)
-- fromTarget is fine, because primal of inputs is a variable,
-- hence it's duplicable
artifactsGradAst =
fst $ revProduceArtifactWithoutInterpretation
IgnoreIncomingCotangent (hAst id id id) ftk
!(!value2Ast, !gradient2Ast) =
revInterpretArtifact7 artifactsGradAst
!(!value2AstS, !gradient2AstS) =
revInterpretArtifact7 (simplifyArtifactRev artifactsGradAst)
artifactsGradAstT =
fst $ revProduceArtifactWithoutInterpretation
UseIncomingCotangent (hAst simplifyUserCode id id) ftk
!(!value2AstST, !gradient2AstST) =
revInterpretArtifact7 (simplifyArtifactRev artifactsGradAstT)
artifactsSimpleAst =
fst $ revProduceArtifactWithoutInterpretation
IgnoreIncomingCotangent (hAst id id simplifyInlineContract) ftk
!(!value3Ast, !gradient3Ast) =
revInterpretArtifact7 artifactsSimpleAst
!(!value3AstS, !gradient3AstS) =
revInterpretArtifact7 (simplifyArtifactRev artifactsSimpleAst)
artifactsGradAstUnSimp =
fst $ revProduceArtifactWithoutInterpretation
IgnoreIncomingCotangent (hAst unAstNoSimplify AstNoSimplify id) ftk
!(!value2AstUnSimp, !gradient2AstUnSimp) =
revInterpretArtifact7 artifactsGradAstUnSimp
!(!value2AstSUnSimp, !gradient2AstSUnSimp) =
revInterpretArtifact7 (simplifyArtifactRev artifactsGradAstUnSimp)
artifactsSimpleAstUnSimp =
fst $ revProduceArtifactWithoutInterpretation
IgnoreIncomingCotangent (hAst unAstNoSimplify AstNoSimplify simplifyInlineContract)
ftk
!(!value3AstUnSimp, !gradient3AstUnSimp) =
revInterpretArtifact7 artifactsSimpleAstUnSimp
!(!value3AstSUnSimp, !gradient3AstSUnSimp) =
revInterpretArtifact7 (simplifyArtifactRev artifactsSimpleAstUnSimp)
artifactsPrimalAst =
fst $ revProduceArtifactWithoutInterpretation
IgnoreIncomingCotangent (hAst unAstNoVectorize AstNoVectorize id) ftk
!(!value4Ast, !gradient4Ast) =
revInterpretArtifact7 artifactsPrimalAst
!(!value4AstS, !gradient4AstS) =
revInterpretArtifact7 (simplifyArtifactRev artifactsPrimalAst)
artifactsPSimpleAst =
fst $ revProduceArtifactWithoutInterpretation
IgnoreIncomingCotangent (hAst unAstNoVectorize AstNoVectorize simplifyInlineContract)
ftk
!(!value5Ast, !gradient5Ast) =
revInterpretArtifact7 artifactsPSimpleAst
-- Due to no vectorization this may result in huge terms,
-- which then take forever to inline into (substitution into huge term)
-- and to simplify, so we ignore this test for now.
(value5AstS, gradient5AstS) =
revInterpretArtifact7 (simplifyArtifactRev artifactsPSimpleAst)
!cderivative = cjvp f vals vals
!derivative = jvp f vals vals
!derivativeRfwd1 = rfwd1ds @Concrete @r @n @m @r f vals
$ toADTensorKindShared ftk vals
return
( value0, value1, value2, value3, value2UnSimp, value3UnSimp
, value4, value5
, gradient1, gradientRrev1, gradient2, gradient3
, gradient2UnSimp, gradientRrev2UnSimp
, gradient3UnSimp, gradientRrev3UnSimp
, gradient4, gradientRrev4, gradient5, gradientRrev5
, astVectSimp, astSimp
, value9, value2Ast, value2AstS, value2AstST, value3Ast, value3AstS
, value2AstUnSimp, value2AstSUnSimp, value3AstUnSimp, value3AstSUnSimp
, value4Ast, value4AstS, value5Ast, value5AstS
, gradient9, gradient2Ast, gradient2AstS, gradient2AstST
, gradient3Ast, gradient3AstS
, gradient2AstUnSimp, gradient2AstSUnSimp
, gradient3AstUnSimp, gradient3AstSUnSimp
, gradient4Ast, gradient4AstS, gradient5Ast, gradient5AstS
, vals, cderivative, derivative, derivativeRfwd1
, secondRrev1, secondRrev2UnSimp, secondRrev3UnSimp
, secondRrev4, secondRrev5
)
assertEqualUpToEpsilon'
:: ( KnownNat n, KnownNat m
, v ~ Concrete (TKR m r)
, w ~ Concrete (ADTensorKind (TKR m r))
, a ~ Concrete (ADTensorKind (TKR n r))
, AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v
, AssertEqualUpToEpsilon (ADTensorScalar r)
, GoodScalar r, NumScalar (ADTensorScalar r), HasCallStack)
=> Rational -- ^ error margin (i.e., the epsilon)
-> Concrete (TKR n r) -- ^ expected reverse derivative value
-> ( ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a
, AstTensor AstMethodLet FullSpan (TKR m r), AstTensor AstMethodLet FullSpan (TKR m r)
, v, v, v, v, v, v, v, v, v, v, v, v, v, v
, a, a, a, a, a, a, a, a, a, a, a, a, a, a
, Concrete (TKR n r), w, w, w
, a, a, a, a, a )
, ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a
, AstTensor AstMethodLet FullSpan (TKR m r), AstTensor AstMethodLet FullSpan (TKR m r)
, v, v, v, v, v, v, v, v, v, v, v, v, v, v
, a, a, a, a, a, a, a, a, a, a, a, a, a, a
, Concrete (TKR n r), w, w, w
, a, a, a, a, a ) )
-> Assertion
{-# INLINE assertEqualUpToEpsilon' #-}
assertEqualUpToEpsilon' errMargin expected' (tup1, tup2) = do
assertEqualUpToEpsilon1 errMargin expected' tup1
assertEqualUpToEpsilon1 errMargin expected' tup2
assertEqualUpToEpsilon1
:: ( KnownNat n, KnownNat m
, v ~ Concrete (TKR m r)
, w ~ Concrete (ADTensorKind (TKR m r))
, a ~ Concrete (ADTensorKind (TKR n r))
, AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v
, AssertEqualUpToEpsilon (ADTensorScalar r)
, GoodScalar r, NumScalar (ADTensorScalar r), HasCallStack)
=> Rational -- ^ error margin (i.e., the epsilon)
-> Concrete (TKR n r) -- ^ expected reverse derivative value
-> ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a
, AstTensor AstMethodLet FullSpan (TKR m r), AstTensor AstMethodLet FullSpan (TKR m r)
, v, v, v, v, v, v, v, v, v, v, v, v, v, v
, a, a, a, a, a, a, a, a, a, a, a, a, a, a
, Concrete (TKR n r), w, w, w
, a, a, a, a, a )
-- ^ actual values
-> Assertion
assertEqualUpToEpsilon1
errMargin expected'
( value0, value1, value2, value3, value2UnSimp, value3UnSimp
, value4, value5
, gradient1, gradientRrev1, gradient2, gradient3
, gradient2UnSimp, gradientRrev2UnSimp
, gradient3UnSimp, gradientRrev3UnSimp
, gradient4, gradientRrev4, gradient5, gradientRrev5
, _astVectSimp, _astSimp
, value9, value2Ast, value2AstS, value2AstST, value3Ast, value3AstS
, value2AstUnSimp, value2AstSUnSimp, value3AstUnSimp, value3AstSUnSimp
, value4Ast, value4AstS, value5Ast, _value5AstS
, gradient9, gradient2Ast, gradient2AstS, gradient2AstST
, gradient3Ast, gradient3AstS
, gradient2AstUnSimp, gradient2AstSUnSimp
, gradient3AstUnSimp, gradient3AstSUnSimp
, gradient4Ast, gradient4AstS, gradient5Ast, _gradient5AstS
, vals, cderivative, derivative, derivativeRfwd1
, secondRrev1, secondRrev2UnSimp, secondRrev3UnSimp
, secondRrev4, secondRrev5 ) = do
let ftk = tftk knownSTK vals
expected = toADTensorKindShared ftk expected'
assertEqualUpToEpsilonWithMark "Val ADVal" errMargin value0 value1
assertEqualUpToEpsilonWithMark "Val Vectorized" errMargin value0 value2
assertEqualUpToEpsilonWithMark "Val Vect+Simp" errMargin value0 value3
assertEqualUpToEpsilonWithMark "Val V UnS" errMargin value0 value2UnSimp
assertEqualUpToEpsilonWithMark "Val V+S UnS" errMargin value0 value3UnSimp
assertEqualUpToEpsilonWithMark "Val NotVect" errMargin value0 value4
assertEqualUpToEpsilonWithMark "Val Simplified" errMargin value0 value5
assertEqualUpToEpsilonWithMark "Grad ADVal" errMargin expected gradient1
assertEqualUpToEpsilonWithMark "Grad ADVal rrev"
errMargin expected gradientRrev1
assertEqualUpToEpsilonWithMark "Grad Vectorized" errMargin expected gradient2
assertEqualUpToEpsilonWithMark "Grad Vect+Simp" errMargin expected gradient3
assertEqualUpToEpsilonWithMark "Grad V UnS" errMargin expected gradient2UnSimp
assertEqualUpToEpsilonWithMark "Grad V UnS rrev2"
errMargin expected gradientRrev2UnSimp
assertEqualUpToEpsilonWithMark "Second V UnS rrev2"
(max 1e-5 errMargin) secondRrev1 secondRrev2UnSimp
assertEqualUpToEpsilonWithMark "Grad V+S UnS"
errMargin expected gradient3UnSimp
assertEqualUpToEpsilonWithMark "Grad V+S UnS rrev"
errMargin expected gradientRrev3UnSimp
assertEqualUpToEpsilonWithMark "Second V+S UnS rrev"
(max 1e-5 errMargin) secondRrev1 secondRrev3UnSimp
assertEqualUpToEpsilonWithMark "Grad NotVect" errMargin expected gradient4
assertEqualUpToEpsilonWithMark "Grad NotVect rrev"
errMargin expected gradientRrev4
assertEqualUpToEpsilonWithMark "Second NotVect rrev"
(max 1e-5 errMargin) secondRrev1 secondRrev4
assertEqualUpToEpsilonWithMark "Grad Simplified" errMargin expected gradient5
assertEqualUpToEpsilonWithMark "Grad Simplified rrev2"
errMargin expected gradientRrev5
assertEqualUpToEpsilonWithMark "Second Simplified rrev2"
(max 1e-5 errMargin) secondRrev1 secondRrev5
assertEqualUpToEpsilonWithMark "Val Ast Vectorized" errMargin value0 value2Ast
assertEqualUpToEpsilonWithMark "Val Ast V S" errMargin value0 value2AstS
assertEqualUpToEpsilonWithMark "Val Ast V ST" errMargin value0 value2AstST
assertEqualUpToEpsilonWithMark "Val Ast Vect+Simp" errMargin value0 value3Ast
assertEqualUpToEpsilonWithMark "Val Ast V+S S" errMargin value0 value3AstS
assertEqualUpToEpsilonWithMark "Val Ast V UnS" errMargin value0
value2AstUnSimp
assertEqualUpToEpsilonWithMark "Val Ast V S UnS" errMargin value0
value2AstSUnSimp
assertEqualUpToEpsilonWithMark "Val Ast Vect+Simp UnS" errMargin value0
value3AstUnSimp
assertEqualUpToEpsilonWithMark "Val Ast V+S S UnS" errMargin value0
value3AstSUnSimp
assertEqualUpToEpsilonWithMark "Val Ast NotVect" errMargin value0 value4Ast
assertEqualUpToEpsilonWithMark "Val Ast NotVect S" errMargin value0 value4AstS
assertEqualUpToEpsilonWithMark "Val Ast Simplified" errMargin value0 value5Ast
-- assertEqualUpToEpsilonWithMark "Val Ast S S" errMargin value0 value5AstS
assertEqualUpToEpsilonWithMark "Grad Ast Vectorized"
errMargin expected gradient2Ast
assertEqualUpToEpsilonWithMark "Grad Ast Vectorized S"
errMargin expected gradient2AstS
assertEqualUpToEpsilonWithMark "Grad Ast Vectorized ST"
errMargin expected gradient2AstST
assertEqualUpToEpsilonWithMark "Grad Ast Vect+Simp"
errMargin expected gradient3Ast
assertEqualUpToEpsilonWithMark "Grad Ast Vect+Simp S"
errMargin expected gradient3AstS
assertEqualUpToEpsilonWithMark "Grad Ast Vectorized UnS"
errMargin expected gradient2AstUnSimp
assertEqualUpToEpsilonWithMark "Grad Ast Vectorized S UnS"
errMargin expected gradient2AstSUnSimp
assertEqualUpToEpsilonWithMark "Grad Ast Vect+Simp UnS"
errMargin expected gradient3AstUnSimp
assertEqualUpToEpsilonWithMark "Grad Ast Vect+Simp S UnS"
errMargin expected gradient3AstSUnSimp
assertEqualUpToEpsilonWithMark "Grad Ast NotVect"
errMargin expected gradient4Ast
assertEqualUpToEpsilonWithMark "Grad Ast NotVect S"
errMargin expected gradient4AstS
assertEqualUpToEpsilonWithMark "Grad Ast Simplified"
errMargin expected gradient5Ast
-- assertEqualUpToEpsilonWithMark "Grad Ast Simplified S"
-- errMargin expected gradient5AstS
assertEqualUpToEpsilonWithMark "Val ADVal Ast" errMargin value0 value9
assertEqualUpToEpsilonWithMark "Grad ADVal Ast" errMargin expected gradient9
assertEqualUpToEpsilonWithMark "Derivatives" errMargin cderivative derivative
assertEqualUpToEpsilonWithMark "Derivatives rjvp"
errMargin cderivative derivativeRfwd1
-- The formula for comparing derivative and gradient is due to @awf
-- at https://github.com/Mikolaj/horde-ad/issues/15#issuecomment-1063251319
-- and a similar property stated mathematically is in Lemma 1 in
-- https://www.microsoft.com/en-us/research/uploads/prod/2021/08/higher-order-ad.pdf
assertEqualUpToEpsilonWithMark "Reverse vs forward"
1e-5 (rfromK $ rsum0 derivative) (rfromK $ rdot0 expected (toADTensorKindShared ftk vals))
{- TODO: this most probably leaks gigabytes of strings from one test case
-- to another in -O0 mode, leading to OOMs, so it's disabled for now.
-- We could also try to stream the strings and compare on the fly.
--
-- No Eq instance, so let's compare the text.
assertEqual "Idempotence of simplification of non-vectorized AST"
(show astSimp)
(show (simplifyInlineContract astSimp))
assertEqual "Idempotence of simplification of vectorized AST"
(show astVectSimp)
(show (simplifyInlineContract astVectSimp))
-}
t16 :: (GoodScalar r, Fractional r) => Concrete (TKR 5 r)
t16 = ringestData (fromList [2, 2, 1, 2, 2]) [5, 2, 6, 1, -2, 0.000001, 0.1, -0.2, 13.1, 9, 8, -4, 34, 2.99432, -33, 26]
t16b :: (GoodScalar r, Fractional r) => Concrete (TKR 4 r)
t16b = ringestData (fromList [2, 2, 2, 2]) [5, 2, 6, 1, -2, 0, 0.1, -0.2, 13.1, 9, 8, -4, 582934, 2.99432, -335, 26]
t48 :: (GoodScalar r, Fractional r) => Concrete (TKR 7 r)
t48 = ringestData (fromList [3, 1, 2, 2, 1, 2, 2]) [18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001, 5, 2, 6, 1, -2, 0.92, 0.1, -0.2, 13.1, 9, 8, -4, 34, 2.99432, -33, 26, 2, 2, 2, 2, -0.2,-0.2,-0.2,-0.2,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003]
t128 :: (GoodScalar r, Fractional r) => Concrete (TKR 10 r)
t128 = ringestData (fromList [1, 2, 2, 1, 2, 2, 2, 2, 2, 1]) [29.1,32.1,40.1,29.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,32.0,53.99432,97.1,25.8943200001, 5, 2, 6, 1, -2, 97.1,58.8943200001,97.1,55.8943200001,97.1,58.8943200001,18.1,29.1,32.1,40.1,32.1,32.1,40.1,53.0,53.99432, -0.00001, 0.1, -0.2, 13.1, 9, 8, -4, 29, 2.99432, -335, 26, 2, 2, 2, 2, -0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,40.1,8.0,11.0,-3.0,25.89432,28.79432,-39.09999999999997,25.8,40.1,8.0,11.0,-3.0,25.89432,28.79432,-19.09999999999997,25.8, 8.1,29.1,32.1,40.1,32.1,40.1,292.0,53.99432,97.1,55.8943200001,97.1,85.8943200001,97.1,85.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,22.0,53.99432,97.1,82.8943200001,97.1,22.8943200001,97.1,58.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,89.0,53.99432,97.1,56.8943200001,97.1,52.8943200001,97.1,55.8943200001]
t128b :: (GoodScalar r, Fractional r) => Concrete (TKR 4 r)
t128b = rreshape (4 :$: 2 :$: 4 :$: 4 :$: ZSR) t128
t128c :: (GoodScalar r, Fractional r) => Concrete (TKR 4 r)
t128c = rreshape (2 :$: 2 :$: 8 :$: 4 :$: ZSR) t128
rrev1 :: forall g r n m r3.
(ADReady g, GoodScalar r, KnownNat n, NumScalar r3)
=> (forall f. ADReady f => f (TKR n r) -> f (TKR m r3)) -> g (TKR n r)
-> g (ADTensorKind (TKR n r))
{-# INLINE rrev1 #-}
rrev1 f u = kgrad (rsum0 . f) (tftk knownSTK u) u
rrevFTK :: forall g x z. (ADReady g, KnownSTK x, TKAllNum z)
=> FullShapeTK z -> (forall f. ADReady f => f x -> f z) -> g x
-> g (ADTensorKind x)
{-# INLINE rrevFTK #-}
rrevFTK ftk f u = kgrad (tsum0Target ftk . f) (tftk knownSTK u) u
rfwd1ds :: forall g r n m r3.
(ADReady g, GoodScalar r, KnownNat n)
=> (forall f. ADReady f => f (TKR n r) -> f (TKR m r3)) -> g (TKR n r)
-> g (ADTensorKind (TKR n r))
-> g (ADTensorKind (TKR m r3))
{-# INLINE rfwd1ds #-}
rfwd1ds f u = rjvp f (tftk knownSTK u) u
rfwd1 :: forall g r n m r3.
(ADReady g, GoodScalar r, NumScalar (ADTensorScalar r), KnownNat n)
=> (forall f. ADReady f => f (TKR n r) -> f (TKR m r3)) -> g (TKR n r)
-> g (ADTensorKind (TKR m r3))
{-# INLINE rfwd1 #-}
rfwd1 f u = rfwd1ds f u (rrepl (rshape u) 1)
srev1 :: forall g r sh sh2 r3.
(ADReady g, GoodScalar r, KnownShS sh, NumScalar r3)
=> (forall f. ADReady f => f (TKS sh r) -> f (TKS sh2 r3)) -> g (TKS sh r)
-> g (ADTensorKind (TKS sh r))
{-# INLINE srev1 #-}
srev1 f u = kgrad (ssum0 . f) (tftk knownSTK u) u
sfwd1 :: forall g r sh sh2 r3.
(ADReady g, GoodScalar r, NumScalar (ADTensorScalar r), KnownShS sh)
=> (forall f. ADReady f => f (TKS sh r) -> f (TKS sh2 r3)) -> g (TKS sh r)
-> g (ADTensorKind (TKS sh2 r3))
{-# INLINE sfwd1 #-}
sfwd1 f u = sjvp f (tftk knownSTK u) u (srepl @_ @(ADTensorScalar r) 1)