packages feed

horde-ad 0.2.0.0 → 0.3.0.0

raw patch · 76 files changed

+27424/−20473 lines, 76 filesdep −indexed-traversabledep ~basedep ~dependent-enummapdep ~ghc-typelits-knownnat

Dependencies removed: indexed-traversable

Dependency ranges changed: base, dependent-enummap, ghc-typelits-knownnat, ghc-typelits-natnormalise, orthotope, ox-arrays, tasty, vector

Files

CHANGELOG.md view
@@ -1,11 +1,17 @@ # Revision history for horde-ad -## [v0.1.0.0](https://github.com/Mikolaj/horde-ad/compare/v0.0.0.0...v0.1.0.0)+## [v0.3.0.0](https://github.com/Mikolaj/horde-ad/compare/v0.2.0.0...v0.3.0.0) -- First version. Released on an unsuspecting world.+- Extensive performance rework+- Update to new versions of ox-arrays, orthotope and dependent-enummap+- Minor polish of the API  ## [v0.2.0.0](https://github.com/Mikolaj/horde-ad/compare/v0.1.0.0...v0.2.0.0)  - Modernize the dep (ilist) that provides imap to make Stackage happy - Make the cabal sublibraries public, as intended - Tweak benchmarks++## [v0.1.0.0](https://github.com/Mikolaj/horde-ad/compare/v0.0.0.0...v0.1.0.0)++- First version. Released on an unsuspecting world.
CREDITS.md view
@@ -6,7 +6,7 @@  Simon Peyton Jones, who sketched in "Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation" talk at https://www.youtube.com/watch?v=qGkgLtpseMs algorithms that this Haskell codebase implements -Faustyna Krawiec, Neel Krishnaswami, Tom Ellis, Andrew Fitzgibbon, Richard Eisenberg, the remaining authors of the paper "Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation" from POPL 2022 that describes and proves correct a comprehensive AD formalism, including some ingenious algorithmic ideas this Haskell codebase implements+Faustyna Krawiec, Neel Krishnaswami, Tom Ellis, Andrew Fitzgibbon, Richard Eisenberg, the remaining authors of the paper "Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation" (https://dl.acm.org/doi/10.1145/3498710) that describes and proves correct a comprehensive AD formalism, including some ingenious algorithmic ideas this Haskell codebase implements  Oleg Grenrus, whose https://hackage.haskell.org/package/overloaded contains inspiring examples related to AD
README.md view
@@ -52,10 +52,10 @@ ```  Instantiated to matrices, `foo` now returns a matrix, not a scalar — but a gradient can only be computed of a function that returns a scalar.-To remediate this, let's sum the whole output of `foo` and only then compute its gradient: (note that `ssum0` returns a zero-dimensional array; `kfromS` extracts the (single) scalar from that)+To remediate this, let's sum the whole output of `foo` and only then compute its gradient: ```hs gradSumFooMatrix :: ThreeMatrices Double -> ThreeMatrices Double-gradSumFooMatrix = cgrad (kfromS . ssum0 . foo)+gradSumFooMatrix = cgrad (ssum0 . foo) ```  This works as well as before:@@ -102,7 +102,7 @@        ((m4 * m5) * dret + m4 * dret) ``` -A concrete value of this symbolic reverse derivative at the same input as before can be obtained by interpreting its program in the context of the operations supplied by the horde-ad library. (Note that the output happens to be the same as `gradSumFooMatrix threeSimpleMatrices` above, which used `cgrad` on `kfromS . ssum0 . foo`; the reason is that `srepl 1.0` happens to be the reverse derivative of `kfromS . ssum0`.)+A concrete value of this symbolic reverse derivative at the same input as before can be obtained by interpreting its program in the context of the operations supplied by the horde-ad library. (Note that the output happens to be the same as `gradSumFooMatrix threeSimpleMatrices` above, which used `cgrad` on `ssum0 . foo`; the reason is that `srepl 1.0` happens to be the reverse derivative of `ssum0`.) ```hs >>> vjpInterpretArtifact artifact (toTarget threeSimpleMatrices) (srepl 1.0) ((sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627]) :: ThreeConcreteMatrices Double)@@ -112,7 +112,7 @@  A shorthand that creates a symbolic gradient program, simplifies it and interprets it with a given input on the default CPU backend is called `grad` and is used exactly the same as (but with often much better performance on the same program than) `cgrad`: ```hs->>> grad (kfromS . ssum0 . fooLet) threeSimpleMatrices+>>> grad (ssum0 . fooLet) threeSimpleMatrices (sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627]) ``` @@ -247,7 +247,7 @@ The tests in this suite also don't contribute big Haskell CAFs, which makes them more reliable micro-benchmarks. Ordinary benchmarks that use package criterion are even more trustworthy,-but they don't have a comparable coverage at this time.+but the included set doesn't have a comparable coverage at this time.   Coding style@@ -273,6 +273,6 @@ Copyright --------- -Copyright 2023--2025 Mikolaj Konarski, Well-Typed LLP and others (see git history)+Copyright 2023--2026 Mikolaj Konarski, Well-Typed LLP and others (see git history)  License: BSD-3-Clause (see file LICENSE)
bench/LongProdBench.hs view
@@ -16,9 +16,10 @@ main =   defaultMain  -- skips the tiny benchmarks     [ bgroup1e4 allxs-{- OOMs, probably for a good reason (huge terms); TODO: diagnose and enable 1e5     , bgroup1e5 allxs     , bgroup1e6 allxs+{- OOMs, probably for a good reason (huge terms); TODO: diagnose and enable 1e7+   and enable more cases also for 1e5 and 1e6     , bgroup1e7 allxs     , bgroup5e7 allxs -}     ]
bench/common/BenchMnistTools.hs view
@@ -10,8 +10,9 @@ import Criterion.Main import Data.Default qualified as Default import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:)) import GHC.Exts (WithDict)-import GHC.TypeLits (KnownNat)+import GHC.TypeLits (KnownNat, type (<=?)) import System.Random import Test.Inspection import Type.Reflection (Typeable)@@ -23,6 +24,7 @@  import Data.Array.Nested qualified as Nested import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)  import MnistData import MnistFcnnRanked1 qualified@@ -47,9 +49,11 @@   withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->   withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+    (stkOfIxR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+    (stkOfIxR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden) :~: True) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden2) :~: True) $   let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters                     Concrete widthHidden widthHidden2 r       valsInit = fst $ randomValue 1 (mkStdGen 44)@@ -83,9 +87,11 @@   withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->   withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+    (stkOfIxR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+    (stkOfIxR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden) :~: True) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden2) :~: True) $   let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters                     Concrete widthHidden widthHidden2 r       valsInit = fst $ randomValue 1 (mkStdGen 44)@@ -141,9 +147,11 @@   withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->   withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+    (stkOfIxR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+    (stkOfIxR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden) :~: True) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden2) :~: True) $   let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters                     Concrete widthHidden widthHidden2 r       valsInit = fst $ randomValue 1 (mkStdGen 44)@@ -179,7 +187,7 @@             widthHiddenSNat widthHidden2SNat             (glyphR, labelR) pars         artRaw = gradArtifact f (valsInit, dataInit)-        art = simplifyArtifactGradient artRaw+        art = simplifyArtifactRev artRaw         go :: [MnistDataLinearR r]            -> Concrete (XParams widthHidden widthHidden2 r)            -> Concrete (XParams widthHidden widthHidden2 r)@@ -187,7 +195,7 @@         go ((glyph, label) : rest) !parameters =           let parametersAndInput =                 tpair parameters (tpair (rconcrete glyph) (rconcrete label))-              gradient = tproject1 $ fst+              gradient = tproject1 $ snd                          $ revInterpretArtifact art parametersAndInput Nothing           in go rest (updateWithGradient gamma knownSTK parameters gradient)         chunk = take batchSize xs@@ -319,7 +327,7 @@   -> Benchmark mnistTrainBench2VTC prefix widthHidden widthHidden2 =   bench prefix-  $ whnf (simplifyArtifactGradient . snd+  $ whnf (simplifyArtifactRev . snd           . MnistFcnnRanked2.mnistTrainBench2VTOGradient               @Double (Proxy @Float) IgnoreIncomingCotangent               1 (mkStdGen 44) widthHidden)@@ -355,7 +363,7 @@         go ((glyph, label) : rest) !parameters =           let parametersAndInput =                 tpair parameters (tpair (rconcrete glyph) (rconcrete label))-              gradient = tproject1 $ fst+              gradient = tproject1 $ snd                          $ revInterpretArtifact art parametersAndInput Nothing           in go rest (updateWithGradient gamma knownSTK parameters gradient)         chunk = take batchSize xs@@ -379,7 +387,7 @@         MnistFcnnRanked2.mnistTrainBench2VTOGradient           @Double (Proxy @Float) IgnoreIncomingCotangent           1 (mkStdGen 44) widthHidden widthHidden2-      !art = simplifyArtifactGradient artRaw+      !art = simplifyArtifactRev artRaw   in mnistTrainBench2VTOO prefix gamma batchSize xs (targetInit, art)  mnistBGroup2VTO :: Int -> Benchmark@@ -388,7 +396,7 @@         MnistFcnnRanked2.mnistTrainBench2VTOGradient           @Double (Proxy @Float) IgnoreIncomingCotangent           1 (mkStdGen 44) 1500 500-      !art = simplifyArtifactGradient artRaw  -- no NFData for AST+      !art = simplifyArtifactRev artRaw  -- no NFData for AST   in env (do     testData0 <- loadMnistData testGlyphsPath testLabelsPath  -- 10k total     let testData = shuffle (mkStdGen 42) testData0@@ -419,7 +427,7 @@         go ((glyph, label) : rest) !parameters =           let parametersAndInput =                 tpair parameters (tpair (rconcrete glyph) (rconcrete label))-              gradient = tproject1 $ fst+              gradient = tproject1 $ snd                          $ revInterpretArtifact art parametersAndInput Nothing           in go rest (updateWithGradient gamma knownSTK parameters gradient)         chunk = take batchSize xs@@ -468,7 +476,7 @@         go ((glyph, label) : rest) !parameters =           let parametersAndInput =                 tpair parameters (tpair (rconcrete glyph) (rconcrete label))-              gradient = tproject1 $ fst+              gradient = tproject1 $ snd                          $ revInterpretArtifact art parametersAndInput Nothing           in go rest (updateWithGradient gamma knownSTK parameters gradient)         chunk = take batchSize xs@@ -504,10 +512,10 @@ -- -- The `Storable` is only needed for overloaded profiling, e.g., with -- cabal bench longMnistBench -ftest_seq -w /home/mikolaj/r/ghc.HEAD/ghc/_build/stage1/bin/ghc --allow-newer --enable-optimization --enable-profiling --profiling-detail=none --ghc-options="-fprof-late-overloaded -fpolymorphic-specialisation" --benchmark-options='-n1 -m pattern "1 VTA MNIST nn with samples: 400/500|150 v" +RTS -pj'-inspect $ hasNoTypeClassesExcept 'mnistTrainBench1VTA [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Storable]-inspect $ hasNoTypeClassesExcept 'mnistTrainBench1VTO [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Storable,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Elt, ''Integral]-inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTA [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Storable]-inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTC [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default]-inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTO [''(~), ''GoodScalar, ''Show, ''Num, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default,       ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''WithDict, ''KnownShS, ''KnownSTK, ''KnownNat, ''Nested.Elt, ''Integral]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench1VTA [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Storable]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench1VTO [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Storable,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Elt, ''Integral]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTA [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Storable]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTC [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTO [''(~), ''GoodScalar, ''Show, ''Num, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default,       ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''WithDict, ''KnownShS, ''KnownSTK, ''KnownNat, ''Nested.Elt, ''Integral] -- inspect $ coreOf 'mnistTrainBench1VTA -}
bench/common/BenchProdTools.hs view
@@ -25,19 +25,19 @@ import HordeAd.Core.Ops  bgroup100, bgroup1000, bgroup1e4, bgroup1e5, bgroup1e6, bgroup1e7, bgroup5e7 :: [Double] -> Benchmark-bgroup100 = envProd 100 $ \args -> bgroup "100" $ benchProd args+bgroup100 = envProd 100 $ bgroup "100" . benchProd -bgroup1000 = envProd 1000 $ \args -> bgroup "1000" $ benchProdShort args+bgroup1000 = envProd 1000 $ bgroup "1000" . benchProdShort -bgroup1e4 = envProd 1e4 $ \args -> bgroup "1e4" $ benchProdShort args+bgroup1e4 = envProd 1e4 $ bgroup "1e4" . benchProdShort -bgroup1e5 = envProd 1e5 $ \args -> bgroup "1e5" $ benchProdShortest args+bgroup1e5 = envProd 1e5 $ bgroup "1e5" . benchProdShortest -bgroup1e6 = envProd 1e6 $ \args -> bgroup "1e6" $ benchProdShortest args+bgroup1e6 = envProd 1e6 $ bgroup "1e6" . benchProdShortest -bgroup1e7 = envProd 1e7 $ \args -> bgroup "1e7" $ benchProdShortest args+bgroup1e7 = envProd 1e7 $ bgroup "1e7" . benchProdShortest -bgroup5e7 = envProd 5e7 $ \args -> bgroup "5e7" $ benchProdShortest args+bgroup5e7 = envProd 5e7 $ bgroup "5e7" . benchProdShortest   -- 5e7 == 5 * 10^7 == 0.5 * 10^8 == 0.5e8  envProd :: r ~ Double@@ -45,8 +45,8 @@         -> (forall n.             ( SNat n             , [Concrete (TKScalar r)]-            , ListR n (Concrete (TKScalar r))-            , ListR n (Concrete (TKS '[] r))+            , IxR n (Concrete (TKScalar r))+            , IxR n (Concrete (TKS '[] r))             , Concrete (TKS '[n] r) )             -> Benchmark)         -> [r]@@ -68,8 +68,8 @@   :: r ~ Double   => ( SNat n      , [Concrete (TKScalar r)]-     , ListR n (Concrete (TKScalar r))-     , ListR n (Concrete (TKS '[] r))+     , IxR n (Concrete (TKScalar r))+     , IxR n (Concrete (TKS '[] r))      , Concrete (TKS '[n] r) )   -> [Benchmark] benchProd ~(snat, list, l, lt, t) = case snat of@@ -96,15 +96,15 @@   :: r ~ Double   => ( SNat n      , [Concrete (TKScalar r)]-     , ListR n (Concrete (TKScalar r))-     , ListR n (Concrete (TKS '[] r))+     , IxR n (Concrete (TKScalar r))+     , IxR n (Concrete (TKS '[] r))      , Concrete (TKS '[n] r) )   -> [Benchmark] benchProdShort ~(snat, list, l, lt, t) = case snat of   SNat ->---    [ bench "cgrad s MapAccum" $ nf (crevSMapAccum snat) t---    , bench "grad s MapAccum" $ nf (revSMapAccum snat) t-    [ bench "cgrad k MapAccum" $ nf (crevScalarMapAccum snat) t+    [ bench "cgrad s MapAccum" $ nf (crevSMapAccum snat) t+    , bench "grad s MapAccum" $ nf (revSMapAccum snat) t+    , bench "cgrad k MapAccum" $ nf (crevScalarMapAccum snat) t     , bench "grad k MapAccum" $ nf (revScalarMapAccum snat) t --    , bench "cgrad k list" $ nf crevScalarList list --    , bench "grad k list" $ nf revScalarList list@@ -124,23 +124,23 @@   :: r ~ Double   => ( SNat n      , [Concrete (TKScalar r)]-     , ListR n (Concrete (TKScalar r))-     , ListR n (Concrete (TKS '[] r))+     , IxR n (Concrete (TKScalar r))+     , IxR n (Concrete (TKS '[] r))      , Concrete (TKS '[n] r) )   -> [Benchmark] benchProdShortest ~(snat, list, l, lt, t) = case snat of   SNat ->---    [ bench "cgrad s MapAccum" $ nf (crevSMapAccum snat) t---    , bench "grad s MapAccum" $ nf (revSMapAccum snat) t-    [ bench "cgrad k MapAccum" $ nf (crevScalarMapAccum snat) t+    [ bench "cgrad s MapAccum" $ nf (crevSMapAccum snat) t+    , bench "grad s MapAccum" $ nf (revSMapAccum snat) t+    , bench "cgrad k MapAccum" $ nf (crevScalarMapAccum snat) t     , bench "grad k MapAccum" $ nf (revScalarMapAccum snat) t --    , bench "cgrad k list" $ nf crevScalarList list --    , bench "grad k list" $ nf revScalarList list-    , bench "cgrad k L" $ nf (crevScalarL snat) l+--    , bench "cgrad k L" $ nf (crevScalarL snat) l --    , bench "grad k L" $ nf (revScalarL snat) l --    , bench "cgrad k R" $ nf (crevScalarR snat) l --    , bench "grad k R" $ nf (revScalarR snat) l-    , bench "cgrad k NotShared" $ nf (crevScalarNotShared snat) l+--    , bench "cgrad k NotShared" $ nf (crevScalarNotShared snat) l --    , bench "cgrad s L" $ nf (crevSL snat) lt --    , bench "grad s L" $ nf (revSL snat) lt --    , bench "cgrad s R" $ nf (crevSR snat) lt@@ -152,7 +152,7 @@ -- We can define sproduct if this benchmark ends up used anywhere, -- because the current codomain of gradientFromDelta rules out -- low-level hacky pipeline tricks that could avoid indexing.-multSMapAccum :: (BaseTensor target, LetTensor target, GoodScalar r)+multSMapAccum :: (BaseTensor target, LetTensor target, NumScalar r)               => SNat n -> target (TKS '[n] r) -> target (TKS '[] r) multSMapAccum SNat = sfold (*) (sscalar 1) {-# SPECIALIZE multSMapAccum :: SNat n -> ADVal Concrete (TKS '[n] Double) -> ADVal Concrete (TKS '[] Double) #-}@@ -160,166 +160,155 @@  crevSMapAccum   :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)-crevSMapAccum snat@SNat = cgrad (kfromS . multSMapAccum snat)+crevSMapAccum snat@SNat =+  cgrad  @_ @_ @_ @Concrete(kfromS . multSMapAccum snat)  revSMapAccum   :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double) revSMapAccum snat@SNat = grad (kfromS . multSMapAccum snat) -multScalarMapAccum :: forall target n r.-                      (BaseTensor target, GoodScalar r)+multScalarMapAccum :: forall target n r. (ADReady target, NumScalar r)                    => SNat n -> target (TKS '[n] r) -> target (TKScalar r)-multScalarMapAccum snat@SNat  =-  tproject1-  . tmapAccumL (Proxy @target)-     snat-     (FTKScalar @r)-     (FTKScalar @Z1)-     (FTKScalar @r)-     (let g :: forall f. ADReady f-            => f (TKScalar r) -> f (TKScalar r)-            -> f (TKProduct (TKScalar r) TKUnit)-          g !acc !e = tpair (acc * e) tunit-      in g)-     1+multScalarMapAccum snat = tfold snat STKScalar STKScalar (*) 1 {-# SPECIALIZE multScalarMapAccum :: SNat n -> ADVal Concrete (TKS '[n] Double) -> ADVal Concrete (TKScalar Double) #-} {-# SPECIALIZE multScalarMapAccum :: SNat n -> AstTensor AstMethodLet FullSpan (TKS '[n] Double) -> AstTensor AstMethodLet FullSpan (TKScalar Double) #-}  crevScalarMapAccum   :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)-crevScalarMapAccum snat@SNat = cgrad (multScalarMapAccum snat)+crevScalarMapAccum snat@SNat =+  cgrad @_ @_ @_ @Concrete (multScalarMapAccum snat)  revScalarMapAccum   :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double) revScalarMapAccum snat@SNat = grad (multScalarMapAccum snat) -multScalarList :: (BaseTensor target, GoodScalar r)+multScalarList :: (BaseTensor target, NumScalar r)                => [target (TKScalar r)] -> target (TKScalar r) multScalarList = foldl1' (*)  crevScalarList   :: [Concrete (TKScalar Double)] -> [Concrete (TKScalar Double)] crevScalarList =-  cgrad multScalarList+  cgrad @_ @_ @_ @Concrete multScalarList  revScalarList   :: [Concrete (TKScalar Double)] -> [Concrete (TKScalar Double)] revScalarList =   grad multScalarList -multScalarL :: (BaseTensor target, GoodScalar r)-            => ListR n (target (TKScalar r)) -> target (TKScalar r)+multScalarL :: (BaseTensor target, NumScalar r)+            => IxR n (target (TKScalar r)) -> target (TKScalar r) multScalarL = foldl1' (*) . Foldable.toList  crevScalarL-  :: SNat n -> ListR n (Concrete (TKScalar Double))-  -> ListR n (Concrete (TKScalar Double))+  :: SNat n -> IxR n (Concrete (TKScalar Double))+  -> IxR n (Concrete (TKScalar Double)) crevScalarL snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $-  cgrad multScalarL+  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $+  cgrad @_ @_ @_ @Concrete multScalarL  revScalarL-  :: SNat n -> ListR n (Concrete (TKScalar Double))-  -> ListR n (Concrete (TKScalar Double))+  :: SNat n -> IxR n (Concrete (TKScalar Double))+  -> IxR n (Concrete (TKScalar Double)) revScalarL snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $+  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $   grad multScalarL -multScalarR :: (BaseTensor target, GoodScalar r)-            => ListR n (target (TKScalar r)) -> target (TKScalar r)+multScalarR :: (BaseTensor target, NumScalar r)+            => IxR n (target (TKScalar r)) -> target (TKScalar r) multScalarR = foldr1 (*)  crevScalarR-  :: SNat n -> ListR n (Concrete (TKScalar Double))-  -> ListR n (Concrete (TKScalar Double))+  :: SNat n -> IxR n (Concrete (TKScalar Double))+  -> IxR n (Concrete (TKScalar Double)) crevScalarR snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $-  cgrad multScalarR+  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $+  cgrad @_ @_ @_ @Concrete multScalarR  revScalarR-  :: SNat n -> ListR n (Concrete (TKScalar Double))-  -> ListR n (Concrete (TKScalar Double))+  :: SNat n -> IxR n (Concrete (TKScalar Double))+  -> IxR n (Concrete (TKScalar Double)) revScalarR snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $+  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $   grad multScalarR -multScalarNotShared :: (BaseTensor target, GoodScalar r)-                    => ListR n (ADVal target (TKScalar r))+multScalarNotShared :: (BaseTensor target, NumScalar r)+                    => IxR n (ADVal target (TKScalar r))                     -> ADVal target (TKScalar r) multScalarNotShared = foldr1 multNotShared  crevScalarNotShared-  :: SNat n -> ListR n (Concrete (TKScalar Double))-  -> ListR n (Concrete (TKScalar Double))+  :: SNat n -> IxR n (Concrete (TKScalar Double))+  -> IxR n (Concrete (TKScalar Double)) crevScalarNotShared snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $-  cgrad multScalarNotShared+  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $+  cgrad @_ @_ @_ @Concrete multScalarNotShared -multSL :: (BaseTensor target, GoodScalar r)-       => ListR n (target (TKS '[] r)) -> target (TKS '[] r)+multSL :: (BaseTensor target, NumScalar r)+       => IxR n (target (TKS '[] r)) -> target (TKS '[] r) multSL = foldl1' (*) . Foldable.toList  crevSL-  :: SNat n -> ListR n (Concrete (TKS '[] Double))-  -> ListR n (Concrete (TKS '[] Double))+  :: SNat n -> IxR n (Concrete (TKS '[] Double))+  -> IxR n (Concrete (TKS '[] Double)) crevSL snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $-  cgrad (kfromS . multSL)+  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $+  cgrad @_ @_ @_ @Concrete (kfromS . multSL)  revSL-  :: SNat n -> ListR n (Concrete (TKS '[] Double))-  -> ListR n (Concrete (TKS '[] Double))+  :: SNat n -> IxR n (Concrete (TKS '[] Double))+  -> IxR n (Concrete (TKS '[] Double)) revSL snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $+  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $   grad (kfromS . multSL) -multSR :: (BaseTensor target, GoodScalar r)-       => ListR n (target (TKS '[] r)) -> target (TKS '[] r)+multSR :: (BaseTensor target, NumScalar r)+       => IxR n (target (TKS '[] r)) -> target (TKS '[] r) multSR = foldr1 (*)  crevSR-  :: SNat n -> ListR n (Concrete (TKS '[] Double))-  -> ListR n (Concrete (TKS '[] Double))+  :: SNat n -> IxR n (Concrete (TKS '[] Double))+  -> IxR n (Concrete (TKS '[] Double)) crevSR snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $-  cgrad (kfromS . multSR)+  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $+  cgrad @_ @_ @_ @Concrete (kfromS . multSR)  revSR-  :: SNat n -> ListR n (Concrete (TKS '[] Double))-  -> ListR n (Concrete (TKS '[] Double))+  :: SNat n -> IxR n (Concrete (TKS '[] Double))+  -> IxR n (Concrete (TKS '[] Double)) revSR snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $+  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $   grad (kfromS . multSR) -multSNotShared :: (BaseTensor target, GoodScalar r)-               => ListR n (ADVal target (TKS '[] r))+multSNotShared :: (BaseTensor target, NumScalar r)+               => IxR n (ADVal target (TKS '[] r))                -> ADVal target (TKS '[] r) multSNotShared = foldr1 multNotShared  crevSNotShared-  :: SNat n -> ListR n (Concrete (TKS '[] Double))-  -> ListR n (Concrete (TKS '[] Double))+  :: SNat n -> IxR n (Concrete (TKS '[] Double))+  -> IxR n (Concrete (TKS '[] Double)) crevSNotShared snat@SNat =-  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $-  cgrad (kfromS . multSNotShared)+  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $+  cgrad @_ @_ @_ @Concrete (kfromS . multSNotShared)  {- TODO: re-enable once -fpolymorphic-specialisation works --- KnownNat and AstSpan are tag types, so it's fine not to specialize+-- KnownNat and KnownSpan are tag types, so it's fine not to specialize -- for them. Some of the other classes come from existential types, -- some of which it's not advantageous to specialize. -- -- This is expected to fail with -O0 and to pass with -O1 -- and -fpolymorphic-specialisation. -- This prevents running benchmarks without optimization, which is a good thing.-inspect $ hasNoTypeClassesExcept 'crevScalarL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt]-inspect $ hasNoTypeClassesExcept 'revScalarL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]-inspect $ hasNoTypeClassesExcept 'crevScalarNotShared [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt]-inspect $ hasNoTypeClassesExcept 'crevSL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''Nested.Storable,    ''ShareTensor]-inspect $ hasNoTypeClassesExcept 'revSL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]-inspect $ hasNoTypeClassesExcept 'crevSMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor]-inspect $ hasNoTypeClassesExcept 'revSMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]-inspect $ hasNoTypeClassesExcept 'crevScalarMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor]-inspect $ hasNoTypeClassesExcept 'revScalarMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]+inspect $ hasNoTypeClassesExcept 'crevScalarL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt]+inspect $ hasNoTypeClassesExcept 'revScalarL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]+inspect $ hasNoTypeClassesExcept 'crevScalarNotShared [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt]+inspect $ hasNoTypeClassesExcept 'crevSL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''Nested.Storable,    ''ShareTensor]+inspect $ hasNoTypeClassesExcept 'revSL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]+inspect $ hasNoTypeClassesExcept 'crevSMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor]+inspect $ hasNoTypeClassesExcept 'revSMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]+inspect $ hasNoTypeClassesExcept 'crevScalarMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor]+inspect $ hasNoTypeClassesExcept 'revScalarMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral] -- inspect $ coreOf 'revScalarL -}
example/MnistCnnRanked2.hs view
@@ -51,14 +51,14 @@  -- | A single convolutional layer with @relu@ and @maxPool@. convMnistLayerR-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: (ADReady target, NumScalar r, Differentiable r)   => target (TKR 4 r)  -- ^ @[c_out, c_in, kh + 1, kw + 1]@   -> target (TKR 4 r)  -- ^ @[batch_size, c_in, h, w]@   -> target (TKR 1 r)  -- ^ @[c_out]@   -> target (TKR 4 r)  -- ^ @[batch_size, c_out, h \`Div\` 2, w \`Div\` 2]@ convMnistLayerR ker input bias =   let (batch_size :$: _ :$: h :$: w :$: ZSR) = rshape input-      yConv = conv2dUnpadded ker input+      yConv = conv2dSame ker input       biasStretched = rtranspose [0, 3, 1, 2]                       $ rreplicate batch_size $ rreplicate h $ rreplicate w bias       yRelu = relu $ yConv + biasStretched@@ -66,7 +66,7 @@  -- | Composition of two convolutional layers. convMnistTwoR-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: (ADReady target, NumScalar r, Differentiable r)   => Int -> Int -> Int   -> PrimalOf target (TKR 4 r)        -- ^ input images @[batch_size, 1, SizeMnistHeight, SizeMnistWidth]@@@ -94,7 +94,7 @@  -- | The neural network composed with the SoftMax-CrossEntropy loss function. convMnistLossFusedR-  :: (ADReady target, ADReady (PrimalOf target), GoodScalar r, Differentiable r)+  :: (ADReady target, ADReady (PrimalOf target), NumScalar r, Differentiable r)   => Int  -- ^ batch_size   -> ( PrimalOf target (TKR 3 r)          -- ^ @[batch_size, SizeMnistHeight, SizeMnistWidth]@@@ -118,7 +118,7 @@ -- and the trained parameters. convMnistTestR   :: forall target r.-     (target ~ Concrete, GoodScalar r, Differentiable r)+     (target ~ Concrete, NumScalar r, Differentiable r)   => Int   -> MnistDataBatchR r   -> ADCnnMnistParameters Concrete r
example/MnistCnnShaped2.hs view
@@ -12,10 +12,11 @@  import Data.Vector.Generic qualified as V import Data.Vector.Storable (Vector)-import GHC.TypeLits (fromSNat, type (*), type (+), type (<=), type Div)+import GHC.TypeLits (type (*), type (+), type (<=), type Div)  import Data.Array.Nested qualified as Nested import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat')  import HordeAd import MnistData@@ -45,7 +46,7 @@   :: forall kh kw h w c_in c_out batch_size target r.      ( 1 <= kh      , 1 <= kw  -- wrongly reported as redundant due to plugins-     , ADReady target, GoodScalar r, Differentiable r )+     , ADReady target, NumScalar r, Differentiable r )   => SNat kh -> SNat kw -> SNat h -> SNat w   -> SNat c_in -> SNat c_out -> SNat batch_size   -> target (TKS '[c_out, c_in, kh + 1, kw + 1] r)@@ -54,13 +55,13 @@   -> target (TKS '[batch_size, c_out, h `Div` 2, w `Div` 2] r) convMnistLayerS SNat SNat SNat SNat SNat SNat SNat                 ker input bias =-  let yConv = conv2dUnpaddedS ker input+  let yConv = conv2dSameS ker input       biasStretched = stranspose @'[0, 3, 1, 2]                       $ sreplicate {-@batch_size-}                       $ sreplicate {-@h-}                       $ sreplicate {-@w-} bias       yRelu = reluS $ yConv + biasStretched-  in maxPool2dUnpaddedS @2 @2 yRelu+  in maxPool2dUnpaddedS (SNat @2) (SNat @2) yRelu  -- | Composition of two convolutional layers. convMnistTwoS@@ -68,7 +69,7 @@        -- @h@ and @w@ are fixed with MNIST data, but not with test data      ( 1 <= kh  -- kernel height is large enough      , 1 <= kw  -- kernel width is large enough-     , ADReady target, GoodScalar r, Differentiable r )+     , ADReady target, NumScalar r, Differentiable r )   => SNat kh -> SNat kw -> SNat h -> SNat w   -> SNat c_out -> SNat n_hidden -> SNat batch_size        -- ^ these boilerplate lines tie type parameters to the corresponding@@ -105,7 +106,7 @@      , 1 <= kh      , 1 <= kw      , ADReady target, ADReady (PrimalOf target)-     , GoodScalar r, Differentiable r )+     , NumScalar r, Differentiable r )   => SNat kh -> SNat kw   -> SNat c_out   -> SNat n_hidden -> SNat batch_size@@ -123,7 +124,7 @@                              input adparameters       targets = str labelS       loss = lossSoftMaxCrossEntropyS targets result-  in kfromPrimal (recip $ kconcrete $ fromInteger $ fromSNat batch_size) * loss+  in kfromPrimal (recip $ kconcrete $ fromIntegral $ fromSNat' batch_size) * loss  -- | A function testing the neural network given testing set of inputs -- and the trained parameters.@@ -133,7 +134,7 @@      , 1 <= kh      , 1 <= kw      , target ~ Concrete-     , GoodScalar r, Differentiable r )+     , NumScalar r, Differentiable r )   => SNat kh -> SNat kw   -> SNat c_out   -> SNat n_hidden -> SNat batch_size@@ -141,7 +142,7 @@   -> ADCnnMnistParametersShaped target h w kh kw c_out n_hidden r   -> r convMnistTestS _ _ _ _ batch_size@SNat _ _-  | sNatValue batch_size == 0 = 0+  | fromSNat' batch_size == 0 = 0 convMnistTestS kh@SNat kw@SNat                c_out@SNat n_hidden@SNat batch_size@SNat                (glyphS, labelS) testParams =@@ -164,4 +165,4 @@       matchesLabels output label | V.maxIndex output == V.maxIndex label = 1                                  | otherwise = 0   in fromIntegral (sum (zipWith matchesLabels outputs labels))-     / fromInteger (fromSNat batch_size)+     / fromIntegral (fromSNat' batch_size)
example/MnistData.hs view
@@ -22,6 +22,7 @@ import Data.Array.Nested qualified as Nested import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat')  import HordeAd @@ -31,7 +32,7 @@ sizeMnistWidth = SNat @SizeMnistWidth  sizeMnistWidthInt :: Int-sizeMnistWidthInt = sNatValue sizeMnistWidth+sizeMnistWidthInt = fromSNat' sizeMnistWidth  type SizeMnistHeight = SizeMnistWidth @@ -52,7 +53,7 @@ sizeMnistLabel = SNat @SizeMnistLabel  sizeMnistLabelInt :: Int-sizeMnistLabelInt = sNatValue sizeMnistLabel+sizeMnistLabelInt = fromSNat' sizeMnistLabel  type LengthTestData = 10000 :: Nat @@ -125,10 +126,8 @@                   => [MnistDataS r] -> MnistDataBatchS batch_size r mkMnistDataBatchS l =   let (inputs, targets) = unzip l-  in ( Nested.sfromListOuter (SNat @batch_size)-       $ NonEmpty.fromList inputs-     , Nested.sfromListOuter (SNat @batch_size)-       $ NonEmpty.fromList targets )+  in ( Nested.sfromListOuter (SNat @batch_size) $ NonEmpty.fromList inputs+     , Nested.sfromListOuter (SNat @batch_size) $ NonEmpty.fromList targets ) {-# SPECIALIZE mkMnistDataBatchS :: forall batch_size. KnownNat batch_size => [MnistDataS Double] -> MnistDataBatchS batch_size Double #-} {-# SPECIALIZE mkMnistDataBatchS :: forall batch_size. KnownNat batch_size => [MnistDataS Float] -> MnistDataBatchS batch_size Float #-}
example/MnistFcnnRanked1.hs view
@@ -9,38 +9,37 @@ import Prelude  import Data.Vector.Generic qualified as V-import GHC.TypeLits (KnownNat, Nat)+import GHC.TypeLits (KnownNat, Nat, type (<=))  import Data.Array.Nested qualified as Nested import Data.Array.Nested.Ranked.Shape  import HordeAd-import HordeAd.Core.Ops (tfromListR)+import HordeAd.Core.Ops (tsfromIxR) import MnistData  -- | The differentiable type of all trainable parameters of this nn. type ADFcnnMnist1Parameters        (target :: Target) (widthHidden :: Nat) (widthHidden2 :: Nat) r =-  ( ( ListR widthHidden (target (TKS '[SizeMnistGlyph] r))+  ( ( IxR widthHidden (target (TKS '[SizeMnistGlyph] r))     , target (TKS '[widthHidden] r) )-  , ( ListR widthHidden2 (target (TKS '[widthHidden] Float))+  , ( IxR widthHidden2 (target (TKS '[widthHidden] Float))     , target (TKS '[widthHidden2] r) )-  , ( ListR SizeMnistLabel (target (TKS '[widthHidden2] r))+  , ( IxR SizeMnistLabel (target (TKS '[widthHidden2] r))     , target (TKS '[SizeMnistLabel] r) )   )  -- | An ad-hoc matrix multiplication analogue for matrices represented -- as lists of vectors. listMatmul1-  :: forall target r w1 w2.-     (ADReady target, GoodScalar r, KnownNat w1)-  => target (TKS '[w1] r) -> ListR w2 (target (TKS '[w1] r))+  :: forall target r w1 w2. (ADReady target, NumScalar r, KnownNat w1, 1 <= w2)+  => target (TKS '[w1] r) -> IxR w2 (target (TKS '[w1] r))   -> target (TKS '[w2] r) {-# INLINE listMatmul1 #-}  -- this doesn't want to specialize listMatmul1 x0 weights = tlet x0 $ \x ->   let f :: target (TKS '[w1] r) -> target (TKS '[] r)-      f v = v `sdot0` x-  in tfromListR knownSTK $ f <$> weights+      f v = sfromK $ v `sdot0` x+  in tsfromIxR $ f <$> weights  -- | Fully connected neural network for the MNIST digit classification task. -- There are two hidden layers and both use the same activation function.@@ -48,7 +47,8 @@ -- The widths of the two hidden layers are @widthHidden@ and @widthHidden2@, -- respectively. afcnnMnist1 :: forall target r widthHidden widthHidden2.-               (ADReady target, GoodScalar r, Differentiable r)+               ( ADReady target, NumScalar r, Differentiable r+               , 1 <= widthHidden, 1 <= widthHidden2 )             => (forall n. KnownNat n                 => target (TKS '[n] r) -> target (TKS '[n] r))             -> (target (TKS '[SizeMnistLabel] r)@@ -71,7 +71,8 @@ -- | The neural network applied to concrete activation functions -- and composed with the appropriate loss function. afcnnMnistLoss1-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: ( ADReady target, NumScalar r, Differentiable r+     , 1 <= widthHidden, 1 <= widthHidden2 )   => SNat widthHidden -> SNat widthHidden2   -> (target (TKR 1 r), target (TKR 1 r))   -> ADFcnnMnist1Parameters target widthHidden widthHidden2 r@@ -80,17 +81,18 @@   let result = afcnnMnist1 logisticS softMax1S                            widthHidden widthHidden2 (sfromR datum) adparams   in lossCrossEntropyV target result--- {-# SPECIALIZE afcnnMnistLoss1 :: (GoodScalar r, Differentiable r) => SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 r), AstTensor AstMethodLet FullSpan (TKR 1 r)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 r -> AstTensor AstMethodLet FullSpan (TKScalar r) #-}-{-# SPECIALIZE afcnnMnistLoss1 :: SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 1 Double)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 Double -> AstTensor AstMethodLet FullSpan (TKScalar Double) #-}-{-# SPECIALIZE afcnnMnistLoss1 :: SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 Float), AstTensor AstMethodLet FullSpan (TKR 1 Float)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 Float -> AstTensor AstMethodLet FullSpan (TKScalar Float) #-}-{-# SPECIALIZE afcnnMnistLoss1 :: SNat widthHidden -> SNat widthHidden2 -> (ADVal Concrete (TKR 1 Double), ADVal Concrete (TKR 1 Double)) -> ADFcnnMnist1Parameters (ADVal Concrete) widthHidden widthHidden2 Double -> ADVal Concrete (TKScalar Double) #-}-{-# SPECIALIZE afcnnMnistLoss1 :: SNat widthHidden -> SNat widthHidden2 -> (ADVal Concrete (TKR 1 Float), ADVal Concrete (TKR 1 Float)) -> ADFcnnMnist1Parameters (ADVal Concrete) widthHidden widthHidden2 Float -> ADVal Concrete (TKScalar Float) #-}+-- {-# SPECIALIZE afcnnMnistLoss1 :: (NumScalar r, Differentiable r, 1 <= widthHidden, 1 <= widthHidden2) => SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 r), AstTensor AstMethodLet FullSpan (TKR 1 r)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 r -> AstTensor AstMethodLet FullSpan (TKScalar r) #-}+{-# SPECIALIZE afcnnMnistLoss1 :: (1 <= widthHidden, 1 <= widthHidden2) => SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 1 Double)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 Double -> AstTensor AstMethodLet FullSpan (TKScalar Double) #-}+{-# SPECIALIZE afcnnMnistLoss1 :: (1 <= widthHidden, 1 <= widthHidden2) => SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 Float), AstTensor AstMethodLet FullSpan (TKR 1 Float)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 Float -> AstTensor AstMethodLet FullSpan (TKScalar Float) #-}+{-# SPECIALIZE afcnnMnistLoss1 :: (1 <= widthHidden, 1 <= widthHidden2) => SNat widthHidden -> SNat widthHidden2 -> (ADVal Concrete (TKR 1 Double), ADVal Concrete (TKR 1 Double)) -> ADFcnnMnist1Parameters (ADVal Concrete) widthHidden widthHidden2 Double -> ADVal Concrete (TKScalar Double) #-}+{-# SPECIALIZE afcnnMnistLoss1 :: (1 <= widthHidden, 1 <= widthHidden2) => SNat widthHidden -> SNat widthHidden2 -> (ADVal Concrete (TKR 1 Float), ADVal Concrete (TKR 1 Float)) -> ADFcnnMnist1Parameters (ADVal Concrete) widthHidden widthHidden2 Float -> ADVal Concrete (TKScalar Float) #-}  -- | A function testing the neural network given testing set of inputs -- and the trained parameters. afcnnMnistTest1   :: forall target widthHidden widthHidden2 r.-     (target ~ Concrete, GoodScalar r, Differentiable r)+     ( target ~ Concrete, NumScalar r, Differentiable r+     , 1 <= widthHidden, 1 <= widthHidden2 )   => SNat widthHidden -> SNat widthHidden2   -> [MnistDataLinearR r]   -> ADFcnnMnist1Parameters target widthHidden widthHidden2 r
example/MnistFcnnRanked2.hs view
@@ -49,8 +49,8 @@ -- The output layer uses a different activation function. -- The widths of the two hidden layers are @widthHidden@ and @widthHidden2@, -- respectively.-afcnnMnist2 :: ( ADReady target, GoodScalar r, Differentiable r-               , GoodScalar q, Differentiable q )+afcnnMnist2 :: ( ADReady target, NumScalar r, Differentiable r+               , NumScalar q, Differentiable q )             => (target (TKR 1 r) -> target (TKR 1 r))             -> (target (TKR 1 r) -> target (TKR 1 r))             -> target (TKR 1 r)@@ -68,8 +68,8 @@ -- | The neural network applied to concrete activation functions -- and composed with the appropriate loss function. afcnnMnistLoss2-  :: ( ADReady target, GoodScalar r, Differentiable r-     , GoodScalar q, Differentiable q )+  :: ( ADReady target, NumScalar r, Differentiable r+     , NumScalar q, Differentiable q )   => (target (TKR 1 r), target (TKR 1 r)) -> ADFcnnMnist2Parameters target r q   -> target (TKScalar r) afcnnMnistLoss2 (datum, target) adparams =@@ -83,8 +83,8 @@ -- and the trained parameters. afcnnMnistTest2   :: forall target r q.-     ( target ~ Concrete, GoodScalar r, Differentiable r-     , GoodScalar q, Differentiable q )+     ( target ~ Concrete, NumScalar r, Differentiable r+     , NumScalar q, Differentiable q )   => [MnistDataLinearR r]   -> ADFcnnMnist2Parameters target r q   -> r@@ -106,8 +106,8 @@ -- between tests and benchmarks and to separate compile-time and run-time -- for benchmarking (this part is considered compile-time). mnistTrainBench2VTOGradient-  :: forall r q. ( GoodScalar r, Differentiable r-                 , GoodScalar q, Differentiable q )+  :: forall r q. ( NumScalar r, Differentiable r+                 , NumScalar q, Differentiable q )   => Proxy q -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int   -> ( Concrete (XParams2 r q)      , AstArtifactRev@@ -116,7 +116,7 @@             (TKProduct (TKR2 1 (TKScalar r))                        (TKR2 1 (TKScalar r))))          (TKScalar r) )-mnistTrainBench2VTOGradient Proxy cotangentHandling range seed widthHidden widthHidden2 =+mnistTrainBench2VTOGradient Proxy cotangentHandling range seed widthHidden widthHidden2 | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat widthHidden $ \(SNat @widthHidden) ->   withSNat widthHidden2 $ \(SNat @widthHidden2) ->   -- Initial parameter generation is counted as part of compilation time.@@ -144,8 +144,8 @@ -- | A version of 'mnistTrainBench2VTOGradient' without any simplification, -- even the AST smart constructors. Intended for benchmarking. mnistTrainBench2VTOGradientX-  :: forall r q. ( GoodScalar r, Differentiable r-                 , GoodScalar q, Differentiable q )+  :: forall r q. ( NumScalar r, Differentiable r+                 , NumScalar q, Differentiable q )   => Proxy q -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int   -> ( Concrete (XParams2 r q)      , AstArtifactRev@@ -154,7 +154,7 @@             (TKProduct (TKR2 1 (TKScalar r))                        (TKR2 1 (TKScalar r))))          (TKScalar r) )-mnistTrainBench2VTOGradientX Proxy cotangentHandling range seed widthHidden widthHidden2 =+mnistTrainBench2VTOGradientX Proxy cotangentHandling range seed widthHidden widthHidden2 | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat widthHidden $ \(SNat @widthHidden) ->   withSNat widthHidden2 $ \(SNat @widthHidden2) ->   -- Initial parameter generation is counted as part of compilation time.
example/MnistRnnRanked2.hs view
@@ -3,10 +3,9 @@ -- for classification of MNIST digits. Sports 2 hidden layers. module MnistRnnRanked2 where -import Prelude hiding (foldl')+import Prelude  import Data.Kind (Type)-import Data.List (foldl') import Data.Vector.Generic qualified as V import Data.Vector.Storable (Vector) import GHC.TypeLits (KnownNat, Nat, type (+))@@ -44,7 +43,7 @@   , target (TKR 1 r) )  zeroStateR-  :: (BaseTensor target, GoodScalar r)+  :: (BaseTensor target, NumScalar r)   => IShR n -> (target (TKR n r)  -- state                     -> a)   -> a@@ -61,7 +60,7 @@  -- | A single recurrent layer with @tanh@ activation function. rnnMnistLayerR-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: (ADReady target, NumScalar r, Differentiable r)   => target (TKR 2 r)  -- ^ in state, @[out_width, batch_size]@   -> target (TKR 2 r)  -- ^ input, @[in_width, batch_size]@   -> LayerWeigthsRNN target r  -- ^ parameters@@ -75,7 +74,7 @@ -- TODO: represent state as a pair to avoid appending; tlet now supports this. -- | Composition of two recurrent layers. rnnMnistTwoR-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: (ADReady target, NumScalar r, Differentiable r)   => target (TKR 2 r)  -- initial state, @[2 * out_width, batch_size]@   -> PrimalOf target (TKR 2 r)  -- @[sizeMnistHeight, batch_size]@   -> ( LayerWeigthsRNN target r  -- sizeMnistHeight out_width@@ -96,7 +95,7 @@ -- | The two-layer recurrent nn with its state initialized to zero -- and the result composed with a fully connected layer. rnnMnistZeroR-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: (ADReady target, NumScalar r, Differentiable r)   => Int  -- ^ batch_size   -> PrimalOf target (TKR 3 r)        -- ^ input data @[sizeMnistWidth, sizeMnistHeight, batch_size]@@@ -112,7 +111,7 @@  -- | The neural network composed with the SoftMax-CrossEntropy loss function. rnnMnistLossFusedR-  :: (ADReady target, ADReady (PrimalOf target), GoodScalar r, Differentiable r)+  :: (ADReady target, ADReady (PrimalOf target), NumScalar r, Differentiable r)   => Int   -> (PrimalOf target (TKR 3 r), PrimalOf target (TKR 2 r))  -- batch_size   -> ADRnnMnistParameters target r  -- SizeMnistHeight out_width@@ -128,7 +127,7 @@ -- and the trained parameters. rnnMnistTestR   :: forall target r.-     (target ~ Concrete, GoodScalar r, Differentiable r)+     (target ~ Concrete, NumScalar r, Differentiable r)   => Int   -> MnistDataBatchR r  -- batch_size   -> ADRnnMnistParameters target r
example/MnistRnnShaped2.hs view
@@ -5,17 +5,17 @@ -- for classification of MNIST digits. Sports 2 hidden layers. module MnistRnnShaped2 where -import Prelude hiding (foldl')+import Prelude  import Data.Kind (Type)-import Data.List (foldl') import Data.Vector.Generic qualified as V import Data.Vector.Storable (Vector)-import GHC.TypeLits (KnownNat, Nat, fromSNat, type (*))+import GHC.TypeLits (KnownNat, Nat, type (*)) -import Data.Array.Nested.Permutation qualified as Permutation import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Permutation qualified as Permutation import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat')  import HordeAd import MnistData@@ -36,7 +36,7 @@   , target (TKS '[out_width] r) )           -- bias  zeroStateS-  :: (BaseTensor target, KnownShS sh, GoodScalar r)+  :: (BaseTensor target, KnownShS sh, NumScalar r)   => (target (TKS sh r)  -- state       -> a)   -> a@@ -53,7 +53,7 @@  -- | A single recurrent layer with @tanh@ activation function. rnnMnistLayerS-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: (ADReady target, NumScalar r, Differentiable r)   => SNat in_width -> SNat out_width -> SNat batch_size        -- ^ these boilerplate lines tie type parameters to the corresponding        -- value parameters (@SNat@ below) denoting basic dimensions@@ -70,7 +70,7 @@ -- TODO: represent state as a pair to avoid appending; tlet now supports this. -- | Composition of two recurrent layers. rnnMnistTwoS-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: (ADReady target, NumScalar r, Differentiable r)   => SNat out_width -> SNat batch_size -> SNat sizeMnistH   -> target (TKS '[2 * out_width, batch_size] r)  -- initial state   -> PrimalOf target (TKS '[sizeMnistH, batch_size] r)@@ -99,7 +99,7 @@ -- | The two-layer recurrent nn with its state initialized to zero -- and the result composed with a fully connected layer. rnnMnistZeroS-  :: (ADReady target, GoodScalar r, Differentiable r)+  :: (ADReady target, NumScalar r, Differentiable r)   => SNat out_width   -> SNat batch_size   -> SNat sizeMnistH -> SNat sizeMnistW@@ -119,7 +119,7 @@ rnnMnistLossFusedS   :: forall target h w out_width batch_size r.      ( h ~ SizeMnistHeight, w ~ SizeMnistWidth, Differentiable r-     , ADReady target, ADReady (PrimalOf target), GoodScalar r)+     , ADReady target, ADReady (PrimalOf target), NumScalar r)   => SNat out_width   -> SNat batch_size   -> ( PrimalOf target (TKS '[batch_size, h, w] r)@@ -136,14 +136,14 @@                              xs adparameters       targets = str labelS       loss = lossSoftMaxCrossEntropyS targets result-  in kfromPrimal (recip $ kconcrete $ fromInteger $ fromSNat batch_size) * loss+  in kfromPrimal (recip $ kconcrete $ fromIntegral $ fromSNat' batch_size) * loss  -- | A function testing the neural network given testing set of inputs -- and the trained parameters. rnnMnistTestS   :: forall target h w out_width batch_size r.      ( h ~ SizeMnistHeight, w ~ SizeMnistWidth-     , target ~ Concrete, Differentiable r, GoodScalar r )+     , target ~ Concrete, Differentiable r, NumScalar r )   => SNat out_width   -> SNat batch_size   -> MnistDataBatchS batch_size r@@ -172,4 +172,4 @@       matchesLabels output label | V.maxIndex output == V.maxIndex label = 1                                  | otherwise = 0   in fromIntegral (sum (zipWith matchesLabels outputs labels))-     / fromInteger (fromSNat batch_size)+     / fromIntegral (fromSNat' batch_size)
horde-ad.cabal view
@@ -20,7 +20,7 @@ -- PVP summary:     +-+------- breaking API changes --                  | | +----- non-breaking API additions --                  | | | +--- code changes with no API change-version:            0.2.0.0+version:            0.3.0.0  -- A short (one-line) description of the package. synopsis:           Higher Order Reverse Derivatives Efficiently - Automatic Differentiation@@ -47,11 +47,11 @@ build-type:         Simple  -- Extra doc files to be distributed with the package, such as a CHANGELOG or a README.-extra-doc-files:    README.md,-                    CHANGELOG.md,+extra-doc-files:    README.md+                    CHANGELOG.md                     CREDITS.md -tested-with:        GHC ==9.10.2 || ==9.12.3+tested-with:        GHC ==9.10.3 || ==9.12.4 || ==9.14.1  bug-reports:        https://github.com/Mikolaj/horde-ad/issues @@ -76,25 +76,44 @@  common options   default-language:   GHC2024-  default-extensions: StrictData, TypeFamilies, TypeFamilyDependencies,-                      FunctionalDependencies, RecordWildCards, MultiWayIf,-                      DefaultSignatures, PatternSynonyms, NoStarIsType,-                      TypeData, TypeAbstractions-  other-extensions:   UnboxedTuples, CPP, ViewPatterns, OverloadedLists,-                      DeriveAnyClass, UndecidableInstances, AllowAmbiguousTypes,-                      QuantifiedConstraints, TemplateHaskell, DerivingVia,-                      ImpredicativeTypes+  default-extensions: StrictData TypeFamilies TypeFamilyDependencies+                      FunctionalDependencies RecordWildCards MultiWayIf+                      DefaultSignatures PatternSynonyms NoStarIsType+                      TypeData TypeAbstractions+                      -- to molify stylish-haskell, until it's fixed+                      LambdaCase+  other-extensions:   UnboxedTuples CPP ViewPatterns OverloadedLists+                      DeriveAnyClass QuantifiedConstraints TemplateHaskell+                      DerivingVia ImpredicativeTypes+                      UndecidableInstances UndecidableSuperClasses+                      AllowAmbiguousTypes+  -- * Warnings and similar.   ghc-options:        -Wall -Wcompat -Wimplicit-prelude -Widentities -Wredundant-constraints -Wmissing-export-lists -Wpartial-fields -Wunused-packages -Wredundant-bang-patterns -Woperator-whitespace -Wredundant-strictness-flags -Wterm-variable-capture -Wmissed-specialisations   ghc-options:        -Wmissing-poly-kind-signatures -Wmissing-role-annotations -Wno-unticked-promoted-constructors-  ghc-options:        -fprint-explicit-kinds+  -- Disables a warning we can't yet fix:+  if impl(ghc >= 9.14)+    ghc-options:      -Wno-pattern-namespace-specifier+  ghc-options:        -fprint-explicit-kinds -fsimpl-tick-factor=200++  -- * Performance.   -- This is supposedly risky, but prevents STG from creating thunks, #23848:   ghc-options:        -fworker-wrapper-cbv-  -- Disabling this breaks performance:-  ghc-options:        -fexpose-all-unfoldings -fspecialise-aggressively -fsimpl-tick-factor=200-  -- Disabling hurts performance to some extent:+  -- These improve performance, but slow down compilation:+  if impl(ghc >= 9.12)+    ghc-options:      -fexpose-overloaded-unfoldings+  ghc-options:        -fspecialise-aggressively+  if impl(ghc >= 9.14)+    ghc-options:      -fpolymorphic-specialisation+  -- This one improves both performance and compilation speed:+  ghc-options:        -fkeep-auto-rules+  -- These improve performance to some extent:   ghc-options:        -fdicts-cheap -flate-dmd-anal-  -- ghc-options: -O2 -flate-specialise+  -- Twice the compilation time and no consistent performance boost:+  -- ghc-options      -flate-specialise+  -- 30% longer compilation for small and inconsistent performance benefits:+  -- ghc-options:     -O2 +  -- * Misc.   if flag(with_expensive_assertions)     cpp-options:      -DWITH_EXPENSIVE_ASSERTIONS @@ -105,36 +124,36 @@   ghc-options:        -fno-ignore-asserts  common exe-options-  ghc-options:        -rtsopts+  ghc-options:        -rtsopts -threaded   -- Make GC more predictable in benchmarks.   ghc-options:        "-with-rtsopts=-H1.5g -A0.75g -I0" -common exe-options-for-tests-that-reset-counters-  ghc-options:        -rtsopts -threaded-  ghc-options:        "-with-rtsopts=-H1.5g -A0.75g -I0"-    -- unconditionally seqential, because counter reset is not thread-safe--common exe-options-test+common exe-options-for-parallel-tests   ghc-options:        -rtsopts -threaded   if flag(test_seq)     ghc-options:      "-with-rtsopts=-H1.5g -A0.75g -I0"   else     -- Beware, -with-rtsopts is not cumulative!-    ghc-options:      "-with-rtsopts=-H1.5g -A0.75g -I0 -N"+    ghc-options:      "-with-rtsopts=-H1.5g -A0.75g -I0 -N8"  library     import:           options      exposed-modules:  HordeAd                       HordeAd.ADEngine-                      HordeAd.AstEngine+                      HordeAd.OpsTensor+                      HordeAd.OpsTensorRanked+                      HordeAd.OpsTensorShaped+                      HordeAd.OpsTensorMixed                       HordeAd.Core.Adaptor                       HordeAd.Core.Ast+                      HordeAd.Core.AstEngine                       HordeAd.Core.AstEnv                       HordeAd.Core.AstFreshId                       HordeAd.Core.AstInline                       HordeAd.Core.AstInterpret-                      HordeAd.Core.AstPrettyPrint+                      HordeAd.Core.AstMethodLet+                      HordeAd.Core.AstMethodShare                       HordeAd.Core.AstSimplify                       HordeAd.Core.AstTools                       HordeAd.Core.AstTraverse@@ -142,6 +161,7 @@                       HordeAd.Core.CarriersADVal                       HordeAd.Core.CarriersAst                       HordeAd.Core.CarriersConcrete+                      HordeAd.Core.Conversion                       HordeAd.Core.ConvertTensor                       HordeAd.Core.Delta                       HordeAd.Core.DeltaEval@@ -150,39 +170,36 @@                       HordeAd.Core.OpsADVal                       HordeAd.Core.OpsAst                       HordeAd.Core.OpsConcrete+                      HordeAd.Core.PPEngine+                      HordeAd.Core.PPTools                       HordeAd.Core.TensorKind                       HordeAd.Core.Types                       HordeAd.Core.Unwind+                      HordeAd.Core.UnwindNum                       HordeAd.External.CommonRankedOps                       HordeAd.External.CommonShapedOps                       HordeAd.External.Optimizer                       HordeAd.External.OptimizerTools-                      HordeAd.OpsTensor-                      HordeAd.OpsTensorRanked-                      HordeAd.OpsTensorShaped-                      HordeAd.OpsTensorMixed      hs-source-dirs:   src      build-depends:         assert-failure < 0.1.4-      , base >= 4.20.1 && < 4.21-          -- backport of https://gitlab.haskell.org/ghc/ghc/-/merge_requests/13498 is required for GHC 9.12; TODO: adjust base bounds once 9.12.3 is out+      , atomic-counter < 0.2+      , base >= 4.20.1 && < 4.23       , Boolean < 0.3       , containers >= 0.6 && < 0.9       , deepseq < 1.7       , data-default < 0.9-      , dependent-enummap < 0.2+      , dependent-enummap >= 0.2 && < 0.3       , dependent-sum >= 0.7 && < 0.8       , enummapset < 0.8-      , ghc-typelits-knownnat < 1-      , ghc-typelits-natnormalise < 1-      , indexed-traversable < 0.2-      , orthotope < 0.2-      , ox-arrays < 0.2+      , ghc-typelits-knownnat >= 0.8.0 && < 1+      , ghc-typelits-natnormalise >= 0.8.1 && < 1+      , ox-arrays >= 0.2 && < 0.3+      , ox-arrays:strided-array-ops       , random >= 1.3.0 && < 1.4       , some >= 1 && < 2-      , atomic-counter < 0.2       , vector >= 0.13.2.0 && < 0.14  library exampleLibrary@@ -349,8 +366,9 @@       , horde-ad:horde-ad        , base+      , orthotope >= 0.1.8 && < 0.1.9       , ox-arrays-      , tasty >= 1.0 && < 1.6+      , tasty >= 1.5.4 && < 1.6       , tasty-hunit < 0.11       , vector @@ -359,6 +377,7 @@     visibility:       public      exposed-modules:  TestAdaptorSimplified+                      TestConvQuickCheck                       TestConvSimplified                       TestGatherSimplified                       TestHighRankSimplified@@ -384,14 +403,15 @@       , ghc-typelits-natnormalise       , ox-arrays       , random-      , tasty >= 1.0 && < 1.6+      , tasty >= 1.5.4 && < 1.6       , tasty-hunit < 0.11       , tasty-quickcheck < 0.12+      , vector  -- The data files for this are not included in the cabal package, -- hence 'buildable: False'. test-suite fullTest-    import:           options, exe-options-for-tests-that-reset-counters+    import:           options, exe-options      if flag(release)       buildable: False@@ -409,10 +429,11 @@       , horde-ad:testToolLibrary        , base-      , tasty >= 1.0+      , tasty +-- This test must be sequential, because counter reset is not thread-safe. test-suite CAFlessTest-    import:           options, exe-options-for-tests-that-reset-counters+    import:           options, exe-options      if flag(release)       buildable: False@@ -430,10 +451,10 @@       , horde-ad:testToolLibrary        , base-      , tasty >= 1.0+      , tasty  test-suite parallelTest-    import:           options, exe-options-test+    import:           options, exe-options-for-parallel-tests      if flag(release)       buildable: False@@ -451,12 +472,34 @@       , horde-ad:testToolLibrary        , base-      , tasty >= 1.0+      , tasty +-- This test can serve as a benchmark when run sequentially.+test-suite sequentialMnistTest+    import:           options, exe-options++    if flag(release)+      buildable: False+    else+      buildable: True++    type:             exitcode-stdio-1.0++    hs-source-dirs:   test++    main-is:          SequentialMnistTest.hs++    build-depends:+      , horde-ad:testCommonLibrary+      , horde-ad:testToolLibrary++      , base+      , tasty+ -- These tests don't require any dataset and are short, so can be enabled -- in the Hackage package. test-suite minimalTest-    import:           options, exe-options-for-tests-that-reset-counters+    import:           options, exe-options      type:             exitcode-stdio-1.0 @@ -469,4 +512,4 @@       , horde-ad:testToolLibrary        , base-      , tasty >= 1.0+      , tasty
src/HordeAd.hs view
@@ -19,7 +19,8 @@   , module HordeAd.Core.ConvertTensor     -- * The main AD API   , module HordeAd.ADEngine-  , module HordeAd.AstEngine+  , module HordeAd.Core.AstEngine+  , module HordeAd.Core.PPEngine     -- * Additional support types and operations   , module HordeAd.Core.Ast   , module HordeAd.Core.CarriersADVal@@ -34,11 +35,12 @@ import Prelude ()  import HordeAd.ADEngine-import HordeAd.AstEngine import HordeAd.Core.Ast+import HordeAd.Core.AstEngine import HordeAd.Core.CarriersADVal import HordeAd.Core.CarriersConcrete import HordeAd.Core.ConvertTensor+import HordeAd.Core.PPEngine import HordeAd.Core.TensorKind import HordeAd.Core.Types import HordeAd.External.CommonRankedOps
src/HordeAd/ADEngine.hs view
@@ -16,36 +16,35 @@ -- of the offered respective APIs. module HordeAd.ADEngine   ( -- * Symbolic reverse derivative adaptors-    grad, vjp+    grad, grad2, vjp, vjp2   , gradArtifact, vjpArtifact   , gradInterpretArtifact, vjpInterpretArtifact     -- * Symbolic forward derivative adaptors-  , jvp, jvpArtifact, jvpInterpretArtifact+  , jvp, jvp2, jvpArtifact, jvpInterpretArtifact     -- * Non-symbolic reverse derivative adaptors-  , cgrad, cvjp+  , cgrad, cgrad2, cvjp, cvjp2     -- * Non-symbolic forward derivative adaptors-  , cjvp+  , cjvp, cjvp2     -- * Internal machinery for symbolic adaptors   , IncomingCotangentHandling(..)-  , revArtifactAdapt, revArtifactDelta+  , revArtifactAdapt, revArtifactAdaptDt, revArtifactDelta   , revProduceArtifactWithoutInterpretation, revInterpretArtifact   , fwdArtifactAdapt, fwdArtifactDelta, fwdInterpretArtifact-    -- * Internal machinery for non-symbolic adaptors-  , cfwdBoth   ) where  import Prelude -import HordeAd.AstEngine+import Data.Proxy (Proxy (Proxy))+ import HordeAd.Core.Adaptor import HordeAd.Core.Ast+import HordeAd.Core.AstEngine import HordeAd.Core.AstEnv import HordeAd.Core.AstInterpret import HordeAd.Core.CarriersADVal import HordeAd.Core.CarriersAst import HordeAd.Core.CarriersConcrete import HordeAd.Core.Delta-import HordeAd.Core.DeltaEval import HordeAd.Core.Ops import HordeAd.Core.OpsADVal import HordeAd.Core.OpsAst@@ -67,7 +66,7 @@ -- the type of concrete contangents to the type of concrete input parameters. grad   :: forall src r tgt.-     ( X src ~ X (Value src), KnownSTK (X src)+     ( GoodScalar r, X src ~ X (Value src), KnownSTK (X src)      , AdaptableTarget (AstTensor AstMethodLet FullSpan) src      , AdaptableTarget Concrete (Value src)      , tgt ~ AstTensor AstMethodLet FullSpan (TKScalar r) )@@ -75,8 +74,24 @@   -> Value src   -> Value src  -- morally Value (ADTensorKind src) {-# INLINE grad #-}-grad f vals = revMaybe f vals Nothing+grad f vals = snd $ grad2 f vals +-- | The @grad2@ operation works just as 'grad', but additionally produces+-- the primal result of the objective function at minimal extra cost.+grad2+  :: forall src r tgt.+     ( GoodScalar r, X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan (TKScalar r) )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> ( Concrete (TKScalar r)+     , Value src )  -- morally Value (ADTensorKind src)+{-# INLINE grad2 #-}+grad2 f vals | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =+  revMaybe f vals Nothing+ -- | This version of the symbolic reverse derivative operation -- explicitly takes the sensitivity parameter (the incoming cotangent). -- It also permits an arbitrary (nested tuple+) type of the domain@@ -98,15 +113,31 @@   -> Concrete (ADTensorKind ztgt)   -> Value src  -- morally Value (ADTensorKind src) {-# INLINE vjp #-}-vjp f vals dt = revMaybe f vals (Just dt)+vjp f vals0 dt = snd $ vjp2 f vals0 dt +-- | The @vjp2@ operation works just as 'vjp', but additionally produces+-- the primal result of the objective function at minimal extra cost.+vjp2+  :: forall src ztgt tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> Concrete (ADTensorKind ztgt)+  -> ( Concrete ztgt+     , Value src )  -- morally Value (ADTensorKind src)+{-# INLINE vjp2 #-}+vjp2 = revMaybeDt+ -- | Compute the reverse derivative not for a specific input, but as symbolic -- function from inputs to the gradient value. -- The function is represented as an "artifact", which is the gradient -- AST term together with the variable corresponding to the input. gradArtifact   :: forall src r tgt.-     ( X src ~ X (Value src), KnownSTK (X src)+     ( GoodScalar r, X src ~ X (Value src), KnownSTK (X src)      , AdaptableTarget (AstTensor AstMethodLet FullSpan) src      , AdaptableTarget Concrete (Value src)      , tgt ~ AstTensor AstMethodLet FullSpan (TKScalar r) )@@ -115,9 +146,10 @@   -> AstArtifactRev (X src) (TKScalar r)        -- ^ the artifact containing the symbolic code of the derivative {-# INLINE gradArtifact #-}-gradArtifact f vals0 =-  let xftk = tftkG (knownSTK @(X src)) $ unConcrete $ toTarget vals0-  in revArtifactAdapt IgnoreIncomingCotangent f xftk+gradArtifact f vals0+  | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =+    let xftk = tftkG (knownSTK @(X src)) $ unConcrete $ toTarget vals0+    in revArtifactAdapt IgnoreIncomingCotangent f xftk  -- | Compute the reverse derivative not for a specific input, but as symbolic -- function from inputs and incoming cotangents to the gradient value.@@ -136,29 +168,32 @@ {-# INLINE vjpArtifact #-} vjpArtifact f vals0 =   let xftk = tftkG (knownSTK @(X src)) $ unConcrete $ toTarget vals0-  in revArtifactAdapt UseIncomingCotangent f xftk+  in revArtifactAdaptDt f xftk  -- | Interpret the "artifact" as a function from a concrete tensor -- to a concrete tensor (possibly adapted, e.g., from horde-ad nested pairs -- to Haskell n-tuples). gradInterpretArtifact   :: forall x r avals.-     (X avals ~ ADTensorKind x, AdaptableTarget Concrete avals)+     (GoodScalar r, X avals ~ ADTensorKind x, AdaptableTarget Concrete avals)   => AstArtifactRev x (TKScalar r)        -- ^ the artifact containing the symbolic code of the derivative   -> Concrete x   -> avals {-# INLINE gradInterpretArtifact #-}-gradInterpretArtifact AstArtifactRev{..} parameters =-  let xftk = varNameToFTK artVarDomainRev-      azftk = varNameToFTK artVarDtRev-                -- STKScalar @(ADTensorScalar r) or STKScalar @Z1-      oneAtF = treplTarget 1 azftk-      env = extendEnv artVarDtRev oneAtF-            $ extendEnv artVarDomainRev parameters emptyEnv-  in if tftkG (ftkToSTK xftk) (unConcrete parameters) == xftk-     then fromTarget $ interpretAstPrimal env artDerivativeRev-     else error "gradInterpretArtifact: reverse derivative parameters must have the same shape as the domain of the objective function"+gradInterpretArtifact AstArtifactRev{..} parameters+  | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =+    let xftk = varNameToFTK artVarDomainRev+        azftk = varNameToFTK artVarDtRev+                  -- STKScalar @(ADTensorScalar r) or STKScalar @Z1+        oneAtF = treplTarget 1 azftk+        env = extendEnv artVarDtRev oneAtF+              $ extendEnv artVarDomainRev parameters emptyEnv+    in if tftkG (ftkToSTK xftk) (unConcrete parameters) == xftk+       then fromTarget $ interpretAstFull env artDerivativeRev+       else error $ "gradInterpretArtifact: reverse derivative parameters must have the same shape as the domain of the objective function: "+                    ++ show ( tftkG (ftkToSTK xftk) (unConcrete parameters)+                            , xftk )  -- | Interpret the "artifact" as a function from concrete tensors -- to a concrete tensor (possibly adapted, e.g., from horde-ad nested pairs@@ -179,23 +214,27 @@             $ extendEnv artVarDomainRev parameters emptyEnv   in if tftkG (ftkToSTK xftk) (unConcrete parameters) == xftk      then if tftkG (ftkToSTK azftk) (unConcrete dt) == azftk-          then fromTarget $ interpretAstPrimal env artDerivativeRev-          else error "vjpInterpretArtifact: reverse derivative incoming cotangent must have the same shape as the codomain of the objective function"-     else error "vjpInterpretArtifact: reverse derivative parameters must have the same shape as the domain of the objective function"+          then fromTarget $ interpretAstFull env artDerivativeRev+          else error $ "vjpInterpretArtifact: reverse derivative incoming cotangent must have the same shape as the codomain of the objective function: "+                       ++ show (tftkG (ftkToSTK azftk) (unConcrete dt), azftk)+     else error $ "vjpInterpretArtifact: reverse derivative parameters must have the same shape as the domain of the objective function: "+                  ++ show (tftkG (ftkToSTK xftk) (unConcrete parameters), xftk)   -- * Symbolic reverse derivative adaptors' internal machinery  revMaybe   :: forall src ztgt tgt.-     ( X src ~ X (Value src), KnownSTK (X src)+     ( TKAllNum (ADTensorKind ztgt)+     , X src ~ X (Value src), KnownSTK (X src)      , AdaptableTarget (AstTensor AstMethodLet FullSpan) src      , AdaptableTarget Concrete (Value src)      , tgt ~ AstTensor AstMethodLet FullSpan ztgt )   => (src -> tgt)  -- ^ the objective function   -> Value src   -> Maybe (Concrete (ADTensorKind ztgt))-  -> Value src  -- morally Value (ADTensorKind src)+  -> ( Concrete ztgt+     , Value src )  -- morally Value (ADTensorKind src) {-# INLINE revMaybe #-} revMaybe f vals0 mdt =   let valsTarget = toTarget vals0@@ -203,13 +242,14 @@       cotangentHandling =         maybe IgnoreIncomingCotangent (const UseIncomingCotangent) mdt       artifactRaw = revArtifactAdapt cotangentHandling f xftk-      artifact = simplifyArtifactGradient artifactRaw-  in fromTarget $ fromADTensorKindShared (ftkToSTK xftk)-     $ fst $ revInterpretArtifact artifact valsTarget mdt+      artifact = simplifyArtifactRev artifactRaw+      (primal, res) = revInterpretArtifact artifact valsTarget mdt+  in (primal, fromTarget $ fromADTensorKindShared xftk res)  revArtifactAdapt   :: forall src ztgt tgt.-     ( AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     ( TKAllNum (ADTensorKind ztgt)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src      , tgt ~ AstTensor AstMethodLet FullSpan ztgt )   => IncomingCotangentHandling   -> (src -> tgt)  -- ^ the objective function@@ -219,17 +259,17 @@ {-# INLINE revArtifactAdapt #-} revArtifactAdapt cotangentHandling f xftk =   let g :: AstTensor AstMethodLet FullSpan (X src) -> tgt-      g !arg = simplifyInline $ ttlet arg $ f . fromTarget+      g !arg = simplifyUserCode $ ttlet arg $ f . fromTarget                                   -- fromTarget requires duplicable   in revProduceArtifact cotangentHandling g emptyEnv xftk  revInterpretArtifact-  :: forall x z.-     AstArtifactRev x z+  :: forall x z. TKAllNum (ADTensorKind z)+  => AstArtifactRev x z        -- ^ the artifact containing the symbolic code of the derivative   -> Concrete x   -> Maybe (Concrete (ADTensorKind z))-  -> (Concrete (ADTensorKind x), Concrete z)+  -> (Concrete z, Concrete (ADTensorKind x)) {-# INLINE revInterpretArtifact #-} revInterpretArtifact AstArtifactRev{..} parameters mdt =   let azftk = varNameToFTK artVarDtRev@@ -241,22 +281,81 @@         Just dt ->           if tftkG (ftkToSTK azftk) (unConcrete dt) == azftk           then extendEnv artVarDtRev dt env-          else error "revInterpretArtifact: reverse derivative incoming cotangent must have the same shape as the codomain of the objective function"-      gradient = interpretAstPrimal envDt artDerivativeRev-      primal = interpretAstPrimal env artPrimalRev-  in (gradient, primal)+          else error $ "revInterpretArtifact: reverse derivative incoming cotangent must have the same shape as the codomain of the objective function: "+                       ++ show (tftkG (ftkToSTK azftk) (unConcrete dt), azftk)+      gradient = interpretAstFull envDt artDerivativeRev+      primal = interpretAstFull env artPrimalRev+  in (primal, gradient) +-- These three functions are as above, but the dt must be provided and so,+-- due to technical reasons, the type is less constrained.+revMaybeDt+  :: forall src ztgt tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> Concrete (ADTensorKind ztgt)+  -> ( Concrete ztgt+     , Value src )  -- morally Value (ADTensorKind src)+{-# INLINE revMaybeDt #-}+revMaybeDt f vals0 dt =+  let valsTarget = toTarget vals0+      xftk = tftkG (knownSTK @(X src)) $ unConcrete valsTarget+      artifactRaw = revArtifactAdaptDt f xftk+      artifact = simplifyArtifactRev artifactRaw+      (primal, res) = revInterpretArtifactDt artifact valsTarget dt+  in (primal, fromTarget $ fromADTensorKindShared xftk res) +revArtifactAdaptDt+  :: forall src ztgt tgt.+     ( AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> FullShapeTK (X src)+  -> AstArtifactRev (X src) ztgt+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE revArtifactAdaptDt #-}+revArtifactAdaptDt f xftk =+  let g :: AstTensor AstMethodLet FullSpan (X src) -> tgt+      g !arg = simplifyUserCode $ ttlet arg $ f . fromTarget+                                  -- fromTarget requires duplicable+  in revProduceArtifactDt g emptyEnv xftk++revInterpretArtifactDt+  :: forall x z.+     AstArtifactRev x z+       -- ^ the artifact containing the symbolic code of the derivative+  -> Concrete x+  -> Concrete (ADTensorKind z)+  -> (Concrete z, Concrete (ADTensorKind x))+{-# INLINE revInterpretArtifactDt #-}+revInterpretArtifactDt AstArtifactRev{..} parameters dt =+  let azftk = varNameToFTK artVarDtRev+      env = extendEnv artVarDomainRev parameters emptyEnv+      envDt =+        if tftkG (ftkToSTK azftk) (unConcrete dt) == azftk+        then extendEnv artVarDtRev dt env+        else error $ "revInterpretArtifactDt: reverse derivative incoming cotangent must have the same shape as the codomain of the objective function: "+                     ++ show (tftkG (ftkToSTK azftk) (unConcrete dt), azftk)+      gradient = interpretAstFull envDt artDerivativeRev+      primal = interpretAstFull env artPrimalRev+  in (primal, gradient)++ -- * Symbolic reverse derivative adaptors' testing-only internal machinery  revArtifactDelta   :: forall src ztgt tgt.-     ( AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     ( TKAllNum (ADTensorKind ztgt)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src      , tgt ~ AstTensor AstMethodLet FullSpan ztgt )   => IncomingCotangentHandling   -> (src -> tgt)  -- ^ the objective function   -> FullShapeTK (X src)-  -> (AstArtifactRev (X src) ztgt, Delta (AstRaw PrimalSpan) ztgt)+  -> (AstArtifactRev (X src) ztgt, Delta (AstRaw FullSpan) ztgt)        -- ^ the artifact containing the symbolic code of the derivative {-# INLINE revArtifactDelta #-} revArtifactDelta cotangentHandling f xftk =@@ -266,26 +365,25 @@                                 (forwardPassByInterpretation g emptyEnv) xftk  revProduceArtifactWithoutInterpretation-  :: forall x z.-     IncomingCotangentHandling-  -> (ADVal (AstRaw PrimalSpan) x -> ADVal (AstRaw PrimalSpan) z)+  :: forall x z. TKAllNum (ADTensorKind z)+  => IncomingCotangentHandling+  -> (ADVal (AstRaw FullSpan) x -> ADVal (AstRaw FullSpan) z)   -> FullShapeTK x-  -> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)+  -> (AstArtifactRev x z, Delta (AstRaw FullSpan) z)        -- ^ the artifact containing the symbolic code of the derivative {-# INLINE revProduceArtifactWithoutInterpretation #-}-revProduceArtifactWithoutInterpretation cotangentHandling f xftk =+revProduceArtifactWithoutInterpretation cotangentHandling f =   -- No simplification performed to let individual tests decide.   revArtifactFromForwardPass cotangentHandling                              (forwardPassByApplication f)-                             xftk  forwardPassByApplication   :: forall x z.-     (ADVal (AstRaw PrimalSpan) x -> ADVal (AstRaw PrimalSpan) z)-  -> AstTensor AstMethodShare PrimalSpan x-  -> AstVarName FullSpan x+     (ADVal (AstRaw FullSpan) x -> ADVal (AstRaw FullSpan) z)+  -> AstTensor AstMethodShare FullSpan x+  -> AstVarName '(FullSpan, x)   -> AstTensor AstMethodLet FullSpan x-  -> ADVal (AstRaw PrimalSpan) z+  -> ADVal (AstRaw FullSpan) z {-# INLINE forwardPassByApplication #-} forwardPassByApplication g astVarPrimal var _astVar =   let deltaInputs = generateDeltaInputs $ varNameToFTK var@@ -311,13 +409,29 @@   -> Value src  -- morally (ADTensorKind src)   -> Concrete (ADTensorKind ztgt) {-# INLINE jvp #-}-jvp f vals0 ds =+jvp f vals0 ds = snd $ jvp2 f vals0 ds++-- | The @jvp2@ operation works just as 'jvp', but additionally produces+-- the primal result of the objective function at the direction parameter+-- (without taking the perturbation into account) at minimal extra cost.+jvp2+  :: forall src ztgt tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> Value src  -- morally (ADTensorKind src)+  -> (Concrete ztgt, Concrete (ADTensorKind ztgt))+{-# INLINE jvp2 #-}+jvp2 f vals0 ds =   let valsTarget = toTarget vals0       xftk = tftkG (knownSTK @(X src)) $ unConcrete valsTarget       artifactRaw = fwdArtifactAdapt f xftk-      artifact = simplifyArtifactDerivative artifactRaw-  in fst $ fwdInterpretArtifact artifact valsTarget-         $ toADTensorKindShared xftk (toTarget ds)+      artifact = simplifyArtifactFwd artifactRaw+  in fwdInterpretArtifact artifact valsTarget+     $ toADTensorKindShared xftk (toTarget ds)        -- the shapes of vals0 vs ds are checked in fwdInterpretArtifact  -- | Compute the forward derivative not for a specific input, but as symbolic@@ -349,7 +463,7 @@   -> Concrete (ADTensorKind x)   -> Concrete (ADTensorKind z) {-# INLINE jvpInterpretArtifact #-}-jvpInterpretArtifact art parameters = fst . fwdInterpretArtifact art parameters+jvpInterpretArtifact art parameters = snd . fwdInterpretArtifact art parameters   -- the shapes of parameters vs ds are checked in fwdInterpretArtifact  @@ -366,7 +480,7 @@ {-# INLINE fwdArtifactAdapt #-} fwdArtifactAdapt f xftk =   let g :: AstTensor AstMethodLet FullSpan (X src) -> tgt-      g !arg = simplifyInline $ ttlet arg $ f . fromTarget+      g !arg = simplifyUserCode $ ttlet arg $ f . fromTarget                                   -- fromTarget requires duplicable   in fwdProduceArtifact g emptyEnv xftk @@ -376,7 +490,7 @@        -- ^ the artifact containing the symbolic code of the derivative   -> Concrete x   -> Concrete (ADTensorKind x)-  -> (Concrete (ADTensorKind z), Concrete z)+  -> (Concrete z, Concrete (ADTensorKind z)) {-# INLINE fwdInterpretArtifact #-} fwdInterpretArtifact AstArtifactFwd{..} parameters ds =   let xftk = varNameToFTK artVarDomainFwd@@ -385,11 +499,13 @@       envD = extendEnv artVarDsFwd ds env   in if tftkG xstk (unConcrete parameters) == xftk      then if tftkG (adSTK xstk) (unConcrete ds) == adFTK xftk-          then let derivative = interpretAstPrimal envD artDerivativeFwd-                   primal = interpretAstPrimal env artPrimalFwd-               in (derivative, primal)-          else error "fwdInterpretArtifact: forward derivative perturbation must have the same shape as the domain of the objective function"-     else error "fwdInterpretArtifact: forward derivative input must have the same shape as the domain of the objective function"+          then let derivative = interpretAstFull envD artDerivativeFwd+                   primal = interpretAstFull env artPrimalFwd+               in (primal, derivative)+          else error $ "fwdInterpretArtifact: forward derivative perturbation must have the same shape as the domain of the objective function: "+                       ++ show (tftkG (adSTK xstk) (unConcrete ds), adFTK xftk)+     else error $ "fwdInterpretArtifact: forward derivative input must have the same shape as the domain of the objective function: "+                  ++ show (tftkG xstk (unConcrete parameters), xftk)   -- * Symbolic forward derivative adaptors' testing-only internal machinery@@ -400,7 +516,7 @@      , tgt ~ AstTensor AstMethodLet FullSpan ztgt )   => (src -> tgt)  -- ^ the objective function   -> FullShapeTK (X src)-  -> (AstArtifactFwd (X src) ztgt, Delta (AstRaw PrimalSpan) ztgt)+  -> (AstArtifactFwd (X src) ztgt, Delta (AstRaw FullSpan) ztgt)        -- ^ the artifact containing the symbolic code of the derivative {-# INLINE fwdArtifactDelta #-} fwdArtifactDelta f xftk =@@ -411,6 +527,13 @@  -- * Non-symbolic reverse derivative adaptors +-- These adaptor are generalized to any suitable target, because+-- it's useful when defining complex nested derivatives.+-- However, it incurs many extra type applications, which is why+-- we refrain from such a generalization for the symbolic adaptors above,+-- which are less likely to be used in nested derivatives+-- and more likely to be used a lot in simple use cases..+ -- We are inlining these functions because they take function arguments -- and are not too large. However, because they are called in many places, -- we break the inline chain not far from the top, to avoid exe blowup.@@ -419,185 +542,155 @@ -- reverse derivative operation sets the incoming cotangent @dt@ to be 1 -- and assumes the codomain of the function to be differentiated is a scalar. cgrad-  :: forall src r tgt.-     ( X src ~ X (DValue src), KnownSTK (X src)-     , AdaptableTarget (ADVal Concrete) src-     , AdaptableTarget Concrete (DValue src)-     , tgt ~ ADVal Concrete (TKScalar r) )+  :: forall src r tgt target.+     ( GoodScalar r, X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal target) src+     , AdaptableTarget target (DValue src)+     , tgt ~ ADVal target (TKScalar r)+     , ADReadyNoLet target, ShareTensor target )   => (src -> tgt)  -- ^ the objective function   -> DValue src   -> DValue src  -- morally DValue (ADTensorKind src) {-# INLINE cgrad #-}-cgrad f vals = crevMaybe f vals Nothing+cgrad f vals = snd $ cgrad2 f vals +cgrad2+  :: forall src r tgt target.+     ( GoodScalar r, X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal target) src+     , AdaptableTarget target (DValue src)+     , tgt ~ ADVal target (TKScalar r)+     , ADReadyNoLet target, ShareTensor target )+  => (src -> tgt)  -- ^ the objective function+  -> DValue src+  -> ( target (TKScalar r)+     , DValue src )  -- morally DValue (ADTensorKind src)+{-# INLINE cgrad2 #-}+cgrad2 f vals | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =+  crevMaybe f vals Nothing+ -- | This more general version of the concrete (non-symbolic) -- reverse derivative operation additionally takes the sensitivity parameter -- (the incoming cotangent). cvjp-  :: forall src ztgt tgt.+  :: forall src ztgt tgt target.      ( X src ~ X (DValue src), KnownSTK (X src)-     , AdaptableTarget (ADVal Concrete) src-     , AdaptableTarget Concrete (DValue src)-     , tgt ~ ADVal Concrete ztgt )+     , AdaptableTarget (ADVal target) src+     , AdaptableTarget target (DValue src)+     , tgt ~ ADVal target ztgt+     , ADReadyNoLet target, ShareTensor target )   => (src -> tgt)  -- ^ the objective function   -> DValue src-  -> Concrete (ADTensorKind ztgt)+  -> target (ADTensorKind ztgt)   -> DValue src  -- morally DValue (ADTensorKind src) {-# INLINE cvjp #-}-cvjp f vals dt = crevMaybe f vals (Just dt)+cvjp f vals0 dt = snd $ cvjp2 f vals0 dt +cvjp2+  :: forall src ztgt tgt target.+     ( X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal target) src+     , AdaptableTarget target (DValue src)+     , tgt ~ ADVal target ztgt+     , ADReadyNoLet target, ShareTensor target )+  => (src -> tgt)  -- ^ the objective function+  -> DValue src+  -> target (ADTensorKind ztgt)+  -> ( target ztgt+     , DValue src )  -- morally DValue (ADTensorKind src)+{-# INLINE cvjp2 #-}+cvjp2 = crevMaybeDt + -- * Non-symbolic reverse derivative adaptors' internal machinery  crevMaybe-  :: forall src ztgt tgt.-     ( X src ~ X (DValue src), KnownSTK (X src)-     , AdaptableTarget (ADVal Concrete) src-     , AdaptableTarget Concrete (DValue src)-     , tgt ~ ADVal Concrete ztgt )+  :: forall src ztgt tgt target.+     ( TKAllNum (ADTensorKind ztgt)+     , X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal target) src+     , AdaptableTarget target (DValue src)+     , tgt ~ ADVal target ztgt+     , ADReadyNoLet target, ShareTensor target )   => (src -> tgt)  -- ^ the objective function   -> DValue src-  -> Maybe (Concrete (ADTensorKind ztgt))-  -> DValue src  -- morally DValue (ADTensorKind src)+  -> Maybe (target (ADTensorKind ztgt))+  -> ( target ztgt+     , DValue src )  -- morally SValue (ADTensorKind src) {-# INLINE crevMaybe #-} crevMaybe f vals0 mdt =   let valsTarget = toTarget vals0-      g :: ADVal Concrete (X src) -> tgt+      g :: ADVal target (X src) -> tgt       g = f . fromTarget-      xftk = tftkG (knownSTK @(X src)) $ unConcrete valsTarget-  in fromTarget $ fromADTensorKindShared (ftkToSTK xftk)-     $ fst $ crevOnParams mdt g xftk valsTarget+      xftk = tftk (knownSTK @(X src)) valsTarget+      (primal, res) = crevOnParams mdt g xftk valsTarget+  in (primal, fromTarget $ fromADTensorKindShared xftk res) +-- This function is as above, but the dt must be provided and so,+-- due to technical reasons, the type is less constrained.+crevMaybeDt+  :: forall src ztgt tgt target.+     ( X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal target) src+     , AdaptableTarget target (DValue src)+     , tgt ~ ADVal target ztgt+     , ADReadyNoLet target, ShareTensor target )+  => (src -> tgt)  -- ^ the objective function+  -> DValue src+  -> target (ADTensorKind ztgt)+  -> ( target ztgt+     , DValue src )  -- morally SValue (ADTensorKind src)+{-# INLINE crevMaybeDt #-}+crevMaybeDt f vals0 dt =+  let valsTarget = toTarget vals0+      g :: ADVal target (X src) -> tgt+      g = f . fromTarget+      xftk = tftk (knownSTK @(X src)) valsTarget+      (primal, res) = crevOnParamsDt dt g xftk valsTarget+  in (primal, fromTarget $ fromADTensorKindShared xftk res) + -- * Non-symbolic forward derivative adaptors  -- | Concrete (non-symbolic) forward derivative operation. It always takes -- the perturbation parameter, by convention. cjvp-  :: forall src ztgt tgt.+  :: forall src ztgt tgt target.      ( X src ~ X (DValue src), KnownSTK (X src)-     , AdaptableTarget (ADVal Concrete) src-     , AdaptableTarget Concrete (DValue src)-     , tgt ~ ADVal Concrete ztgt )+     , AdaptableTarget (ADVal target) src+     , AdaptableTarget target (DValue src)+     , tgt ~ ADVal target ztgt+     , ADReadyNoLet target, ShareTensor target )   => (src -> tgt)  -- ^ the objective function   -> DValue src   -> DValue src  -- morally (ADTensorKind src)-  -> Concrete (ADTensorKind ztgt)+  -> target (ADTensorKind ztgt) {-# INLINE cjvp #-}-cjvp f vals ds = fst $ cfwdBoth f vals ds----- * Non-symbolic forward derivative adaptors' internal machinery+cjvp f vals ds = snd $ cjvp2 f vals ds -cfwdBoth-  :: forall src ztgt tgt.+-- | The @cjvp2@ operation works just as 'cjvp', but additionally produces+-- the primal result of the objective function at the direction parameter+-- (without taking the perturbation into account) at no extra cost.+cjvp2+  :: forall src ztgt tgt target.      ( X src ~ X (DValue src), KnownSTK (X src)-     , AdaptableTarget (ADVal Concrete) src-     , AdaptableTarget Concrete (DValue src)-     , tgt ~ ADVal Concrete ztgt )+     , AdaptableTarget (ADVal target) src+     , AdaptableTarget target (DValue src)+     , tgt ~ ADVal target ztgt+     , ADReadyNoLet target, ShareTensor target )   => (src -> tgt)  -- ^ the objective function   -> DValue src   -> DValue src  -- morally (ADTensorKind src)-  -> (Concrete (ADTensorKind ztgt), Concrete ztgt)-{-# INLINE cfwdBoth #-}-cfwdBoth f vals0 ds =+  -> (target ztgt, target (ADTensorKind ztgt))+{-# INLINE cjvp2 #-}+cjvp2 f vals0 ds =   let valsTarget = toTarget vals0-      xftk = tftkG (knownSTK @(X src)) $ unConcrete valsTarget-      g :: ADVal Concrete (X src) -> tgt+      xftk = tftk (knownSTK @(X src)) valsTarget+      g :: ADVal target (X src) -> tgt       g = f . fromTarget       dsTarget = toTarget ds-  in if tftkG (ftkToSTK xftk) (unConcrete dsTarget) == xftk+  in if tftk (ftkToSTK xftk) dsTarget == xftk      then cfwdOnParams xftk valsTarget g           $ toADTensorKindShared xftk dsTarget-     else error "cfwdBoth: forward derivative input must have the same shape as the perturbation argument"-------- This specialization is not possible where the functions are defined,--- due to dependency cycles, but it's possible here:-{-# SPECIALIZE gradientFromDelta :: FullShapeTK x -> FullShapeTK z -> Concrete (ADTensorKind z) -> Delta Concrete z -> Concrete (ADTensorKind x) #-}-{-# SPECIALIZE evalRev :: FullShapeTK y -> EvalState Concrete -> Concrete (ADTensorKind y) -> Delta Concrete y -> EvalState Concrete #-}-{-# SPECIALIZE evalRevFTK :: EvalState Concrete -> Concrete (ADTensorKind y) -> Delta Concrete y -> EvalState Concrete #-}--- RULE left-hand side too complicated to desugar:--- {-# SPECIALIZE evalRevSame :: y ~ ADTensorKind y => EvalState Concrete -> Concrete (ADTensorKind y) -> Delta Concrete y -> EvalState Concrete #-}-{-# SPECIALIZE evalRevFromnMap :: EvalState Concrete -> EvalState Concrete #-}--{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKScalar Double) -> Delta Concrete (TKScalar Double) -> EvalState Concrete #-}-{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKScalar Float) -> Delta Concrete (TKScalar Float) -> EvalState Concrete #-}-{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKR n Double) -> Delta Concrete (TKR n Double) -> EvalState Concrete #-}-{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKR n Float) -> Delta Concrete (TKR n Float) -> EvalState Concrete #-}-{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKS sh Double) -> Delta Concrete (TKS sh Double) -> EvalState Concrete #-}-{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKS sh Float) -> Delta Concrete (TKS sh Float) -> EvalState Concrete #-}-{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKX sh Double) -> Delta Concrete (TKX sh Double) -> EvalState Concrete #-}-{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKX sh Float) -> Delta Concrete (TKX sh Float) -> EvalState Concrete #-}----- These and all other SPECIALIZE pragmas are needed due to the already--- mostly fixed issues #21286 and others, even just to compare--- the output with them and without.--- This is needed for all three AstSpan values, to handle recursive calls--- from interpretAstDual, etc.-{-# SPECIALIZE interpretAst-  :: AstEnv (ADVal Concrete)-  -> AstTensor AstMethodLet PrimalSpan y-  -> ADVal Concrete y #-}-{-# SPECIALIZE interpretAst-  :: AstEnv (ADVal (AstRaw PrimalSpan))-  -> AstTensor AstMethodLet PrimalSpan y-  -> ADVal (AstRaw PrimalSpan) y #-}-{-# SPECIALIZE interpretAst-  :: AstEnv Concrete-  -> AstTensor AstMethodLet PrimalSpan y-  -> Concrete y #-}-{-# SPECIALIZE interpretAst-  :: AstEnv (ADVal Concrete)-  -> AstTensor AstMethodLet DualSpan y-  -> ADVal Concrete y #-}-{-# SPECIALIZE interpretAst-  :: AstEnv (ADVal (AstRaw PrimalSpan))-  -> AstTensor AstMethodLet DualSpan y-  -> ADVal (AstRaw PrimalSpan) y #-}-{-# SPECIALIZE interpretAst-  :: AstEnv Concrete-  -> AstTensor AstMethodLet DualSpan y-  -> Concrete y #-}-{-# SPECIALIZE interpretAst-  :: AstEnv (ADVal Concrete)-  -> AstTensor AstMethodLet FullSpan y-  -> ADVal Concrete y #-}-{-# SPECIALIZE interpretAst-  :: AstEnv (ADVal (AstRaw PrimalSpan))-  -> AstTensor AstMethodLet FullSpan y-  -> ADVal (AstRaw PrimalSpan) y #-}-{-# SPECIALIZE interpretAst-  :: AstEnv Concrete-  -> AstTensor AstMethodLet FullSpan y-  -> Concrete y #-}--{-# SPECIALIZE interpretAstPrimal-  :: AstEnv (ADVal Concrete)-  -> AstTensor AstMethodLet PrimalSpan y-  -> Concrete y #-}-{-# SPECIALIZE interpretAstPrimal-  :: AstEnv (ADVal (AstRaw PrimalSpan))-  -> AstTensor AstMethodLet PrimalSpan y-  -> AstRaw PrimalSpan y #-}-{-# SPECIALIZE interpretAstPrimal-  :: AstEnv Concrete-  -> AstTensor AstMethodLet PrimalSpan y-  -> Concrete y #-}--{-# SPECIALIZE interpretAstBool-  :: AstEnv (ADVal Concrete)-  -> AstBool AstMethodLet-  -> Bool #-}-{-# SPECIALIZE interpretAstBool-  :: AstEnv (ADVal (AstRaw PrimalSpan))-  -> AstBool AstMethodLet-  -> AstBool AstMethodShare #-}-{-# SPECIALIZE interpretAstBool-  :: AstEnv Concrete-  -> AstBool AstMethodLet-  -> Bool #-}+     else error $ "cjvp2: forward derivative input must have the same shape as the perturbation argument: "+                  ++ show (tftk (ftkToSTK xftk) dsTarget, xftk)
− src/HordeAd/AstEngine.hs 
@@ -1,149 +0,0 @@--- | Predefined common functions for simplification and pretty-printing of AST.-module HordeAd.AstEngine-  ( -- * The joint inlining and simplification term transformations-    simplifyArtifact, simplifyArtifactGradient, simplifyArtifactDerivative-  , simplifyInline, simplifyInlineContract, simplifyInlineContractNoExpand-    -- * Pretty-printing terms in a few useful configurations-  , printAstVarName-  , printAstSimple, printAstPretty, printAstPrettyButNested-  , printArtifactSimple, printArtifactPretty-  , printArtifactPrimalSimple, printArtifactPrimalPretty-  ) where--import Prelude--import Data.EnumMap.Strict qualified as EM-import Data.IntMap.Strict (IntMap)-import Data.IntMap.Strict qualified as IM--import HordeAd.Core.Ast-import HordeAd.Core.AstInline-import HordeAd.Core.AstPrettyPrint-import HordeAd.Core.AstTraverse---- * The joint inlining and simplification term transformation---- | Simplify the whole reverse derivative artifact (which includes--- also the primal value computed during the differentiation process).-simplifyArtifact :: forall x z.-                    AstArtifactRev x z -> AstArtifactRev x z-simplifyArtifact art =-  let !der = simplifyInlineContract $ artDerivativeRev art in-  let !prim = simplifyInlineContract $ artPrimalRev art-  in art {artDerivativeRev = der, artPrimalRev = prim}---- | Simplify only the gradient in the reverse derivative artifact.-simplifyArtifactGradient :: forall x z.-                            AstArtifactRev x z -> AstArtifactRev x z-simplifyArtifactGradient art =-  art { artDerivativeRev =-        simplifyInlineContract $ artDerivativeRev art }---- | Simplify only the derivative in the forward derivative artifact.-simplifyArtifactDerivative :: forall x z.-                              AstArtifactFwd x z -> AstArtifactFwd x z-simplifyArtifactDerivative art =-  art { artDerivativeFwd =-        simplifyInlineContract $ artDerivativeFwd art }---- | A mixture of simplification and inlining to use when the resultng--- term is not yet supposed to be interpreted using a computational backed,--- but rather to be stored and later composed with other terms.-simplifyInline-  :: forall z s. AstSpan s-  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z-simplifyInline =-  simplifyAst . expandAst . snd . inlineAst EM.empty-  . simplifyAst . expandAst . snd . inlineAst EM.empty-  . simplifyAst---- | A mixture of simplification, inlining and recognition of additional--- backend-specific primitives, to be used just before a term--- is interpreted as a value in the computational backend.-simplifyInlineContract-  :: forall z s. AstSpan s-  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z-simplifyInlineContract =-  contractAst . expandAst . snd . inlineAst EM.empty-  . simplifyAst . expandAst . snd . inlineAst EM.empty-  . simplifyAst--simplifyInlineContractNoExpand-  :: forall z s. AstSpan s-  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z-simplifyInlineContractNoExpand =-  contractAst . simplifyAst . snd . inlineAst EM.empty-  . simplifyAst . snd . inlineAst EM.empty-  . simplifyAst----- * Pretty-printing terms in a few useful configurations--printAstVarRename :: AstSpan s-                  => IntMap String -> AstVarName s y -> String-printAstVarRename renames var =-  printAstVar (defaulPrintConfig {varRenames = renames}) var ""--printAstSimpleRename :: AstSpan s-                     => IntMap String -> AstTensor ms s y -> String-printAstSimpleRename renames t =-  printAst-    (defaulPrintConfig {loseRoudtrip = False, varRenames = renames}) 0 t ""--printAstPrettyRename :: AstSpan s-                     => IntMap String -> AstTensor ms s y -> String-printAstPrettyRename renames t =-  printAst (defaulPrintConfig {varRenames = renames}) 0 t ""--printAstVarName :: AstSpan s-                => AstVarName s y -> String-printAstVarName var =-  printAstVar defaulPrintConfig var ""---- | Print an AST term in a form close to being able to roundtrip,--- including explicit sharing preservation.-printAstSimple :: AstSpan s-               => AstTensor ms s y -> String-printAstSimple t =-  printAst (defaulPrintConfig {loseRoudtrip = False}) 0 t ""---- | Print an AST term in a readable form that does not roundtrip,--- and where Haskell @let@ (sharing on Haskell heap) is used instead--- of explicit sharing of subterms.-printAstPretty :: AstSpan s-               => AstTensor ms s y -> String-printAstPretty t =-  printAst defaulPrintConfig 0 t ""--printAstPrettyButNested :: AstSpan s-                        => AstTensor ms s y -> String-printAstPrettyButNested t =-  printAst (defaulPrintConfig {ignoreNestedLambdas = False}) 0 t ""--printArtifactSimple :: AstArtifactRev x z -> String-printArtifactSimple AstArtifactRev{..} =-  let nDt = fromEnum (varNameToAstVarId artVarDtRev) - 100000000-      renames = IM.singleton nDt "dret"-      varsPP = [ printAstVarRename renames artVarDtRev-               , printAstVarRename renames artVarDomainRev ]-  in "\\" ++ unwords varsPP-          ++ " -> " ++ printAstSimpleRename renames artDerivativeRev--printArtifactPretty :: AstArtifactRev x z -> String-printArtifactPretty AstArtifactRev{..} =-  let nDt = fromEnum (varNameToAstVarId artVarDtRev) - 100000000-      renames = IM.singleton nDt "dret"-      varsPP = [ printAstVarRename renames artVarDtRev-               , printAstVarRename renames artVarDomainRev ]-  in "\\" ++ unwords varsPP-          ++ " -> " ++ printAstPrettyRename renames artDerivativeRev--printArtifactPrimalSimple :: AstArtifactRev x z -> String-printArtifactPrimalSimple AstArtifactRev{..} =-  "\\" ++ printAstVarName artVarDomainRev-       ++ " -> " ++ printAstSimple artPrimalRev--printArtifactPrimalPretty :: AstArtifactRev x z -> String-printArtifactPrimalPretty AstArtifactRev{..} =-  "\\" ++ printAstVarName artVarDomainRev-       ++ " -> " ++ printAstPretty artPrimalRev
src/HordeAd/Core/Adaptor.hs view
@@ -11,7 +11,7 @@ module HordeAd.Core.Adaptor   ( AdaptableTarget(..), TermValue(..), DualNumberValue(..)   , ForgetShape(..), RandomValue(..)-  , stkOfListR+  , stkOfIxR     -- * Helper classes and types   , Tups, NoShapeTensorKind   ) where@@ -30,13 +30,14 @@ import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (unsafeCoerceRefl)+import Data.Array.Nested.Types (pattern SZ, unsafeCoerceRefl)  import HordeAd.Core.Ast import HordeAd.Core.CarriersADVal import HordeAd.Core.CarriersConcrete import HordeAd.Core.ConvertTensor import HordeAd.Core.Ops+import HordeAd.Core.OpsADVal () import HordeAd.Core.OpsAst () import HordeAd.Core.TensorKind import HordeAd.Core.Types@@ -100,21 +101,21 @@   type X (target y) = y   toTarget = id   fromTarget t = t-  {-# SPECIALIZE instance AdaptableTarget Concrete (Concrete (TKS sh Double)) #-}-  {-# SPECIALIZE instance AdaptableTarget Concrete (Concrete (TKS sh Float)) #-}-    -- a failed attempt to specialize without -fpolymorphic-specialisation  instance KnownSTK y          => TermValue (AstTensor AstMethodLet FullSpan y) where   type Value (AstTensor AstMethodLet FullSpan y) = Concrete y   fromValue t = tconcrete (tftkG (knownSTK @y) $ unConcrete t) t -instance (BaseTensor target, BaseTensor (PrimalOf target), KnownSTK y)-         => DualNumberValue (target y) where-  type DValue (target y) = Concrete y-  fromDValue t = tfromPrimal (knownSTK @y)-                 $ tconcrete (tftkG (knownSTK @y) $ unConcrete t) t+instance DualNumberValue (Concrete y) where+  type DValue (Concrete y) = Concrete y+  fromDValue = id +instance (BaseTensor target, KnownSTK y)+         => DualNumberValue (ADVal target y) where+  type DValue (ADVal target y) = target y+  fromDValue t = fromPrimalFTK (tftk (knownSTK @y) t) t+ instance ForgetShape (target (TKScalar r)) where   type NoShape (target (TKScalar r)) = target (TKScalar r)   forgetShape = id@@ -176,7 +177,6 @@            arr = createRandomVector (shsSize (knownShS @sh)) g1        in (arr, g2))       (srepl def, g)-   where srepl = tsconcrete . Nested.sreplicateScal knownShS   -- {-# SPECIALIZE instance (KnownShS sh, GoodScalar r, Fractional r, Random r) => RandomValue (Concrete (TKS sh r)) #-}   {-# SPECIALIZE instance KnownShS sh => RandomValue (Concrete (TKS sh Double)) #-}   {-# SPECIALIZE instance KnownShS sh => RandomValue (Concrete (TKS sh Float)) #-}@@ -208,60 +208,59 @@  -- * Compound instances -instance (BaseTensor target, ConvertTensor target, GoodScalar r)+instance (BaseTensor target, GoodScalar r)          => AdaptableTarget target [target (TKScalar r)] where   type X [target (TKScalar r)] = TKR 1 r   toTarget l = if null l                then trconcrete Nested.remptyArray-               else trfromVector $ V.fromList $ map rfromK l-  fromTarget = map kfromR . trunravelToList-                              -- inefficient, but we probably can't do better+               else let v = V.fromList l+                    in trfromVectorLinear (V.length v :$: ZSR) v+  fromTarget = trtoListLinear  -- inefficient, but we probably can't do better -instance (BaseTensor target, ConvertTensor target, GoodScalar r)+instance (BaseTensor target, GoodScalar r)          => AdaptableTarget target                             (Data.Vector.Vector (target (TKScalar r))) where   type X (Data.Vector.Vector (target (TKScalar r))) = TKR 1 r   toTarget v = if V.null v                then trconcrete Nested.remptyArray-               else trfromVector $ V.map rfromK v-  fromTarget =-    V.fromList . map kfromR . trunravelToList-                                -- inefficient, but we probably can't do better+               else trfromVectorLinear (V.length v :$: ZSR) v+  fromTarget t = V.fromListN (rwidth t) . trtoListLinear $ t+    -- inefficient, but we probably can't do better  type family Tups n t where   Tups 0 t = TKUnit   Tups n t = TKProduct t (Tups (n - 1) t) -stkOfListR :: forall t n.-              SingletonTK t -> SNat n -> SingletonTK (Tups n t)-stkOfListR _ (SNat' @0) = stkUnit-stkOfListR stk SNat =+stkOfIxR :: forall t n.+            SingletonTK t -> SNat n -> SingletonTK (Tups n t)+stkOfIxR _ SZ = stkUnit+stkOfIxR stk SNat =   gcastWith (unsafeCoerceRefl :: (1 <=? n) :~: True) $   gcastWith (unsafeCoerceRefl :: Tups n t :~: TKProduct t (Tups (n - 1) t)) $-  STKProduct stk (stkOfListR stk (SNat @(n - 1)))+  STKProduct stk (stkOfIxR stk (SNat @(n - 1)))  instance (BaseTensor target, KnownNat n, AdaptableTarget target a)-         => AdaptableTarget target (ListR n a) where-  type X (ListR n a) = Tups n (X a)-  toTarget ZR = tkconcrete Z1-  toTarget ((:::) @n1 a rest) =+         => AdaptableTarget target (IxR n a) where+  type X (IxR n a) = Tups n (X a)+  toTarget ZIR = tkconcrete Z1+  toTarget ((:.:) @n1 a rest) =     gcastWith (unsafeCoerceRefl-               :: X (ListR n a) :~: TKProduct (X a) (X (ListR n1 a))) $+               :: X (IxR n a) :~: TKProduct (X a) (X (IxR n1 a))) $     let a1 = toTarget a         rest1 = toTarget rest     in tpair a1 rest1   fromTarget tups = case SNat @n of-    SNat' @0 -> ZR+    SZ -> ZIR     _ ->       gcastWith (unsafeCoerceRefl :: (1 <=? n) :~: True) $       gcastWith (unsafeCoerceRefl-                 :: X (ListR n a) :~: TKProduct (X a) (X (ListR (n - 1) a))) $+                 :: X (IxR n a) :~: TKProduct (X a) (X (IxR (n - 1) a))) $       let (a1, rest1) = (tproject1 tups, tproject2 tups)           a = fromTarget a1-          rest = fromTarget rest1-      in (a ::: rest)-  {-# SPECIALIZE instance (KnownNat n, AdaptableTarget (AstTensor AstMethodLet FullSpan) a) => AdaptableTarget (AstTensor AstMethodLet FullSpan) (ListR n a) #-}-  {-# SPECIALIZE instance (KnownNat n, AdaptableTarget (ADVal Concrete) a) => AdaptableTarget (ADVal Concrete) (ListR n a) #-}+          rest = fromTarget @_ @(IxR (n - 1) a) rest1+      in (a :.: rest)+  {-# SPECIALIZE instance (KnownNat n, AdaptableTarget (AstTensor AstMethodLet FullSpan) a) => AdaptableTarget (AstTensor AstMethodLet FullSpan) (IxR n a) #-}+  {-# SPECIALIZE instance (KnownNat n, AdaptableTarget (ADVal Concrete) a) => AdaptableTarget (ADVal Concrete) (IxR n a) #-}  instance TermValue a => TermValue [a] where   type Value [a] = [Value a]@@ -271,10 +270,10 @@   type Value (Data.Vector.Vector a) = Data.Vector.Vector (Value a)   fromValue = V.map fromValue -instance TermValue a => TermValue (ListR n a) where-  type Value (ListR n a) = ListR n (Value a)-  fromValue ZR = ZR-  fromValue (a ::: rest) = fromValue a ::: fromValue rest+instance TermValue a => TermValue (IxR n a) where+  type Value (IxR n a) = IxR n (Value a)+  fromValue ZIR = ZIR+  fromValue (a :.: rest) = fromValue a :.: fromValue rest  instance DualNumberValue a => DualNumberValue [a] where   type DValue [a] = [DValue a]@@ -284,10 +283,10 @@   type DValue (Data.Vector.Vector a) = Data.Vector.Vector (DValue a)   fromDValue = V.map fromDValue -instance DualNumberValue a => DualNumberValue (ListR n a) where-  type DValue (ListR n a) = ListR n (DValue a)-  fromDValue ZR = ZR-  fromDValue (a ::: rest) = fromDValue a ::: fromDValue rest+instance DualNumberValue a => DualNumberValue (IxR n a) where+  type DValue (IxR n a) = IxR n (DValue a)+  fromDValue ZIR = ZIR+  fromDValue (a :.: rest) = fromDValue a :.: fromDValue rest  instance ForgetShape [a] where   type NoShape [a] = [a]@@ -297,19 +296,19 @@   type NoShape (Data.Vector.Vector a) = Data.Vector.Vector a   forgetShape = id -instance ForgetShape a => ForgetShape (ListR n a) where-  type NoShape (ListR n a) = ListR n (NoShape a)-  forgetShape ZR = ZR-  forgetShape (a ::: rest) = forgetShape a ::: forgetShape rest+instance ForgetShape a => ForgetShape (IxR n a) where+  type NoShape (IxR n a) = IxR n (NoShape a)+  forgetShape ZIR = ZIR+  forgetShape (a :.: rest) = forgetShape a :.: forgetShape rest -instance (RandomValue a, KnownNat n) => RandomValue (ListR n a) where+instance (RandomValue a, KnownNat n) => RandomValue (IxR n a) where   randomValue range g = case cmpNat (Proxy @n) (Proxy @0)  of     LTI -> error "randomValue: impossible"-    EQI -> (ZR, g)+    EQI -> (ZIR, g)     GTI -> gcastWith (unsafeCoerceRefl :: (1 <=? n) :~: True) $            let (v, g1) = randomValue range g-               (rest, g2) = randomValue @(ListR (n - 1) a) range g1-           in (v ::: rest, g2)+               (rest, g2) = randomValue @(IxR (n - 1) a) range g1+           in (v :.: rest, g2)   -- * Tuple instances
src/HordeAd/Core/Ast.hs view
@@ -20,32 +20,38 @@ -- -- (A copy of the text above is in "HordeAd.Core.Ops".) module HordeAd.Core.Ast-  ( -- * The AstSpan tags and class-    AstSpanType(..), AstSpan(..), sameAstSpan+  ( -- * The AstSpan tags, singletons and operations+    AstSpan(..), PrimalSpan, SAstSpan(..), KnownSpan(..), withKnownSpan+  , primalPart, dualPart, plainPart, fromPrimal, fromDual, fromPlain     -- * Variables and related types   , AstVarId, intToAstVarId-  , AstInt, IntVarName, pattern AstIntVar-  , AstVarName, mkAstVarName, varNameToAstVarId, varNameToFTK, varNameToBounds-  , astVar+  , AstInt, IntVarName, pattern AstIntVar, AstBool+  , AstVarName(..), FtkAndBounds(..)+  , mkAstVarName, mkAstVarNameBounds+  , reshapeVarName, respanVarName, reboundsVarName+  , varNameToAstVarId, varNameToSpan, varNameToFTK, varNameToBounds   , AstArtifactRev(..), AstArtifactFwd(..)-  , AstIxS, AstVarListS, pattern AstLeqInt+  , AstIxS, AstVarListS+  , ListS(..)  -- TODO: ListS should be hidden and its constructor+               -- exposed as AstVarListS constructor+  , pattern ZS, pattern (::$), pattern AstLeqInt     -- * AST-  , AstMethodOfSharing(..), AstTensor(..)-  , AstHFun(..)-  , AstBool(..), OpCodeNum1(..), OpCode1(..), OpCode2(..), OpCodeIntegral2(..)+  , AstMethodOfSharing(..), AstTensor(..), AstHFun(..)+  , OpCodeNum1(..), OpCode1(..), OpCode2(..), OpCodeIntegral2(..)   ) where -import Prelude hiding (foldl')+import Prelude  import Data.Dependent.EnumMap.Strict qualified as DMap-import Data.Functor.Const-import Data.Int (Int64)+import Data.Foldable qualified as Foldable import Data.Kind (Type)-import Data.Some import Data.Type.Equality (TestEquality (..), (:~:) (Refl)) import Data.Vector.Strict qualified as Data.Vector+import GHC.Exts (withDict)+import GHC.IsList (IsList)+import GHC.IsList qualified as IsList import GHC.TypeLits (type (+), type (<=))-import Type.Reflection (Typeable, typeRep)+import Type.Reflection (typeRep)  import Data.Array.Nested (type (++)) import Data.Array.Nested qualified as Nested@@ -54,49 +60,157 @@ import Data.Array.Nested.Shaped.Shape import Data.Array.Nested.Types (Init) +import HordeAd.Core.Conversion import HordeAd.Core.TensorKind import HordeAd.Core.Types --- * The AstSpan tags and class+-- * The AstSpan tags, singletons and operations --- | A kind (a type intended to be promoted) marking whether an AST term--- is supposed to denote the primal part of a dual number, the dual part--- or the whole dual number. It's mainly used to index the terms--- of the AstTensor type  and related GADTs.-type data AstSpanType = PrimalSpan | DualSpan | FullSpan+-- | A type intended to be promoted that marks whether an AST term+-- is supposed to denote the (n-th iteration of taking the) primal part+-- of a dual number, the dual part, the whole dual number or a plain value.+-- It's mainly used to index the terms of the AstTensor type+-- and related GADTs.+type data AstSpan =+  FullSpan | PrimalStepSpan AstSpan | DualSpan | PlainSpan+type PrimalSpan = PrimalStepSpan FullSpan -class Typeable s => AstSpan (s :: AstSpanType) where-  fromPrimal :: AstTensor ms PrimalSpan y -> AstTensor ms s y-  fromDual :: AstTensor ms DualSpan y -> AstTensor ms s y-  primalPart :: AstTensor ms s y -> AstTensor ms PrimalSpan y-  dualPart :: AstTensor ms s y -> AstTensor ms DualSpan y+-- | The singleton type for `AstSpan`.+type role SAstSpan nominal+data SAstSpan s where+  SFullSpan :: SAstSpan FullSpan+  SPrimalStepSpan :: SAstSpan s -> SAstSpan (PrimalStepSpan s)+  SDualSpan :: SAstSpan DualSpan+  SPlainSpan :: SAstSpan PlainSpan --- These are weak instance and we can't move them to AstSimplify,--- because it's too late and also astPrimalPart only works on AstMethodLet.-instance AstSpan PrimalSpan where-  fromPrimal = id-  fromDual t = AstPrimalPart $ AstFromDual t  -- this is primal zero-  primalPart t = t-  dualPart t = AstDualPart $ AstFromPrimal t  -- this is dual zero+instance TestEquality SAstSpan where+  testEquality SFullSpan SFullSpan = Just Refl+  testEquality (SPrimalStepSpan s1) (SPrimalStepSpan s2)+    | Just Refl <- testEquality s1 s2 = Just Refl+  testEquality SDualSpan SDualSpan = Just Refl+  testEquality SPlainSpan SPlainSpan = Just Refl+  testEquality _ _ = Nothing -instance AstSpan DualSpan where-  fromPrimal t = AstDualPart $ AstFromPrimal t  -- this is dual zero-  fromDual = id-  primalPart t = AstPrimalPart $ AstFromDual t  -- this is primal zero-  dualPart t = t+-- These are weak definitions rewriting-wise and we can't move them+-- to AstSimplify to improve this, because it's too late+-- and also astPrimalPart only works on AstMethodLet.+primalPart :: forall s ms y. KnownSpan s+           => AstTensor ms s y -> AstTensor ms (PrimalStepSpan s) y+primalPart t = case knownSpan @s of+  SFullSpan -> cAstPrimalPart t+  SPrimalStepSpan{} -> cAstPrimalPart t+  SDualSpan -> fullSpanToStep knownSpan $ AstFromDual t  -- this is primal zero+  SPlainSpan -> cAstPrimalPart t -instance AstSpan FullSpan where-  fromPrimal = AstFromPrimal-  fromDual = AstFromDual-  primalPart (AstFromPrimal t) = t-  primalPart t = AstPrimalPart t-  dualPart (AstFromDual t) = t-  dualPart t = AstDualPart t+dualPart :: forall s ms y. KnownSpan s+         => AstTensor ms s y -> AstTensor ms DualSpan y+dualPart t = case knownSpan @s of+  SFullSpan -> cAstDualPart t+  SPrimalStepSpan s -> cAstDualPart $ stepToFullSpan s t  -- this is dual zero+  SDualSpan -> t+  SPlainSpan -> AstDualPart $ AstFromPlain t  -- this is dual zero -sameAstSpan :: forall s1 s2. (AstSpan s1, AstSpan s2) => Maybe (s1 :~: s2)-sameAstSpan = testEquality (typeRep @s1) (typeRep @s2)+plainPart :: forall s ms y. KnownSpan s+          => AstTensor ms s y -> AstTensor ms PlainSpan y+plainPart t = case knownSpan @s of+  SFullSpan -> cAstPlainPart t+  SPrimalStepSpan{} -> cAstPlainPart t+  SDualSpan -> AstPlainPart $ AstFromDual t  -- this is plain zero+  SPlainSpan -> t +fromPrimal :: forall s ms y. KnownSpan s+           => AstTensor ms (PrimalStepSpan s) y -> AstTensor ms s y+fromPrimal t = case knownSpan @s of+  SFullSpan -> cAstFromPrimal t+  SPrimalStepSpan{} -> cAstFromPrimal t+  SDualSpan -> cAstDualPart $ stepToFullSpan knownSpan t  -- this is dual zero+  SPlainSpan -> cAstPlainPart t +fromDual :: forall s ms y. KnownSpan s+         => AstTensor ms DualSpan y -> AstTensor ms s y+fromDual t = case knownSpan @s of+  SFullSpan -> AstFromDual t+  SPrimalStepSpan s -> fullSpanToStep s $ AstFromDual t  -- this is primal zero+  SDualSpan -> t+  SPlainSpan -> AstPlainPart $ AstFromDual t  -- this is plain zero++fromPlain :: forall s ms y. KnownSpan s+          => AstTensor ms PlainSpan y -> AstTensor ms s y+fromPlain t = case knownSpan @s of+  SFullSpan -> AstFromPlain t+  SPrimalStepSpan{} -> AstFromPlain t+  SDualSpan -> AstDualPart $ AstFromPlain t  -- this is dual zero+  SPlainSpan -> t++class KnownSpan (s :: AstSpan) where+  knownSpan :: SAstSpan s++instance KnownSpan FullSpan where+  knownSpan = SFullSpan++instance KnownSpan s => KnownSpan (PrimalStepSpan s) where+  knownSpan = SPrimalStepSpan (knownSpan @s)++instance KnownSpan DualSpan where+  knownSpan = SDualSpan++instance KnownSpan PlainSpan where+  knownSpan = SPlainSpan++-- | Turn a singleton into a constraint via a continuation.+withKnownSpan :: forall s r. SAstSpan s -> (KnownSpan s => r) -> r+withKnownSpan = withDict @(KnownSpan s)++fullSpanToStep :: SAstSpan s+               -> AstTensor ms FullSpan y+               -> AstTensor ms (PrimalStepSpan s) y+fullSpanToStep = \case+  SFullSpan -> cAstPrimalPart+  SPrimalStepSpan sspan -> withKnownSpan sspan+                           $ cAstPrimalPart . fullSpanToStep sspan+  SDualSpan -> cAstPrimalPart . cAstDualPart+  SPlainSpan -> cAstPrimalPart . cAstPlainPart++stepToFullSpan :: SAstSpan s+               -> AstTensor ms (PrimalStepSpan s) y+               -> AstTensor ms FullSpan y+stepToFullSpan = \case+  SFullSpan -> cAstFromPrimal+  SPrimalStepSpan sspan -> withKnownSpan sspan+                           $ stepToFullSpan sspan . cAstFromPrimal+  SDualSpan -> AstFromDual . cAstFromPrimal+  SPlainSpan -> AstFromPlain . cAstFromPrimal++cAstPrimalPart :: forall y s ms. KnownSpan s+               => AstTensor ms s y -> AstTensor ms (PrimalStepSpan s) y+cAstPrimalPart (AstFromPrimal t) = t+cAstPrimalPart (AstFromPlain t) = AstFromPlain t+cAstPrimalPart t = AstPrimalPart t++cAstDualPart :: forall y ms.+                AstTensor ms FullSpan y -> AstTensor ms DualSpan y+cAstDualPart (AstFromDual t) = t+cAstDualPart t = AstDualPart t++cAstPlainPart :: forall y s ms. KnownSpan s+              => AstTensor ms s y -> AstTensor ms PlainSpan y+cAstPlainPart (AstFromPlain v) = v+cAstPlainPart (AstPrimalPart v) = cAstPlainPart v+cAstPlainPart (AstFromPrimal v) = cAstPlainPart v+cAstPlainPart t | SPlainSpan <- knownSpan @s = t+cAstPlainPart t = AstPlainPart t++cAstFromPrimal :: forall y s ms. KnownSpan s+               => AstTensor ms (PrimalStepSpan s) y -> AstTensor ms s y+cAstFromPrimal (AstFromPlain t) = cAstFromPlain t+cAstFromPrimal t = AstFromPrimal t++cAstFromPlain :: forall y s ms. KnownSpan s+              => AstTensor ms PlainSpan y -> AstTensor ms s y+cAstFromPlain t | SPlainSpan <- knownSpan @s = t+cAstFromPlain t = AstFromPlain t++ -- * Variables and related types  newtype AstVarId = AstVarId Int@@ -105,92 +219,185 @@ intToAstVarId :: Int -> AstVarId intToAstVarId = AstVarId -type role AstVarName phantom nominal-data AstVarName :: AstSpanType -> TK -> Type where-  AstVarName :: forall s y.-                FullShapeTK y -> Int64 -> Int64 -> AstVarId-             -> AstVarName s y+type role AstVarName nominal+data AstVarName :: (AstSpan, TK) -> Type where+  AstVarName :: AstVarId -> FtkAndBounds s_y -> AstVarName s_y -instance Eq (AstVarName s y) where-  AstVarName _ _ _ varId1 == AstVarName _ _ _ varId2 = varId1 == varId2+instance Eq (AstVarName '(s, y)) where+  AstVarName varId1 _ == AstVarName varId2 _ = varId1 == varId2 -instance Show (AstVarName s y) where-  showsPrec d (AstVarName _ _ _ varId) =+instance Show (AstVarName '(s, y)) where+  showsPrec d (AstVarName varId _) =     showsPrec d varId  -- less verbose, more readable -instance DMap.Enum1 (AstVarName s) where-  type Enum1Info (AstVarName s) = Some FtkAndBounds-  fromEnum1 (AstVarName ftk minb maxb varId) =-    (fromEnum varId, Some (FtkAndBounds ftk minb maxb))-  toEnum1 varIdInt (Some (FtkAndBounds ftk minb maxb)) =-    Some $ AstVarName ftk minb maxb $ toEnum varIdInt+instance TestEquality AstVarName where+  testEquality (AstVarName _ ftkBounds1) (AstVarName _ ftkBounds2)+    | Just Refl <- testEquality ftkBounds1 ftkBounds2 =+      Just Refl+  testEquality _ _ = Nothing +instance DMap.Enum1 AstVarName where+  type Enum1Info AstVarName = FtkAndBounds+  fromEnum1 (AstVarName varId ftkBounds) = (fromEnum varId, ftkBounds)+  toEnum1 varIdInt = AstVarName (toEnum varIdInt)+ type role FtkAndBounds nominal-data FtkAndBounds y = FtkAndBounds (FullShapeTK y) Int64 Int64+data FtkAndBounds :: (AstSpan, TK) -> Type where+  FtkAndBoundsFull :: FullShapeTK y+                   -> FtkAndBounds '(FullSpan, y)+  FtkAndBoundsPrimal :: FullShapeTK y -> SAstSpan s+                     -> FtkAndBounds '(PrimalStepSpan s, y)+  FtkAndBoundsDual :: FullShapeTK y+                   -> FtkAndBounds '(DualSpan, y)+  FtkAndBoundsPlain :: FullShapeTK y+                    -> FtkAndBounds '(PlainSpan, y)+  FtkAndBoundsBounds :: Int -> Int+                     -> FtkAndBounds '(PlainSpan, TKScalar Int) -instance TestEquality (AstVarName s) where-  testEquality (AstVarName ftk1 _ _ _) (AstVarName ftk2 _ _ _) =-    matchingFTK ftk1 ftk2+instance TestEquality FtkAndBounds where+  testEquality ftkBounds1 ftkBounds2 = case (ftkBounds1, ftkBounds2) of+    (FtkAndBoundsFull ftk1, FtkAndBoundsFull ftk2)+      | Just Refl <- matchingFTK ftk1 ftk2 ->+        Just Refl+    (FtkAndBoundsPrimal ftk1 sspan1, FtkAndBoundsPrimal ftk2 sspan2)+      | Just Refl <- testEquality sspan1 sspan2+      , Just Refl <- matchingFTK ftk1 ftk2 ->+        Just Refl+    (FtkAndBoundsDual ftk1, FtkAndBoundsDual ftk2)+      | Just Refl <- matchingFTK ftk1 ftk2 ->+        Just Refl+    (FtkAndBoundsPlain ftk1, FtkAndBoundsPlain ftk2)+      | Just Refl <- matchingFTK ftk1 ftk2 ->+        Just Refl+    (FtkAndBoundsBounds _ _ , FtkAndBoundsBounds _ _) ->+      Just Refl+    _ -> Nothing -mkAstVarName :: forall s y.-                FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId-             -> AstVarName s y-mkAstVarName ftk Nothing = AstVarName ftk (-1000000000) 1000000000-mkAstVarName ftk (Just (minb, maxb)) = AstVarName ftk minb maxb+mkAstVarName :: forall s y. KnownSpan s+             => FullShapeTK y -> AstVarId -> AstVarName '(s, y)+mkAstVarName ftk varId =+  let ftkBounds = case knownSpan @s of+        SFullSpan -> FtkAndBoundsFull ftk+        SPrimalStepSpan sspan -> FtkAndBoundsPrimal ftk sspan+        SDualSpan -> FtkAndBoundsDual ftk+        SPlainSpan -> FtkAndBoundsPlain ftk+  in AstVarName varId ftkBounds -varNameToAstVarId :: AstVarName s y -> AstVarId-varNameToAstVarId (AstVarName _ _ _ varId) = varId+mkAstVarNameBounds :: (Int, Int) -> AstVarId+                   -> AstVarName '(PlainSpan, TKScalar Int)+{-# INLINE mkAstVarNameBounds #-}+mkAstVarNameBounds (lb, ub) varId = AstVarName varId $ FtkAndBoundsBounds lb ub -varNameToFTK :: AstVarName s y -> FullShapeTK y-varNameToFTK (AstVarName ftk _ _ _) = ftk+reshapeVarName :: FullShapeTK z -> AstVarName '(s, y) -> AstVarName '(s, z)+reshapeVarName ftk (AstVarName varId ftkBounds) =+  AstVarName varId $ case ftkBounds of+    FtkAndBoundsFull{} -> FtkAndBoundsFull ftk+    (FtkAndBoundsPrimal _ sspan) -> FtkAndBoundsPrimal ftk sspan+    FtkAndBoundsDual{} -> FtkAndBoundsDual ftk+    FtkAndBoundsPlain{} -> FtkAndBoundsPlain ftk+    FtkAndBoundsBounds{}+      | FTKScalar @r <- ftk+      , Just Refl <- testEquality (typeRep @r) (typeRep @Int) -> ftkBounds+    FtkAndBoundsBounds{} -> FtkAndBoundsPlain ftk -varNameToBounds :: AstVarName s y -> Maybe (Int64, Int64)-varNameToBounds (AstVarName _ minb maxb _) =-  if minb == -1000000000 && maxb == 1000000000-  then Nothing-  else Just (minb, maxb)+-- | This fails if the variable had bounds (that would be now lost,+-- unless the new span is the same as old, which is just as irregular).+respanVarName :: forall s s2 y. KnownSpan s2+              => AstVarName '(s, y) -> AstVarName '(s2, y)+respanVarName var@(AstVarName varId ftkBounds) = case ftkBounds of+  FtkAndBoundsBounds{} -> error "respanVarName: bounds lost"+  _ -> mkAstVarName (varNameToFTK var) varId -astVar :: AstSpan s-       => AstVarName s y -> AstTensor ms s y-astVar (AstVarName (FTKScalar @r) lb ub _)-  | lb == ub-  , Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =-    fromPrimal $ AstConcreteK lb-astVar varName = AstVar varName+reboundsVarName :: (Int, Int) -> AstVarName '(PlainSpan, TKScalar Int)+                -> AstVarName '(PlainSpan, TKScalar Int)+reboundsVarName (lb, ub) (AstVarName varId _) =+  mkAstVarNameBounds (lb, ub) varId +varNameToAstVarId :: AstVarName s_y -> AstVarId+varNameToAstVarId (AstVarName varId _) = varId++varNameToSpan :: AstVarName '(s, y) -> SAstSpan s+varNameToSpan (AstVarName _ ftkBounds) = case ftkBounds of+  FtkAndBoundsFull{} -> SFullSpan+  (FtkAndBoundsPrimal _ sspan) -> SPrimalStepSpan sspan+  FtkAndBoundsDual{} -> SDualSpan+  FtkAndBoundsPlain{} -> SPlainSpan+  FtkAndBoundsBounds{} -> SPlainSpan++varNameToFTK :: AstVarName '(s, y) -> FullShapeTK y+varNameToFTK (AstVarName _ ftkBounds) = case ftkBounds of+  (FtkAndBoundsFull ftk) -> ftk+  (FtkAndBoundsPrimal ftk _) -> ftk+  (FtkAndBoundsDual ftk) -> ftk+  (FtkAndBoundsPlain ftk) -> ftk+  FtkAndBoundsBounds{} -> FTKScalar++varNameToBounds :: AstVarName '(s, y) -> Maybe (Int, Int)+varNameToBounds (AstVarName _ (FtkAndBoundsBounds lb ub)) = Just (lb, ub)+varNameToBounds _ = Nothing+ -- | The reverse derivative artifact. type role AstArtifactRev nominal nominal data AstArtifactRev x z = AstArtifactRev-  { artVarDtRev      :: AstVarName PrimalSpan (ADTensorKind z)-  , artVarDomainRev  :: AstVarName PrimalSpan x-  , artDerivativeRev :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)-  , artPrimalRev     :: AstTensor AstMethodLet PrimalSpan z+  { artVarDtRev      :: AstVarName '(FullSpan, ADTensorKind z)+  , artVarDomainRev  :: AstVarName '(FullSpan, x)+  , artDerivativeRev :: AstTensor AstMethodLet FullSpan (ADTensorKind x)+  , artPrimalRev     :: ~(AstTensor AstMethodLet FullSpan z)+      -- rarely used, so not forced   }  deriving Show  -- | The forward derivative artifact. type role AstArtifactFwd nominal nominal data AstArtifactFwd x z = AstArtifactFwd-  { artVarDsFwd      :: AstVarName PrimalSpan (ADTensorKind x)-  , artVarDomainFwd  :: AstVarName PrimalSpan x-  , artDerivativeFwd :: AstTensor AstMethodLet PrimalSpan (ADTensorKind z)-  , artPrimalFwd     :: AstTensor AstMethodLet PrimalSpan z+  { artVarDsFwd      :: AstVarName '(FullSpan, ADTensorKind x)+  , artVarDomainFwd  :: AstVarName '(FullSpan, x)+  , artDerivativeFwd :: AstTensor AstMethodLet FullSpan (ADTensorKind z)+  , artPrimalFwd     :: ~(AstTensor AstMethodLet FullSpan z)+      -- rarely used, so not forced   }  deriving Show  -- | This is the (arbitrarily) chosen representation of terms denoting -- integers in the indexes of tensor operations.-type AstInt ms = AstTensor ms PrimalSpan (TKScalar Int64)--- ~ IntOf (AstTensor ms FullSpan)+type AstInt ms = AstTensor ms PlainSpan (TKScalar Int)+-- ~ IntOf (AstTensor ms s) -type IntVarName = AstVarName PrimalSpan (TKScalar Int64)+type IntVarName = AstVarName '(PlainSpan, TKScalar Int)  pattern AstIntVar :: IntVarName -> AstInt ms pattern AstIntVar var <- AstVar var +-- This is needed only for AstVarListS. Parameter i is needed only+-- to derive Foldable.+type role ListS nominal representational+newtype ListS sh i = AstVarListS {unAstVarListS :: IxS sh i}+  deriving (Foldable)++instance Show (AstVarListS sh) where+  showsPrec d (AstVarListS ix) = showsPrec d ix++instance KnownShS sh => IsList (ListS sh i) where+  type Item (ListS sh i) = i+  fromList = AstVarListS . IsList.fromList+  toList = Foldable.toList++pattern ZS :: forall sh i. () => sh ~ '[] => ListS sh i+pattern ZS = AstVarListS ZIS++pattern (::$)+  :: forall {sh1} {i}.+     forall n sh. (n : sh ~ sh1)+  => i -> ListS sh i -> ListS sh1 i+pattern i ::$ l <- AstVarListS (i :.$ (AstVarListS -> l))+  where i ::$ AstVarListS l = AstVarListS (i :.$ l)+infixr 3 ::$++{-# COMPLETE ZS, (::$) #-}+ -- Data invariant: the var names have bounds of the form (0, k - 1), -- where the corresponding dimension in sh is k. This is never checked.-type AstVarListS sh = ListS sh (Const IntVarName)+type AstVarListS sh = ListS sh IntVarName  -- There's no data invariant here. The shape matches rather the argument -- of indexing (or gather) than the indexes.@@ -202,18 +409,23 @@  matchAstLeqInt :: AstBool ms -> Maybe (AstInt ms, AstInt ms) matchAstLeqInt (AstLeqK @r t u)-  | Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =+  | Just Refl <- testEquality (typeRep @r) (typeRep @Int) =       Just (t, u) matchAstLeqInt _ = Nothing +type AstBool ms = AstTensor ms PlainSpan (TKScalar Bool)+-- ~ BoolOf (AstTensor ms s) + -- * AST  type data AstMethodOfSharing = AstMethodShare | AstMethodLet  -- | AST for tensors that are meant to be differentiated.+--+-- Some terms have no semantics, e.g., currently terms with nested primal span. type role AstTensor nominal nominal nominal-data AstTensor :: AstMethodOfSharing -> AstSpanType -> Target where+data AstTensor :: AstMethodOfSharing -> AstSpan -> Target where   -- General operations, for scalar, ranked, shared and other tensors at once   AstPair :: forall y z ms s.              AstTensor ms s y -> AstTensor ms s z@@ -222,178 +434,187 @@                  AstTensor ms s (TKProduct y z) -> AstTensor ms s y   AstProject2 :: forall y z ms s.                  AstTensor ms s (TKProduct y z) -> AstTensor ms s z-  AstFromVector :: forall y k ms s.-                   SNat k -> SingletonTK y-                -> Data.Vector.Vector (AstTensor ms s y)-                -> AstTensor ms s (BuildTensorKind k y)-  AstSum :: forall y k ms s.-            SNat k -> SingletonTK y-         -> AstTensor ms s (BuildTensorKind k y)-         -> AstTensor ms s y-  AstReplicate :: forall y k ms s.-                  SNat k -> SingletonTK y-               -> AstTensor ms s y-               -> AstTensor ms s (BuildTensorKind k y)-  AstMapAccumRDer-    :: forall accy by ey k ms s.-       SNat k-    -> FullShapeTK by-    -> FullShapeTK ey-    -> AstHFun s s-               (TKProduct accy ey) (TKProduct accy by)-    -> AstHFun s s-               (TKProduct (ADTensorKind (TKProduct accy ey))-                          (TKProduct accy ey))-               (ADTensorKind (TKProduct accy by))-    -> AstHFun s s-               (TKProduct (ADTensorKind (TKProduct accy by))-                          (TKProduct accy ey))-               (ADTensorKind (TKProduct accy ey))-    -> AstTensor ms s accy-    -> AstTensor ms s (BuildTensorKind k ey)-    -> AstTensor ms s (TKProduct accy (BuildTensorKind k by))   AstMapAccumLDer     :: forall accy by ey k ms s.        SNat k     -> FullShapeTK by     -> FullShapeTK ey-    -> AstHFun s s+    -> AstHFun s                (TKProduct accy ey) (TKProduct accy by)-    -> AstHFun s s+    -> AstHFun s                (TKProduct (ADTensorKind (TKProduct accy ey))                           (TKProduct accy ey))                (ADTensorKind (TKProduct accy by))-    -> AstHFun s s+    -> AstHFun s                (TKProduct (ADTensorKind (TKProduct accy by))                           (TKProduct accy ey))                (ADTensorKind (TKProduct accy ey))     -> AstTensor ms s accy     -> AstTensor ms s (BuildTensorKind k ey)     -> AstTensor ms s (TKProduct accy (BuildTensorKind k by))-  AstApply :: (AstSpan s1, AstSpan s)-           => AstHFun s1 s x z -> AstTensor ms s1 x -> AstTensor ms s z-  AstVar :: AstVarName s y -> AstTensor ms s y-  AstCond :: forall y ms s.-             AstBool ms -> AstTensor ms s y -> AstTensor ms s y-          -> AstTensor ms s y+  AstApply :: AstHFun s x z -> AstTensor ms s x -> AstTensor ms s z+  AstVar :: AstVarName '(s, y) -> AstTensor ms s y   AstBuild1 :: forall y k ms s.                SNat k -> SingletonTK y             -> (IntVarName, AstTensor ms s y)             -> AstTensor ms s (BuildTensorKind k y)    -- Sharing-related operations, mutually exclusive via AstMethodOfSharing-  AstLet :: forall y z s s2. AstSpan s-         => AstVarName s y -> AstTensor AstMethodLet s y+  AstLet :: forall y z s s2.+            AstVarName '(s, y) -> AstTensor AstMethodLet s y          -> AstTensor AstMethodLet s2 z          -> AstTensor AstMethodLet s2 z-  AstShare :: AstVarName s y -> AstTensor AstMethodShare s y+  AstShare :: AstVarName '(s, y) -> AstTensor AstMethodShare s y            -> AstTensor AstMethodShare s y   AstToShare :: AstTensor AstMethodLet s y              -> AstTensor AstMethodShare s y    -- Explicit dual numbers handling, eliminated in interpretation to ADVal-  AstPrimalPart :: forall y ms.-                   AstTensor ms FullSpan y -> AstTensor ms PrimalSpan y+  AstPrimalPart :: forall y s ms. KnownSpan s+                => AstTensor ms s y -> AstTensor ms (PrimalStepSpan s) y   AstDualPart :: forall y ms.                  AstTensor ms FullSpan y -> AstTensor ms DualSpan y-  AstFromPrimal :: forall y ms.-                   AstTensor ms PrimalSpan y -> AstTensor ms FullSpan y+  AstPlainPart :: forall y s ms. KnownSpan s+               => AstTensor ms s y -> AstTensor ms PlainSpan y+  AstFromPrimal :: forall y s ms.+                   AstTensor ms (PrimalStepSpan s) y -> AstTensor ms s y   AstFromDual :: forall y ms.                  AstTensor ms DualSpan y -> AstTensor ms FullSpan y+  AstFromPlain :: forall y s ms.+                  AstTensor ms PlainSpan y -> AstTensor ms s y    -- Scalar arithmetic (to avoid the slowness of indexes as 1-element tensors)-  AstPlusK :: GoodScalar r+  AstPlusK :: NumScalar r            => AstTensor ms s (TKScalar r)            -> AstTensor ms s (TKScalar r)            -> AstTensor ms s (TKScalar r)-  AstTimesK :: GoodScalar r+  AstTimesK :: NumScalar r             => AstTensor ms s (TKScalar r)             -> AstTensor ms s (TKScalar r)             -> AstTensor ms s (TKScalar r)-  AstN1K :: GoodScalar r+  AstN1K :: NumScalar r          => OpCodeNum1 -> AstTensor ms s (TKScalar r)          -> AstTensor ms s (TKScalar r)-  AstR1K :: (RealFloatH r, Nested.FloatElt r, GoodScalar r)+  AstR1K :: (NumScalar r, Differentiable r)          => OpCode1 -> AstTensor ms s (TKScalar r)          -> AstTensor ms s (TKScalar r)-  AstR2K :: (RealFloatH r, Nested.FloatElt r, GoodScalar r)+  AstR2K :: (NumScalar r, Differentiable r)          => OpCode2 -> AstTensor ms s (TKScalar r)          -> AstTensor ms s (TKScalar r)          -> AstTensor ms s (TKScalar r)-  AstI2K :: (IntegralH r, Nested.IntElt r, GoodScalar r)+  AstI2K :: (NumScalar r, IntegralH r, Nested.IntElt r)          => OpCodeIntegral2 -> AstTensor ms s (TKScalar r)          -> AstTensor ms s (TKScalar r)          -> AstTensor ms s (TKScalar r)   AstConcreteK :: GoodScalar r-               => r -> AstTensor ms PrimalSpan (TKScalar r)-  AstFloorK :: (GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2)-            => AstTensor ms PrimalSpan (TKScalar r1)-            -> AstTensor ms PrimalSpan (TKScalar r2)-  AstFromIntegralK :: (GoodScalar r1, Integral r1, GoodScalar r2)-                   => AstTensor ms PrimalSpan (TKScalar r1)-                   -> AstTensor ms PrimalSpan (TKScalar r2)-  AstCastK :: (GoodScalar r1, RealFrac r1, RealFrac r2, GoodScalar r2)+               => r -> AstTensor ms PlainSpan (TKScalar r)+  AstFloorK :: (NumScalar r1, Differentiable r1, NumScalar r2, Integral r2)+            => AstTensor ms PlainSpan (TKScalar r1)+            -> AstTensor ms PlainSpan (TKScalar r2)+  AstFromIntegralK :: (NumScalar r1, Integral r1, NumScalar r2)+                   => AstTensor ms PlainSpan (TKScalar r1)+                   -> AstTensor ms PlainSpan (TKScalar r2)+  AstCastK :: (NumScalar r1, Differentiable r1, NumScalar r2, Differentiable r2)            => AstTensor ms s (TKScalar r1) -> AstTensor ms s (TKScalar r2)+  AstArgMinK :: forall n r ms. NumScalar r+             => AstTensor ms PlainSpan (TKS '[n] r)+             -> AstTensor ms PlainSpan (TKScalar Int)+  AstArgMaxK :: forall n r ms. NumScalar r+             => AstTensor ms PlainSpan (TKS '[n] r)+             -> AstTensor ms PlainSpan (TKScalar Int)+  AstIndexK :: forall shm r s ms. GoodScalar r+            => AstTensor ms s (TKS shm r) -> AstIxS ms shm+            -> AstTensor ms s (TKScalar r)+    -- out of bounds indexing is permitted and the results is def (==0)    -- Shaped arithmetic-  AstPlusS :: GoodScalar r+  AstPlusS :: NumScalar r            => AstTensor ms s (TKS sh r)            -> AstTensor ms s (TKS sh r)            -> AstTensor ms s (TKS sh r)-  AstTimesS :: GoodScalar r+  AstTimesS :: NumScalar r             => AstTensor ms s (TKS sh r)             -> AstTensor ms s (TKS sh r)             -> AstTensor ms s (TKS sh r)-  AstN1S :: GoodScalar r+  AstN1S :: NumScalar r          => OpCodeNum1 -> AstTensor ms s (TKS sh r)          -> AstTensor ms s (TKS sh r)-  AstR1S :: (RealFloatH r, Nested.FloatElt r, GoodScalar r)+  AstR1S :: (NumScalar r, Differentiable r)          => OpCode1 -> AstTensor ms s (TKS sh r)          -> AstTensor ms s (TKS sh r)-  AstR2S :: (RealFloatH r, Nested.FloatElt r, GoodScalar r)+  AstR2S :: (NumScalar r, Differentiable r)          => OpCode2 -> AstTensor ms s (TKS sh r)          -> AstTensor ms s (TKS sh r)          -> AstTensor ms s (TKS sh r)-  AstI2S :: (IntegralH r, Nested.IntElt r, GoodScalar r)+  AstI2S :: (NumScalar r, IntegralH r, Nested.IntElt r)          => OpCodeIntegral2 -> AstTensor ms s (TKS sh r)          -> AstTensor ms s (TKS sh r)          -> AstTensor ms s (TKS sh r)   AstConcreteS :: GoodScalar r-               => Nested.Shaped sh r -> AstTensor ms PrimalSpan (TKS sh r)-  AstFloorS :: (GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2)-            => AstTensor ms PrimalSpan (TKS sh r1)-            -> AstTensor ms PrimalSpan (TKS sh r2)-  AstFromIntegralS :: (GoodScalar r1, Integral r1, GoodScalar r2)-                   => AstTensor ms PrimalSpan (TKS sh r1)-                   -> AstTensor ms PrimalSpan (TKS sh r2)-  AstCastS :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+               => Nested.Shaped sh r -> AstTensor ms PlainSpan (TKS sh r)+  AstFloorS :: (NumScalar r1, Differentiable r1, NumScalar r2, Integral r2)+            => AstTensor ms PlainSpan (TKS sh r1)+            -> AstTensor ms PlainSpan (TKS sh r2)+  AstFromIntegralS :: (NumScalar r1, Integral r1, NumScalar r2)+                   => AstTensor ms PlainSpan (TKS sh r1)+                   -> AstTensor ms PlainSpan (TKS sh r2)+  AstCastS :: (NumScalar r1, Differentiable r1, NumScalar r2, Differentiable r2)            => AstTensor ms s (TKS sh r1)            -> AstTensor ms s (TKS sh r2)--  -- Shaped tensor operations+  AstArgMinS :: forall k sh r ms. NumScalar r+             => AstTensor ms PlainSpan (TKS (k ': sh) r)+             -> AstTensor ms PlainSpan (TKS (Init (k ': sh)) Int)+  AstArgMaxS :: forall k sh r ms. NumScalar r+             => AstTensor ms PlainSpan (TKS (k ': sh) r)+             -> AstTensor ms PlainSpan (TKS (Init (k ': sh)) Int)   AstIndexS :: forall shm shn x s ms.                ShS shn             -> AstTensor ms s (TKS2 (shm ++ shn) x) -> AstIxS ms shm             -> AstTensor ms s (TKS2 shn x)-    -- out of bounds indexing is permitted and the results is def (==0)-  AstScatterS :: forall shm shn shp x s ms.-                 ShS shn -> AstTensor ms s (TKS2 (shm ++ shn) x)++  AstCondK :: forall r ms s. GoodScalar r+           => AstBool ms+           -> AstTensor ms s (TKScalar r)+           -> AstTensor ms s (TKScalar r)+           -> AstTensor ms s (TKScalar r)+  AstCondS :: forall sh x ms s.+              AstBool ms+           -> AstTensor ms s (TKS2 sh x)+           -> AstTensor ms s (TKS2 sh x)+           -> AstTensor ms s (TKS2 sh x)+  AstFromVectorK :: forall shm r s ms. GoodScalar r+                 => ShS shm -> Data.Vector.Vector (AstTensor ms s (TKScalar r))+                 -> AstTensor ms s (TKS shm r)+  AstFromVectorS :: forall shm shn x s ms.+                    ShS shm -> Data.Vector.Vector (AstTensor ms s (TKS2 shn x))+                 -> AstTensor ms s (TKS2 (shm ++ shn) x)+  AstSumK :: forall shm r s ms. NumScalar r+          => AstTensor ms s (TKS shm r)+          -> AstTensor ms s (TKScalar r)+  AstSumS :: forall shm shn x s ms. TKAllNum x+          => ShS shm -> AstTensor ms s (TKS2 (shm ++ shn) x)+          -> AstTensor ms s (TKS2 shn x)+  -- Shaped tensor operations+  AstScatterS :: forall shm shn shp x s ms. TKAllNum x+              => ShS shm -> ShS shn -> ShS shp+              -> AstTensor ms s (TKS2 (shm ++ shn) x)               -> (AstVarListS shm, AstIxS ms shp)               -> AstTensor ms s (TKS2 (shp ++ shn) x)     -- out of bounds indexing is permitted and the results is def (==0)+  AstReplicateK :: forall shm r s ms. GoodScalar r+                => ShS shm -> AstTensor ms s (TKScalar r)+                -> AstTensor ms s (TKS shm r)+  AstReplicateS :: forall shm shn x s ms.+                   ShS shm -> AstTensor ms s (TKS2 shn x)+                -> AstTensor ms s (TKS2 (shm ++ shn) x)   AstGatherS :: forall shm shn shp x s ms.-                ShS shn -> AstTensor ms s (TKS2 (shp ++ shn) x)+                ShS shm -> ShS shn -> ShS shp+             -> AstTensor ms s (TKS2 (shp ++ shn) x)              -> (AstVarListS shm, AstIxS ms shp)              -> AstTensor ms s (TKS2 (shm ++ shn) x)     -- out of bounds indexing is permitted and the results is def (==0)-  AstMinIndexS :: forall n sh r r2 ms. (GoodScalar r, GoodScalar r2)-               => AstTensor ms PrimalSpan (TKS (n ': sh) r)-               -> AstTensor ms PrimalSpan (TKS (Init (n ': sh)) r2)-  AstMaxIndexS :: forall n sh r r2 ms. (GoodScalar r, GoodScalar r2)-               => AstTensor ms PrimalSpan (TKS (n ': sh) r)-               -> AstTensor ms PrimalSpan (TKS (Init (n ': sh)) r2)-  AstIotaS :: forall n r ms. GoodScalar r-           => SNat n -> AstTensor ms PrimalSpan (TKS '[n] r)+  AstIotaS :: forall n r ms. NumScalar r+           => SNat n -> AstTensor ms PlainSpan (TKS '[n] r)   AstAppendS :: forall m n sh x ms s.                 AstTensor ms s (TKS2 (m ': sh) x)              -> AstTensor ms s (TKS2 (n ': sh) x)@@ -415,32 +636,53 @@   AstConvert :: TKConversion a b -> AstTensor ms s a -> AstTensor ms s b    -- Backend-specific primitives-  AstSum0S :: AstTensor ms s (TKS2 sh x)-           -> AstTensor ms s (TKS2 '[] x)-  AstDot0S :: GoodScalar r-           => AstTensor ms s (TKS sh r) -> AstTensor ms s (TKS sh r)-           -> AstTensor ms s (TKS '[] r)-  AstDot1InS :: forall sh n r ms s. GoodScalar r+  AstDot0 :: NumScalar r+          => AstTensor ms s (TKS sh r) -> AstTensor ms s (TKS sh r)+          -> AstTensor ms s (TKScalar r)+  AstDot1InS :: forall sh n r ms s. NumScalar r              => ShS sh -> SNat n              -> AstTensor ms s (TKS (sh ++ '[n]) r)              -> AstTensor ms s (TKS (sh ++ '[n]) r)              -> AstTensor ms s (TKS sh r)-  AstMatmul2S :: GoodScalar r+  AstMatmul2S :: NumScalar r               => SNat m -> SNat n -> SNat p               -> AstTensor ms s (TKS '[m, n] r)               -> AstTensor ms s (TKS '[n, p] r)               -> AstTensor ms s (TKS '[m, p] r) +  -- Booleans+  AstBoolNotK :: AstBool ms -> AstBool ms+  AstBoolNotS :: AstTensor ms PlainSpan (TKS sh Bool)+              -> AstTensor ms PlainSpan (TKS sh Bool)+  AstBoolAndK :: AstBool ms -> AstBool ms -> AstBool ms+  AstBoolAndS :: AstTensor ms PlainSpan (TKS sh Bool)+              -> AstTensor ms PlainSpan (TKS sh Bool)+              -> AstTensor ms PlainSpan (TKS sh Bool)+  -- There are existential variables here.+  AstLeqK :: forall r ms. NumScalar r+          => AstTensor ms PlainSpan (TKScalar r)+          -> AstTensor ms PlainSpan (TKScalar r)+          -> AstBool ms+  AstLeq :: forall sh r ms. NumScalar r+         => AstTensor ms PlainSpan (TKS sh r)+         -> AstTensor ms PlainSpan (TKS sh r)+         -> AstBool ms+  AstLeqS :: forall shb sh r ms. NumScalar r+          => ShS shb -> ShS sh+          -> AstTensor ms PlainSpan (TKS (shb ++ sh) r)+          -> AstTensor ms PlainSpan (TKS (shb ++ sh) r)+          -> AstTensor ms PlainSpan (TKS shb Bool)+ deriving instance Show (AstTensor ms s y)   -- for this to work, AstConcreteS can't take a Concrete;   -- an alternative might be @Has Show (AstTensor ms s)@, but then we'd need   -- to write @has@ before we apply @show@ and we'd weaken @AllTargetShow@ -type role AstHFun nominal nominal nominal nominal-data AstHFun s s2 x z where-  AstLambda :: ~(AstVarName s x)-            -> ~(AstTensor AstMethodLet s2 z)-            -> AstHFun s s2 x z+type role AstHFun nominal nominal nominal+data AstHFun s x z where+  AstLambda :: ~(AstVarName '(s, x))+            -> ~(AstTensor AstMethodLet s z)+            -> AstHFun s x z     -- ^ The function body can't have any free variables outside those     -- listed in the first component of the pair; this reflects     -- the quantification in 'HordeAd.Core.Ops.rrev'@@ -449,32 +691,18 @@     -- The constructor is non-strict in order not to pre-compute     -- higher derivatives (e.g., inside folds) that are never going to be used.     -- As a side effect, all lambdas (closed functions) are processed-    -- lazily, which makes no harm, since they have no outside free variables+    -- lazily, which causes no harm, since they have no outside free variables     -- and so can't easiliy induce leaks by retaining outside values (e.g.,     -- big environments from which values for the variables would be drawn).     -- The cost of computing a reverse derivative of a fold nested inside     -- the function argument n times is reduced by the laziness from 20^n     -- to under 2^n (old experimental results). Note, however,     -- that if the n-th forward and reverse derivative is taken,-    -- the laziness is defeated.--deriving instance Show (AstHFun s s2 x z)+    -- the laziness is defeated. To make the variable argument strict+    -- we'd need to modify some other code fragments, while the performance+    -- impact seems mixed. -type role AstBool nominal-data AstBool ms where-  AstBoolConst :: Bool -> AstBool ms-  AstBoolNot :: AstBool ms -> AstBool ms-  AstBoolAnd :: AstBool ms -> AstBool ms -> AstBool ms-  -- There are existential variables here.-  AstLeqK :: forall r ms. GoodScalar r-          => AstTensor ms PrimalSpan (TKScalar r)-          -> AstTensor ms PrimalSpan (TKScalar r)-          -> AstBool ms-  AstLeqS :: forall sh r ms. GoodScalar r-          => AstTensor ms PrimalSpan (TKS sh r)-          -> AstTensor ms PrimalSpan (TKS sh r)-          -> AstBool ms-deriving instance Show (AstBool ms)+deriving instance Show (AstHFun s x z)  data OpCodeNum1 =     NegateOp | AbsOp | SignumOp
+ src/HordeAd/Core/AstEngine.hs view
@@ -0,0 +1,60 @@+-- | Predefined common functions for the joint inlining and simplification+-- of AST terms.+module HordeAd.Core.AstEngine+  ( simplifyArtifactRev, simplifyArtifactFwd+  , simplifyUserCode, simplifyInlineContract+  ) where++import Prelude++import HordeAd.Core.Ast+import HordeAd.Core.AstInline+import HordeAd.Core.AstTraverse++-- | Simplify the whole reverse derivative artifact (which includes+-- also the primal value computed during the differentiation process).+{-# INLINE simplifyArtifactRev #-}+simplifyArtifactRev :: AstArtifactRev x z -> AstArtifactRev x z+simplifyArtifactRev art =+  let !der = simplifyInlineContract $ artDerivativeRev art in+  let prim = simplifyInlineContract $ artPrimalRev art+  in art {artDerivativeRev = der, artPrimalRev = prim}++-- | Simplify the whole forward derivative artifact (which includes+-- also the primal value computed during the differentiation process).+{-# INLINE simplifyArtifactFwd #-}+simplifyArtifactFwd :: AstArtifactFwd x z -> AstArtifactFwd x z+simplifyArtifactFwd art =+  let !der = simplifyInlineContract $ artDerivativeFwd art in+  let prim = simplifyInlineContract $ artPrimalFwd art+  in art {artDerivativeFwd = der, artPrimalFwd = prim}++-- | A mixture of simplification and inlining to use when the resultng+-- term is not yet supposed to be interpreted using a computational backed,+-- but rather to be stored and later composed with other terms.+-- A typical example is user code to be differentiated afterwards.+{-# INLINE simplifyUserCode #-}+simplifyUserCode+  :: forall z s. KnownSpan s+  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z+simplifyUserCode =+  simplifyAst+  . expandAst . inlineAstTensor+  . simplifyAst . expandAst . inlineAstTensor+  . simplifyAst+    -- no need to start with letDown, because either the user writes+    -- good lets or unsharing happens just before this simplification++-- | A mixture of simplification, inlining and recognition of additional+-- backend-specific primitives, to be used just before a term+-- is interpreted as a value in the computational backend.+{-# INLINE simplifyInlineContract #-}+simplifyInlineContract+  :: forall z s. KnownSpan s+  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z+simplifyInlineContract =+  letDownAst . contractAst . contractAst+  . expandAst . inlineAstTensor+  . simplifyAst . expandAst . inlineAstTensor+  . simplifyAst+    -- this usually starts with letDown from unsharing
src/HordeAd/Core/AstEnv.hs view
@@ -1,21 +1,22 @@-{-# LANGUAGE QuantifiedConstraints #-}+{-# LANGUAGE UndecidableInstances #-} -- | The environment datatype and operations for creating and accessing it. module HordeAd.Core.AstEnv-  ( AstEnv, emptyEnv, showsPrecAstEnv+  ( AstEnv, SpanTargetFam, SpanTarget(..)+  , lemPlainOfSpan, dictSpanFam, toFromFullSpan, toFullSpan, fromFullSpan+  , emptyEnv  -- TODO: showsPrecAstEnv   , extendEnv, extendEnvI, extendEnvVarsS   ) where  import Prelude -import Data.Coerce (coerce) import Data.Dependent.EnumMap.Strict (DEnumMap) import Data.Dependent.EnumMap.Strict qualified as DMap-import Data.Dependent.Sum-import Data.Foldable qualified as Foldable import Data.Kind (Type)-import Text.Show (showListWith)+import Data.Proxy (Proxy)+import Data.Type.Equality ((:~:) (Refl))  import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)  import HordeAd.Core.Ast import HordeAd.Core.Ops@@ -25,44 +26,133 @@ -- | The environment that keeps values assigned to variables -- during interpretation. type AstEnv :: Target -> Type-type AstEnv target = DEnumMap (AstVarName FullSpan) target-  -- We can't easily index over span and tensor kind at once,-  -- so instead we represent PrimalSpan values as FullSpan-  -- (dual number) values with zero dual component and DualSpan values-  -- as FullSpan values with zero primal component.+type AstEnv target = DEnumMap AstVarName (SpanTarget target) +type family SpanTargetFam target s :: Target where+  -- This variant causes typing problems;+--SpanTargetFam target (PrimalStepSpan s2) = SpanTargetFam (PrimalOf target) s2+  SpanTargetFam target (PrimalStepSpan s2) = PrimalOf (SpanTargetFam target s2)+  SpanTargetFam target DualSpan = target  -- !!! not: DualOf target, see below+  SpanTargetFam target FullSpan = target+  SpanTargetFam target PlainSpan = PlainOf target+  -- The equation for DualSpan is hacky to prevent interpreter crashing fast+  -- due to no tensor operations defined for Delta expressions. Instead of+  -- producing the Delta expressions (in non-symbolic instances), we produce+  -- full dual numbers for which tensor operations are always defined.+  -- To keep the artificial representation sound, we try to manually ensure+  -- in the interpreter code that the primal part of these dual numbers+  -- is always zero, so only the dual part (the Delta expression) matters.++-- This is needed, because type families can't yet be partially applied.+type role SpanTarget nominal nominal+newtype SpanTarget target s_y = SpanTarget (SpanTargetFamUncurried target s_y)+type family SpanTargetFamUncurried target s_y :: Type where+  SpanTargetFamUncurried target '(s, y) = SpanTargetFam target s y++lemPlainOfSpan :: ADReady target+               => Proxy target -> SAstSpan s+               -> PlainOf (SpanTargetFam target s) :~: PlainOf target+{-# INLINE lemPlainOfSpan #-}+lemPlainOfSpan _ = \case+  SFullSpan -> Refl+  SPrimalStepSpan SFullSpan -> Refl+  -- This is true morally and in all instances, even though it's+  -- not derivable.+  SPrimalStepSpan _ -> unsafeCoerceRefl+  SDualSpan -> Refl  -- due to the dual hack+  SPlainSpan -> Refl++dictSpanFam :: ADReady target+            => Proxy target -> SAstSpan s+            -> Dict0 (ADReadyClasses (SpanTargetFam target s))+{-# INLINE dictSpanFam #-}+dictSpanFam _ = \case+  SFullSpan -> Dict0+  SPrimalStepSpan SFullSpan -> Dict0+  SPrimalStepSpan _ ->+    error "dictSpanFam: these operations on nested primal terms are illegal"+  SDualSpan -> Dict0  -- due to the dual hack+  SPlainSpan -> Dict0++toFromFullSpan+  :: BaseTensor target+  => SingletonTK y -> SAstSpan s+  -> ( SpanTargetFam target s y -> SpanTargetFam target FullSpan y+     , SpanTargetFam target FullSpan y -> SpanTargetFam target s y )+{-# INLINE toFromFullSpan #-}+toFromFullSpan stk = \case+  SFullSpan -> (id, id)+  SPrimalStepSpan SFullSpan -> (tfromPrimal stk, tprimalPart)+  SPrimalStepSpan _ ->+    error "toFromFullSpan: nested primal numbers are not converted to full dual numbers"+  {- This would require an arbitrarily large number of dictionaries+     for PrimalOf (PrimalOf (... (PrimalOf target)))+     or equating the type of primal numbers and dual numbers with+     a zero dual part, after a finite number of steps, which is troublesome,+     or at step zero, which is crude and gives too few typing hints.+  SPrimalStepSpan s4 ->+    let (toFull, fromFull) = toFromFull stk s4+    in (toFull . tfromPrimal stk, tprimalPart . fromFull)+  -}+  SDualSpan -> (id, id)  -- due to the dual hack+  SPlainSpan -> (tfromPlain stk, tplainPart)++toFullSpan :: BaseTensor target+           => SingletonTK y -> SAstSpan s+           -> SpanTargetFam target s y -> SpanTargetFam target FullSpan y+{-# INLINE toFullSpan #-}+toFullSpan stk = \case+  SFullSpan -> id+  SPrimalStepSpan SFullSpan -> tfromPrimal stk+  SPrimalStepSpan _ ->+    error "toFullSpan: nested primal numbers are not converted to full dual numbers"+  SDualSpan -> id  -- due to the dual hack+  SPlainSpan -> tfromPlain stk++fromFullSpan :: BaseTensor target+             => SAstSpan s+             -> SpanTargetFam target FullSpan y -> SpanTargetFam target s y+{-# INLINE fromFullSpan #-}+fromFullSpan = \case+  SFullSpan -> id+  SPrimalStepSpan SFullSpan -> tprimalPart+  SPrimalStepSpan _ ->+    error "fromFullSpan: nested primal numbers are not converted to full dual numbers"+  SDualSpan -> id  -- due to the dual hack+  SPlainSpan -> tplainPart+ emptyEnv :: AstEnv target emptyEnv = DMap.empty +{- TODO: showsPrecAstEnv-  :: (forall y. KnownSTK y => Show (target y))+  :: AllTargetShow target   => Int -> AstEnv target -> ShowS showsPrecAstEnv d demap =   showParen (d > 10) $     showString "fromList "     . showListWith-        (\(k :=> target) ->+        (\(k :=> SpanTarget target) ->            withKnownSTK (ftkToSTK $ varNameToFTK k) $            showsPrec 2 k . showString " :=> " . showsPrec 1 target)         (DMap.toList demap)+-}  extendEnv :: forall target s y.-             AstVarName s y -> target y -> AstEnv target+             AstVarName '(s, y) -> SpanTargetFam target s y -> AstEnv target           -> AstEnv target-extendEnv var !t !env =-  let var2 :: AstVarName FullSpan y-      var2 = coerce var  -- only FullSpan variables permitted in env; see above-  in DMap.insertWithKey (\_ _ _ -> error $ "extendEnv: duplicate " ++ show var)-                        var2 t env+extendEnv !var !t !env =+  if DMap.member var env+  then error $ "extendEnv: duplicate " ++ show var+  else DMap.insert var (SpanTarget t) env -extendEnvI :: BaseTensor target-           => IntVarName -> IntOf target -> AstEnv target+extendEnvI :: IntVarName -> IntOf target -> AstEnv target            -> AstEnv target-extendEnvI var !i !env = extendEnv var (tfromPrimal STKScalar i) env+extendEnvI !var !i !env = extendEnv var i env -extendEnvVarsS :: forall target sh. BaseTensor target-               => AstVarListS sh -> IxSOf target sh -> AstEnv target+extendEnvVarsS :: forall target sh.+                  AstVarListS sh -> IxSOf target sh -> AstEnv target                -> AstEnv target-extendEnvVarsS vars !ix !env =-  let assocs = zip (listsToList vars) (Foldable.toList ix)-  in foldr (uncurry extendEnvI) env assocs+extendEnvVarsS ZS ZIS !env = env+extendEnvVarsS (var ::$ vars) (i :.$ ix) env =+  extendEnvVarsS vars ix (extendEnvI var i env)
src/HordeAd/Core/AstFreshId.hs view
@@ -5,10 +5,12 @@ -- with @unsafePerformIO@ outside, so some of the impurity escapes -- and is encapsulated elsewhere. module HordeAd.Core.AstFreshId-  ( funToAstIO, funToAst, funToAst2, fun1ToAst+  ( funToAstIO, funToAst+  , funToAstIntIO, funToAstInt+  , funToAstIntMaybeIO, funToAstIntMaybe+  , funToAstAutoBoundsIO, funToAstNoBoundsIO   , funToAstRevIO, funToAstFwdIO-  , funToAstIntVarIO, funToAstIntVar, funToAstI-  , funToVarsIxS, funToAstIxS+  , funToVarsIxS     -- * Low level counter manipulation to be used only in sequential tests   , resetVarCounter   ) where@@ -16,13 +18,15 @@ import Prelude  import Control.Concurrent.Counter (Counter, add, new, set)-import Data.Int (Int64)+import Data.Type.Equality (testEquality, (:~:) (Refl)) import GHC.Exts (IsList (..)) import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)  import Data.Array.Nested.Shaped.Shape  import HordeAd.Core.Ast+import HordeAd.Core.AstTools import HordeAd.Core.TensorKind import HordeAd.Core.Types @@ -43,142 +47,126 @@ unsafeGetFreshAstVarId =   intToAstVarId <$> add unsafeAstVarCounter 1 -unsafeGetFreshAstVarName :: FullShapeTK y -> Maybe (Int64, Int64)-                         -> IO (AstVarName s y)-{-# INLINE unsafeGetFreshAstVarName #-}-unsafeGetFreshAstVarName ftk bounds =-  mkAstVarName ftk bounds-  . intToAstVarId <$> add unsafeAstVarCounter 1+funToAstIO :: KnownSpan s+           => FullShapeTK y -> (AstTensor ms s y -> AstTensor ms s2 z)+           -> IO (AstVarName '(s, y), AstTensor ms s2 z)+{-# INLINE funToAstIO  #-}+funToAstIO ftk f = do+  !freshId <- unsafeGetFreshAstVarId+  let !var = mkAstVarName ftk freshId+      x = f $ astVar var+  return (var, x) -funToAstIO2 :: forall y z s s2 ms. AstSpan s-            => FullShapeTK y -> Maybe (Int64, Int64)-            -> (AstTensor ms s y -> AstTensor ms s2 z)-            -> IO (AstVarName s y, AstTensor ms s2 z)-{-# INLINE funToAstIO2 #-}-funToAstIO2 ftk bounds f = do-  freshId <- unsafeGetFreshAstVarName ftk bounds-  let !x = f (astVar freshId)-  return (freshId, x)--- Warning: adding a bang before freshId breaks fragile tests.--- Probably GHC then optimizes differently and less predictably--- and so changes results between -O0 vs -O1 and possibly also--- between different GHC versions and between local vs CI setup.+funToAst :: KnownSpan s+         => FullShapeTK y -> (AstTensor ms s y -> AstTensor ms s2 z)+         -> (AstVarName '(s, y), AstTensor ms s2 z)+{-# NOINLINE funToAst #-}+funToAst ftk = unsafePerformIO . funToAstIO ftk -funToAst2 :: AstSpan s-          => FullShapeTK y -> Maybe (Int64, Int64)-          -> (AstTensor ms s y -> AstTensor ms s2 z)-          -> (AstVarName s y, AstTensor ms s2 z)-{-# NOINLINE funToAst2 #-}-funToAst2 ftk bounds = unsafePerformIO . funToAstIO2 ftk bounds+funToAstIntIO :: (Int, Int) -> (AstInt ms -> AstTensor ms s2 z)+              -> IO (IntVarName, AstTensor ms s2 z)+{-# INLINE funToAstIntIO #-}+funToAstIntIO bds f = do+  !freshId <- unsafeGetFreshAstVarId+  let !var = mkAstVarNameBounds bds freshId+      x = f $ astVar var+  return (var, x) -funToAstIO :: forall y z s ms. AstSpan s-           => FullShapeTK y-           -> (AstTensor ms s y -> AstTensor ms s z)-           -> IO (AstVarName s y, AstTensor ms s z)-{-# INLINE funToAstIO #-}-funToAstIO ftk = funToAstIO2 ftk Nothing+funToAstInt :: (Int, Int) -> (AstInt ms -> AstTensor ms s2 z)+            -> (IntVarName, AstTensor ms s2 z)+{-# NOINLINE funToAstInt #-}+funToAstInt bds = unsafePerformIO . funToAstIntIO bds -funToAst :: AstSpan s-         => FullShapeTK y -> Maybe (Int64, Int64)-         -> (AstTensor ms s y -> AstTensor ms s z)-         -> (AstVarName s y, AstTensor ms s z)-{-# NOINLINE funToAst #-}-funToAst ftk bounds = unsafePerformIO . funToAstIO2 ftk bounds+funToAstIntMaybeIO :: Maybe (Int, Int) -> ((IntVarName, AstInt ms) -> a)+                   -> IO a+{-# INLINE funToAstIntMaybeIO #-}+funToAstIntMaybeIO mbounds f = do+  !freshId <- unsafeGetFreshAstVarId+  let !var = case mbounds of+        Nothing -> mkAstVarName FTKScalar freshId+        Just bds -> mkAstVarNameBounds bds freshId+      x = astVar var+  return $! f (var, x) -fun1ToAstIO :: FullShapeTK y -> (AstVarName s y -> AstTensor ms s y)-            -> IO (AstTensor ms s y)-{-# INLINE fun1ToAstIO #-}-fun1ToAstIO ftk f = do-  !freshId <- unsafeGetFreshAstVarName ftk Nothing-  return $! f freshId+funToAstIntMaybe :: Maybe (Int, Int) -> ((IntVarName, AstInt ms) -> a) -> a+{-# NOINLINE funToAstIntMaybe #-}+funToAstIntMaybe mbounds = unsafePerformIO . funToAstIntMaybeIO mbounds -fun1ToAst :: FullShapeTK y -> (AstVarName s y -> AstTensor ms s y)-          -> AstTensor ms s y-{-# NOINLINE fun1ToAst #-}-fun1ToAst ftk = unsafePerformIO . fun1ToAstIO ftk+funToAstAutoBoundsIO :: forall r s ms. KnownSpan s+                     => FullShapeTK (TKScalar r) -> AstTensor ms s (TKScalar r)+                     -> IO (AstVarName '(s, TKScalar r))+{-# INLINE funToAstAutoBoundsIO #-}+funToAstAutoBoundsIO ftk@FTKScalar a = do+  !freshId <- unsafeGetFreshAstVarId+  case knownSpan @s of+    SPlainSpan | Just Refl <- testEquality (typeRep @r) (typeRep @Int)+               , Just bds <- intBounds a ->+      pure $! mkAstVarNameBounds bds freshId+    _ -> pure $! mkAstVarName ftk freshId +funToAstNoBoundsIO :: KnownSpan s+                   => FullShapeTK y -> IO (AstVarName '(s, y))+{-# INLINE funToAstNoBoundsIO #-}+funToAstNoBoundsIO ftk = do+  !freshId <- unsafeGetFreshAstVarId+  pure $! mkAstVarName ftk freshId+ funToAstRevIO :: forall x.                  FullShapeTK x-              -> IO ( AstVarName PrimalSpan x-                    , AstTensor AstMethodShare PrimalSpan x-                    , AstVarName FullSpan x+              -> IO ( AstTensor AstMethodShare FullSpan x+                    , AstVarName '(FullSpan, x)                     , AstTensor AstMethodLet FullSpan x ) {-# INLINE funToAstRevIO #-} funToAstRevIO ftk = do   !freshId <- unsafeGetFreshAstVarId-  let varPrimal :: AstVarName PrimalSpan x-      varPrimal = mkAstVarName ftk Nothing freshId-      var :: AstVarName FullSpan x-      var = mkAstVarName ftk Nothing freshId-      astVarPrimal :: AstTensor AstMethodShare PrimalSpan x-      !astVarPrimal = astVar varPrimal+  let var :: AstVarName '(FullSpan, x)+      var = mkAstVarName ftk freshId+      astVarPrimal :: AstTensor AstMethodShare FullSpan x+      !astVarPrimal = astVar var       astVarD :: AstTensor AstMethodLet FullSpan x       !astVarD = astVar var-  return (varPrimal, astVarPrimal, var, astVarD)+  return (astVarPrimal, var, astVarD)  funToAstFwdIO :: forall x.                  FullShapeTK x-              -> IO ( AstVarName PrimalSpan (ADTensorKind x)-                    , AstTensor AstMethodShare PrimalSpan (ADTensorKind x)-                    , AstVarName PrimalSpan x-                    , AstTensor AstMethodShare PrimalSpan x-                    , AstVarName FullSpan x+              -> IO ( AstVarName '(FullSpan, ADTensorKind x)+                    , AstTensor AstMethodShare FullSpan (ADTensorKind x)+                    , AstTensor AstMethodShare FullSpan x+                    , AstVarName '(FullSpan, x)                     , AstTensor AstMethodLet FullSpan x ) {-# INLINE funToAstFwdIO #-} funToAstFwdIO ftk = do-  !freshIdDs <- unsafeGetFreshAstVarId+  !freshIdD <- unsafeGetFreshAstVarId   !freshId <- unsafeGetFreshAstVarId-  let varPrimalD :: AstVarName PrimalSpan (ADTensorKind x)-      varPrimalD = mkAstVarName (adFTK ftk) Nothing freshIdDs-      varPrimal :: AstVarName PrimalSpan x-      varPrimal = mkAstVarName ftk Nothing freshId-      var :: AstVarName FullSpan x-      var = mkAstVarName ftk Nothing freshId-      astVarPrimalD :: AstTensor AstMethodShare PrimalSpan (ADTensorKind x)+  let varPrimalD :: AstVarName '(FullSpan, ADTensorKind x)+      varPrimalD = mkAstVarName (adFTK ftk) freshIdD+      var :: AstVarName '(FullSpan, x)+      var = mkAstVarName ftk freshId+      astVarPrimalD :: AstTensor AstMethodShare FullSpan (ADTensorKind x)       !astVarPrimalD = astVar varPrimalD-      astVarPrimal :: AstTensor AstMethodShare PrimalSpan x-      !astVarPrimal = astVar varPrimal+      astVarPrimal :: AstTensor AstMethodShare FullSpan x+      !astVarPrimal = astVar var       astVarD :: AstTensor AstMethodLet FullSpan x       !astVarD = astVar var-  return (varPrimalD, astVarPrimalD, varPrimal, astVarPrimal, var, astVarD)--funToAstIntVarIO :: Maybe (Int64, Int64) -> ((IntVarName, AstInt ms) -> a)-                 -> IO a-{-# INLINE funToAstIntVarIO #-}-funToAstIntVarIO bounds f = do-  !varName <- unsafeGetFreshAstVarName (FTKScalar @Int64) bounds-  return $! f (varName, astVar varName)--funToAstIntVar :: Maybe (Int64, Int64) -> ((IntVarName, AstInt ms) -> a) -> a-{-# NOINLINE funToAstIntVar #-}-funToAstIntVar bounds = unsafePerformIO . funToAstIntVarIO bounds--funToAstI :: Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)-{-# NOINLINE funToAstI #-}-funToAstI bounds f = unsafePerformIO . funToAstIntVarIO bounds-                     $ \ (!var, !i) -> let !x = f i in (var, x)+  return (varPrimalD, astVarPrimalD, astVarPrimal, var, astVarD)  funToVarsIxIOS-  :: forall sh a ms.-     ShS sh -> ((AstVarListS sh, AstIxS ms sh) -> a) -> IO a+  :: ShS sh -> (AstVarListS sh -> AstIxS ms sh -> AstTensor ms s2 z)+  -> IO (AstTensor ms s2 z) {-# INLINE funToVarsIxIOS #-} funToVarsIxIOS sh f = withKnownShS sh $ do-  let freshBound n =-        unsafeGetFreshAstVarName (FTKScalar @Int64)-                                 (Just (0, fromIntegral n - 1))-  !varList <- mapM freshBound $ shsToList sh-  let !vars = fromList varList-  let !ix = fromList $ map astVar varList-  return $! f (vars, ix)+  let unsafeGetFreshIntVarName :: Int -> IO IntVarName+      unsafeGetFreshIntVarName n = do+        freshId <- unsafeGetFreshAstVarId+        return $! mkAstVarNameBounds (0, n - 1) freshId+  varList <- mapM unsafeGetFreshIntVarName $ shsToList sh+  let !ix = fromList varList+      vars = AstVarListS ix+      asts = fmap astVar ix+  return $! f vars asts  funToVarsIxS-  :: ShS sh -> ((AstVarListS sh, AstIxS ms sh) -> a) -> a+  :: ShS sh -> (AstVarListS sh -> AstIxS ms sh -> AstTensor ms s2 z)+  -> AstTensor ms s2 z {-# NOINLINE funToVarsIxS #-} funToVarsIxS sh = unsafePerformIO . funToVarsIxIOS sh--funToAstIxS-  :: ShS sh -> (AstIxS ms sh -> AstIxS ms sh2)-  -> (AstVarListS sh, AstIxS ms sh2)-{-# NOINLINE funToAstIxS #-}-funToAstIxS sh f = unsafePerformIO $ funToVarsIxIOS sh-                   $ \ (!vars, !ix) -> let !x = f ix in (vars, x)
src/HordeAd/Core/AstInline.hs view
@@ -1,11 +1,13 @@+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif -- | Inlining and global sharing elimination. module HordeAd.Core.AstInline   ( -- * Inlining-    inlineAst+    inlineAstTensor     -- * Translation of global sharing to local lets   , unshareAstTensor-    -- * Helper classes and types-  , AstMemo   ) where  import Prelude@@ -16,83 +18,60 @@ import Data.Dependent.Sum (DSum (..)) import Data.EnumMap.Strict qualified as EM import Data.Foldable qualified as Foldable-import Data.List (mapAccumR)+import Data.Kind (Type)+import Data.List (sortOn)+import Data.Ord qualified import Data.Some-import Data.Type.Equality ((:~:) (Refl))-import GHC.Exts (IsList (..))+import Data.Vector.Generic qualified as V -import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX) import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat') -import HordeAd.Core.Ast (AstBool, AstTensor)-import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))+import HordeAd.Core.Ast (AstTensor)+import HordeAd.Core.Ast hiding (AstTensor (..)) import HordeAd.Core.Ast qualified as Ast import HordeAd.Core.AstSimplify (substituteAst) import HordeAd.Core.AstTools-import HordeAd.Core.TensorKind-import HordeAd.Core.Types+import HordeAd.Core.Types (TK, mapAccumL')  -- * The pass that inlines lets with the bottom-up strategy  type AstMemo = EM.EnumMap AstVarId Int --- | This inlines 'HordeAd.Core.Ast.AstLet', traversing the term bottom-up.+-- | This inlines occurences of 'HordeAd.Core.Ast.AstLet', traversing+-- the term bottom-up.+inlineAstTensor+  :: forall s y. KnownSpan s+  => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y+inlineAstTensor = snd . inlineAst 1 EM.empty++-- | This inlines occurences of 'HordeAd.Core.Ast.AstLet', traversing+-- the term bottom-up. inlineAst-  :: forall s y. AstSpan s-  => AstMemo -> AstTensor AstMethodLet s y+  :: forall s y. KnownSpan s+  => Int -> AstMemo -> AstTensor AstMethodLet s y   -> (AstMemo, AstTensor AstMethodLet s y)-inlineAst memo v0 = case v0 of+inlineAst c !memo v0 = case v0 of   Ast.AstPair t1 t2 ->-    let (memo2, v1) = inlineAst memo t1-        (memo3, v2) = inlineAst memo2 t2+    let (memo2, v1) = inlineAst c memo t1+        (memo3, v2) = inlineAst c memo2 t2     in (memo3, Ast.AstPair v1 v2)-  Ast.AstProject1 t -> second Ast.AstProject1 (inlineAst memo t)-  Ast.AstProject2 t -> second Ast.AstProject2 (inlineAst memo t)-  Ast.AstFromVector snat stk l ->-    let (memo2, l2) = mapAccumR inlineAst memo l-    in (memo2, Ast.AstFromVector snat stk l2)-  Ast.AstSum snat stk v -> second (Ast.AstSum snat stk) (inlineAst memo v)-  Ast.AstReplicate snat stk v ->-    second (Ast.AstReplicate snat stk) (inlineAst memo v)-  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->-    let (memo1, f2) = inlineAstHFun memo f-        (memo2, df2) = inlineAstHFun memo1 df-        (memo3, rf2) = inlineAstHFun memo2 rf-        (memo4, acc02) = inlineAst memo3 acc0-        (memo5, es2) = inlineAst memo4 es-    in (memo5, Ast.AstMapAccumRDer k bftk eftk f2 df2 rf2 acc02 es2)+  Ast.AstProject1 t -> second Ast.AstProject1 (inlineAst c memo t)+  Ast.AstProject2 t -> second Ast.AstProject2 (inlineAst c memo t)   Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->-    let (memo1, f2) = inlineAstHFun memo f-        (memo2, df2) = inlineAstHFun memo1 df-        (memo3, rf2) = inlineAstHFun memo2 rf-        (memo4, acc02) = inlineAst memo3 acc0-        (memo5, es2) = inlineAst memo4 es+    let (memo1, f2) = inlineAstHFun c memo f+        (memo2, df2) = inlineAstHFun c memo1 df+        (memo3, rf2) = inlineAstHFun c memo2 rf+        (memo4, acc02) = inlineAst c memo3 acc0+        (memo5, es2) = inlineAst c memo4 es     in (memo5, Ast.AstMapAccumLDer k bftk eftk f2 df2 rf2 acc02 es2)   Ast.AstApply t ll ->-    let (memo1, t2) = inlineAstHFun memo t-        (memo2, ll2) = inlineAst memo1 ll+    let (memo1, t2) = inlineAstHFun c memo t+        (memo2, ll2) = inlineAst c memo1 ll     in (memo2, Ast.AstApply t2 ll2)-  Ast.AstVar var ->-    let f Nothing = Just 1-        f (Just count) = Just $ count + 1-    in (EM.alter f (varNameToAstVarId var) memo, v0)-  Ast.AstCond b a2 a3 ->-    -- This is a place where our inlining may increase code size-    -- by enlarging both branches due to not considering number of syntactic-    -- occurrences, but only dynamic occurrences. Tensor expressions-    -- in conditionals are problematic and special enough-    -- that we can let it be until problems are encountered in the wild.-    -- See https://github.com/VMatthijs/CHAD/blob/main/src/Count.hs#L88-L152.-    let (memo1, b1) = inlineAstBool memo b-        (memoA2, t2) = inlineAst EM.empty a2-        (memoA3, t3) = inlineAst EM.empty a3-        memo4 = EM.unionWith max memoA2 memoA3-        memo5 = EM.unionWith (+) memo1 memo4-    in (memo5, Ast.AstCond b1 t2 t3)+  Ast.AstVar var -> (EM.adjust (+ c) (varNameToAstVarId var) memo, v0)   Ast.AstBuild1 k stk (var, v) ->-    let (memoV0, !v2) = inlineAst EM.empty v-        memo1 = EM.unionWith-                  (\c1 c0 -> c1 + sNatValue k * c0) memo memoV0+    let (memo1, !v2) = inlineAst (fromSNat' k * c) memo v     in (memo1, Ast.AstBuild1 k stk (var, v2))    Ast.AstLet var u v ->@@ -100,202 +79,253 @@     -- the delete and hence var couldn't appear in memo, so we can make     -- the recursive call for v with memo intact, to record extra occurrences     -- of other variables without the costly summing of maps.-    let vv = varNameToAstVarId var-        (memo1, v2) = inlineAst memo v-        memo1NoVar = EM.delete vv memo1-        (memo2, u2) = inlineAst memo1NoVar u-    in case EM.findWithDefault 0 vv memo1 of-      0 -> (memo1, v2)-      1 -> (memo2, substituteAst u2 var v2)-      count | astIsSmall (count < 10) u ->-        let (memoU0, u0) = inlineAst EM.empty u-            memo3 = EM.unionWith (\c1 c0 -> c1 + count * c0) memo1NoVar memoU0-                      -- u is small, so the union is fast+    withKnownSpan (varNameToSpan var) $+    let (memo1, v2) = inlineAst c (EM.insert (varNameToAstVarId var) 0 memo) v+        memo1NoVar = EM.delete (varNameToAstVarId var) memo1+    in case memo1 EM.! varNameToAstVarId var of+      0 -> (memo1NoVar, v2)+      count | count <= c  -- occurs once and in the same nesting situation+              || astIsSmall (count < 100) u ->+        let (memo3, u0) = inlineAst count memo1NoVar u         in (memo3, substituteAst u0 var v2)-      _ -> (memo2, Ast.AstLet var u2 v2)+      _ ->+        let (memo2, u2) = inlineAst c memo1NoVar u+        in (memo2, Ast.AstLet var u2 v2) -  Ast.AstPrimalPart a -> second Ast.AstPrimalPart $ inlineAst memo a-  Ast.AstDualPart a -> second Ast.AstDualPart $ inlineAst memo a-  Ast.AstFromPrimal a -> second Ast.AstFromPrimal $ inlineAst memo a-  Ast.AstFromDual a -> second Ast.AstFromDual $ inlineAst memo a+  Ast.AstPrimalPart a -> second Ast.AstPrimalPart $ inlineAst c memo a+  Ast.AstDualPart a -> second Ast.AstDualPart $ inlineAst c memo a+  Ast.AstPlainPart a -> second Ast.AstPlainPart $ inlineAst c memo a+  Ast.AstFromPrimal a -> second Ast.AstFromPrimal $ inlineAst c memo a+  Ast.AstFromDual a -> second Ast.AstFromDual $ inlineAst c memo a+  Ast.AstFromPlain a -> second Ast.AstFromPlain $ inlineAst c memo a    Ast.AstPlusK u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstPlusK u2 v3)   Ast.AstTimesK u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstTimesK u2 v3)   Ast.AstN1K opCode u ->-    let (memo2, u2) = inlineAst memo u+    let (memo2, u2) = inlineAst c memo u     in (memo2, Ast.AstN1K opCode u2)   Ast.AstR1K opCode u ->-    let (memo2, u2) = inlineAst memo u+    let (memo2, u2) = inlineAst c memo u     in (memo2, Ast.AstR1K opCode u2)   Ast.AstR2K opCode u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstR2K opCode u2 v3)   Ast.AstI2K opCode u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstI2K opCode u2 v3)   Ast.AstConcreteK{} -> (memo, v0)-  Ast.AstFloorK a -> second Ast.AstFloorK $ inlineAst memo a-  Ast.AstFromIntegralK a -> second Ast.AstFromIntegralK $ inlineAst memo a-  Ast.AstCastK a -> second Ast.AstCastK $ inlineAst memo a+  Ast.AstFloorK a -> second Ast.AstFloorK $ inlineAst c memo a+  Ast.AstFromIntegralK a -> second Ast.AstFromIntegralK $ inlineAst c memo a+  Ast.AstCastK a -> second Ast.AstCastK $ inlineAst c memo a+  Ast.AstArgMinK a -> second Ast.AstArgMinK $ inlineAst c memo a+  Ast.AstArgMaxK a -> second Ast.AstArgMaxK $ inlineAst c memo a+  Ast.AstIndexK v ix ->+    let (memo1, v2) = inlineAst c memo v+        (memo2, ix2) = mapAccumL' (inlineAst c) memo1 (Foldable.toList ix)+    in (memo2, Ast.AstIndexK v2 (ixsFromIxS ix ix2))    Ast.AstPlusS u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstPlusS u2 v3)   Ast.AstTimesS u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstTimesS u2 v3)   Ast.AstN1S opCode u ->-    let (memo2, u2) = inlineAst memo u+    let (memo2, u2) = inlineAst c memo u     in (memo2, Ast.AstN1S opCode u2)   Ast.AstR1S opCode u ->-    let (memo2, u2) = inlineAst memo u+    let (memo2, u2) = inlineAst c memo u     in (memo2, Ast.AstR1S opCode u2)   Ast.AstR2S opCode u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstR2S opCode u2 v3)   Ast.AstI2S opCode u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstI2S opCode u2 v3)   Ast.AstConcreteS{} -> (memo, v0)-  Ast.AstFloorS a -> second Ast.AstFloorS $ inlineAst memo a-  Ast.AstFromIntegralS v -> second Ast.AstFromIntegralS $ inlineAst memo v-  Ast.AstCastS v -> second Ast.AstCastS $ inlineAst memo v+  Ast.AstFloorS a -> second Ast.AstFloorS $ inlineAst c memo a+  Ast.AstFromIntegralS v -> second Ast.AstFromIntegralS $ inlineAst c memo v+  Ast.AstCastS v -> second Ast.AstCastS $ inlineAst c memo v+  Ast.AstArgMinS a -> second Ast.AstArgMinS $ inlineAst c memo a+  Ast.AstArgMaxS a -> second Ast.AstArgMaxS $ inlineAst c memo a+  Ast.AstIndexS @shm shn v ix ->+    let (memo1, v2) = inlineAst c memo v+        (memo2, ix2) = mapAccumL' (inlineAst c) memo1 (Foldable.toList ix)+    in (memo2, Ast.AstIndexS @shm shn v2 (ixsFromIxS ix ix2)) -  Ast.AstIndexS @sh1 shn v ix ->-    let (memo1, v2) = inlineAst memo v-        (memo2, ix2) = mapAccumR inlineAst memo1 (Foldable.toList ix)-    in withKnownShS (shsFromIxS ix)-       (memo2, Ast.AstIndexS @sh1 shn v2 (fromList ix2))-  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->-    let (memo1, v2) = inlineAst memo v-        (memoI0, ix2) = mapAccumR inlineAst EM.empty (Foldable.toList ix)-        count = shsSize (shsFromIxS ix) * shsSize shn-        memo2 = EM.unionWith (\c1 c0 -> c1 + count * c0) memo1 memoI0-        !ix3 = withKnownShS (shsFromIxS ix) $ fromList ix2-    in (memo2, Ast.AstScatterS @shm @shn @shp shn v2 (vars, ix3))-  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->-    let (memo1, v2) = inlineAst memo v-        (memoI0, ix2) = mapAccumR inlineAst EM.empty (Foldable.toList ix)-        count = shsSize (shsFromListS vars) * shsSize shn-        memo2 = EM.unionWith (\c1 c0 -> c1 + count * c0) memo1 memoI0-        !ix3 = withKnownShS (shsFromIxS ix) $ fromList ix2-    in (memo2, Ast.AstGatherS @shm @shn @shp shn v2 (vars, ix3))-  Ast.AstMinIndexS a -> second Ast.AstMinIndexS $ inlineAst memo a-  Ast.AstMaxIndexS a -> second Ast.AstMaxIndexS $ inlineAst memo a+  -- This is a place where our inlining may increase code size+  -- by enlarging both branches due to not considering number of syntactic+  -- occurrences, but only dynamic occurrences. Tensor expressions+  -- in conditionals are problematic and special enough+  -- that we can let it be until problems are encountered in the wild.+  -- See https://github.com/VMatthijs/CHAD/blob/main/src/Count.hs#L88-L152.+  Ast.AstCondK b a2 a3 ->+    let (memoA2, t2) = inlineAst c memo a2+        (memoA3, t3) = inlineAst c memo a3+        memo4 = EM.unionWith max memoA2 memoA3+        (memo5, b1) = inlineAst c memo4 b+    in (memo5, Ast.AstCondK b1 t2 t3)+  Ast.AstCondS b a2 a3 ->+    let (memoA2, t2) = inlineAst c memo a2+        (memoA3, t3) = inlineAst c memo a3+        memo4 = EM.unionWith max memoA2 memoA3+        (memo5, b1) = inlineAst c memo4 b+    in (memo5, Ast.AstCondS b1 t2 t3)+  Ast.AstFromVectorK shm l ->+    let (memo2, l2) = mapAccumL' (inlineAst c) memo $ V.toList l+    in (memo2, Ast.AstFromVectorK shm $ V.fromListN (V.length l) l2)+  Ast.AstFromVectorS shm l ->+    let (memo2, l2) = mapAccumL' (inlineAst c) memo $ V.toList l+    in (memo2, Ast.AstFromVectorS shm $ V.fromListN (V.length l) l2)+  Ast.AstSumK v -> second Ast.AstSumK (inlineAst c memo v)+  Ast.AstSumS shm v -> second (Ast.AstSumS shm) (inlineAst c memo v)+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    let (memo1, v2) = inlineAst c memo v+        (memo2, ix2) = mapAccumL' (inlineAst (shsSize shp * c)) memo1+                                  (Foldable.toList ix)+        !ix3 = ixsFromIxS ix ix2+    in (memo2, Ast.AstScatterS shm shn shp v2 (vars, ix3))+  Ast.AstReplicateK shm v -> second (Ast.AstReplicateK shm) (inlineAst c memo v)+  Ast.AstReplicateS shm v -> second (Ast.AstReplicateS shm) (inlineAst c memo v)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    let (memo1, v2) = inlineAst c memo v+        (memo2, ix2) = mapAccumL' (inlineAst (shsSize shp * c)) memo1+                                  (Foldable.toList ix)+        !ix3 = ixsFromIxS ix ix2+    in (memo2, Ast.AstGatherS shm shn shp v2 (vars, ix3))   Ast.AstIotaS{} -> (memo, v0)   Ast.AstAppendS x y ->-    let (memo1, t1) = inlineAst memo x-        (memo2, t2) = inlineAst memo1 y+    let (memo1, t1) = inlineAst c memo x+        (memo2, t2) = inlineAst c memo1 y     in (memo2, Ast.AstAppendS t1 t2)-  Ast.AstSliceS i n k v -> second (Ast.AstSliceS i n k) (inlineAst memo v)-  Ast.AstReverseS v -> second Ast.AstReverseS (inlineAst memo v)-  Ast.AstTransposeS perm v -> second (Ast.AstTransposeS perm) $ inlineAst memo v-  Ast.AstReshapeS sh v -> second (Ast.AstReshapeS sh) (inlineAst memo v)+  Ast.AstSliceS i n k v -> second (Ast.AstSliceS i n k) (inlineAst c memo v)+  Ast.AstReverseS v -> second Ast.AstReverseS (inlineAst c memo v)+  Ast.AstTransposeS perm v ->+    second (Ast.AstTransposeS perm) $ inlineAst c memo v+  Ast.AstReshapeS sh v -> second (Ast.AstReshapeS sh) (inlineAst c memo v) -  Ast.AstConvert c v -> second (Ast.AstConvert c) $ inlineAst memo v+  Ast.AstConvert c1 v -> second (Ast.AstConvert c1) $ inlineAst c memo v -  Ast.AstSum0S v -> second Ast.AstSum0S (inlineAst memo v)-  Ast.AstDot0S u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v-    in (memo3, Ast.AstDot0S u2 v3)+  Ast.AstDot0 u v ->+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v+    in (memo3, Ast.AstDot0 u2 v3)   Ast.AstDot1InS m n u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstDot1InS m n u2 v3)   Ast.AstMatmul2S m n p u v ->-    let (memo2, u2) = inlineAst memo u-        (memo3, v3) = inlineAst memo2 v+    let (memo2, u2) = inlineAst c memo u+        (memo3, v3) = inlineAst c memo2 v     in (memo3, Ast.AstMatmul2S m n p u2 v3) +  Ast.AstBoolNotK arg ->+    let (memo2, arg2) = inlineAst c memo arg+    in (memo2, Ast.AstBoolNotK arg2)+  Ast.AstBoolNotS arg ->+    let (memo2, arg2) = inlineAst c memo arg+    in (memo2, Ast.AstBoolNotS arg2)+  Ast.AstBoolAndK arg1 arg2 ->+    let (memo1, b1) = inlineAst c memo arg1+        (memo2, b2) = inlineAst c memo1 arg2+    in (memo2, Ast.AstBoolAndK b1 b2)+  Ast.AstBoolAndS arg1 arg2 ->+    let (memo1, b1) = inlineAst c memo arg1+        (memo2, b2) = inlineAst c memo1 arg2+    in (memo2, Ast.AstBoolAndS b1 b2)+  Ast.AstLeqK arg1 arg2 ->+    let (memo1, r1) = inlineAst c memo arg1+        (memo2, r2) = inlineAst c memo1 arg2+    in (memo2, Ast.AstLeqK r1 r2)+  Ast.AstLeq arg1 arg2 ->+    let (memo1, r1) = inlineAst c memo arg1+        (memo2, r2) = inlineAst c memo1 arg2+    in (memo2, Ast.AstLeq r1 r2)+  Ast.AstLeqS shb sh arg1 arg2 ->+    let (memo1, r1) = inlineAst c memo arg1+        (memo2, r2) = inlineAst c memo1 arg2+    in (memo2, Ast.AstLeqS shb sh r1 r2)+ inlineAstHFun-  :: AstSpan s2-  => AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)-inlineAstHFun memo v0 = case v0 of+  :: KnownSpan s+  => Int -> AstMemo -> AstHFun s x y -> (AstMemo, AstHFun s x y)+inlineAstHFun c !memo v0 = case v0 of   Ast.AstLambda var l ->     -- No other free variables in l, so no outside lets can reach there,     -- so we don't need to pass the information from v upwards.-    (memo, Ast.AstLambda var (snd $ inlineAst EM.empty l))--inlineAstBool :: AstMemo -> AstBool AstMethodLet-              -> (AstMemo, AstBool AstMethodLet)-inlineAstBool memo v0 = case v0 of-  Ast.AstBoolConst{} -> (memo, v0)-  Ast.AstBoolNot arg ->-    let (memo2, arg2) = inlineAstBool memo arg-    in (memo2, Ast.AstBoolNot arg2)-  Ast.AstBoolAnd arg1 arg2 ->-    let (memo1, b1) = inlineAstBool memo arg1-        (memo2, b2) = inlineAstBool memo1 arg2-    in (memo2, Ast.AstBoolAnd b1 b2)-  Ast.AstLeqK arg1 arg2 ->-    let (memo1, r1) = inlineAst memo arg1-        (memo2, r2) = inlineAst memo1 arg2-    in (memo2, Ast.AstLeqK r1 r2)-  Ast.AstLeqS arg1 arg2 ->-    let (memo1, r1) = inlineAst memo arg1-        (memo2, r2) = inlineAst memo1 arg2-    in (memo2, Ast.AstLeqS r1 r2)+    -- However, how many times this code is copied by the context matters+    -- so @c@ is taken into account instead of setting 1.+    (memo, Ast.AstLambda var (snd $ inlineAst c EM.empty l))   -- * Translation of global sharing to normal lets -type AstBindings = DEnumMap (AstVarName PrimalSpan)-                            (AstTensor AstMethodLet PrimalSpan)+type AstBindings = DEnumMap AstVarName SpanTarget -bindsToLet :: forall s y.-              AstTensor AstMethodLet s y -> AstBindings+type role SpanTarget nominal+data SpanTarget :: (AstSpan, TK) -> Type where+  SpanTarget :: AstTensor AstMethodLet s y -> SpanTarget '(s, y)++bindsToLet :: forall s y. KnownSpan s+           => AstTensor AstMethodLet s y -> AstBindings            -> AstTensor AstMethodLet s y-{-# INLINE bindsToLet #-}  -- help list fusion-bindsToLet u0 bs = foldl' bindToLet u0 (DMap.toDescList bs)+bindsToLet u0 !memo = foldl' bindToLet u0 l  where+  varFromDSum :: DSum AstVarName SpanTarget -> AstVarId+  varFromDSum (var :=> _) = varNameToAstVarId var+  l :: [DSum AstVarName SpanTarget]+  l = sortOn (Data.Ord.Down . varFromDSum) (DMap.toList memo)+  -- Lets are immediately pushed down before other rewrites block+  -- some opportunities.   bindToLet :: AstTensor AstMethodLet s y-            -> DSum (AstVarName PrimalSpan)-                    (AstTensor AstMethodLet PrimalSpan)+            -> DSum AstVarName SpanTarget             -> AstTensor AstMethodLet s y-  bindToLet !u (var :=> w) = Ast.AstLet var w u+  bindToLet !u (var :=> SpanTarget w) = astLetDown var w u  -- | This replaces 'HordeAd.Core.Ast.AstShare' with 'HordeAd.Core.Ast.AstLet', -- traversing the term bottom-up.-unshareAstTensor :: AstTensor AstMethodShare PrimalSpan y-                 -> AstTensor AstMethodLet PrimalSpan y+unshareAstTensor :: AstTensor AstMethodShare FullSpan y+                 -> AstTensor AstMethodLet FullSpan y unshareAstTensor tShare =   let (memoOut, tLet) = unshareAst DMap.empty tShare   in bindsToLet tLet memoOut +-- Splitting the variable list to make it more typed complicates+-- and slows down the code, so let's keep it just [AstVarId].+closeOccurs :: [AstVarId] -> AstBindings -> (AstBindings, AstBindings)+closeOccurs vars !memo =+  let varsOccur :: SpanTarget s_y -> Bool+      varsOccur (SpanTarget t) = any (`varInAst` t) vars+      (memoLocal, memoGlobal) = DMap.partition varsOccur memo+  in if DMap.null memoLocal+     then (memoLocal, memoGlobal)+     else let vars2 = map (\(Some var) -> varNameToAstVarId var)+                          (DMap.keys memoLocal)+              (memoLocal2, memoGlobal2) = closeOccurs vars2 memoGlobal+          in (DMap.union memoLocal memoLocal2, memoGlobal2)+ -- This works only because the other code never inserts the same rshare -- into more than one index element, with the share containing -- the gather/scatter/build variables corresponding to the index. unshareAstScoped-  :: forall z s. AstSpan s+  :: forall z s. KnownSpan s   => [IntVarName] -> AstBindings -> AstTensor AstMethodShare s z   -> (AstBindings, AstTensor AstMethodLet s z)-unshareAstScoped vars0 memo0 v0 =+unshareAstScoped vars0 !memo0 v0 =   let (memo1, v1) = unshareAst memo0 v0       memoDiff = DMap.difference memo1 memo0-      varsOccur :: [AstVarId] -> AstTensor AstMethodLet PrimalSpan y -> Bool-      varsOccur vs d = any (`varInAst` d) vs-      closeOccurs :: [AstVarId] -> AstBindings -> (AstBindings, AstBindings)-      closeOccurs vars memo =-        let (memoLocal, memoGlobal) = DMap.partition (varsOccur vars) memo-        in if DMap.null memoLocal-           then (memoLocal, memoGlobal)-           else let vars2 = map (\(Some var) -> varNameToAstVarId var)-                                (DMap.keys memoLocal)-                    (memoLocal2, memoGlobal2) = closeOccurs vars2 memoGlobal-                in (DMap.union memoLocal memoLocal2, memoGlobal2)       (memoLocal1, memoGlobal1) =         closeOccurs (map varNameToAstVarId vars0) memoDiff   in (DMap.union memo0 memoGlobal1, bindsToLet v1 memoLocal1)@@ -304,26 +334,16 @@ -- but we mark it as potentially containing lets, because in the future -- we may optimize this by inserting some lets not at the top-level. unshareAst-  :: forall s y. AstSpan s+  :: forall s y. KnownSpan s   => AstBindings -> AstTensor AstMethodShare s y   -> (AstBindings, AstTensor AstMethodLet s y)-unshareAst memo = \case+unshareAst !memo = \case   Ast.AstPair t1 t2 ->     let (memo1, v1) = unshareAst memo t1         (memo2, v2) = unshareAst memo1 t2     in (memo2, Ast.AstPair v1 v2)   Ast.AstProject1 t -> second Ast.AstProject1 (unshareAst memo t)   Ast.AstProject2 t -> second Ast.AstProject2 (unshareAst memo t)-  Ast.AstFromVector snat stk l ->-    let (memo2, l2) = mapAccumR unshareAst memo l-    in (memo2, Ast.AstFromVector snat stk l2)-  Ast.AstSum snat stk v -> second (Ast.AstSum snat stk) (unshareAst memo v)-  Ast.AstReplicate snat stk v ->-    second (Ast.AstReplicate snat stk) (unshareAst memo v)-  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->-    let (memo1, acc02) = unshareAst memo acc0-        (memo2, es2) = unshareAst memo1 es-    in (memo2, Ast.AstMapAccumRDer k bftk eftk f df rf acc02 es2)   Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->     let (memo1, acc02) = unshareAst memo acc0         (memo2, es2) = unshareAst memo1 es@@ -333,62 +353,25 @@         (memo2, ll2) = unshareAst memo1 ll     in (memo2, Ast.AstApply t2 ll2)   Ast.AstVar v -> (memo, Ast.AstVar v)-  Ast.AstCond b a2 a3 ->-    let (memo1, b1) = unshareAstBool memo b-        (memo2, t2) = unshareAst memo1 a2-        (memo3, t3) = unshareAst memo2 a3-    in (memo3, Ast.AstCond b1 t2 t3)   Ast.AstBuild1 snat stk (var, v) ->     let (memo1, !v2) = unshareAstScoped [var] memo v     in (memo1, Ast.AstBuild1 snat stk (var, v2))    -- We assume v is the same if var is the same.-  Ast.AstShare varRaw a | Just Refl <- sameAstSpan @s @PrimalSpan -> case a of-    AstFromS' @y2 ftkz v ->-      let var = mkAstVarName-                  (ftkAst v) (varNameToBounds varRaw) (varNameToAstVarId varRaw)-          astVar0 = cAstFromS @y2 ftkz $ Ast.AstVar var-      in if var `DMap.member` memo-         then (memo, astVar0)-         else let (memo1, !a2) = unshareAst memo v-                    -- DMap is strict, but let's be paranoid-              in (DMap.insert var a2 memo1, astVar0)-    -- The PrimalSpan check ensures there's no need to match for-    -- Ast.AstFromPrimal (Ast.AstFromS).-    _ -> case varNameToFTK varRaw of-      ftk@(FTKR @_ @x sh' x) ->-        withShsFromShR sh' $ \(sh :: ShS sh) ->-          let var = mkAstVarName (FTKS sh x) (varNameToBounds varRaw)-                                 (varNameToAstVarId varRaw)-              astVar0 = cAstFromS @(TKS2 sh x) ftk $ Ast.AstVar var-          in if var `DMap.member` memo-             then (memo, astVar0)-             else let (memo1, !a2) = unshareAst memo (cAstSFromR @sh sh a)-                  in (DMap.insert var a2 memo1, astVar0)-      ftk@(FTKX @_ @x sh' x) ->-        withShsFromShX sh' $ \(sh :: ShS sh) ->-          let var = mkAstVarName (FTKS sh x) (varNameToBounds varRaw)-                                 (varNameToAstVarId varRaw)-              astVar0 = cAstFromS @(TKS2 sh x) ftk $ Ast.AstVar var-          in if var `DMap.member` memo-             then (memo, astVar0)-             else let (memo1, !a2) = unshareAst memo (cAstSFromX @sh sh a)-                  in (DMap.insert var a2 memo1, astVar0)-      -- it maybe not be worth it to recursively convert product-      -- so let's not do that until profiling shows we need it-      _ -> let var = varRaw-               astVar0 = Ast.AstVar var-           in if var `DMap.member` memo-              then (memo, astVar0)-              else let (memo1, !a2) = unshareAst memo a-                   in (DMap.insert var a2 memo1, astVar0)-  Ast.AstShare{} -> error "unshareAst: AstShare not in PrimalSpan"+  Ast.AstShare var a ->+    let astVar0 = Ast.AstVar var+    in if var `DMap.member` memo+       then (memo, astVar0)+       else let (memo1, a2) = unshareAst memo a+            in (DMap.insert var (SpanTarget a2) memo1, astVar0)   Ast.AstToShare v -> (memo, v)  -- nothing to unshare in this subtree    Ast.AstPrimalPart a -> second Ast.AstPrimalPart $ unshareAst memo a   Ast.AstDualPart a -> second Ast.AstDualPart $ unshareAst memo a+  Ast.AstPlainPart a -> second Ast.AstPlainPart $ unshareAst memo a   Ast.AstFromPrimal a -> second Ast.AstFromPrimal $ unshareAst memo a   Ast.AstFromDual a -> second Ast.AstFromDual $ unshareAst memo a+  Ast.AstFromPlain a -> second Ast.AstFromPlain $ unshareAst memo a    Ast.AstPlusK u v ->     let (memo2, u2) = unshareAst memo u@@ -416,6 +399,12 @@   Ast.AstFloorK a -> second Ast.AstFloorK $ unshareAst memo a   Ast.AstFromIntegralK v -> second Ast.AstFromIntegralK $ unshareAst memo v   Ast.AstCastK v -> second Ast.AstCastK $ unshareAst memo v+  Ast.AstArgMinK v -> second Ast.AstArgMinK $ unshareAst memo v+  Ast.AstArgMaxK v -> second Ast.AstArgMaxK $ unshareAst memo v+  Ast.AstIndexK v ix ->+    let (memo1, v2) = unshareAst memo v+        (memo2, ix2) = mapAccumL' unshareAst memo1 (Foldable.toList ix)+    in (memo2, Ast.AstIndexK v2 (ixsFromIxS ix ix2))    Ast.AstPlusS u v ->     let (memo2, u2) = unshareAst memo u@@ -443,28 +432,45 @@   Ast.AstFloorS a -> second Ast.AstFloorS $ unshareAst memo a   Ast.AstFromIntegralS v -> second Ast.AstFromIntegralS $ unshareAst memo v   Ast.AstCastS v -> second Ast.AstCastS $ unshareAst memo v--  Ast.AstIndexS @sh1 shn v ix ->+  Ast.AstArgMinS a -> second Ast.AstArgMinS $ unshareAst memo a+  Ast.AstArgMaxS a -> second Ast.AstArgMaxS $ unshareAst memo a+  Ast.AstIndexS @shm shn v ix ->     let (memo1, v2) = unshareAst memo v-        (memo2, ix2) = mapAccumR unshareAst memo1 (Foldable.toList ix)-    in withKnownShS (shsFromIxS ix)-       (memo2, Ast.AstIndexS @sh1 shn v2 (fromList ix2))-  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->-    let (memo1, ix2) =-          mapAccumR (unshareAstScoped $ listsToList vars)-                    memo (Foldable.toList ix)+        (memo2, ix2) = mapAccumL' unshareAst memo1 (Foldable.toList ix)+    in (memo2, Ast.AstIndexS @shm shn v2 (ixsFromIxS ix ix2))++  Ast.AstCondK b a2 a3 ->+    let (memo1, b1) = unshareAst memo b+        (memo2, t2) = unshareAst memo1 a2+        (memo3, t3) = unshareAst memo2 a3+    in (memo3, Ast.AstCondK b1 t2 t3)+  Ast.AstCondS b a2 a3 ->+    let (memo1, b1) = unshareAst memo b+        (memo2, t2) = unshareAst memo1 a2+        (memo3, t3) = unshareAst memo2 a3+    in (memo3, Ast.AstCondS b1 t2 t3)+  Ast.AstFromVectorK shm l ->+    let (memo2, l2) = mapAccumL' unshareAst memo $ V.toList l+    in (memo2, Ast.AstFromVectorK shm $ V.fromListN (V.length l) l2)+  Ast.AstFromVectorS shm l ->+    let (memo2, l2) = mapAccumL' unshareAst memo $ V.toList l+    in (memo2, Ast.AstFromVectorS shm $ V.fromListN (V.length l) l2)+  Ast.AstSumK v -> second Ast.AstSumK (unshareAst memo v)+  Ast.AstSumS shm v -> second (Ast.AstSumS shm) (unshareAst memo v)+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    let (memo1, ix2) = mapAccumL' (unshareAstScoped $ Foldable.toList vars)+                                  memo (Foldable.toList ix)         (memo2, v2) = unshareAst memo1 v-        !ix3 = withKnownShS (shsFromIxS ix) $ fromList ix2-    in (memo2, Ast.AstScatterS @shm @shn @shp shn v2 (vars, ix3))-  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->-    let (memo1, ix2) =-          mapAccumR (unshareAstScoped $ listsToList vars)-                    memo (Foldable.toList ix)+        !ix3 = ixsFromIxS ix ix2+    in (memo2, Ast.AstScatterS shm shn shp v2 (vars, ix3))+  Ast.AstReplicateK shm v -> second (Ast.AstReplicateK shm) (unshareAst memo v)+  Ast.AstReplicateS shm v -> second (Ast.AstReplicateS shm) (unshareAst memo v)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    let (memo1, ix2) = mapAccumL' (unshareAstScoped $ Foldable.toList vars)+                                  memo (Foldable.toList ix)         (memo2, v2) = unshareAst memo1 v-        !ix3 = withKnownShS (shsFromIxS ix) $ fromList ix2-    in (memo2, Ast.AstGatherS @shm @shn @shp shn v2 (vars, ix3))-  Ast.AstMinIndexS a -> second Ast.AstMinIndexS $ unshareAst memo a-  Ast.AstMaxIndexS a -> second Ast.AstMaxIndexS $ unshareAst memo a+        !ix3 = ixsFromIxS ix ix2+    in (memo2, Ast.AstGatherS shm shn shp v2 (vars, ix3))   Ast.AstIotaS snat -> (memo, Ast.AstIotaS snat)   Ast.AstAppendS x y ->     let (memo1, t1) = unshareAst memo x@@ -478,11 +484,10 @@    Ast.AstConvert c v -> second (Ast.AstConvert c) $ unshareAst memo v -  Ast.AstSum0S v -> second Ast.AstSum0S (unshareAst memo v)-  Ast.AstDot0S u v ->+  Ast.AstDot0 u v ->     let (memo2, u2) = unshareAst memo u         (memo3, v3) = unshareAst memo2 v-    in (memo3, Ast.AstDot0S u2 v3)+    in (memo3, Ast.AstDot0 u2 v3)   Ast.AstDot1InS m n u v ->     let (memo2, u2) = unshareAst memo u         (memo3, v3) = unshareAst memo2 v@@ -492,31 +497,38 @@         (memo3, v3) = unshareAst memo2 v     in (memo3, Ast.AstMatmul2S m n p u2 v3) +  Ast.AstBoolNotK arg ->+    let (memo2, arg2) = unshareAst memo arg+    in (memo2, Ast.AstBoolNotK arg2)+  Ast.AstBoolNotS arg ->+    let (memo2, arg2) = unshareAst memo arg+    in (memo2, Ast.AstBoolNotS arg2)+  Ast.AstBoolAndK arg1 arg2 ->+    let (memo1, b1) = unshareAst memo arg1+        (memo2, b2) = unshareAst memo1 arg2+    in (memo2, Ast.AstBoolAndK b1 b2)+  Ast.AstBoolAndS arg1 arg2 ->+    let (memo1, b1) = unshareAst memo arg1+        (memo2, b2) = unshareAst memo1 arg2+    in (memo2, Ast.AstBoolAndS b1 b2)+  Ast.AstLeqK arg1 arg2 ->+    let (memo1, r1) = unshareAst memo arg1+        (memo2, r2) = unshareAst memo1 arg2+    in (memo2, Ast.AstLeqK r1 r2)+  Ast.AstLeq arg1 arg2 ->+    let (memo1, r1) = unshareAst memo arg1+        (memo2, r2) = unshareAst memo1 arg2+    in (memo2, Ast.AstLeq r1 r2)+  Ast.AstLeqS shb sh arg1 arg2 ->+    let (memo1, r1) = unshareAst memo arg1+        (memo2, r2) = unshareAst memo1 arg2+    in (memo2, Ast.AstLeqS shb sh r1 r2)+ unshareAstHFun-  :: AstBindings -> AstHFun s s2 x y -> (AstBindings, AstHFun s s2 x y)+  :: AstBindings -> AstHFun s x y -> (AstBindings, AstHFun s x y) unshareAstHFun memo v0 = case v0 of   Ast.AstLambda{} ->     -- No other free variables in l, so no outside lets can reach there,     -- so we don't need to pass the information from v upwards     -- nor remove the Share constructors.     (memo, v0)--unshareAstBool :: AstBindings -> AstBool AstMethodShare-               -> (AstBindings, AstBool AstMethodLet)-unshareAstBool memo = \case-  Ast.AstBoolConst t -> (memo, Ast.AstBoolConst t)-  Ast.AstBoolNot arg ->-    let (memo2, arg2) = unshareAstBool memo arg-    in (memo2, Ast.AstBoolNot arg2)-  Ast.AstBoolAnd arg1 arg2 ->-    let (memo1, b1) = unshareAstBool memo arg1-        (memo2, b2) = unshareAstBool memo1 arg2-    in (memo2, Ast.AstBoolAnd b1 b2)-  Ast.AstLeqK arg1 arg2 ->-    let (memo1, r1) = unshareAst memo arg1-        (memo2, r2) = unshareAst memo1 arg2-    in (memo2, Ast.AstLeqK r1 r2)-  Ast.AstLeqS arg1 arg2 ->-    let (memo1, r1) = unshareAst memo arg1-        (memo2, r2) = unshareAst memo1 arg2-    in (memo2, Ast.AstLeqS r1 r2)
src/HordeAd/Core/AstInterpret.hs view
@@ -1,4 +1,6 @@ {-# LANGUAGE CPP #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} -- | Interpretation of AST terms in an arbitrary tensor operations -- class instance. With the exception of the the interpretation -- of the sharing mechanisms and any other performance tweaks,@@ -6,25 +8,26 @@ -- The sharing mechanisms are translated so as to preserve sharing in case -- the instance is a term algebra as well. module HordeAd.Core.AstInterpret-  ( interpretAstFull, interpretAstPrimal, interpretAstDual, interpretAst-    -- * Exported only to specialize elsewhere-  , interpretAstBool+  ( interpretAstFull, interpretAstPrimal, interpretAstDual, interpretAstPlain   ) where  import Prelude -import Data.Coerce (coerce) import Data.Dependent.EnumMap.Strict qualified as DMap+import Data.List.NonEmpty (NonEmpty (..))+import Data.List.NonEmpty qualified as NonEmpty import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (testEquality, (:~:) (Refl))+import Data.Type.Equality ((:~:) (Refl)) import Data.Vector.Generic qualified as V-import Type.Reflection (typeRep) +import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Shaped.Shape  import HordeAd.Core.Ast import HordeAd.Core.AstEnv import HordeAd.Core.AstTools+import HordeAd.Core.Conversion import HordeAd.Core.ConvertTensor import HordeAd.Core.Ops import HordeAd.Core.TensorKind@@ -41,116 +44,53 @@ {-# INLINE interpretAstFull #-} interpretAstFull = interpretAst --- Strict environment and strict ADVal and Delta make this hard to optimize.--- Either the environment has to be traversed to remove the dual parts or--- the dual part needs to be potentially needlessly computed.--- However, with correct sharing and large tensors, the overall cost--- is negligible, so we optimize only minimally.--- It helps that usually the dual part is either trivially computed--- to be zero or is used elsewhere. It's rarely really lost and forgotten. interpretAstPrimal   :: forall target y. ADReady target   => AstEnv target -> AstTensor AstMethodLet PrimalSpan y   -> PrimalOf target y-interpretAstPrimal !env v1 = case v1 of-  -- This prevents computing the complex dual parts for mapAccum in ADVal.-  AstMapAccumRDer k bftk eftk f0 df0 rf0 acc0 es ->-    let f = interpretAstHFunPrimal env f0-        df = interpretAstHFunPrimal env df0-        rf = interpretAstHFunPrimal env rf0-        acc02 = interpretAstPrimal env acc0-        es2 = interpretAstPrimal env es-    in tmapAccumRDer (Proxy @(PrimalOf target))-                     k (ftkAst acc0) bftk eftk f df rf acc02 es2-  AstMapAccumLDer k bftk eftk f0 df0 rf0 acc0 es ->-    let f = interpretAstHFunPrimal env f0-        df = interpretAstHFunPrimal env df0-        rf = interpretAstHFunPrimal env rf0-        acc02 = interpretAstPrimal env acc0-        es2 = interpretAstPrimal env es-    in tmapAccumLDer (Proxy @(PrimalOf target))-                     k (ftkAst acc0) bftk eftk f df rf acc02 es2-  -- This prevents multiple ifH expansions in ADVal.-  AstCond b a1 a2 ->-    let c = interpretAstBool env b-    in tcond (ftkToSTK $ ftkAst a1) c-             (interpretAstPrimal env a1) (interpretAstPrimal env a2)-  _ -> tprimalPart (interpretAst env v1)+{-# INLINE interpretAstPrimal #-}+interpretAstPrimal = interpretAst  interpretAstDual   :: forall target y. ADReady target   => AstEnv target -> AstTensor AstMethodLet DualSpan y   -> DualOf target y {-# INLINE interpretAstDual #-}-interpretAstDual !env v1 =-  tdualPart (ftkToSTK $ ftkAst v1) (interpretAst env v1)+interpretAstDual env a = tdualPart (ftkToSTK (ftkAst a)) $ interpretAst env a --- A more precise type signature would result in @PrimalOf target@--- whenever @s@ is @PrimalSpan@, but this would complicate things,--- e.g., we'd need an extra type family------ type family SpanTarget s target :: Target where---   SpanTarget PrimalSpan target = PrimalOf target---   SpanTarget DualSpan target = DualOf target---   SpanTarget FullSpan target = target------ to be used in AstEnv and the codomain of interpretAst and a lot of other--- code would need to be changed. So instead we promote results to @target@--- similarly as in AstEnv and simiarly as we omit @PrimalOf@ in the signatures--- of most "Ops" methods.---+interpretAstPlain+  :: forall target y. ADReady target+  => AstEnv target -> AstTensor AstMethodLet PlainSpan y+  -> PlainOf target y+{-# INLINE interpretAstPlain #-}+interpretAstPlain = interpretAst+ -- | Interpret a term in an environment.------ Note that for 'PrimalSpan' term, the results of this function--- land in @target y@ and not in @PrimalOf target y@.--- To make it sound nevertheless, we maintain an invariant that a value--- of interpretation of a term with 'PrimalSpan' has zero dual part--- and of a term with 'DualSpan' has zero primal part.--- The invariants holds by the properties of instances of @Ops@--- (see especially the ADVal instance, which zeroes dual part of many ops)--- and structural induction on Ast, inspecting spans of constructors.--- This promotion from @PrimalOf target y@ to @target y@ coincides--- with how most operations that in Ast have 'PrimalSpan',--- don't have 'PrimalOf' (but have full target instead)--- in their method signatures in @Ops@, for user convenience.--- See, e.g., 'AstConcreteS' vs 'tsconcrete' and 'AstFloorS' vs 'tsfloor'. interpretAst-  :: forall target s y. (ADReady target, AstSpan s)+  :: forall target s y. (ADReady target, KnownSpan s)   => AstEnv target -> AstTensor AstMethodLet s y-  -> target y-interpretAst !env = \case+  -> SpanTargetFam target s y+{-# INLINEABLE interpretAst #-}+interpretAst !env | Refl <- lemPlainOfSpan (Proxy @target) (knownSpan @s)+                  , Dict0 <- dictSpanFam (Proxy @target) (knownSpan @s) = \case   AstPair t1 t2 -> tpair (interpretAst env t1) (interpretAst env t2)   AstProject1 t -> tproject1 (interpretAst env t)   AstProject2 t -> tproject2 (interpretAst env t)-  AstFromVector snat stk l ->-    let l2 = V.map (interpretAst env) l-    in tfromVector snat stk l2-  AstSum snat stk v -> tsum snat stk $ interpretAst env v-  AstReplicate snat stk v ->-    treplicate snat stk (interpretAst env v)-  AstMapAccumRDer k bftk eftk f0 df0 rf0 acc0 es ->-    let f = interpretAstHFun env f0-        df = interpretAstHFun env df0-        rf = interpretAstHFun env rf0-        acc02 = interpretAst env acc0-        es2 = interpretAst env es-    in tmapAccumRDer (Proxy @target) k (ftkAst acc0) bftk eftk f df rf acc02 es2   AstMapAccumLDer k bftk eftk f0 df0 rf0 acc0 es ->     let f = interpretAstHFun env f0         df = interpretAstHFun env df0         rf = interpretAstHFun env rf0         acc02 = interpretAst env acc0         es2 = interpretAst env es-    in tmapAccumLDer (Proxy @target) k (ftkAst acc0) bftk eftk f df rf acc02 es2-  AstApply t ll ->-    let t2 = interpretAstHFun env t-        ll2 = interpretAst env ll-    in tApply t2 ll2+    in tmapAccumLDer (Proxy @(SpanTargetFam target s))+                     k (ftkAst acc0) bftk eftk f df rf acc02 es2+  AstApply f t ->+    let f2 = interpretAstHFun env f+        t2 = interpretAst env t+    in tapply f2 t2   AstVar var ->-    let var2 :: AstVarName FullSpan y-        var2 = coerce var  -- only FullSpan variables permitted in env-    in case DMap.lookup var2 env of-      Just t ->+    case DMap.lookup var env of+      Just (SpanTarget t) -> #ifdef WITH_EXPENSIVE_ASSERTIONS         withKnownSTK (ftkToSTK $ varNameToFTK var) $         -- We can't assert anything about bounds, because values can be@@ -162,57 +102,60 @@         t       _ -> error $ "interpretAst: unknown AstVar " ++ show var                    -- ++ " in environment " ++ showsPrecAstEnv 0 env ""-  AstCond b a1 a2 ->-    let c = interpretAstBool env b-    in tcond (ftkToSTK (ftkAst a1)) c-             (interpretAst env a1) (interpretAst env a2)+  -- TODO: recognize nested builds and, for Concrete, call tbuild instead;+  -- also recognize map and zipWith in nested builds and call these   AstBuild1 snat stk (var, v) ->     let f i = interpretAst (extendEnvI var i env) v     in tbuild1 snat stk f    -- We assume there are no nested lets with the same variable.-  ---  -- Note that without the second sameAstSpan check, AstLet with both PrimalSpan-  -- would get translated to a composition of ttletPrimal and tfromPrimal,-  -- which doesn't make a difference in a translation from PrimalSpan-  -- terms to PrimalSpan terms, but does in a translation from PrimalSpan-  -- terms to FullSpan terms, causing a loss of a dual part.-  ---  -- However, right now this whole code fragment is disabled, because-  -- it increases the allocation in testsuites by ~3% and slows down the VTO1-  -- benchmark 5 times. To be re-evaluated when rewriting is changed-  -- and also more examples are available.-  AstLet {-@_ @_ @s1-} var u v -> {- case ( sameAstSpan @s1 @PrimalSpan-                                          , sameAstSpan @s @FullSpan ) of-    (Just Refl, Just Refl) ->-      let t = interpretAstPrimal env u-          stk = ftkToSTK (ftkAst u)-          env2 wPrimal = extendEnv var (tfromPrimal stk wPrimal) env-      in ttletPrimal t (\wPrimal -> interpretAst (env2 wPrimal) v)-        -- @ttletPrimal@ can be more frugal in some targets, though we pay-        -- for it with @ftkAst@-    _ -> -}-      let t = interpretAst env u-          env2 w = extendEnv var w env-      in ttlet t (\w -> interpretAst (env2 w) v)--  AstPrimalPart a ->-    tfromPrimal (ftkToSTK (ftkAst a)) (tprimalPart $ interpretAstFull env a)+  AstLet @_ @_ @_ @s2 var u v ->+    let (toFull, fromFull) =+          toFromFullSpan @target (ftkToSTK (ftkAst v)) (knownSpan @s2)+        t = withKnownSpan (varNameToSpan var) $ interpretAst env u+        env2 w = extendEnv var w env+    in case varNameToSpan var of+         SFullSpan ->+           fromFull $ ttlet t (\w -> toFull $ interpretAst (env2 w) v)+         SPrimalStepSpan SFullSpan ->+           fromFull $ ttletPrimal t (\w -> toFull $ interpretAst (env2 w) v)+         SPrimalStepSpan _ ->+           error "interpretAst: can't store a nested primal value"+             -- actually, we could, by substituting into v0, etc.+         SDualSpan ->+           fromFull $ ttlet t (\w -> toFull $ interpretAst (env2 w) v)+             -- due to the dual hack+         SPlainSpan | SPlainSpan <- knownSpan @s2 ->  -- a speedup+           ttlet t (\w -> interpretAst (env2 w) v)+         SPlainSpan ->+           fromFull $ ttletPlain t (\w -> toFull $ interpretAst (env2 w) v)+  AstPrimalPart @_ @s2 a -> case knownSpan @s2 of+    SFullSpan -> tprimalPart $ interpretAst env a+    SPrimalStepSpan SFullSpan -> tprimalPart $ interpretAst env a+    SPrimalStepSpan _ ->+      error "interpretAst: can't convert a nested primal value"  -- (... easily)+    SDualSpan -> tdefTarget (ftkAst a)  -- primal zero+    SPlainSpan -> tprimalPart $ interpretAst env a   AstDualPart a ->-    tfromDual (tdualPart (ftkToSTK (ftkAst a)) $ interpretAstFull env a)-  AstFromPrimal a ->-    -- By the invariant, interpretation of @a@ has zero dual part,-    -- so we don't have to do the following to remove the dual part,-    -- but we still do, because there's almost no rewriting of delta-    -- expressions, so even though they are semantically zero, they'd build-    -- up considerably if not wiped out regularly. By constrast, operations-    -- on AstConstant are rewritten eagerly to AstConstant, so for AstFromDual-    -- we really don't need to do anything.-    tfromPrimal (ftkToSTK (ftkAst a)) (interpretAstPrimal env a)-  AstFromDual a -> interpretAst env a-    -- By the invariant, interpretation of @a@ has zero primal part,-    -- so we don't have to do the following to remove the primal part:-    --   tfromDual (interpretAstDual env a)+    -- We zero the primal part, but keep it a dual number, not its second+    -- component, that is a Delta expression (in non-symbolic instances).+    tfromDual $ tdualPart (ftkToSTK (ftkAst a)) $ interpretAst env a+  AstPlainPart @_ @s2 a -> case knownSpan @s2 of+    SFullSpan -> tplainPart $ interpretAst env a+    SPrimalStepSpan SFullSpan -> tplainPart $ interpretAst env a+    SPrimalStepSpan _ ->+      error "interpretAst: can't convert a nested primal value"  -- (... easily)+    SDualSpan -> tdefTarget (ftkAst a)  -- plain zero+    SPlainSpan -> interpretAst env a+  AstFromPrimal a -> tfromPrimal (ftkToSTK (ftkAst a)) $ interpretAst env a+  AstFromDual a ->+    -- Not @tfromDual $ interpretAst env a@, because dual parts are represented+    -- as dual numbers with zero primal parts, so nothing needs to be done here,+    -- because the part is already zeroed by inductive assumption, so this+    -- also works as a representation of a dual number with zero primal part,+    -- which is the semantics of `AstFromDual`.+    interpretAst env a+  AstFromPlain a -> tfromPlain (ftkToSTK (ftkAst a)) $ interpretAst env a    AstPlusK u v -> interpretAst env u + interpretAst env v   AstTimesK u v -> interpretAst env u * interpretAst env v@@ -230,18 +173,16 @@     let u2 = interpretAst env u         v2 = interpretAst env v     in interpretAstI2 opCode u2 v2-  AstConcreteK k ->-    tkconcrete @target k-      -- this is equal to the following-      -- (and similarly for tsconcretet and tsiota below):-      -- tfromPrimal @target STKScalar $ tkconcrete @(PrimalOf target) k-  AstFloorK v ->-    -- By the invariant v has zero dual part, so the following suffices:-    tkfloor $ interpretAst env v-  AstFromIntegralK v ->-    -- By the invariant v has zero dual part, so the following suffices:-    tkfromIntegral $ interpretAst env v+  AstConcreteK k -> tkconcrete k+  AstFloorK v -> tkfloor $ interpretAst env v+  AstFromIntegralK v -> tkfromIntegral $ interpretAst env v   AstCastK v -> tkcast $ interpretAst env v+  AstArgMinK v -> tkargMin $ interpretAst env v+  AstArgMaxK v -> tkargMax $ interpretAst env v+  AstIndexK v ix ->+    let v2 = interpretAst env v+        ix3 = interpretAst env <$> ix+    in tsindex0 v2 ix3    AstPlusS u v -> interpretAst env u + interpretAst env v   AstTimesS u v -> interpretAst env u * interpretAst env v@@ -252,106 +193,130 @@   AstI2S opCode u v ->     interpretAstI2 opCode (interpretAst env u) (interpretAst env v)   AstConcreteS a -> tsconcrete a-  AstFloorS v ->-    -- By the invariant v has zero dual part, so the following suffices:-    tsfloor $ interpretAst env v-  AstFromIntegralS v ->-    -- By the invariant v has zero dual part, so the following suffices:-    tsfromIntegral $ interpretAst env v-  AstCastS @r1 @r2 v ->-    -- Specializing for the cases covered by rules in GHC.Internal.Float.-    case testEquality (typeRep @r1) (typeRep @Double) of-      Just Refl -> case testEquality (typeRep @r2) (typeRep @Float) of-        Just Refl -> tscast @_ @Double @Float $ interpretAst env v-        _ -> tscast @_ @Double $ interpretAst env v-      _ -> case testEquality (typeRep @r1) (typeRep @Float) of-        Just Refl -> case testEquality (typeRep @r2) (typeRep @Double) of-          Just Refl -> tscast @_ @Float @Double $ interpretAst env v-          _ -> tscast @_ @Float $ interpretAst env v-        _ -> tscast $ interpretAst env v+  AstFloorS v -> tsfloor $ interpretAst env v+  AstFromIntegralS v -> tsfromIntegral $ interpretAst env v+  AstCastS v -> tscast $ interpretAst env v+  AstArgMinS v -> tsargMin $ interpretAst env v+  AstArgMaxS v -> tsargMax $ interpretAst env v+  AstIndexS @shm shn v ix ->+    withKnownShS shn $+    withKnownSTK (stkAstX v) $+    let v2 = interpretAst env v+        ix3 = interpretAst env <$> ix+    in tsindex @_ @shm v2 ix3 -  AstIndexS @sh1 sh2 v ix -> case ftkToSTK (ftkAst v) of-    STKS _ x ->-      withKnownShS (shsFromIxS ix) $-      withKnownShS sh2 $-      withKnownSTK x $-      let v2 = interpretAst env v-          ix3 = interpretAstPrimal env <$> ix-      in tsindex @target @sh1 v2 ix3-  {- TODO: this breaks specialization:-  AstScatterS shn v (ZS, ix) -> case ftkToSTK (ftkAst v) of-    STKS _ x ->-      withKnownShS shn $-      withKnownShS (shsFromIxS ix) $-      withKnownSTK x $-      tsoneHot (interpretAst env v) (interpretAstPrimal env <$> ix) -}-  AstScatterS @shm @shn @shp-              shn v (vars, ix) -> case ftkToSTK (ftkAst v) of-    STKS _ x ->-      withKnownShS (shsFromListS vars) $+  AstCondK b a1 a2 ->+    let c = interpretAst env b+    in kcond c (interpretAst env a1) (interpretAst env a2)+  AstCondS b a1 a2 | FTKS sh x <- ftkAst a1 ->+    withKnownShS sh $+    withKnownSTK (ftkToSTK x) $+    let c = interpretAst env b+    in scond c (interpretAst env a1) (interpretAst env a2)+  AstFromVectorK shm l ->+    let l2 = V.map (interpretAst env) l+    in tsfromVectorLinear shm l2+  AstFromVectorS shm l -> case V.uncons l of+    Just (v, _) | FTKS shn x <- ftkAst v ->       withKnownShS shn $-      withKnownShS (shsFromIxS ix) $-      withKnownSTK x $-      let t1 = interpretAst env v-          f2 :: IxSOf target shm -> IxSOf target shp-          f2 !ix2 = interpretAstPrimal (extendEnvVarsS vars ix2 env) <$> ix-      in tsscatter @_ @shm @shn @shp t1 f2-  AstGatherS shn v (ZS, ix) -> case ftkToSTK (ftkAst v) of-    STKS _ x ->+      withKnownSTK (ftkToSTK x) $+      let l2 = V.map (interpretAst env) l+      in tsfromVectorN shm l2+    Nothing -> error "interpretAst: empty vector in AstFromVectorS"+  AstSumK v -> case ftkAst v of+    FTKS shm _ ->+      withKnownShS shm $+      tssum0 (interpretAst env v)+  t@(AstSumS @shm shm v) -> case ftkAst t of+    FTKS shn x ->+      withKnownShS shm $       withKnownShS shn $-      withKnownShS (shsFromIxS ix) $-      withKnownSTK x $-      tsindex (interpretAst env v) (interpretAstPrimal env <$> ix)-  AstGatherS @shm @shn @shp-             shn v (vars, ix) -> case ftkToSTK (ftkAst v) of-    STKS _ x ->-      withKnownShS (shsFromListS vars) $+      withKnownSTK (ftkToSTK x) $+      tssumN @_ @shm (interpretAst env v)+  -- TODO: this breaks specialization:+  AstScatterS _ shn shp v (ZS, ix) ->+    withKnownShS shn $+    withKnownShS shp $+    withKnownSTK (stkAstX v) $+    tsoneHot (interpretAst env v) (interpretAst env <$> ix)+  AstScatterS @_ @_ @shp (SNat :$$ _) shn shp v (var ::$ ZS, ix) ->+    withKnownShS shn $+    withKnownShS shp $+    withKnownSTK (stkAstX v) $+    let t1 = interpretAst env v+        f2 :: IntOf target -> IxSOf target shp+        f2 !i2 = interpretAst (extendEnvI var i2 env) <$> ix+    in tsscatter1 t1 f2+  AstScatterS @shm @shn @shp shm shn shp v (vars, ix) ->+    withKnownShS shm $+    withKnownShS shn $+    withKnownShS shp $+    withKnownSTK (stkAstX v) $+    let t1 = interpretAst env v+        f2 :: IxSOf target shm -> IxSOf target shp+        f2 !ix2 = interpretAst (extendEnvVarsS vars ix2 env) <$> ix+    in tsscatter @_ @_ @shn t1 f2+  AstReplicateK shm v ->+    tsreplicate0N shm (interpretAst env v)+  AstReplicateS shm v -> case ftkAst v of+    FTKS shn x ->       withKnownShS shn $-      withKnownShS (shsFromIxS ix) $-      withKnownSTK x $-      let t1 = interpretAst env v-          f2 :: IxSOf target shm -> IxSOf target shp-          f2 !ix2 = interpretAstPrimal (extendEnvVarsS vars ix2 env) <$> ix-      in tsgather @_ @shm @shn @shp t1 f2-  AstMinIndexS v ->-    -- By the invariant v has zero dual part, so the following suffices:-    tsminIndex $ interpretAst env v-  AstMaxIndexS v ->-    -- By the invariant v has zero dual part, so the following suffices:-    tsmaxIndex $ interpretAst env v+      withKnownSTK (ftkToSTK x) $+      tsreplicateN shm (interpretAst env v)+  AstGatherS _ shn _ v (ZS, ix) -> interpretAst env (AstIndexS shn v ix)+  AstGatherS @_ @_ @shp (SNat :$$ _) shn shp v (var ::$ ZS, ix) ->+    withKnownShS shn $+    withKnownShS shp $+    withKnownSTK (stkAstX v) $+    let t1 = interpretAst env v+        f2 :: IntOf target -> IxSOf target shp+        f2 !i2 = interpretAst (extendEnvI var i2 env) <$> ix+    in tsgather1 t1 f2+  AstGatherS @shm @shn @shp shm shn shp v (vars, ix) ->+    withKnownShS shm $+    withKnownShS shn $+    withKnownShS shp $+    withKnownSTK (stkAstX v) $+    let t1 = interpretAst env v+        f2 :: IxSOf target shm -> IxSOf target shp+        f2 !ix2 = interpretAst (extendEnvVarsS vars ix2 env) <$> ix+    in tsgather @_ @_ @shn t1 f2   AstIotaS SNat -> tsiota-  AstAppendS a b -> case ftkToSTK (ftkAst a) of-    STKS _ x ->-      withKnownSTK x $-      let t1 = interpretAst env a-          t2 = interpretAst env b-      in tsappend t1 t2-  AstSliceS i n k v -> case ftkToSTK (ftkAst v) of-    STKS _ x ->-      withKnownSTK x $-      tsslice i n k $ interpretAst env v-  AstReverseS v -> case ftkToSTK (ftkAst v) of-    STKS _ x ->-      withKnownSTK x $-      tsreverse (interpretAst env v)-  AstTransposeS perm v -> case ftkToSTK (ftkAst v) of-    STKS _ x ->-      withKnownSTK x $-      tstranspose perm $ interpretAst env v-  AstReshapeS sh2 v -> case ftkToSTK (ftkAst v) of-    STKS _ x ->-      withKnownSTK x $-      tsreshape sh2 (interpretAst env v)+  t@(AstAppendS @_ @_ @shRest @x a b) -> case ftkAst t of+    FTKS (SNat :$$ shRest) x ->+      withKnownShS shRest $+      withKnownSTK (ftkToSTK x) $+      let flattenAppend+            :: AstTensor AstMethodLet s (TKS2 (m ': shRest) x)+            -> NonEmpty (SpanTargetFam target s (TKR2 (Rank (m ': shRest)) x))+          flattenAppend (AstAppendS c d) = flattenAppend c <> flattenAppend d+          flattenAppend u | (FTKS (SNat :$$ _) _) <- ftkAst u =+            NonEmpty.singleton (rfromS $ interpretAst env u)+      in case flattenAppend t of+        -- At least 3 elements and a concrete instance:+        ne@(_ :| _ : _ : _) | isConcreteInstance @target ->+          sfromR $ trconcat ne+        _ ->+          let t1 = interpretAst env a+              t2 = interpretAst env b+          in tsappend t1 t2+  AstSliceS i n k v ->+    withKnownSTK (stkAstX v) $+    tsslice i n k $ interpretAst env v+  AstReverseS v ->+    withKnownSTK (stkAstX v) $+    tsreverse (interpretAst env v)+  AstTransposeS perm v ->+    withKnownSTK (stkAstX v) $+    tstranspose perm $ interpretAst env v+  AstReshapeS sh2 v ->+    withKnownSTK (stkAstX v) $+    tsreshape sh2 (interpretAst env v)    AstConvert c a ->     tconvert c (ftkToSTK (ftkAst a)) (interpretAst env a) -  AstSum0S v -> case ftkToSTK (ftkAst v) of-    STKS sh x ->-      withKnownShS sh $-      withKnownSTK x $-      tssum0 (interpretAst env v)-  AstDot0S u v -> case ftkAst u of+  AstDot0 u v -> case ftkAst u of     FTKS sh _ ->       withKnownShS sh $       tsdot0 (interpretAst env u) (interpretAst env v)@@ -361,50 +326,69 @@   AstMatmul2S SNat SNat SNat u v ->     tsmatmul2 (interpretAst env u) (interpretAst env v) +  AstBoolNotK arg ->+    tfromPlain STKScalar $ notB $ interpretAst env arg+  AstBoolNotS arg | FTKS sh _ <- ftkAst arg ->+    withKnownShS sh $+    tfromPlain (ftkToSTK $ ftkAst arg)+    $ tsmap0N notB $ interpretAst env arg+  AstBoolAndK arg1 arg2 ->+    let b1 = interpretAst env arg1+        b2 = interpretAst env arg2+    in tfromPlain STKScalar $ b1 &&* b2+  AstBoolAndS arg1 arg2 | FTKS sh _ <- ftkAst arg1 ->+    withKnownShS sh $+    let b1 = interpretAst env arg1+        b2 = interpretAst env arg2+    in tfromPlain (ftkToSTK $ ftkAst arg1)+       $ tszipWith0N (&&*) b1 b2+  AstLeqK arg1 arg2 ->+    let r1 = interpretAst env arg1+        r2 = interpretAst env arg2+    in tfromPlain STKScalar $ r1 <=. r2+  AstLeq arg1 arg2 ->+    let r1 = interpretAst env arg1+        r2 = interpretAst env arg2+    in tfromPlain STKScalar $ r1 <=. r2+  AstLeqS @shb @sh shb sh arg1 arg2 | Refl <- lemAppNil @shb ->+    withKnownShS shb $+    withKnownShS sh $+    let r1 = interpretAst env arg1+        r2 = interpretAst env arg2+    in tfromPlain (STKS shb STKScalar)+       $ sunNest+       $ szipWithNested+           (\a1 a2 ->+             let c = convCmp ConvXS+                             (convCmp (Conv0X (STKS ZSS STKScalar))+                                      (convCmp ConvXS (Conv0X STKScalar)))+             in tconvert c STKScalar $ sunNest a1 <=. sunNest a2)+           (snest @_ @_ @sh shb r1) (snest shb r2)+ interpretAstHFun-  :: forall target x y s s2. (AstSpan s2, BaseTensor target)-  => AstEnv target -> AstHFun s s2 x y-  -> HFunOf target x y+  :: forall target x y s. (KnownSpan s, BaseTensor (SpanTargetFam target s))+  => AstEnv target -> AstHFun s x y+  -> HFunOf (SpanTargetFam target s) x y {-# INLINE interpretAstHFun #-}-interpretAstHFun _env (AstLambda var l) =-  tlambda @target (varNameToFTK var)-  $ HFun $ \ws -> interpretAst (extendEnv var ws emptyEnv) l+interpretAstHFun _env (AstLambda var t) =+  tlambda @(SpanTargetFam target s) (varNameToFTK var)+  $ HFun $ \ (ws :: f x) ->+              toFullSpan @f (ftkToSTK (ftkAst t)) (knownSpan @s)+              $ interpretAst @f+                  (extendEnv var (fromFullSpan (knownSpan @s) ws) emptyEnv) t       -- Interpretation in empty environment makes sense here, because       -- there are no free variables except for the one declared. -interpretAstHFunPrimal-  :: forall target x y. ADReady target-  => AstEnv target -> AstHFun PrimalSpan PrimalSpan x y-  -> HFunOf (PrimalOf target) x y-{-# INLINE interpretAstHFunPrimal #-}-interpretAstHFunPrimal _env (AstLambda var l) =-  tlambda @(PrimalOf target) (varNameToFTK var)-  $ HFun $ \ws -> interpretAst (extendEnv var ws emptyEnv) l-      -- This is probably optimized as much as possible, because-      -- thanks to the invariant, we get zero dual part from this-      -- PrimalSpan term and so interpretAstPrimal and tfromPrimal-      -- is not needed (and would not be possible, because we lack-      -- FullShapeTK y). From the other end, due to (PrimalOf target),-      -- there won't be any dual part coming from an argument.--interpretAstBool :: ADReady target-                 => AstEnv target -> AstBool AstMethodLet-                 -> BoolOf target-interpretAstBool !env = \case-  AstBoolConst a -> if a then true else false-  AstBoolNot arg -> notB $ interpretAstBool env arg-  AstBoolAnd arg1 arg2 ->-    let b1 = interpretAstBool env arg1-        b2 = interpretAstBool env arg2-    in b1 &&* b2-  AstLeqK arg1 arg2 ->-    let r1 = interpretAstPrimal env arg1-        r2 = interpretAstPrimal env arg2-    in r1 <=. r2-  AstLeqS arg1 arg2 ->-    let r1 = interpretAstPrimal env arg1-        r2 = interpretAstPrimal env arg2-    in r1 <=. r2+-- This version accepts nested arrays, because they are needed here.+szipWithNested :: ( KnownShS sh, KnownSTK x, KnownSTK x1, KnownSTK x2+                  , BaseTensor target )+               => (target (TKS2 '[] x1) -> target (TKS2 '[] x2)+                   -> target (TKS2 '[] x))+               -> target (TKS2 sh x1) -> target (TKS2 sh x2)+               -> target (TKS2 sh x)+{-# INLINE szipWithNested #-}+szipWithNested @sh f u v | Refl <- lemAppNil @sh =+  tsbuild @_ @sh (\ix -> f (tsindex u ix) (tsindex v ix))   -- * Interpretation of arithmetic, boolean and relation operations
+ src/HordeAd/Core/AstMethodLet.hs view
@@ -0,0 +1,292 @@+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Arithmetic instances for AST with sharing method AstMethodLet.+-- The bulk of the instances is defined in "AstSimplify", the remaining+-- instances are here.+module HordeAd.Core.AstMethodLet+  (+  ) where++import Prelude++import Data.Type.Equality (testEquality, (:~:) (Refl))++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)+import Data.Array.Nested.Shaped.Shape++import HordeAd.Core.Ast+import HordeAd.Core.AstSimplify+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersAst+import HordeAd.Core.OpsConcrete ()+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++liftRFromS1 :: forall n x s. KnownSpan s+            => (forall sh.+                   AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKS2 sh x))+            -> AstTensor AstMethodLet s (TKR2 n x)+            -> AstTensor AstMethodLet s (TKR2 n x)+{-# INLINE liftRFromS1 #-}+liftRFromS1 f a = case ftkAst a of+  FTKR sh' x ->+    withShsFromShR sh' $ \(sh :: ShS sh) ->+      astConvUpRFromS sh x $ f (astConvDownSFromR sh x a)++liftRFromS2 :: forall n x s. KnownSpan s+            => (forall sh.+                   AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKS2 sh x))+            -> AstTensor AstMethodLet s (TKR2 n x)+            -> AstTensor AstMethodLet s (TKR2 n x)+            -> AstTensor AstMethodLet s (TKR2 n x)+{-# INLINE liftRFromS2 #-}+liftRFromS2 f a b  = case ftkAst a of+  FTKR sh' x ->+    withShsFromShR sh' $ \(sh :: ShS sh) ->+      astConvUpRFromS sh x+      $ f (astConvDownSFromR sh x a) (astConvDownSFromR sh x b)++liftXFromS1 :: forall sh' x s. KnownSpan s+            => (forall sh.+                   AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKS2 sh x))+            -> AstTensor AstMethodLet s (TKX2 sh' x)+            -> AstTensor AstMethodLet s (TKX2 sh' x)+{-# INLINE liftXFromS1 #-}+liftXFromS1 f a = case ftkAst a of+  FTKX sh' x ->+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      astConvUpXFromS sh' x $ f (astConvDownSFromX sh x a)++liftXFromS2 :: forall sh' x s. KnownSpan s+            => (forall sh.+                   AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKS2 sh x))+            -> AstTensor AstMethodLet s (TKX2 sh' x)+            -> AstTensor AstMethodLet s (TKX2 sh' x)+            -> AstTensor AstMethodLet s (TKX2 sh' x)+{-# INLINE liftXFromS2 #-}+liftXFromS2 f a b = case ftkAst a of+  FTKX sh' x ->+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      astConvUpXFromS sh' x+      $ f (astConvDownSFromX sh x a) (astConvDownSFromX sh x b)+++-- * Unlawful numeric instances for ranked AST; lawful modulo evaluation++instance (NumScalar r, KnownSpan s)+         => Num (AstTensor AstMethodLet s (TKR n r)) where+  (+) = liftRFromS2 (+)+  (-) = liftRFromS2 (-)+  (*) = liftRFromS2 (*)+  negate = liftRFromS1 negate+  abs = liftRFromS1 abs+  signum = liftRFromS1 signum+  fromInteger i = error $ "fromInteger is not defined for ranked tensors: "+                          ++ show i++instance (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)+         => IntegralH (AstTensor AstMethodLet s (TKR n r)) where+  quotH = liftRFromS2 quotH+  remH = liftRFromS2 remH++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Fractional (AstTensor AstMethodLet s (TKR n r)) where+  (/) = liftRFromS2 (/)+  recip = liftRFromS1 recip+  fromRational r = error $ "fromRational is not defined for ranked tensors: "+                           ++ show r++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Floating (AstTensor AstMethodLet s (TKR n r)) where+  pi = error "pi is not defined for tensors"+  exp = liftRFromS1 exp+  log = liftRFromS1 log+  sqrt = liftRFromS1 sqrt+  (**) = liftRFromS2 (**)+  logBase = liftRFromS2 logBase+  sin = liftRFromS1 sin+  cos = liftRFromS1 cos+  tan = liftRFromS1 tan+  asin = liftRFromS1 asin+  acos = liftRFromS1 acos+  atan = liftRFromS1 atan+  sinh = liftRFromS1 sinh+  cosh = liftRFromS1 cosh+  tanh = liftRFromS1 tanh+  asinh = liftRFromS1 asinh+  acosh = liftRFromS1 acosh+  atanh = liftRFromS1 atanh++instance (NumScalar r, Differentiable r, KnownSpan s)+         => RealFloatH (AstTensor AstMethodLet s (TKR n r)) where+  atan2H = liftRFromS2 atan2H++-- TODO: refactor with something like liftRFromS2+instance (KnownSpan s, NumScalar r)+         => EqH (AstTensor AstMethodLet s) (TKR n r) where+  v ==. u = case ftkAst v of+    FTKR shv' x -> case ftkAst u of+      FTKR shu' _ ->+        withShsFromShR shv' $ \shv ->+          withShsFromShR shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                astConvDownSFromR shu x v ==. astConvDownSFromR shv x u+              _ -> error $ "(==.): shapes don't match: "+                           ++ show (shu, shv)++instance (KnownSpan s, NumScalar r)+         => OrdH (AstTensor AstMethodLet s) (TKR n r) where+  v <=. u = case ftkAst v of+    FTKR shv' x -> case ftkAst u of+      FTKR shu' _ ->+        withShsFromShR shv' $ \shv ->+          withShsFromShR shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                astConvDownSFromR shu x v <=. astConvDownSFromR shv x u+              _ -> error $ "(<=.): shapes don't match: "+                           ++ show (shu, shv)+++-- * Unlawful numeric instances for mixed AST; lawful modulo evaluation++instance (NumScalar r, KnownSpan s)+         => Num (AstTensor AstMethodLet s (TKX sh r)) where+  (+) = liftXFromS2 (+)+  (-) = liftXFromS2 (-)+  (*) = liftXFromS2 (*)+  negate = liftXFromS1 negate+  abs = liftXFromS1 abs+  signum = liftXFromS1 signum+  fromInteger i = error $ "fromInteger is not defined for mixed tensors: "+                          ++ show i++instance (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)+         => IntegralH (AstTensor AstMethodLet s (TKX sh r)) where+  quotH = liftXFromS2 quotH+  remH = liftXFromS2 remH++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Fractional (AstTensor AstMethodLet s (TKX sh r)) where+  (/) = liftXFromS2 (/)+  recip = liftXFromS1 recip+  fromRational r = error $ "fromRational is not defined for mixed tensors: "+                           ++ show r++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Floating (AstTensor AstMethodLet s (TKX sh r)) where+  pi = error "pi is not defined for tensors"+  exp = liftXFromS1 exp+  log = liftXFromS1 log+  sqrt = liftXFromS1 sqrt+  (**) = liftXFromS2 (**)+  logBase = liftXFromS2 logBase+  sin = liftXFromS1 sin+  cos = liftXFromS1 cos+  tan = liftXFromS1 tan+  asin = liftXFromS1 asin+  acos = liftXFromS1 acos+  atan = liftXFromS1 atan+  sinh = liftXFromS1 sinh+  cosh = liftXFromS1 cosh+  tanh = liftXFromS1 tanh+  asinh = liftXFromS1 asinh+  acosh = liftXFromS1 acosh+  atanh = liftXFromS1 atanh++instance (NumScalar r, Differentiable r, KnownSpan s)+         => RealFloatH (AstTensor AstMethodLet s (TKX sh r)) where+  atan2H = liftXFromS2 atan2H++instance (KnownSpan s, NumScalar r)+         => EqH (AstTensor AstMethodLet s) (TKX sh r) where+  v ==. u = case ftkAst v of+    FTKX shv' x -> case ftkAst u of+      FTKX shu' _ ->+        withShsFromShX shv' $ \shv ->+          withShsFromShX shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                astConvDownSFromX shu x v ==. astConvDownSFromX shv x u+              _ -> error $ "(==.): shapes don't match: "+                           ++ show (shu, shv)++instance (KnownSpan s, NumScalar r)+         => OrdH (AstTensor AstMethodLet s) (TKX sh r) where+  v <=. u = case ftkAst v of+    FTKX shv' x -> case ftkAst u of+      FTKX shu' _ ->+        withShsFromShX shv' $ \shv ->+          withShsFromShX shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                astConvDownSFromX shu x v <=. astConvDownSFromX shv x u+              _ -> error $ "(<=.): shapes don't match: "+                           ++ show (shu, shv)+++--- * AstNoVectorize and AstNoSimplify instances++instance Boolean (AstNoVectorize PlainSpan (TKScalar Bool)) where+  true = AstNoVectorize true+  false = AstNoVectorize false+  notB b = AstNoVectorize (notB $ unAstNoVectorize b)+  b &&* c = AstNoVectorize (unAstNoVectorize b &&* unAstNoVectorize c)+  b ||* c = AstNoVectorize (unAstNoVectorize b ||* unAstNoVectorize c)++instance (EqH (AstTensor AstMethodLet s) y)+         => EqH (AstNoVectorize s) y where+  AstNoVectorize v ==. AstNoVectorize u = AstNoVectorize $ v ==. u+instance (OrdH (AstTensor AstMethodLet s) y)+         => OrdH (AstNoVectorize s) y where+  AstNoVectorize v <=. AstNoVectorize u = AstNoVectorize $ v <=. u++deriving instance Eq (AstNoVectorize s y)+deriving instance Ord (AstNoVectorize s y)+deriving instance Num (AstTensor AstMethodLet s y) => Num (AstNoVectorize s y)+deriving instance (IntegralH (AstTensor AstMethodLet s y))+                  => IntegralH (AstNoVectorize s y)+deriving instance Fractional (AstTensor AstMethodLet s y)+                  => Fractional (AstNoVectorize s y)+deriving instance Floating (AstTensor AstMethodLet s y)+                  => Floating (AstNoVectorize s y)+deriving instance (RealFloatH (AstTensor AstMethodLet s y))+                  => RealFloatH (AstNoVectorize s y)++instance Boolean (AstNoSimplify PlainSpan (TKScalar Bool)) where+  true = AstNoSimplify true+  false = AstNoSimplify false+  notB b = AstNoSimplify (notB $ unAstNoSimplify b)+  b &&* c = AstNoSimplify (unAstNoSimplify b &&* unAstNoSimplify c)+  b ||* c = AstNoSimplify (unAstNoSimplify b ||* unAstNoSimplify c)++instance (EqH (AstTensor AstMethodLet s) y)+         => EqH (AstNoSimplify s) y where+  AstNoSimplify v ==. AstNoSimplify u = AstNoSimplify $ v ==. u+instance (OrdH (AstTensor AstMethodLet s) y)+         => OrdH (AstNoSimplify s) y where+  AstNoSimplify v <=. AstNoSimplify u = AstNoSimplify $ v <=. u++deriving instance Eq (AstNoSimplify s y)+deriving instance Ord (AstNoSimplify s y)+deriving instance Num (AstTensor AstMethodLet s y) => Num (AstNoSimplify s y)+deriving instance (IntegralH (AstTensor AstMethodLet s y))+                  => IntegralH (AstNoSimplify s y)+deriving instance Fractional (AstTensor AstMethodLet s y)+                  => Fractional (AstNoSimplify s y)+deriving instance Floating (AstTensor AstMethodLet s y)+                  => Floating (AstNoSimplify s y)+deriving instance (RealFloatH (AstTensor AstMethodLet s y))+                  => RealFloatH (AstNoSimplify s y)
+ src/HordeAd/Core/AstMethodShare.hs view
@@ -0,0 +1,497 @@+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Arithmetic instances for AST with sharing method AstMethodShare.+module HordeAd.Core.AstMethodShare+  ( cAstConvDownKFromS, cAstConvDownSFromR, cAstConvDownSFromX+  , cAstConvUpSFromK, cAstConvUpRFromS, cAstConvUpXFromS+  , astShareNoSimplify+  ) where++import Prelude++import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import System.IO.Unsafe (unsafePerformIO)++import Data.Array.Nested (MapJust)+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.Ast+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersAst+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Conversion+import HordeAd.Core.OpsConcrete ()+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++cAstConvDownKFromS :: forall r s.+                      AstTensor AstMethodShare s (TKS '[] r)+                   -> AstTensor AstMethodShare s (TKScalar r)+cAstConvDownKFromS (AstConvert (ConvCmp ConvXS (Conv0X STKScalar)) a) = a+cAstConvDownKFromS a = AstConvert (ConvCmp ConvX0 ConvSX) a++cAstConvDownSFromR :: forall sh x s.+                      ShS sh -> FullShapeTK x+                   -> AstTensor AstMethodShare s (TKR2 (Rank sh) x)+                   -> AstTensor AstMethodShare s (TKS2 sh x)+cAstConvDownSFromR sh _ (AstConvert (ConvCmp (ConvXR _x) ConvSX) a)+  | let FTKS sh2 _ = ftkAst a+  , Just Refl <- testEquality sh sh2 = a+cAstConvDownSFromR sh x t | Refl <- lemRankReplicate (Proxy @(Rank sh)) =+  AstConvert (ConvCmp (ConvXS' (FTKS sh x)) ConvRX) t++cAstConvDownSFromX :: forall sh sh' x s. Rank sh ~ Rank sh'+                   => ShS sh -> FullShapeTK x+                   -> AstTensor AstMethodShare s (TKX2 sh' x)+                   -> AstTensor AstMethodShare s (TKS2 sh x)+cAstConvDownSFromX sh _ (AstConvert (ConvCmp (ConvXX' (FTKX _sh' _x)) ConvSX) a)+  | let FTKS sh2 _ = ftkAst a+  , Just Refl <- testEquality sh sh2 = a+cAstConvDownSFromX sh x t = AstConvert (ConvXS' (FTKS sh x)) t++cAstConvUpSFromK :: forall r s. GoodScalar r+                 => AstTensor AstMethodShare s (TKScalar r)+                 -> AstTensor AstMethodShare s (TKS '[] r)+cAstConvUpSFromK = AstConvert (ConvCmp ConvXS (Conv0X STKScalar))++cAstConvUpRFromS :: forall sh x s.+                    ShS sh -> FullShapeTK x+                 -> AstTensor AstMethodShare s (TKS2 sh x)+                 -> AstTensor AstMethodShare s (TKR2 (Rank sh) x)+cAstConvUpRFromS sh x | Refl <- lemRankMapJust sh =+  AstConvert (ConvCmp (ConvXR (ftkToSTK x)) ConvSX)++cAstConvUpXFromS :: forall sh sh' x s. Rank sh ~ Rank sh'+                 => IShX sh' -> FullShapeTK x+                 -> AstTensor AstMethodShare s (TKS2 sh x)+                 -> AstTensor AstMethodShare s (TKX2 sh' x)+cAstConvUpXFromS sh' x =+  gcastWith (unsafeCoerceRefl :: Rank (MapJust sh) :~: Rank sh) $+  AstConvert (ConvCmp (ConvXX' (FTKX sh' x)) ConvSX)++liftRFromS1 :: forall n x s.+               (forall sh.+                   AstTensor AstMethodShare s (TKS2 sh x)+                -> AstTensor AstMethodShare s (TKS2 sh x))+            -> AstTensor AstMethodShare s (TKR2 n x)+            -> AstTensor AstMethodShare s (TKR2 n x)+{-# INLINE liftRFromS1 #-}+liftRFromS1 f a = case ftkAst a of+  FTKR sh' x ->+    withShsFromShR sh' $ \(sh :: ShS sh) ->+      cAstConvUpRFromS sh x+      $ f (cAstConvDownSFromR sh x a)++liftRFromS2 :: forall n x s.+               (forall sh.+                   AstTensor AstMethodShare s (TKS2 sh x)+                -> AstTensor AstMethodShare s (TKS2 sh x)+                -> AstTensor AstMethodShare s (TKS2 sh x))+            -> AstTensor AstMethodShare s (TKR2 n x)+            -> AstTensor AstMethodShare s (TKR2 n x)+            -> AstTensor AstMethodShare s (TKR2 n x)+{-# INLINE liftRFromS2 #-}+liftRFromS2 f a b  = case ftkAst a of+  FTKR sh' x ->+    withShsFromShR sh' $ \(sh :: ShS sh) ->+      cAstConvUpRFromS sh x+      $ f (cAstConvDownSFromR sh x a) (cAstConvDownSFromR sh x b)++liftXFromS1 :: forall sh' x s.+               (forall sh.+                   AstTensor AstMethodShare s (TKS2 sh x)+                -> AstTensor AstMethodShare s (TKS2 sh x))+            -> AstTensor AstMethodShare s (TKX2 sh' x)+            -> AstTensor AstMethodShare s (TKX2 sh' x)+{-# INLINE liftXFromS1 #-}+liftXFromS1 f a = case ftkAst a of+  FTKX sh' x ->+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      cAstConvUpXFromS sh' x+      $ f (cAstConvDownSFromX sh x a)++liftXFromS2 :: forall sh' x s.+               (forall sh.+                   AstTensor AstMethodShare s (TKS2 sh x)+                -> AstTensor AstMethodShare s (TKS2 sh x)+                -> AstTensor AstMethodShare s (TKS2 sh x))+            -> AstTensor AstMethodShare s (TKX2 sh' x)+            -> AstTensor AstMethodShare s (TKX2 sh' x)+            -> AstTensor AstMethodShare s (TKX2 sh' x)+{-# INLINE liftXFromS2 #-}+liftXFromS2 f a b = case ftkAst a of+  FTKX sh' x ->+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      cAstConvUpXFromS sh' x+      $ f (cAstConvDownSFromX sh x a) (cAstConvDownSFromX sh x b)+++-- * Unlawful numeric instances for AST scalars; they are lawful modulo evaluation++-- The normal form has AstConcreteK or AstFromPlain (AstConcreteK),+-- if any, as the first argument of the constructor.+-- No flattening is performed beyond that.+instance (NumScalar r, KnownSpan s)+         => Num (AstTensor AstMethodShare s (TKScalar r)) where+  (+) = AstPlusK+  (*) = AstTimesK+  negate = AstN1K NegateOp+  abs = AstN1K AbsOp+  signum = AstN1K SignumOp+  {-# INLINE fromInteger #-}+  fromInteger i = fromPlain $ AstConcreteK (fromInteger i)++instance (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)+         => IntegralH (AstTensor AstMethodShare s (TKScalar r)) where+  quotH = AstI2K QuotOp+  remH = AstI2K RemOp++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Fractional (AstTensor AstMethodShare s (TKScalar r)) where+  (/) = AstR2K DivideOp+  recip = AstR1K RecipOp+  {-# INLINE fromRational #-}+  fromRational r = fromPlain $ AstConcreteK (fromRational r)++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Floating (AstTensor AstMethodShare s (TKScalar r)) where+  pi = error "pi is not defined for tensors"+  exp = AstR1K ExpOp+  log = AstR1K LogOp+  sqrt = AstR1K SqrtOp+  (**) = AstR2K PowerOp+  logBase = AstR2K LogBaseOp+  sin = AstR1K SinOp+  cos = AstR1K CosOp+  tan = AstR1K TanOp+  asin = AstR1K AsinOp+  acos = AstR1K AcosOp+  atan = AstR1K AtanOp+  sinh = AstR1K SinhOp+  cosh = AstR1K CoshOp+  tanh = AstR1K TanhOp+  asinh = AstR1K AsinhOp+  acosh = AstR1K AcoshOp+  atanh = AstR1K AtanhOp++instance (NumScalar r, Differentiable r, KnownSpan s)+         => RealFloatH (AstTensor AstMethodShare s (TKScalar r)) where+  atan2H = AstR2K Atan2Op+++-- * Unlawful numeric instances for shaped AST; lawful modulo evaluation++instance NumScalar r+         => Num (AstTensor AstMethodShare s (TKS sh r)) where+  (+) = AstPlusS+  (*) = AstTimesS+  negate = AstN1S NegateOp+  abs = AstN1S AbsOp+  signum = AstN1S SignumOp+  fromInteger i = error $ "fromInteger is not defined for shaped tensors: "+                          ++ show i++instance (NumScalar r, IntegralH r, Nested.IntElt r)+         => IntegralH (AstTensor AstMethodShare s (TKS sh r)) where+  quotH = AstI2S QuotOp+  remH = AstI2S RemOp++instance (NumScalar r, Differentiable r)+         => Fractional (AstTensor AstMethodShare s (TKS sh r)) where+  (/) = AstR2S DivideOp+  recip = AstR1S RecipOp+  fromRational r = error $ "fromRational is not defined for shaped tensors: "+                           ++ show r++instance (NumScalar r, Differentiable r)+         => Floating (AstTensor AstMethodShare s (TKS sh r)) where+  pi = error "pi is not defined for tensors"+  exp = AstR1S ExpOp+  log = AstR1S LogOp+  sqrt = AstR1S SqrtOp+  (**) = AstR2S PowerOp+  logBase = AstR2S LogBaseOp+  sin = AstR1S SinOp+  cos = AstR1S CosOp+  tan = AstR1S TanOp+  asin = AstR1S AsinOp+  acos = AstR1S AcosOp+  atan = AstR1S AtanOp+  sinh = AstR1S SinhOp+  cosh = AstR1S CoshOp+  tanh = AstR1S TanhOp+  asinh = AstR1S AsinhOp+  acosh = AstR1S AcoshOp+  atanh = AstR1S AtanhOp++instance (NumScalar r, Differentiable r)+         => RealFloatH (AstTensor AstMethodShare s (TKS sh r)) where+  atan2H = AstR2S Atan2Op+++-- * Unlawful instances of AST for bool; they are lawful modulo evaluation++instance Boolean (AstBool AstMethodShare) where+  true = AstConcreteK True+  false = AstConcreteK False+  notB = AstBoolNotK+  (&&*) = AstBoolAndK+  b ||* c = notB (notB b &&* notB c)++-- Since u and v are duplicated here, they need to be shared.+-- We share their difference, which would most likely appear in the+-- inequalities once they are rewritten, to ensure it's shared and whatever+-- vectorization substitutes into it is shared as well.+-- Otherwise, if u and v are variables, the sharing would vanish+-- before vectoriation complicates the expression a bit, making it+-- worth sharing.+instance (KnownSpan s, NumScalar r)+         => EqH (AstTensor AstMethodShare s) (TKScalar r) where+  vUnshared ==. uUnshared =+    let uv = astShareNoSimplify (uUnshared - vUnshared)+    in 0 <=. uv &&* uv <=. 0++instance (KnownSpan s, NumScalar r)+         => EqH (AstTensor AstMethodShare s) (TKS sh r) where+  vUnshared ==. uUnshared =+    let uv = astShareNoSimplify (uUnshared - vUnshared)+        zero = fromPlain $ AstConcreteS $ defTargetRep $ ftkAst vUnshared+    in zero <=. uv &&* uv <=. zero++instance (KnownSpan s, NumScalar r)+         => OrdH (AstTensor AstMethodShare s) (TKScalar r) where+  v <=. u = AstLeqK (plainPart v) (plainPart u)++instance (KnownSpan s, NumScalar r)+         => OrdH (AstTensor AstMethodShare s) (TKS sh r) where+  v <=. u = AstLeq (plainPart v) (plainPart u)+++-- * Unlawful numeric instances for ranked AST; lawful modulo evaluation++instance NumScalar r+         => Num (AstTensor AstMethodShare s (TKR n r)) where+  (+) = liftRFromS2 (+)+  (-) = liftRFromS2 (-)+  (*) = liftRFromS2 (*)+  negate = liftRFromS1 negate+  abs = liftRFromS1 abs+  signum = liftRFromS1 signum+  fromInteger i = error $ "fromInteger is not defined for ranked tensors: "+                          ++ show i++instance (NumScalar r, IntegralH r, Nested.IntElt r)+         => IntegralH (AstTensor AstMethodShare s (TKR n r)) where+  quotH = liftRFromS2 quotH+  remH = liftRFromS2 remH++instance (NumScalar r, Differentiable r)+         => Fractional (AstTensor AstMethodShare s (TKR n r)) where+  (/) = liftRFromS2 (/)+  recip = liftRFromS1 recip+  fromRational r = error $ "fromRational is not defined for ranked tensors: "+                           ++ show r++instance (NumScalar r, Differentiable r)+         => Floating (AstTensor AstMethodShare s (TKR n r)) where+  pi = error "pi is not defined for tensors"+  exp = liftRFromS1 exp+  log = liftRFromS1 log+  sqrt = liftRFromS1 sqrt+  (**) = liftRFromS2 (**)+  logBase = liftRFromS2 logBase+  sin = liftRFromS1 sin+  cos = liftRFromS1 cos+  tan = liftRFromS1 tan+  asin = liftRFromS1 asin+  acos = liftRFromS1 acos+  atan = liftRFromS1 atan+  sinh = liftRFromS1 sinh+  cosh = liftRFromS1 cosh+  tanh = liftRFromS1 tanh+  asinh = liftRFromS1 asinh+  acosh = liftRFromS1 acosh+  atanh = liftRFromS1 atanh++instance (NumScalar r, Differentiable r)+         => RealFloatH (AstTensor AstMethodShare s (TKR n r)) where+  atan2H = liftRFromS2 atan2H++instance (KnownSpan s, NumScalar r)+         => EqH (AstTensor AstMethodShare s) (TKR n r) where+  v ==. u = case ftkAst v of+    FTKR shv' x -> case ftkAst u of+      FTKR shu' _ ->+        withShsFromShR shv' $ \shv ->+          withShsFromShR shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                cAstConvDownSFromR shu x v ==. cAstConvDownSFromR shv x u+              _ -> error $ "(==.): shapes don't match: "+                           ++ show (shu, shv)++instance (KnownSpan s, NumScalar r)+         => OrdH (AstTensor AstMethodShare s) (TKR n r) where+  v <=. u = case ftkAst v of+    FTKR shv' x -> case ftkAst u of+      FTKR shu' _ ->+        withShsFromShR shv' $ \shv ->+          withShsFromShR shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                cAstConvDownSFromR shu x v <=. cAstConvDownSFromR shv x u+              _ -> error $ "(<=.): shapes don't match: "+                           ++ show (shu, shv)+++-- * Unlawful numeric instances for mixed AST; lawful modulo evaluation++instance NumScalar r+         => Num (AstTensor AstMethodShare s (TKX sh r)) where+  (+) = liftXFromS2 (+)+  (-) = liftXFromS2 (-)+  (*) = liftXFromS2 (*)+  negate = liftXFromS1 negate+  abs = liftXFromS1 abs+  signum = liftXFromS1 signum+  fromInteger i = error $ "fromInteger is not defined for mixed tensors: "+                          ++ show i++instance (NumScalar r, IntegralH r, Nested.IntElt r)+         => IntegralH (AstTensor AstMethodShare s (TKX sh r)) where+  quotH = liftXFromS2 quotH+  remH = liftXFromS2 remH++instance (NumScalar r, Differentiable r)+         => Fractional (AstTensor AstMethodShare s (TKX sh r)) where+  (/) = liftXFromS2 (/)+  recip = liftXFromS1 recip+  fromRational r = error $ "fromRational is not defined for mixed tensors: "+                           ++ show r++instance (NumScalar r, Differentiable r)+         => Floating (AstTensor AstMethodShare s (TKX sh r)) where+  pi = error "pi is not defined for tensors"+  exp = liftXFromS1 exp+  log = liftXFromS1 log+  sqrt = liftXFromS1 sqrt+  (**) = liftXFromS2 (**)+  logBase = liftXFromS2 logBase+  sin = liftXFromS1 sin+  cos = liftXFromS1 cos+  tan = liftXFromS1 tan+  asin = liftXFromS1 asin+  acos = liftXFromS1 acos+  atan = liftXFromS1 atan+  sinh = liftXFromS1 sinh+  cosh = liftXFromS1 cosh+  tanh = liftXFromS1 tanh+  asinh = liftXFromS1 asinh+  acosh = liftXFromS1 acosh+  atanh = liftXFromS1 atanh++instance (NumScalar r, Differentiable r)+         => RealFloatH (AstTensor AstMethodShare s (TKX sh r)) where+  atan2H = liftXFromS2 atan2H++instance (KnownSpan s, NumScalar r)+         => EqH (AstTensor AstMethodShare s) (TKX sh r) where+  v ==. u = case ftkAst v of+    FTKX shv' x -> case ftkAst u of+      FTKX shu' _ ->+        withShsFromShX shv' $ \shv ->+          withShsFromShX shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                cAstConvDownSFromX shu x v ==. cAstConvDownSFromX shv x u+              _ -> error $ "(==.): shapes don't match: "+                           ++ show (shu, shv)++instance (KnownSpan s, NumScalar r)+         => OrdH (AstTensor AstMethodShare s) (TKX sh r) where+  v <=. u = case ftkAst v of+    FTKX shv' x -> case ftkAst u of+      FTKX shu' _ ->+        withShsFromShX shv' $ \shv ->+          withShsFromShX shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                cAstConvDownSFromX shu x v <=. cAstConvDownSFromX shv x u+              _ -> error $ "(<=.): shapes don't match: "+                           ++ show (shu, shv)+++-- * AstRaw instances++instance Boolean (AstRaw PlainSpan (TKScalar Bool)) where+  true = AstRaw true+  false = AstRaw false+  notB b = AstRaw (notB $ unAstRaw b)+  b &&* c = AstRaw (unAstRaw b &&* unAstRaw c)+  b ||* c = AstRaw (unAstRaw b ||* unAstRaw c)++instance (EqH (AstTensor AstMethodShare s) y)+         => EqH (AstRaw s) y where+  AstRaw v ==. AstRaw u = AstRaw $ v ==. u+instance (OrdH (AstTensor AstMethodShare s) y)+         => OrdH (AstRaw s) y where+  AstRaw v <=. AstRaw u = AstRaw $ v <=. u++deriving instance Eq (AstRaw s y)+deriving instance Ord (AstRaw s y)+deriving instance Num (AstTensor AstMethodShare s y) => Num (AstRaw s y)+deriving instance IntegralH (AstTensor AstMethodShare s y)+                  => IntegralH (AstRaw s y)+deriving instance Fractional (AstTensor AstMethodShare s y)+                  => Fractional (AstRaw s y)+deriving instance Floating (AstTensor AstMethodShare s y)+                  => Floating (AstRaw s y)+deriving instance RealFloatH (AstTensor AstMethodShare s y)+                  => RealFloatH (AstRaw s y)+++-- * Misc++astShareNoSimplify :: KnownSpan s+                   => AstTensor AstMethodShare s y+                   -> AstTensor AstMethodShare s y+{-# NOINLINE astShareNoSimplify #-}+astShareNoSimplify a | astIsSmall True a = a+                         -- too important an optimization to skip+astShareNoSimplify a = case a of+  AstFromPrimal v -> fromPrimal $ astShareNoSimplify v+  AstFromDual v -> fromDual $ astShareNoSimplify v+  AstFromPlain v -> fromPlain $ astShareNoSimplify v+  _ -> unsafePerformIO $ case ftkAst a of+    ftk@FTKScalar -> do+        var <- funToAstAutoBoundsIO ftk a+        pure $! astShare var a+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> do+        let v = cAstConvDownSFromR sh x a+        var <- funToAstNoBoundsIO (FTKS sh x)+        pure $! cAstConvUpRFromS sh x $ astShare var v+    FTKX sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> do+        let v = cAstConvDownSFromX sh x a+        var <- funToAstNoBoundsIO (FTKS sh x)+        pure $! cAstConvUpXFromS sh' x $ astShare var v+    FTKS ZSS x@FTKScalar -> do+        let v = cAstConvDownKFromS a+        var <- funToAstAutoBoundsIO x v+        pure $! cAstConvUpSFromK $ astShare var v+    -- calling recursively for product may be not worth it+    ftk -> do+        var <- funToAstNoBoundsIO ftk+        pure $! astShare var a
− src/HordeAd/Core/AstPrettyPrint.hs 
@@ -1,584 +0,0 @@--- | Pretty-printing of the AST. Some of the variants of pretty-printing--- almost roundtrip, while others are more readable but less faithful.-module HordeAd.Core.AstPrettyPrint-  ( PrintConfig(..), defaulPrintConfig-  , printAstVar, printAst-  ) where--import Prelude--import Data.Foldable qualified as Foldable-import Data.IntMap.Strict (IntMap)-import Data.IntMap.Strict qualified as IM-import Data.List (intersperse)-import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality ((:~:) (Refl))-import Data.Vector.Generic qualified as V--import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Permutation (Perm (..), permToList)-import Data.Array.Nested.Shaped.Shape--import HordeAd.Core.Ast-import HordeAd.Core.AstTools-import HordeAd.Core.TensorKind-import HordeAd.Core.Types---- * Pretty-printing config---- Modeled after https://github.com/VMatthijs/CHAD/blob/755fc47e1f8d1c3d91455f123338f44a353fc265/src/TargetLanguage.hs#L335.---- TODO: ensure that terms roundtrip if neither loseRoudtrip--- nor ignoreNestedLambdas is set and that explicit sharing is then preserved--- as opposed to displaying sharing as Haskell lets.--- Note that disabling ignoreNestedLambdas causes derivatives to be computed,--- so pretty-printing in this way can be very expensive.-data PrintConfig = PrintConfig-  { loseRoudtrip        :: Bool-  , ignoreNestedLambdas :: Bool-  , varRenames          :: IntMap String-  }--defaulPrintConfig :: PrintConfig-defaulPrintConfig = PrintConfig-  { loseRoudtrip        = True-  , ignoreNestedLambdas = True-  , varRenames          = IM.empty-  }----- * Pretty-printing of variables--printAstVarId :: String -> PrintConfig -> AstVarId -> ShowS-printAstVarId prefix cfg var =-  let n = fromEnum var - 100000000-  in showString $ case IM.lookup n (varRenames cfg) of-    Just name | name /= "" -> name-    _ -> prefix ++ show n--printAstVar :: forall s y. AstSpan s-            => PrintConfig -> AstVarName s y -> ShowS-printAstVar cfg var = case isTensorInt (Proxy @s) (varNameToFTK var) of-  Just Refl -> printAstIntVar cfg var-  _ -> let prefix = case lengthSTK (ftkToSTK $ varNameToFTK var) of-             0 -> "x"-             1 -> "v"-             2 -> "m"-             3 -> "t"-             4 -> "u"-             _ -> "w"-       in printAstVarId prefix cfg (varNameToAstVarId var)--printAstIntVar :: PrintConfig -> IntVarName -> ShowS-printAstIntVar cfg var = printAstVarId "i" cfg (varNameToAstVarId var)----- * Pretty-printing of AST terms---- Precedences used are as in Haskell.-printAst :: forall s y ms. AstSpan s-         => PrintConfig -> Int -> AstTensor ms s y -> ShowS-printAst cfg d = \case-  AstPair t1 t2 ->-    showParen (d > 10)-    $ showString "tpair "-      . printAst cfg 11 t1-      . showString " "-      . printAst cfg 11 t2-  AstProject1 t -> printPrefixOp printAst cfg d "tproject1" [t]-  AstProject2 t -> printPrefixOp printAst cfg d "tproject2" [t]-  AstFromVector snat stk l ->-   if loseRoudtrip cfg-   then case stk of-    STKR{} ->-      showParen (d > 10)-      $ showString "rfromVector "-        . (showParen True-           $ showString "fromList "-             . showListWith (printAst cfg 0) (V.toList l))-    STKS{} ->-      showParen (d > 10)-      $ showString "sfromVector "-        . (showParen True-           $ showString "fromList "-             . showListWith (printAst cfg 0) (V.toList l))-    STKX{} ->-      showParen (d > 10)-      $ showString "xfromVector "-        . (showParen True-           $ showString "fromList "-             . showListWith (printAst cfg 0) (V.toList l))-    _ ->  -- scalar and product-      showParen (d > 10)-      $ showString "tfromVector "-        . (showParen True-           $ showString "fromList "-             . showListWith (printAst cfg 0) (V.toList l))-   else showParen (d > 10)-        $ showString ("tfromVector (" ++ show snat ++ ") (" ++ show stk ++ ") ")-          . (showParen True-             $ showString "fromList "-               . showListWith (printAst cfg 0) (V.toList l))-  -- This is too common to be verbose even in no loseRoudtrip mode.-  AstSum snat stk v -> case stk of-    STKR{} -> printPrefixOp printAst cfg d "rsum" [v]-    STKS{} -> printPrefixOp printAst cfg d-                            ("ssum @" ++ show (sNatValue snat)) [v]-    STKX{} -> printPrefixOp printAst cfg d-                            ("xsum @" ++ show (sNatValue snat)) [v]-    _ ->  -- scalar and product-      printPrefixOp printAst cfg d-                    ("tsum (" ++ show snat ++ ") (" ++ show stk ++ ")") [v]-  -- This is too common to be verbose even in no loseRoudtrip mode.-  AstReplicate snat stk v -> case stk of-    STKR{} -> printPrefixOp printAst cfg d-                            ("rreplicate " ++ show (sNatValue snat)) [v]-    STKS{} -> printPrefixOp printAst cfg d-                            ("sreplicate @" ++ show (sNatValue snat)) [v]-    STKX{} -> printPrefixOp printAst cfg d-                            ("xreplicate @" ++ show (sNatValue snat)) [v]-    _ ->  -- scalar and product-      printPrefixOp-        printAst cfg d-        ("treplicate (" ++ show snat ++ ") (" ++ show stk ++ ")") [v]-  AstMapAccumRDer k bftk eftk f df rf acc0 es ->-   if loseRoudtrip cfg-   then-    showParen (d > 10)-    $ showString "tmapAccumRDer "-      . showParen True (shows k)-      . showString " "-      . printAstHFun cfg 10 f-      . showString " "-      . printAstHFun cfg 10 df-      . showString " "-      . printAstHFun cfg 10 rf-      . showString " "-      . printAst cfg 11 acc0-      . showString " "-      . printAst cfg 11 es-   else-    showParen (d > 10)-    $ showString "tmapAccumRDer "-      . showParen True (shows k)-      . showString " "-      . showParen True (shows bftk)-      . showString " "-      . showParen True (shows eftk)-      . showString " "-      . printAstHFun cfg 10 f-      . showString " "-      . printAstHFun cfg 10 df-      . showString " "-      . printAstHFun cfg 10 rf-      . showString " "-      . printAst cfg 11 acc0-      . showString " "-      . printAst cfg 11 es-  AstMapAccumLDer k bftk eftk f df rf acc0 es ->-   if loseRoudtrip cfg-   then-    showParen (d > 10)-    $ showString "tmapAccumLDer "-      . showParen True (shows k)-      . showString " "-      . printAstHFun cfg 10 f-      . showString " "-      . printAstHFun cfg 10 df-      . showString " "-      . printAstHFun cfg 10 rf-      . showString " "-      . printAst cfg 11 acc0-      . showString " "-      . printAst cfg 11 es-   else-    showParen (d > 10)-    $ showString "tmapAccumLDer "-      . showParen True (shows k)-      . showString " "-      . showParen True (shows bftk)-      . showString " "-      . showParen True (shows eftk)-      . showString " "-      . printAstHFun cfg 10 f-      . showString " "-      . printAstHFun cfg 10 df-      . showString " "-      . printAstHFun cfg 10 rf-      . showString " "-      . printAst cfg 11 acc0-      . showString " "-      . printAst cfg 11 es-  AstApply t ll -> showParen (d > 10)-                   $ showString "tApply "-                     . printAstHFunOneUnignore cfg 10 t-                         -- this is a lambda, but not nested, so always printed-                     . showString " "-                     . printAst cfg 11 ll-  AstVar var -> printAstVar cfg var-  AstCond b a1 a2 ->-    showParen (d > 10)-    $ showString "ifH "-      . printAstBool cfg 11 b-      . showString " "-      . printAst cfg 11 a1-      . showString " "-      . printAst cfg 11 a2-  AstBuild1 k stk (var, v) ->-   if loseRoudtrip cfg-   then-    showParen (d > 10)-    $ showString "tbuild1 ("-      . shows k-      . showString ") "-      . (showParen True-         $ showString "\\"-           . printAstIntVar cfg var-           . showString " -> "-           . printAst cfg 0 v)-   else-    showParen (d > 10)-    $ showString "tbuild1 ("-      . shows k-      . showString ") "-      . showParen True (shows stk)-      . showString " "-      . (showParen True-         $ showString "\\"-           . printAstIntVar cfg var-           . showString " -> "-           . printAst cfg 0 v)--  t@(AstLet @_ @_ @s1 var0 u0 v0) ->-    if loseRoudtrip cfg-    then let collect :: AstTensor AstMethodLet s y -> ([(ShowS, ShowS)], ShowS)-             collect (AstLet var u v) =-               let name = printAstVar cfg var-                   uPP = printAst cfg 0 u-                   (rest, corePP) = collect v-               in ((name, uPP) : rest, corePP)-             collect v = ([], printAst cfg 0 v)-             (pairs, core) = collect t-         in showParen (d > 0)-            $ showString "let "-              . foldr (.) id (intersperse (showString " ; ")-                  [name . showString " = " . uPP | (name, uPP) <- pairs])-              . showString " in "-              . core-    else let keyword = case ( sameAstSpan @s1 @PrimalSpan-                            , sameAstSpan @s @FullSpan ) of-               (Just Refl, Just Refl) -> "ttletPrimal "-               _ -> "tlet "-         in showParen (d > 10)-            $ showString keyword-              . printAst cfg 11 u0-              . showString " "-              . (showParen True-                 $ showString "\\"-                   . printAstVar cfg var0-                   . showString " -> "-                   . printAst cfg 0 v0)-  AstShare _var v -> printPrefixOp printAst cfg d "tshare" [v]-  AstToShare v -> printPrefixOp printAst cfg d "toShare" [v]--  AstPrimalPart a ->-    if loseRoudtrip cfg-    then case ftkAst a of-      FTKR{} -> printPrefixOp printAst cfg d "rprimalPart" [a]-      FTKS{} -> printPrefixOp printAst cfg d "sprimalPart" [a]-      FTKX{} -> printPrefixOp printAst cfg d "xprimalPart" [a]-      _      -> printPrefixOp printAst cfg d "tprimalPart" [a]-    else printPrefixOp printAst cfg d "tprimalPart" [a]-  AstDualPart a ->-    if loseRoudtrip cfg-    then case ftkAst a of-      FTKR{} -> printPrefixOp printAst cfg d "rdualPart" [a]-      FTKS{} -> printPrefixOp printAst cfg d "sdualPart" [a]-      FTKX{} -> printPrefixOp printAst cfg d "xdualPart" [a]-      _      -> printPrefixOp printAst cfg d "tdualPart" [a]-    else printPrefixOp printAst cfg d-                       ("tdualPart (" ++ show (ftkToSTK (ftkAst a)) ++ ")") [a]-  AstFromPrimal a ->-    if loseRoudtrip cfg-    then printAst cfg d a-    else printPrefixOp-           printAst cfg d-           ("tfromPrimal (" ++ show (ftkToSTK (ftkAst a)) ++ ")") [a]-  AstFromDual a ->-    if loseRoudtrip cfg-    then printAst cfg d a-    else printPrefixOp printAst cfg d "tfromDual" [a]--  AstPlusK u v -> printBinaryOp printAst cfg d u (6, "+") v-  AstTimesK u v -> printBinaryOp printAst cfg d u (7, "*") v-  AstN1K opCode u -> printAstN1R printAst cfg d opCode u-  AstR1K opCode u -> printAstR1R printAst cfg d opCode u-  AstR2K opCode u v -> printAstR2R printAst cfg d opCode u v-  AstI2K opCode u v -> printAstI2R printAst cfg d opCode u v-  AstConcreteK k -> showNumber k-  AstFloorK v ->-    printPrefixOp printAst cfg d "kfloor" [v]-  AstFromIntegralK v ->-    printPrefixOp printAst cfg d "kfromIntegral" [v]-  AstCastK v ->-    printPrefixOp printAst cfg d "kcast" [v]--  AstPlusS u v -> printBinaryOp printAst cfg d u (6, "+") v-  AstTimesS u v -> printBinaryOp printAst cfg d u (7, "*") v-  AstN1S opCode u -> printAstN1R printAst cfg d opCode u-  AstR1S opCode u -> printAstR1R printAst cfg d opCode u-  AstR2S opCode u v -> printAstR2R printAst cfg d opCode u v-  AstI2S opCode u v -> printAstI2R printAst cfg d opCode u v-  AstConcreteS a -> case Nested.sshape a of-    ZSS -> showParen (d > 10)-           $ showString "sscalar "-             . showNumber (Nested.sunScalar a)-    _ -> showParen (d > 10)-         $ showString "sconcrete "-           . (showParen True-              $ shows a)-  AstFloorS a ->-    printPrefixOp printAst cfg d "sfloor" [a]-  AstFromIntegralS a ->-    printPrefixOp printAst cfg d "sfromIntegral" [a]-  AstCastS a ->-    printPrefixOp printAst cfg d "scast" [a]--  AstIndexS _ v ix ->-    showParen (d > 9)-    $ printAst cfg 10 v-      . showString " !$ "-      . showListWith (printAst cfg 0) (Foldable.toList ix)-  AstScatterS _ v (ZS, ix) ->-    showParen (d > 9)-    $ showString "soneHot "-      . printAst cfg 11 v-      . showString " "-      . showListWith (printAst cfg 0) (Foldable.toList ix)-  AstScatterS sh v (vars, ix) ->-   if loseRoudtrip cfg-   then-    showParen (d > 10)-    $ showString "sscatter "-      . printAst cfg 11 v-      . showString " "-      . (showParen True-         $ showString "\\"-           . showListWith (printAstIntVar cfg)-                          (listsToList vars)-           . showString " -> "-           . showListWith (printAst cfg 0) (Foldable.toList ix))-   else-    showParen (d > 10)-    $ showString ("sscatter " ++ show sh ++ " ")-      . printAst cfg 11 v-      . showString " "-      . (showParen True-         $ showString "\\"-           . showListWith (printAstIntVar cfg)-                          (listsToList vars)-           . showString " -> "-           . showListWith (printAst cfg 0) (Foldable.toList ix))-  {- Let's re-enable this when/if we remove AstIndexS altogether-     or at least stop rewriting this to AstIndexS but instead optimize-     the instances for this case:-  AstGatherS _ v (ZS, ix) ->-    showParen (d > 9)-    $ printAst cfg 10 v-      . showString " !$ "-      . showListWith (printAst cfg 0) (Foldable.toList ix) -}-  AstGatherS sh v (vars, ix) ->-   if loseRoudtrip cfg-   then-    showParen (d > 10)-    $ showString "sgather "-      . printAst cfg 11 v-      . showString " "-      . (showParen True-         $ showString "\\"-           . showListWith (printAstIntVar cfg)-                          (listsToList vars)-           . showString " -> "-           . showListWith (printAst cfg 0) (Foldable.toList ix))-   else-    showParen (d > 10)-    $ showString ("sgather " ++ show sh ++ " ")-      . printAst cfg 11 v-      . showString " "-      . (showParen True-         $ showString "\\"-           . showListWith (printAstIntVar cfg)-                          (listsToList vars)-           . showString " -> "-           . showListWith (printAst cfg 0) (Foldable.toList ix))-  AstMinIndexS a -> printPrefixOp printAst cfg d "sminIndex" [a]-  AstMaxIndexS a -> printPrefixOp printAst cfg d "smaxIndex" [a]-  AstIotaS snat ->-    showParen (d > 10)-    $ showString ("siota (" ++ show snat ++ ")")-  AstAppendS t1 t2 ->-    showParen (d > 10)-    $ showString "sappend "-      . printAst cfg 11 t1-      . showString " "-      . printAst cfg 11 t2-  AstSliceS i n _k v ->-    printPrefixOp printAst cfg d-                  ("sslice (" ++ show i ++ ") (" ++ show n ++ ")") [v]-  AstReverseS v -> printPrefixOp printAst cfg d "sreverse" [v]-  AstTransposeS (SNat' @1 `PCons` SNat' @0 `PCons` PNil) v ->-    printPrefixOp printAst cfg d "str" [v]-  AstTransposeS perm v ->-   if loseRoudtrip cfg-   then printPrefixOp printAst cfg d-                      ("stranspose @" ++ show (permToList perm)) [v]-   else printPrefixOp-          printAst cfg d-          ("ttranspose (makePerm @" ++ show (permToList perm) ++ ")") [v]-  AstReshapeS sh2 v ->-    printPrefixOp printAst cfg d ("sreshape @" ++ show (shsToList sh2)) [v]--  -- TODO: pretty-print correctly szip, sunzip, snestS, sunNestS-  -- or at least make sure they get printed as tconvert, not as the others-  AstConvert c t -> case (ftkToSTK (ftkAst t), convertFTK c (ftkAst t)) of-    (STKScalar, FTKS{}) -> printPrefixOp printAst cfg d "sfromK" [t]-    (STKR{}, FTKS{}) -> printPrefixOp printAst cfg d "sfromR" [t]-    (STKX{}, FTKS{}) -> printPrefixOp printAst cfg d "sfromX" [t]-    (STKS{}, FTKScalar) -> printPrefixOp printAst cfg d "kfromS" [t]-    (STKS{}, FTKR{}) -> printPrefixOp printAst cfg d "rfromS" [t]-    (STKS{}, FTKX{}) -> printPrefixOp printAst cfg d "xfromS" [t]-    (ystk, _) -> let s = "tconvert (" ++ show c ++ ") (" ++ show ystk ++ ")"-                 in printPrefixOp printAst cfg d s [t]--  AstSum0S v ->-    printPrefixOp printAst cfg d "ssum0" [v]-  AstDot0S u v ->-    printPrefixOp printAst cfg d "sdot0" [u, v]-  AstDot1InS _ _ u v ->-    printPrefixOp printAst cfg d "sdot1In" [u, v]-  AstMatmul2S SNat SNat SNat u v ->-    showParen (d > 10)-    $ showString "smatmul2 "-      . printAst cfg 11 u-      . showString " "-      . printAst cfg 11 v--showNumber :: (Ord a, Num a, Show a) => a -> ShowS-{-# INLINE showNumber #-}-showNumber a = showParen (a < 0) $ shows a---- Differs from standard only in the space after comma.-showListWith :: (a -> ShowS) -> [a] -> ShowS-{-# INLINE showListWith #-}-showListWith = showCollectionWith "[" ", " "]"--showCollectionWith :: String -> String -> String -> (a -> ShowS) -> [a] -> ShowS-{-# INLINE showCollectionWith #-}-showCollectionWith start _   end _     []     s = start ++ end ++ s-showCollectionWith start sep end showx (x:xs) s = start ++ showx x (showl xs)- where-  showl []     = end ++ s-  showl (y:ys) = sep ++ showx y (showl ys)--printAstHFun :: (AstSpan s, AstSpan s2)-             => PrintConfig -> Int -> AstHFun s s2 x y -> ShowS-printAstHFun cfg d = \case-  AstLambda var l ->-    if loseRoudtrip cfg-    then if ignoreNestedLambdas cfg-         then showString "<lambda>"-         else showParen (d > 0)-              $ showString "\\"-                . printAstVar cfg var-                . showString " -> "-                . printAst cfg 0 l-    else showParen (d > 0)-         $ showString "tlambda $ \\"-           . printAstVar cfg var-           . showString " -> "-           . printAst cfg 0 l--printAstHFunOneUnignore :: (AstSpan s, AstSpan s2)-                        => PrintConfig -> Int -> AstHFun s s2 x y -> ShowS-printAstHFunOneUnignore cfg d = \case-  AstLambda var l ->-    if loseRoudtrip cfg-    then showParen (d > 0)-         $ showString "\\"-           . printAstVar cfg var-           . showString " -> "-           . printAst cfg 0 l-    else showParen (d > 0)-         $ showString "tlambda $ \\"-           . printAstVar cfg var-           . showString " -> "-           . printAst cfg 0 l--printAstBool :: PrintConfig -> Int -> AstBool ms -> ShowS-printAstBool cfg d = \case-  AstBoolConst b -> showString $ if b then "true" else "false"-  AstBoolNot u -> printPrefixOp printAstBool cfg d "notB" [u]-  AstBoolAnd u v -> printBinaryOp printAstBool cfg d u (3, "&&*") v-  AstLeqK u v -> printBinaryOp printAst cfg d u (4, "<=.") v-  AstLeqS u v -> printBinaryOp printAst cfg d u (4, "<=.") v--printAstN1R :: (PrintConfig -> Int -> a -> ShowS)-           -> PrintConfig -> Int -> OpCodeNum1 -> a -> ShowS-printAstN1R pr cfg d opCode u = case opCode of-  NegateOp -> printPrefixOp pr cfg d "negate" [u]-  AbsOp -> printPrefixOp pr cfg d "abs" [u]-  SignumOp -> printPrefixOp pr cfg d "signum" [u]--printAstR1R :: (PrintConfig -> Int -> a -> ShowS)-           -> PrintConfig -> Int -> OpCode1 -> a -> ShowS-printAstR1R pr cfg d opCode u = case opCode of-  RecipOp -> printPrefixOp pr cfg d "recip" [u]-  ExpOp -> printPrefixOp pr cfg d "exp" [u]-  LogOp -> printPrefixOp pr cfg d "log" [u]-  SqrtOp -> printPrefixOp pr cfg d "sqrt" [u]-  SinOp -> printPrefixOp pr cfg d "sin" [u]-  CosOp -> printPrefixOp pr cfg d "cos" [u]-  TanOp -> printPrefixOp pr cfg d "tan" [u]-  AsinOp -> printPrefixOp pr cfg d "asin" [u]-  AcosOp -> printPrefixOp pr cfg d "acos" [u]-  AtanOp -> printPrefixOp pr cfg d "atan" [u]-  SinhOp -> printPrefixOp pr cfg d "sinh" [u]-  CoshOp -> printPrefixOp pr cfg d "cosh" [u]-  TanhOp -> printPrefixOp pr cfg d "tanh" [u]-  AsinhOp -> printPrefixOp pr cfg d "asinh" [u]-  AcoshOp -> printPrefixOp pr cfg d "acosh" [u]-  AtanhOp -> printPrefixOp pr cfg d "atanh" [u]--printAstR2R :: (PrintConfig -> Int -> a -> ShowS)-           -> PrintConfig -> Int -> OpCode2 -> a -> a -> ShowS-printAstR2R pr cfg d opCode u v = case opCode of-  DivideOp -> printBinaryOp pr cfg d u (7, "/") v-  PowerOp -> printBinaryOp pr cfg d u (8, "**") v-  LogBaseOp -> printPrefixOp pr cfg d "logBase" [u, v]-  Atan2Op -> printPrefixOp pr cfg d "atan2H" [u, v]--printAstI2R :: (PrintConfig -> Int -> a -> ShowS)-           -> PrintConfig -> Int -> OpCodeIntegral2 -> a -> a -> ShowS-printAstI2R pr cfg d opCode u v = case opCode of-  QuotOp -> printPrefixOp pr cfg d "quotH" [u, v]-  RemOp -> printPrefixOp pr cfg d "remH" [u, v]--printPrefixOp :: (PrintConfig -> Int -> a -> ShowS)-              -> PrintConfig -> Int -> String -> [a]-              -> ShowS-{-# INLINE printPrefixOp #-}-printPrefixOp pr cfg d funcname args =-  let rs = map (\arg -> showString " " . pr cfg 11 arg) args-  in showParen (d > 10)-     $ showString funcname-       . foldr (.) id rs--printBinaryOp :: (PrintConfig -> Int -> a -> ShowS)-              -> PrintConfig -> Int -> a -> (Int, String) -> a-              -> ShowS-{-# INLINE printBinaryOp #-}-printBinaryOp pr cfg d left (prec, opstr) right =-  showParen (d > prec)-  $ pr cfg (prec + 1) left-    . showString (" " ++ opstr ++ " ")-    . pr cfg (prec + 1) right
src/HordeAd/Core/AstSimplify.hs view
@@ -1,3881 +1,5746 @@-{-# LANGUAGE AllowAmbiguousTypes, ViewPatterns #-}-{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}-{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}-{-# OPTIONS_GHC -Wno-orphans #-}--- {-# OPTIONS_GHC -fmax-pmcheck-models=10000 #-}--- {-# OPTIONS_GHC -freduction-depth=10000 #-}--- {-# OPTIONS_GHC -fconstraint-solver-iterations=10000 #-}--- | This module holds smart constructors for AST, that is,--- term-simplifying combinators corresponding to the Ast constructors.--- The combinators simplify only on the basis of inspecting the roots of their--- argument term trees. If the arguments get modified,--- the modified forms are again inspected and potentially simplified again.------ The limited simplification via combinators is enough to uncover redexes--- for the factorization rules to fire and to undo some of the complication--- introduced by vectorization. The intention is to leave intact as much--- as possible of the original terms provided by the user while making--- sure subterms introduced by vectorization are maximally simplified.-module HordeAd.Core.AstSimplify-  ( RewritePhase(..), SimplifyKnobs (..), defaultKnobs-  , -- * The simplifying combinators, one for almost each AST constructor-    astPair, astProject1, astProject2, astFromVector, astSum, astReplicate-  , astMapAccumRDer, astMapAccumLDer, astApply, astCond-  , astConcrete, astConcreteK, astConcreteS--  , astLet--  , astPrimalPart, astDualPart--  , astFloorK, astFromIntegralK, astCastK--  , astFloorS, astFromIntegralS, astCastS--  , astIndexS, astIndexKnobsS, astScatterS, astGatherS, astGatherKnobsS-  , astAppendS, astSliceS, astReverseS, astTransposeS, astReshapeS--  , astConvert-  , astFromS', astKFromS', astRFromS', astXFromS'-  , astSFromK', astSFromR', astSFromX'-  , astSum0S, astDot0S, astDot1InS, astMatmul2S--    -- * Helper combinators-  , astLetFun-    -- * Substitution operations-  , substituteAst, substituteAstIxS, substituteAstBool-  ) where--import Prelude--import Control.Exception.Assert.Sugar-import Control.Monad (mapAndUnzipM, mplus)-import Data.Foldable qualified as Foldable-import Data.Functor.Const-import Data.Functor.Product qualified as Fun-import Data.GADT.Compare-import Data.Int (Int64)-import Data.List (findIndex)-import Data.Maybe (catMaybes, fromMaybe, isJust)-import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))-import Data.Type.Ord (Compare)-import Data.Vector.Generic qualified as V-import Data.Vector.Strict qualified as Data.Vector-import GHC.Exts (IsList (..))-import GHC.TypeLits-  ( Nat-  , OrderingI (..)-  , cmpNat-  , fromSNat-  , sameNat-  , type (+)-  , type (-)-  , type (<=)-  , type (<=?)-  )-import System.IO.Unsafe (unsafePerformIO)-import Type.Reflection (typeRep)-import Unsafe.Coerce (unsafeCoerce)--import Data.Array.Nested (Replicate, type (++))-import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Convert-  (shrFromShS, shxFromShS, withShsFromShR, withShsFromShX)-import Data.Array.Nested.Lemmas-import Data.Array.Nested.Mixed.Shape-import Data.Array.Nested.Permutation (DropLen, Perm (..), TakeLen, permInverse)-import Data.Array.Nested.Permutation qualified as Permutation-import Data.Array.Nested.Ranked.Shape-import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types-  (Head, Init, Last, Tail, snatPlus, unsafeCoerceRefl)--import HordeAd.Core.Ast-  ( AstBool (AstBoolConst)-  , AstTensor (AstConcreteK, AstConcreteS, AstPlusK, AstPlusS, AstTimesK, AstTimesS)-  )-import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))-import HordeAd.Core.Ast qualified as Ast-import HordeAd.Core.AstFreshId-import HordeAd.Core.AstTools-import HordeAd.Core.CarriersAst (sunReplicate1, sunReplicateN, sunReplicateScal)-import HordeAd.Core.CarriersConcrete-import HordeAd.Core.ConvertTensor-import HordeAd.Core.Ops-import HordeAd.Core.TensorKind-import HordeAd.Core.Types-import HordeAd.Core.Unwind--data RewritePhase =-    PhaseUnspecified-  | PhaseVectorization-  | PhaseSimplification-  | PhaseExpansion-  | PhaseContraction- deriving (Show, Eq)--data SimplifyKnobs = SimplifyKnobs-  { knobPhase :: RewritePhase-  }- deriving Show--defaultKnobs :: SimplifyKnobs-defaultKnobs = SimplifyKnobs PhaseUnspecified---- @PhaseVectorization@ should only affect the topmost redex.-deVect :: SimplifyKnobs -> SimplifyKnobs-deVect (SimplifyKnobs PhaseVectorization) = SimplifyKnobs PhaseUnspecified-deVect knobs = knobs----- * Expressing operations as Gather; introduces new variable names---- | We keep AstTranspose terms for as long as possible, because--- they are small and fuse nicely in many cases. For some forms of indexing--- and nesting with reshape and gather they don't fuse, which is when--- this function is invoked.-astTransposeAsGatherS-  :: forall perm sh s r. AstSpan s-  => SimplifyKnobs -> Permutation.Perm perm-  -> AstTensor AstMethodLet s (TKS2 sh r)-  -> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)-{-# NOINLINE astTransposeAsGatherS #-}-astTransposeAsGatherS knobs perm v =-  let FTKS shn _ = ftkAst v-      shnPermuted = shsPermute perm (shsTakeLen perm shn)-  in funToVarsIxS @_ @AstMethodLet shnPermuted $ \ (!vars, !ix) ->-    -- See astGatherCase.AstTransposeS for similar code with more comments.-    gcastWith (lemRankMapJust $ shsTakeLen perm shn) $-    gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $-    permInverse perm $ \(invperm :: Nested.Perm invperm) proof ->-      case proof (ssxFromShX $ shxFromShS $ shsTakeLen perm shn) of-        Refl ->-          gcastWith (unsafeCoerceRefl-                     :: DropLen invperm-                          (Permutation.Permute perm (TakeLen perm sh))-                        :~: '[]) $-          gcastWith (lemAppNil-                       @(Permutation.Permute invperm-                           (Permutation.Permute perm (TakeLen perm sh)))) $-          -- Seriously? This should be deduced from the above:-          gcastWith (unsafeCoerceRefl-                     :: Permutation.PermutePrefix invperm-                          (Permutation.Permute perm (TakeLen perm sh))-                        :~: Permutation.Permute invperm-                          (Permutation.Permute perm (TakeLen perm sh))) $-          -- This should follow from @proof@, if not for MapJust:-          gcastWith (unsafeCoerceRefl-                     :: Permutation.Permute invperm-                          (Permutation.Permute perm (TakeLen perm sh))-                        :~: TakeLen perm sh) $-          let asts :: AstIxS AstMethodLet (TakeLen perm sh)-              asts = ixsPermutePrefix invperm ix-          in gcastWith (unsafeCoerceRefl-                        :: TakeLen perm sh ++ DropLen perm sh :~: sh) $-             astGatherKnobsS @(Permutation.Permute perm (TakeLen perm sh))-                             @(DropLen perm sh)-                             @(TakeLen perm sh)-                             knobs (shsDropLen perm shn) v (vars, asts)---- TODO: To make this less disastrous, we need to add an extra constructor--- to AstIndex with the semantics "this index reshaped from shIn to shOut"--- that fuses perfectly with itself and absorbs normal indexes--- by substitution. Or perhaps make this the only constructor, with normal--- indexes represented as "this index reshaped from sh to sh".--- Or only extend AstGather and possibly also AstIndex with the extra--- shIn and shOut arguments. This complicates any code related to--- AstGather and AstIndex, but often prevents nested reshapes from affecting--- term size in any way. But we'd need to be careful to avoid breaking such--- an index into components, because that forces index normalization,--- e.g., index(gather) can no longer simplify recursively by one index--- component at a time (probably possible only if the index is shorter--- that the list of variables fo the gather). There are probably bad cases--- where term size blowup can't be avoided, because the index has to be--- normalized between each reshape.------ | This generates big terms that don't simplify well,--- so we keep the AstReshape form until simplification gets stuck.--- In fact, to simplify the terms we'd need advanced solving of equations--- in integer arithmetic modulo.-astReshapeAsGatherS-  :: forall sh sh2 r s. AstSpan s-  => SimplifyKnobs -> ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh r)-  -> AstTensor AstMethodLet s (TKS2 sh2 r)-{-# NOINLINE astReshapeAsGatherS #-}-astReshapeAsGatherS knobs shOut v | Refl <- lemAppNil @sh2-                                  , Refl <- lemAppNil @sh-                                  , FTKS shIn _ <- ftkAst v =-  funToVarsIxS shOut $ \ (!vars, !ix) ->-    let fromInt :: Int -> AstInt AstMethodLet-        fromInt i = AstConcreteK (fromIntegral i)-        iUnshared :: AstInt AstMethodLet-        iUnshared = toLinearIdxS @sh2 @'[] fromInt shOut ix-        asts :: AstInt AstMethodLet -> AstIxS AstMethodLet sh-        asts i = fromLinearIdxS fromInt shIn i-    in gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh) $-       gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[]) $--- This can't be done, because i depends on vars:---     astLetFunB iUnshared $ \i ->-       let i = iUnshared  -- sharing broken-       in astGatherKnobsS @sh2 @'[] @sh knobs ZSS v (vars, asts i)----- * The simplifying combinators, one for almost each AST constructor--astPair :: AstSpan s-        => AstTensor AstMethodLet s x -> AstTensor AstMethodLet s y-        -> AstTensor AstMethodLet s (TKProduct x y)-astPair (Ast.AstFromPrimal v1) (Ast.AstFromPrimal v2) =-  Ast.AstFromPrimal $ astPair v1 v2-astPair (Ast.AstFromDual v1) (Ast.AstFromDual v2) =-  Ast.AstFromDual $ astPair v1 v2-astPair (Ast.AstConvert c1 v1) (Ast.AstConvert c2 v2)-  | checkAstFromS c1 v1 && checkAstFromS c2 v2 =-    astConvert (ConvT2 c1 c2) $ astPair v1 v2-astPair (Ast.AstConvert c1 v1) v2-  | checkAstFromS c1 v1 =-    astConvert (ConvT2 c1 ConvId) $ astPair v1 v2-astPair v1 (Ast.AstConvert c2 v2)-  | checkAstFromS c2 v2 =-    astConvert (ConvT2 ConvId c2) $ astPair v1 v2-astPair v1 v2 = Ast.AstPair v1 v2--astProject1-  :: forall x z s. AstSpan s-  => AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s x-astProject1 u = case u of-  Ast.AstPair x _z -> x-  Ast.AstCond b v1 v2 -> astCond b (astProject1 v1) (astProject1 v2)-  Ast.AstLet var t v -> astLet var t (astProject1 v)-  Ast.AstFromPrimal u1 -> Ast.AstFromPrimal $ astProject1 u1-  Ast.AstFromDual u1 -> Ast.AstFromDual $ astProject1 u1-  Ast.AstConvert c t | FTKProduct yftk _ <- ftkAst t-                     , FTKProduct zftk _ <- convertFTK c (ftkAst t)-                     , Just Refl <- matchingFTK yftk zftk -> astProject1 t-  -- TODO: generalize this somehow to arbitrary Conversions of the right type.-  -- At worst, just generate the canonical (?) c1 for the types at hand.-  Ast.AstConvert c@(ConvT2 c1 _c2) t-    | checkAstFromS c t ->-      astConvert c1 $ astProject1 t-  _ -> Ast.AstProject1 u--astProject2-  :: forall x z s. AstSpan s-  => AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s z-astProject2 u = case u of-  Ast.AstPair _x z -> z-  Ast.AstCond b v1 v2 -> astCond b (astProject2 v1) (astProject2 v2)-  Ast.AstLet var t v -> astLet var t (astProject2 v)-  Ast.AstFromPrimal u1 -> Ast.AstFromPrimal $ astProject2 u1-  Ast.AstFromDual u1 -> Ast.AstFromDual $ astProject2 u1-  Ast.AstConvert c t | FTKProduct _ yftk <- ftkAst t-                     , FTKProduct _ zftk <- convertFTK c (ftkAst t)-                     , Just Refl <- matchingFTK yftk zftk -> astProject2 t-  Ast.AstConvert c@(ConvT2 _c1 c2) t-    | checkAstFromS c t ->-      astConvert c2 $ astProject2 t-  _ -> Ast.AstProject2 u--astFromVector :: forall y k s. AstSpan s-              => SNat k -> SingletonTK y-              -> Data.Vector.Vector (AstTensor AstMethodLet s y)-              -> AstTensor AstMethodLet s (BuildTensorKind k y)-astFromVector (SNat' @1) stk v = astReplicate (SNat @1) stk (v V.! 0)-astFromVector snat@SNat stk l = fromMaybe (Ast.AstFromVector snat stk l) $-  -- This disables some rules, e.g., indexing or summing of fromVector-  -- of concrete arrays, but allocating an extra array of the same size-  -- as the fromVector is not a big deal and early rules are better-  -- then the same rules in contraction phase.-  (case (sameAstSpan @s @PrimalSpan, stk) of-     (Just Refl, STKScalar) ->-       let unConc :: AstTensor AstMethodLet PrimalSpan y-                  -> Maybe (Concrete y)-           unConc (AstConcreteK a) = Just $ Concrete a-           unConc _ = Nothing-       in case V.mapM unConc l of-         Just l4 | V.null l4 -> error "astFromVector: empty vector"-         Just l4 -> Just $ astConcreteS (tfromVector snat stk l4)-         Nothing -> Nothing-     (Just Refl, STKS _ STKScalar) ->-       let unConc :: AstTensor AstMethodLet PrimalSpan y-                  -> Maybe (Concrete y)-           unConc (AstConcreteS a) = Just $ Concrete a-           unConc _ = Nothing-       in case V.mapM unConc l of-         Just l4 | V.null l4 -> error "astFromVector: empty vector"-         Just l4 -> Just $ astConcreteS (tfromVector snat stk l4)-         Nothing -> Nothing-     _ -> Nothing)-  `mplus`-  (case sameAstSpan @s @FullSpan of-     Just Refl ->-       let unFromPrimal :: AstTensor AstMethodLet FullSpan y-                        -> Maybe (AstTensor AstMethodLet PrimalSpan y)-           unFromPrimal (Ast.AstFromPrimal t) = Just t-           unFromPrimal _ = Nothing-       in case V.mapM unFromPrimal l of-         Just l2 | V.null l2 -> error "astFromVector: empty vector"-         Just l2 -> Just $ Ast.AstFromPrimal $ astFromVector snat stk l2-         Nothing -> Nothing-     _ -> Nothing)-  `mplus`-  (case sameAstSpan @s @FullSpan of-     Just Refl ->-       let unFromDual :: AstTensor AstMethodLet FullSpan y-                        -> Maybe (AstTensor AstMethodLet DualSpan y)-           unFromDual (Ast.AstFromDual t) = Just t-           unFromDual _ = Nothing-       in case V.mapM unFromDual l of-         Just l2 | V.null l2 -> error "astFromVector: empty vector"-         Just l2 -> Just $ Ast.AstFromDual $ astFromVector snat stk l2-         Nothing -> Nothing-     _ -> Nothing)-  `mplus`-  (let unFrom :: FullShapeTK x-              -> AstTensor AstMethodLet s y-              -> Maybe (AstTensor AstMethodLet s x)-       unFrom xftk (AstFromS' _ t) =-         case matchingFTK (ftkAst t) xftk of-           Just Refl -> Just t-           Nothing -> error "astFromVector: impossible shape"-       unFrom _ _ = Nothing-   in case V.uncons l of-     Just (Ast.AstConvert c t, _) ->-       let xftk = ftkAst t-       in case V.mapM (unFrom xftk) l of-         Just l2 ->-           -- Here we heavily depend on c being semantically determined-           -- by the domain and codomain. We choose one such c,-           -- not necessarily the most efficient of them all.-           Just $ astConvert (buildTKConversion snat xftk c)-                $ astFromVector snat (ftkToSTK xftk) l2-         Nothing -> Nothing-     Just{} -> Nothing-     Nothing -> error "astFromVector: empty vector")--astSum :: forall y k s. AstSpan s-       => SNat k -> SingletonTK y-       -> AstTensor AstMethodLet s (BuildTensorKind k y)-       -> AstTensor AstMethodLet s y-astSum snat@SNat stk t0 = case t0 of-  _ | Just Refl <- testEquality snat (SNat @0) ->-    let ftk = razeFTK snat stk (ftkAst t0)-    in fromPrimal $ astConcrete ftk (tdefTarget ftk)-  AstConcreteS @_ @sh2 t -> case stk of-    STKS @sh _ STKScalar ->-      gcastWith (unsafeCoerceRefl :: k ': sh :~: sh2) $-      astConcreteS (tsum snat stk $ Concrete t)-    STKScalar ->-      gcastWith (unsafeCoerceRefl :: '[k] :~: sh2) $-      astConcreteK (tsum snat stk $ Concrete t)-  Ast.AstIotaS @_ @r (SNat @n) ->-    let i :: r-        i = fromInteger $ valueOf @n * (valueOf @n - 1) `div` 2-    in case stk of-      STKScalar -> AstConcreteK i-      STKS ZSS STKScalar -> AstConcreteS $ Nested.sscalar i-  Ast.AstReverseS v -> astSum snat stk v-  _ | Just Refl <- testEquality snat (SNat @1)-    , STKScalar <- stk ->-      astKFromS' $ astIndexS ZSS t0 (0 :.$ ZIS)-  _ | Just Refl <- testEquality snat (SNat @1)-    , STKS sh _  <- stk ->  -- other cases too rare-      astIndexS sh t0 (0 :.$ ZIS)  -- astReshape slows down the CNNO test-  Ast.AstFromVector @y2 _ _ l ->-    gcastWith (unsafeCoerceRefl :: y2 :~: y) $-    case stk of-      STKScalar -> foldr1 (+) l-      STKR _ STKScalar -> foldr1 (+) l-      STKS _ STKScalar -> foldr1 (+) l-      STKX _ STKScalar -> foldr1 (+) l-      _ -> Ast.AstSum snat stk t0-  -- See the analogous astSliceS rule.-  Ast.AstTransposeS (SNat' @1 `PCons` SNat' @0 `PCons` PNil) t-    | FTKS (_ :$$ _ :$$ _) _ <- ftkAst t-    , STKS (snat1 :$$ sh3) x <- stk-    , Just u <- unRepl1 t ->-      astReplicate snat1 (STKS sh3 x)-      $ astSum snat (STKS sh3 x) u-  Ast.AstReplicate _ STKScalar v | STKScalar <- stk ->-    v * fromPrimal (AstConcreteK $ fromInteger $ fromSNat snat)-  Ast.AstReplicate _ _ v | STKR _ (STKScalar @r) <- stk ->-    case ftkAst v of-      ftk@(FTKR sh' FTKScalar) ->-        withShsFromShR sh' $ \(sh :: ShS sh) ->-          v * astFromS'-                ftk (fromPrimal $ AstConcreteS @r-                     $ Nested.sreplicateScal sh $ fromInteger $ fromSNat snat)-  Ast.AstReplicate _ _ v | STKX _ (STKScalar @r) <- stk ->-    case ftkAst v of-      ftk@(FTKX sh' FTKScalar) ->-        withShsFromShX sh' $ \(sh :: ShS sh) ->-          v * astFromS'-                ftk (fromPrimal $ AstConcreteS @r-                     $ Nested.sreplicateScal sh $ fromInteger $ fromSNat snat)-  Ast.AstReplicate _ STKS{} v | STKS sh (STKScalar @r) <- stk ->-    case ftkAst v of-      ftk ->-          v * astFromS'-                ftk (fromPrimal $ AstConcreteS @r-                     $ Nested.sreplicateScal sh $ fromInteger $ fromSNat snat)-  -- This keeps tensors alive for longer, but it enables new simplifications,-  -- while hiding a sum inside let not often prevents other simplifications,-  -- because there are few redexes with sum but not at the top.-  Ast.AstLet var u v -> astLet var u (astSum snat stk v)-  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astSum snat stk v-  Ast.AstFromDual v -> Ast.AstFromDual $ astSum snat stk v-  Ast.AstScatterS @shm @shn @shp shn v (vars, (:.$) @k2 i1 rest)-    | STKS{} <- stk ->-      -- This boolean term may have free variables that act as universally-      -- quantified.-      case 0 <=. i1 &&* i1 <=. valueOf @k2 - 1 of-        AstBoolConst True ->-          astScatterS @shm @shn @(Tail shp) shn v (vars, rest)-        _ -> Ast.AstSum snat stk t0-  Ast.AstConvert c t | checkAstFromS c t -> case ftkAst t of-    FTKS (snat2 :$$ rest) x ->-      -- Here we'd need to change the types inside c, so instead we construct-      -- a new conversion based on the domain and codomain.-      astFromS' (razeFTK snat stk (ftkAst t0))-      $ astSum snat2 (STKS rest (ftkToSTK x)) t-    _ -> Ast.AstSum snat stk t0  -- products probably not worth the effort-  _ -> Ast.AstSum snat stk t0--astReplicate :: forall y k s. AstSpan s-             => SNat k -> SingletonTK y-             -> AstTensor AstMethodLet s y-             -> AstTensor AstMethodLet s (BuildTensorKind k y)-astReplicate snat stk t0 = case t0 of-  Ast.AstPair t1 t2 | STKProduct stk1 stk2 <- stk ->-    astPair (astReplicate snat stk1 t1) (astReplicate snat stk2 t2)-  -- This doesn't prevent indexing of replicate, because indexing goes inside-  -- the conditional, but it prevents the sum(replicate(cond)) simplification,-  -- because sum can't go inside, because it's costly and cond is eager.-  -- Ast.AstCond b v1 v2 ->-  --  astCond b (astReplicate snat stk v1) (astReplicate snat stk v2)-  -- TODO: This rules is, in principle, very good, because it permits many other-  -- rules to fire. However, one of these other rules is indexing of transpose-  -- that in some cases complicates terms and causes OOM in CNNI tests.-  -- We need to restrict the indexing rule more effectively first.-  -- Ast.AstLet var t v -> astLet var t (astReplicate snat stk v)-  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astReplicate snat stk v-  Ast.AstFromDual v -> Ast.AstFromDual $ astReplicate snat stk v-  AstConcreteK t -> astConcreteS $ treplicate snat stk $ Concrete t-  AstConcreteS t -> astConcreteS $ treplicate snat stk $ Concrete t-  Ast.AstConvert c t | checkAstFromS c t ->-    let xftk = ftkAst t-    in astConvert (buildTKConversion snat xftk c)-                  (astReplicate snat (ftkToSTK xftk) t)-  _ -> Ast.AstReplicate snat stk t0-  -- TODO: maybe add a rule and then generalize:-  -- replicate n1 (str (replicate n2 u))-  -- ~> transpose [0, 2, 1] (replicate n1 (replicate n2 u))-  -- but the reverse rule is already in astTransposeS---- TODO: also pull up AstFromPrimal, etc.-astMapAccumRDer-  :: forall accy by ey k s. AstSpan s-  => SNat k-  -> FullShapeTK by-  -> FullShapeTK ey-  -> AstHFun s s-             (TKProduct accy ey) (TKProduct accy by)-  -> AstHFun s s-             (TKProduct (ADTensorKind (TKProduct accy ey))-                        (TKProduct accy ey))-             (ADTensorKind (TKProduct accy by))-  -> AstHFun s s-             (TKProduct (ADTensorKind (TKProduct accy by))-                        (TKProduct accy ey))-             (ADTensorKind (TKProduct accy ey))-  -> AstTensor AstMethodLet s accy-  -> AstTensor AstMethodLet s (BuildTensorKind k ey)-  -> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))-astMapAccumRDer k bftk eftk (AstLambda varf vf)-                            (AstLambda vard vd)-                            (AstLambda varr vr)-                (AstFromS' @accyFrom accftk acc0From) es =-  let accftkFrom = ftkAst acc0From-      accFromSTK = ftkToSTK accftkFrom-      ftkf2 = FTKProduct accftkFrom eftk-      varf2 = mkAstVarName ftkf2 (varNameToBounds varf) (varNameToAstVarId varf)-      astf2 = astVar varf2-      vf2 =-        let subbed =-              substituteAst-                (astPair (astFromS' @accyFrom accftk (astProject1 astf2))-                         (astProject2 astf2))-                varf vf-        in astSFrom' @(TKProduct accy by)-                    (STKProduct accFromSTK (ftkToSTK bftk))-                    subbed-      ftkd2 = FTKProduct-                (adFTK $ FTKProduct accftkFrom eftk)-                (FTKProduct accftkFrom eftk)-      vard2 = mkAstVarName ftkd2 (varNameToBounds vard) (varNameToAstVarId vard)-      astd2 = astVar vard2-      vd2 =-        let subbed =-              substituteAst-                (astPair (astPair (astFromS' @(ADTensorKind accyFrom)-                                     (adFTK accftk)-                                     (astProject1 (astProject1 astd2)))-                                  (astProject2 (astProject1 astd2)))-                         (astPair (astFromS' @accyFrom accftk-                                     (astProject1 (astProject2 astd2)))-                                  (astProject2 (astProject2 astd2))))-                vard vd-        in astSFrom' @(ADTensorKind (TKProduct accy by))-                    (adSTK $ STKProduct accFromSTK (ftkToSTK bftk))-                    subbed-      ftkr2 = FTKProduct-                (adFTK $ FTKProduct accftkFrom bftk)-                (FTKProduct accftkFrom eftk)-      varr2 = mkAstVarName ftkr2 (varNameToBounds varr) (varNameToAstVarId varr)-      astr2 = astVar varr2-      vr2 =-        let subbed =-              substituteAst-                (astPair (astPair (astFromS' @(ADTensorKind accyFrom)-                                     (adFTK accftk)-                                     (astProject1 (astProject1 astr2)))-                                  (astProject2 (astProject1 astr2)))-                         (astPair (astFromS' @accyFrom accftk-                                     (astProject1 (astProject2 astr2)))-                                  (astProject2 (astProject2 astr2))))-                varr vr-        in astSFrom' @(ADTensorKind (TKProduct accy ey))-                    (adSTK $ STKProduct accFromSTK (ftkToSTK eftk))-                    subbed-  in astFromS' @(TKProduct accyFrom (BuildTensorKind k by))-               (FTKProduct accftk (buildFTK k bftk))-     $ astMapAccumRDer k bftk eftk (AstLambda varf2 vf2)-                                   (AstLambda vard2 vd2)-                                   (AstLambda varr2 vr2)-                                   acc0From es-astMapAccumRDer k bftk eftk (AstLambda varf vf)-                            (AstLambda vard vd)-                            (AstLambda varr vr)-                acc0 (AstFromS' @esShsFrom _esShsFTK esFrom) =-  let accftk = ftkAst acc0-      accstk = ftkToSTK accftk-      esShsFrom = ftkAst esFrom-      esShsFromSTK = ftkToSTK esShsFrom-  in case razeSTK esShsFromSTK of-    (eftkFromSTK :: SingletonTK eyFrom) ->-      gcastWith (unsafeCoerceRefl-                 :: BuildTensorKind k eyFrom :~: esShsFrom) $-      let eftkFrom = razeFTK k eftkFromSTK esShsFrom-          ftkf2 = FTKProduct accftk eftkFrom-          varf2 =-            mkAstVarName ftkf2 (varNameToBounds varf) (varNameToAstVarId varf)-          astf2 = astVar varf2-          vf2 =-            let subbed =-                  substituteAst-                    (astPair (astProject1 astf2)-                             (astFromS' @eyFrom eftk (astProject2 astf2)))-                    varf vf-            in subbed-          ftkd2 = FTKProduct-                    (adFTK $ FTKProduct accftk eftkFrom)-                    (FTKProduct accftk eftkFrom)-          vard2 =-            mkAstVarName ftkd2 (varNameToBounds vard) (varNameToAstVarId vard)-          astd2 = astVar vard2-          vd2 =-            let subbed =-                  substituteAst-                    (astPair (astPair (astProject1 (astProject1 astd2))-                                      (astFromS' @(ADTensorKind eyFrom)-                                         (adFTK eftk)-                                         (astProject2 (astProject1 astd2))))-                             (astPair (astProject1 (astProject2 astd2))-                                      (astFromS' @eyFrom eftk-                                         (astProject2 (astProject2 astd2)))))-                    vard vd-            in subbed-          ftkr2 = FTKProduct-                    (adFTK $ FTKProduct accftk bftk)-                    (FTKProduct accftk eftkFrom)-          varr2 =-            mkAstVarName ftkr2 (varNameToBounds varr) (varNameToAstVarId varr)-          astr2 = astVar varr2-          vr2 =-            let subbed =-                  substituteAst-                    (astPair (astProject1 astr2)-                             (astPair (astProject1 (astProject2 astr2))-                                      (astFromS' @eyFrom eftk-                                         (astProject2 (astProject2 astr2)))))-                    varr vr-            in astSFrom' @(ADTensorKind (TKProduct accy ey))-                        (adSTK $ STKProduct accstk eftkFromSTK)-                        subbed-      in astMapAccumRDer k bftk eftkFrom (AstLambda varf2 vf2)-                                         (AstLambda vard2 vd2)-                                         (AstLambda varr2 vr2)-                                         acc0 esFrom-astMapAccumRDer k bftk eftk f df rf acc0 es =-  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es--astMapAccumLDer-  :: forall accy by ey k s. AstSpan s-  => SNat k-  -> FullShapeTK by-  -> FullShapeTK ey-  -> AstHFun s s-             (TKProduct accy ey) (TKProduct accy by)-  -> AstHFun s s-             (TKProduct (ADTensorKind (TKProduct accy ey))-                        (TKProduct accy ey))-             (ADTensorKind (TKProduct accy by))-  -> AstHFun s s-             (TKProduct (ADTensorKind (TKProduct accy by))-                        (TKProduct accy ey))-             (ADTensorKind (TKProduct accy ey))-  -> AstTensor AstMethodLet s accy-  -> AstTensor AstMethodLet s (BuildTensorKind k ey)-  -> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))-astMapAccumLDer k bftk eftk (AstLambda varf vf)-                            (AstLambda vard vd)-                            (AstLambda varr vr)-                (AstFromS' @accyFrom accftk acc0From) es =-  let accftkFrom = ftkAst acc0From-      accFromSTK = ftkToSTK accftkFrom-      ftkf2 = FTKProduct accftkFrom eftk-      varf2 = mkAstVarName ftkf2 (varNameToBounds varf) (varNameToAstVarId varf)-      astf2 = astVar varf2-      vf2 =-        let subbed =-              substituteAst-                (astPair (astFromS' @accyFrom accftk (astProject1 astf2))-                         (astProject2 astf2))-                varf vf-        in astSFrom' @(TKProduct accy by)-                    (STKProduct accFromSTK (ftkToSTK bftk))-                    subbed-      ftkd2 = FTKProduct-                (adFTK $ FTKProduct accftkFrom eftk)-                (FTKProduct accftkFrom eftk)-      vard2 = mkAstVarName ftkd2 (varNameToBounds vard) (varNameToAstVarId vard)-      astd2 = astVar vard2-      vd2 =-        let subbed =-              substituteAst-                (astPair (astPair (astFromS' @(ADTensorKind accyFrom)-                                     (adFTK accftk)-                                     (astProject1 (astProject1 astd2)))-                                  (astProject2 (astProject1 astd2)))-                         (astPair (astFromS' @accyFrom accftk-                                     (astProject1 (astProject2 astd2)))-                                  (astProject2 (astProject2 astd2))))-                vard vd-        in astSFrom' @(ADTensorKind (TKProduct accy by))-                    (adSTK $ STKProduct accFromSTK (ftkToSTK bftk))-                    subbed-      ftkr2 = FTKProduct-                (adFTK $ FTKProduct accftkFrom bftk)-                (FTKProduct accftkFrom eftk)-      varr2 = mkAstVarName ftkr2 (varNameToBounds varr) (varNameToAstVarId varr)-      astr2 = astVar varr2-      vr2 =-        let subbed =-              substituteAst-                (astPair (astPair (astFromS' @(ADTensorKind accyFrom)-                                     (adFTK accftk)-                                     (astProject1 (astProject1 astr2)))-                                  (astProject2 (astProject1 astr2)))-                         (astPair (astFromS' @accyFrom accftk-                                     (astProject1 (astProject2 astr2)))-                                  (astProject2 (astProject2 astr2))))-                varr vr-        in astSFrom' @(ADTensorKind (TKProduct accy ey))-                    (adSTK $ STKProduct accFromSTK (ftkToSTK eftk))-                    subbed-  in astFromS' @(TKProduct accyFrom (BuildTensorKind k by))-               (FTKProduct accftk (buildFTK k bftk))-     $ astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)-                                   (AstLambda vard2 vd2)-                                   (AstLambda varr2 vr2)-                                   acc0From es-astMapAccumLDer k bftk eftk (AstLambda varf vf)-                            (AstLambda vard vd)-                            (AstLambda varr vr)-                acc0 (AstFromS' @esShsFrom _esShsFTK esFrom) =-  let accftk = ftkAst acc0-      accstk = ftkToSTK accftk-      esShsFrom = ftkAst esFrom-      esShsFromSTK = ftkToSTK esShsFrom-  in case razeSTK esShsFromSTK of-    (eftkFromSTK :: SingletonTK eyFrom) ->-      gcastWith (unsafeCoerceRefl-                 :: BuildTensorKind k eyFrom :~: esShsFrom) $-      let eftkFrom = razeFTK k eftkFromSTK esShsFrom-          ftkf2 = FTKProduct accftk eftkFrom-          varf2 =-            mkAstVarName ftkf2 (varNameToBounds varf) (varNameToAstVarId varf)-          astf2 = astVar varf2-          vf2 =-            let subbed =-                  substituteAst-                    (astPair (astProject1 astf2)-                             (astFromS' @eyFrom eftk (astProject2 astf2)))-                    varf vf-            in subbed-          ftkd2 = FTKProduct-                    (adFTK $ FTKProduct accftk eftkFrom)-                    (FTKProduct accftk eftkFrom)-          vard2 =-            mkAstVarName ftkd2 (varNameToBounds vard) (varNameToAstVarId vard)-          astd2 = astVar vard2-          vd2 =-            let subbed =-                  substituteAst-                    (astPair (astPair (astProject1 (astProject1 astd2))-                                      (astFromS' @(ADTensorKind eyFrom)-                                         (adFTK eftk)-                                         (astProject2 (astProject1 astd2))))-                             (astPair (astProject1 (astProject2 astd2))-                                      (astFromS' @eyFrom eftk-                                         (astProject2 (astProject2 astd2)))))-                    vard vd-            in subbed-          ftkr2 = FTKProduct-                    (adFTK $ FTKProduct accftk bftk)-                    (FTKProduct accftk eftkFrom)-          varr2 =-            mkAstVarName ftkr2 (varNameToBounds varr) (varNameToAstVarId varr)-          astr2 = astVar varr2-          vr2 =-            let subbed =-                  substituteAst-                    (astPair (astProject1 astr2)-                             (astPair (astProject1 (astProject2 astr2))-                                      (astFromS' @eyFrom eftk-                                         (astProject2 (astProject2 astr2)))))-                    varr vr-            in astSFrom' @(ADTensorKind (TKProduct accy ey))-                        (adSTK $ STKProduct accstk eftkFromSTK)-                        subbed-      in astMapAccumLDer k bftk eftkFrom (AstLambda varf2 vf2)-                                         (AstLambda vard2 vd2)-                                         (AstLambda varr2 vr2)-                                         acc0 esFrom-astMapAccumLDer k bftk eftk f df rf acc0 es =-  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es--astApply :: forall x z s1 s. (AstSpan s1, AstSpan s)-         => AstHFun s1 s x z -> AstTensor AstMethodLet s1 x-         -> AstTensor AstMethodLet s z-astApply (AstLambda !var !v) u = astLet var u v--astCond :: AstSpan s-        => AstBool AstMethodLet-        -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y-        -> AstTensor AstMethodLet s y-astCond (AstBoolConst b) v w = if b then v else w-astCond (Ast.AstBoolNot b) v w = astCond b w v-astCond b (Ast.AstFromPrimal v) (Ast.AstFromPrimal w) =-  Ast.AstFromPrimal $ astCond b v w-astCond b (Ast.AstFromDual v) (Ast.AstFromDual w) =-  Ast.AstFromDual $ astCond b v w-astCond b v@(AstFromS' FTKScalar _) w = Ast.AstCond b v w--- We rely here on c and the other conversion being semantically equal.-astCond b (Ast.AstConvert c v) (AstFromS' _ w) =-  case matchingFTK (ftkAst v) (ftkAst w) of-    Just Refl -> astConvert c $ astCond b v w-    Nothing -> error "astCond: shapes don't match"-astCond b v w = Ast.AstCond b v w---- Invariant: if the variable has bounds, the expression can only have--- values within the bounds (regardless of what the `bounds` call would say).-astLet :: forall y z s s2. (AstSpan s, AstSpan s2)-       => AstVarName s y -> AstTensor AstMethodLet s y-       -> AstTensor AstMethodLet s2 z-       -> AstTensor AstMethodLet s2 z-astLet _var _u v@Ast.AstConcreteK{} = v-astLet _var _u v@Ast.AstConcreteS{} = v-astLet _var _u v@Ast.AstIotaS{} = v-astLet var u v@(Ast.AstVar var2) =-  if varNameToAstVarId var2 == varNameToAstVarId var-  then case sameAstSpan @s @s2 of-    Just Refl -> case testEquality var var2 of-      Just Refl -> u-      _ -> error "astLet: wrong variable types at AstVar"-    _ -> error "astLet: wrong span at AstVar"-  else v-astLet var u v@(Ast.AstPrimalPart (Ast.AstVar var2)) =  -- a common noop-  if varNameToAstVarId var2 == varNameToAstVarId var-  then case sameAstSpan @s @FullSpan of-    Just Refl -> case testEquality var var2 of-      Just Refl -> astPrimalPart u-      _ -> error "astLet: wrong variable types at AstPrimalPart"-    _ -> error "astLet: wrong span at AstPrimalPart"-  else v-astLet var u v@(Ast.AstDualPart (Ast.AstVar var2)) =  -- a noop-  if varNameToAstVarId var2 == varNameToAstVarId var-  then case sameAstSpan @s @FullSpan of-    Just Refl -> case testEquality var var2 of-      Just Refl -> astDualPart u-      _ -> error "astLet: wrong variable types at AstDualPart"-    _ -> error "astLet: wrong span at AstDualPart"-  else v-astLet var u v | astIsSmall True u =-  substituteAst u var v-astLet var (Ast.AstPair u1 u2) v =-  astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->-    substituteAst (Ast.AstPair ast1 ast2) var v-astLet var (Ast.AstFromPrimal (Ast.AstPair u1 u2)) v =-  astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->-    substituteAst (Ast.AstFromPrimal (Ast.AstPair ast1 ast2)) var v-astLet var (Ast.AstFromDual (Ast.AstPair u1 u2)) v =-  astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->-    substituteAst (Ast.AstFromDual (Ast.AstPair ast1 ast2)) var v-astLet var (Ast.AstLet varN uN (Ast.AstPair u1 u2)) v =-  astLet varN uN-  $ astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->-      substituteAst (Ast.AstPair ast1 ast2) var v-astLet var (Ast.AstFromPrimal (Ast.AstLet varN uN (Ast.AstPair u1 u2))) v =-  astLet varN uN-  $ astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->-      substituteAst (Ast.AstFromPrimal (Ast.AstPair ast1 ast2)) var v-astLet var (Ast.AstFromDual (Ast.AstLet varN uN (Ast.AstPair u1 u2))) v =-  astLet varN uN-  $ astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->-      substituteAst (Ast.AstFromDual (Ast.AstPair ast1 ast2)) var v--- This is a common case, e.g., from representing conditionals.-astLet var (Ast.AstFromVector snat stk u) v | V.length u == 2 =-  astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->-    substituteAst (Ast.AstFromVector snat stk $ fromList [ast1, ast2]) var v-astLet var (Ast.AstFromPrimal-              (Ast.AstFromVector snat stk u)) v | V.length u == 2 =-  astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->-    substituteAst (Ast.AstFromPrimal (Ast.AstFromVector snat stk-                                      $ fromList [ast1, ast2])) var v-astLet var (Ast.AstFromDual-              (Ast.AstFromVector snat stk u)) v | V.length u == 2 =-  astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->-    substituteAst (Ast.AstFromDual (Ast.AstFromVector snat stk-                                    $ fromList [ast1, ast2])) var v-astLet var (Ast.AstLet varN uN-              (Ast.AstFromVector snat stk u)) v | V.length u == 2 =-  astLet varN uN-  $ astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->-      substituteAst (Ast.AstFromVector snat stk $ fromList [ast1, ast2]) var v-astLet var (Ast.AstFromPrimal-              (Ast.AstLet varN uN-                 (Ast.AstFromVector snat stk u))) v | V.length u == 2 =-  astLet varN uN-  $ astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->-      substituteAst (Ast.AstFromPrimal (Ast.AstFromVector snat stk-                                        $ fromList [ast1, ast2])) var v-astLet var (Ast.AstFromDual-              (Ast.AstLet varN uN-                 (Ast.AstFromVector snat stk u))) v | V.length u == 2 =-  astLet varN uN-  $ astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->-      substituteAst (Ast.AstFromDual (Ast.AstFromVector snat stk-                                      $ fromList [ast1, ast2])) var v-astLet var (Ast.AstReplicate snat stk a) v =-  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)-      ast = Ast.AstReplicate snat stk $ astVar var2-  in astLet var2 a (substituteAst ast var v)-astLet var (Ast.AstFromPrimal (Ast.AstReplicate snat stk a)) v =-  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)-      ast = Ast.AstFromPrimal (Ast.AstReplicate snat stk $ astVar var2)-  in astLet var2 a (substituteAst ast var v)-astLet var (Ast.AstFromDual (Ast.AstReplicate snat stk a)) v =-  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)-      ast = Ast.AstFromDual (Ast.AstReplicate snat stk $ astVar var2)-  in astLet var2 a (substituteAst ast var v)-astLet var (Ast.AstTransposeS perm a) v =-  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)-      ast = Ast.AstTransposeS perm $ astVar var2-  in astLet var2 a (substituteAst ast var v)-astLet var (Ast.AstFromPrimal (Ast.AstTransposeS perm a)) v =-  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)-      ast = Ast.AstFromPrimal (Ast.AstTransposeS perm $ astVar var2)-  in astLet var2 a (substituteAst ast var v)-astLet var (Ast.AstFromDual (Ast.AstTransposeS perm a)) v =-  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)-      ast = Ast.AstFromDual (Ast.AstTransposeS perm $ astVar var2)-  in astLet var2 a (substituteAst ast var v)-astLet var u@(AstFromS' FTKScalar _) v = Ast.AstLet var u v-astLet var (Ast.AstConvert c a) v | checkAstFromS c a =-  let var2 =-        mkAstVarName (ftkAst a) (varNameToBounds var) (varNameToAstVarId var)-      ast = astConvert c $ astVar var2-  in astLet var2 a (substituteAst ast var v)-astLet var u (Ast.AstFromPrimal v0) = Ast.AstFromPrimal $ astLet var u v0-astLet var u (Ast.AstFromDual v0) = Ast.AstFromDual $ astLet var u v0-astLet var u v@(AstFromS' FTKScalar _) = Ast.AstLet var u v-astLet var u (Ast.AstConvert c v) | checkAstFromS c v =-  astConvert c $ astLet var u v-astLet var u v = Ast.AstLet var u v--astPrimalPart :: AstTensor AstMethodLet FullSpan y-              -> AstTensor AstMethodLet PrimalSpan y-astPrimalPart t = case t of-  Ast.AstPair t1 t2 -> astPair (astPrimalPart t1) (astPrimalPart t2)-  Ast.AstProject1 v -> astProject1 (astPrimalPart v)-  Ast.AstProject2 v -> astProject2 (astPrimalPart v)-  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map astPrimalPart l)-  Ast.AstSum snat stk v ->-    astSum snat stk (astPrimalPart v)-  Ast.AstReplicate snat stk v ->-    astReplicate snat stk (astPrimalPart v)-  {- Ast.AstMapAccumRDer k bftk eftk (AstLambda varf vf)-                                  (AstLambda vard vd)-                                  (AstLambda varr vr) acc0 es ->-    astMapAccumRDer k bftk eftk f df rf-                    (astPrimalPart acc0) (astPrimalPart es)-  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->-    astMapAccumLDer k bftk eftk f df rf-                    (astPrimalPart acc0) (astPrimalPart es) -}-  Ast.AstMapAccumRDer{} -> Ast.AstPrimalPart t  -- TODO-  Ast.AstMapAccumLDer{} -> Ast.AstPrimalPart t-  Ast.AstApply (AstLambda !var !v) ll ->-    astApply (AstLambda var (astPrimalPart v)) ll-  Ast.AstVar{} -> Ast.AstPrimalPart t  -- the only normal form-  Ast.AstCond b a2 a3 -> astCond b (astPrimalPart a2) (astPrimalPart a3)-  Ast.AstBuild1 k stk (var, v) ->-    let !v2 = astPrimalPart v-    in Ast.AstBuild1 k stk (var, v2)--  Ast.AstLet var u v -> astLet var u (astPrimalPart v)--  Ast.AstFromPrimal v -> v-  Ast.AstFromDual v ->-    let ftk = ftkAst v-    in astConcrete ftk (tdefTarget ftk)--  AstPlusK u v -> astPrimalPart u + astPrimalPart v-  AstTimesK u v -> astPrimalPart u * astPrimalPart v-  Ast.AstN1K NegateOp u -> negate (astPrimalPart u)-  Ast.AstN1K AbsOp u -> abs (astPrimalPart u)-  Ast.AstN1K SignumOp u -> signum (astPrimalPart u)-  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astPrimalPart u)-  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (astPrimalPart u) (astPrimalPart v)-  Ast.AstI2K QuotOp u v -> quotH (astPrimalPart u) (astPrimalPart v)-  Ast.AstI2K RemOp u v -> remH (astPrimalPart u) (astPrimalPart v)-  Ast.AstCastK v -> astCastK $ astPrimalPart v--  AstPlusS u v -> astPrimalPart u + astPrimalPart v-  AstTimesS u v -> astPrimalPart u * astPrimalPart v-  Ast.AstN1S NegateOp u -> negate (astPrimalPart u)-  Ast.AstN1S AbsOp u -> abs (astPrimalPart u)-  Ast.AstN1S SignumOp u -> signum (astPrimalPart u)-  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astPrimalPart u)-  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astPrimalPart u)-                                             (astPrimalPart v)-  Ast.AstI2S QuotOp u v -> quotH (astPrimalPart u) (astPrimalPart v)-  Ast.AstI2S RemOp u v -> remH (astPrimalPart u) (astPrimalPart v)-  Ast.AstCastS v -> astCastS $ astPrimalPart v--  Ast.AstIndexS shn v ix ->-    astIndexS shn (astPrimalPart v) ix-  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->-    astScatterS @shm @shn @shp shn (astPrimalPart v) (vars, ix)-  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->-    astGatherS @shm @shn @shp shn (astPrimalPart v) (vars, ix)-  Ast.AstAppendS x y -> astAppendS (astPrimalPart x) (astPrimalPart y)-  Ast.AstSliceS i n k v -> astSliceS i n k (astPrimalPart v)-  Ast.AstReverseS v -> astReverseS (astPrimalPart v)-  Ast.AstTransposeS perm v -> astTransposeS perm (astPrimalPart v)-  Ast.AstReshapeS sh v -> astReshapeS sh (astPrimalPart v)--  -- All conversions need to stay down here to cancel out.-  Ast.AstConvert{} -> Ast.AstPrimalPart t--  -- These should not appear in this context unless via wacky tests.-  Ast.AstSum0S{} -> Ast.AstPrimalPart t-  Ast.AstDot0S{} -> Ast.AstPrimalPart t-  Ast.AstDot1InS{} -> Ast.AstPrimalPart t-  Ast.AstMatmul2S{} -> Ast.AstPrimalPart t---- Note how this can't be pushed down into, say, multiplication, because it--- multiplies the dual part by the primal part. Addition is fine, though.-astDualPart :: AstTensor AstMethodLet FullSpan y-            -> AstTensor AstMethodLet DualSpan y-astDualPart t = case t of-  Ast.AstPair t1 t2 -> astPair (astDualPart t1) (astDualPart t2)-  Ast.AstProject1 v -> astProject1 (astDualPart v)-  Ast.AstProject2 v -> astProject2 (astDualPart v)-  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map astDualPart l)-  Ast.AstSum snat stk v ->-    astSum snat stk (astDualPart v)-  Ast.AstReplicate snat stk v ->-    astReplicate snat stk (astDualPart v)-  {- Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->-    astMapAccumRDer k bftk eftk f df rf-                    (astDualPart acc0) (astDualPart es)-  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->-    astMapAccumLDer k bftk eftk f df rf-                    (astDualPart acc0) (astDualPart es) -}-  Ast.AstMapAccumRDer{} -> Ast.AstDualPart t  -- TODO-  Ast.AstMapAccumLDer{} -> Ast.AstDualPart t-  Ast.AstApply (AstLambda !var !v) ll ->-    astApply (AstLambda var (astDualPart v)) ll-  Ast.AstVar{} -> Ast.AstDualPart t-  Ast.AstCond b a2 a3 -> astCond b (astDualPart a2) (astDualPart a3)-  Ast.AstBuild1 k stk (var, v) ->-    let !v2 = astDualPart v-    in Ast.AstBuild1 k stk (var, v2)--  Ast.AstLet var u v -> astLet var u (astDualPart v)--  Ast.AstFromPrimal v ->-    let ftk = ftkAst v-    in Ast.AstDualPart $ Ast.AstFromPrimal-       $ astConcrete ftk (tdefTarget ftk)-           -- let's hope this is smaller than v-  Ast.AstFromDual v -> v--  AstPlusK u v -> astDualPart u + astDualPart v-  -- This one is mathematically wrong, dual numbers don't mult like that:-  -- AstTimesK u v -> astDualPart u * astDualPart v-  Ast.AstN1K NegateOp u -> negate (astDualPart u)-  {- Some of these are wrong, so let's be conservative:-  Ast.AstN1K AbsOp u -> abs (astDualPart u)-  Ast.AstN1K SignumOp u -> signum (astDualPart u)-  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astDualPart u)-  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (astDualPart u) (astDualPart v)-  Ast.AstI2K QuotOp u v -> quotH (astDualPart u) (astDualPart v)-  Ast.AstI2K RemOp u v -> remH (astDualPart u) (astDualPart v)-  -}-  Ast.AstCastK v -> astCastK $ astDualPart v--  AstPlusS u v -> astDualPart u + astDualPart v-  -- This one is mathematically wrong, dual numbers don't mult like that:-  -- AstTimesS u v -> astDualPart u * astDualPart v-  Ast.AstN1S NegateOp u -> negate (astDualPart u)-  {- Some of these are wrong, so let's be conservative:-  Ast.AstN1S AbsOp u -> abs (astDualPart u)-  Ast.AstN1S SignumOp u -> signum (astDualPart u)-  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astDualPart u)-  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astDualPart u)-                                             (astDualPart v)-  Ast.AstI2S opCode u v -> Ast.AstI2S opCode (astDualPart u)-                                             (astDualPart v)-  -}-  Ast.AstCastS v -> astCastS $ astDualPart v--  Ast.AstIndexS shn v ix ->-    astIndexS shn (astDualPart v) ix-  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->-    astScatterS @shm @shn @shp shn (astDualPart v) (vars, ix)-  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->-    astGatherS @shm @shn @shp shn (astDualPart v) (vars, ix)-  Ast.AstAppendS x y -> astAppendS (astDualPart x) (astDualPart y)-  Ast.AstSliceS i n k v -> astSliceS i n k (astDualPart v)-  Ast.AstReverseS v -> astReverseS (astDualPart v)-  Ast.AstTransposeS perm v -> astTransposeS perm (astDualPart v)-  Ast.AstReshapeS sh v -> astReshapeS sh (astDualPart v)--  -- All conversions need to stay down here to cancel out.-  Ast.AstConvert{} -> Ast.AstDualPart t--  -- These should not appear in this context unless via wacky tests.-  Ast.AstSum0S{} -> Ast.AstDualPart t-  Ast.AstDot0S{} -> Ast.AstDualPart t-  Ast.AstDot1InS{} -> Ast.AstDualPart t-  Ast.AstMatmul2S{} -> Ast.AstDualPart t--  _ -> Ast.AstDualPart t--astConcreteK :: GoodScalar r-             => Concrete (TKScalar r)-             -> AstTensor AstMethodLet PrimalSpan (TKScalar r)-astConcreteK = AstConcreteK . unConcrete--astFloorK :: (GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2)-          => AstTensor AstMethodLet PrimalSpan (TKScalar r1)-          -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)-astFloorK t = case t of-  Ast.AstLet var u v -> astLet var u (astFloorK v)-  -- This increases work and term size, because conditional is eager.-  -- Ast.AstCond b a2 a3 -> Ast.AstCond b (astFloorK a2) (astFloorK a3)-  -- These values are small, so we can simplify them ASAP.-  AstConcreteK k -> astConcreteK (tkfloor $ Concrete k)-  Ast.AstFloorK v -> astFloorK v-  Ast.AstFromIntegralK v -> astFromIntegralK v-  Ast.AstCastK v -> astFloorK v-  _ -> Ast.AstFloorK t---- Beware that increasing the number of calls to this constructor--- sometimes increases runtime, because not enough copies cancel out.--- Hence the commented out rules below.-astFromIntegralK :: forall r1 r2. (GoodScalar r1, GoodScalar r2, Integral r1)-                 => AstTensor AstMethodLet PrimalSpan (TKScalar r1)-                 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)-astFromIntegralK t = case t of-  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t-  Ast.AstLet var u v -> astLet var u (astFromIntegralK v)-  AstConcreteK k -> astConcreteK (tkfromIntegral $ Concrete k)-  Ast.AstN1K NegateOp u -> negate (astFromIntegralK u)-  Ast.AstN1K AbsOp u -> abs (astFromIntegralK u)-  Ast.AstN1K SignumOp u -> signum (astFromIntegralK u)-  Ast.AstFromIntegralK v -> astFromIntegralK v-  _ -> Ast.AstFromIntegralK t--astCastK :: forall r1 r2 s.-            (GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2, AstSpan s)-         => AstTensor AstMethodLet s (TKScalar r1)-         -> AstTensor AstMethodLet s (TKScalar r2)-astCastK t = case t of-  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t-  Ast.AstLet var u v -> astLet var u (astCastK v)-  AstConcreteK k -> astConcreteK (tkcast $ Concrete k)-  -- TODO: which should go deeper, casts or fromPrimal? Or maybe alternate-  -- in different phases to make sure both can cancel out?-  -- Rethink. For now, astFromPrimalis not called, to avoid loops.-  -- The same with many others-  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astCastK v-  Ast.AstFromDual v -> Ast.AstFromDual $ astCastK v-  Ast.AstN1K NegateOp u -> negate (astCastK u)-  Ast.AstN1K AbsOp u -> abs (astCastK u)-  Ast.AstN1K SignumOp u -> signum (astCastK u)---  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astCastK u)-  Ast.AstFromIntegralK v -> astFromIntegralK v-  Ast.AstCastK v -> astCastK v-  _ -> Ast.AstCastK t--astConcreteS :: GoodScalar r-             => Concrete (TKS sh r)-             -> AstTensor AstMethodLet PrimalSpan (TKS sh r)-astConcreteS = AstConcreteS . unConcrete--astFloorS :: forall r1 r2 sh.-             (GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2)-          => AstTensor AstMethodLet PrimalSpan (TKS sh r1)-          -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)-astFloorS t = case t of-  _ | FTKS (snat :$$ sh2) _ <- ftkAst t-    , Just u <- unRepl1 t ->-      astReplicate snat (STKS sh2 STKScalar) (astFloorS u)-  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->-    Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, astFloorS v)-  Ast.AstBuild1 snat STKScalar (var, v) ->-    Ast.AstBuild1 snat STKScalar (var, astFloorK v)-  Ast.AstLet var u v -> astLet var u (astFloorS v)-  Ast.AstScatterS shn v (vars, ix) ->-    astScatterS shn (astFloorS v) (vars, ix)-  Ast.AstGatherS shn v (vars, ix) ->-    astGatherS shn (astFloorS v) (vars, ix)-  Ast.AstIotaS snat -> Ast.AstIotaS snat-  Ast.AstReverseS v -> astReverseS (astFloorS v)-  Ast.AstTransposeS perm v -> astTransposeS perm (astFloorS v)-  Ast.AstReshapeS sh v -> astReshapeS sh (astFloorS v)-  AstSFromK' a -> astSFromK' (astFloorK a)-  Ast.AstFloorS v -> astFloorS v-  Ast.AstFromIntegralS v -> astFromIntegralS v-  Ast.AstCastS v -> astFloorS v-  _ -> Ast.AstFloorS t--astFromIntegralS :: forall r1 r2 sh. (GoodScalar r1, GoodScalar r2, Integral r1)-                 => AstTensor AstMethodLet PrimalSpan (TKS sh r1)-                 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)-astFromIntegralS t = case t of-  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t-  _ | FTKS (snat :$$ sh2) _ <- ftkAst t-    , Just u <- unRepl1 t ->-      astReplicate snat (STKS sh2 STKScalar) (astFromIntegralS u)---  Ast.AstFromVector snat (STKS sh STKScalar) l ->---   astFromVector snat (STKS sh STKScalar) (V.map astFromIntegralS l)---  Ast.AstFromVector snat STKScalar l ->---   astFromVector snat STKScalar (V.map astFromIntegralK l)-  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->-    Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, astFromIntegralS v)-  Ast.AstBuild1 snat STKScalar (var, v) ->-    Ast.AstBuild1 snat STKScalar (var, astFromIntegralK v)-  Ast.AstLet var u v -> astLet var u (astFromIntegralS v)-  Ast.AstN1S NegateOp u -> negate (astFromIntegralS u)-  Ast.AstN1S AbsOp u -> abs (astFromIntegralS u)-  Ast.AstN1S SignumOp u -> signum (astFromIntegralS u)-  Ast.AstFromIntegralS v -> astFromIntegralS v---  Ast.AstIndexS shn v ix -> astIndexS shn (astFromIntegralS v) ix-    -- increases work; also index goes into fromIntegral, so we'd loop-  Ast.AstScatterS shn v (vars, ix) ->-    astScatterS shn (astFromIntegralS v) (vars, ix)-  Ast.AstGatherS shn v (vars, ix) ->-    astGatherS shn (astFromIntegralS v) (vars, ix)-  Ast.AstIotaS snat -> Ast.AstIotaS snat---  Ast.AstSliceS i n k v -> astSliceS i n k (astFromIntegralS v)-  Ast.AstReverseS v -> astReverseS (astFromIntegralS v)-  Ast.AstTransposeS perm v -> astTransposeS perm (astFromIntegralS v)-  Ast.AstReshapeS sh v -> astReshapeS sh (astFromIntegralS v)-  AstSFromK' a -> astSFromK' (astFromIntegralK a)-  _ -> Ast.AstFromIntegralS t--astCastS :: forall r1 r2 s sh.-            (GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2, AstSpan s)-         => AstTensor AstMethodLet s (TKS sh r1)-         -> AstTensor AstMethodLet s (TKS sh r2)-astCastS t = case t of-  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t-  _ | FTKS (snat :$$ sh2) _ <- ftkAst t-    , Just u <- unRepl1 t ->-      astReplicate snat (STKS sh2 STKScalar) (astCastS u)---  Ast.AstFromVector snat (STKS sh STKScalar) l ->---   astFromVector snat (STKS sh STKScalar) (V.map astCastS l)---  Ast.AstFromVector snat STKScalar l ->---   astFromVector snat STKScalar (V.map astCastK l)-  {- This (and other similar rules) is bad, because it changes semantics-     and also impacts performance negatively (a is larger than sum a):-  Ast.AstSum snat (STKS sh STKScalar) a ->-    astSum snat (STKS sh STKScalar) (astCastS a) -}-  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->-    Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, astCastS v)-  Ast.AstBuild1 snat STKScalar (var, v) ->-    Ast.AstBuild1 snat STKScalar (var, astCastK v)-  Ast.AstLet var u v -> astLet var u (astCastS v)-  Ast.AstPrimalPart a -> Ast.AstPrimalPart $ astCastS a-  Ast.AstDualPart a -> Ast.AstDualPart $ astCastS a-  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astCastS v-  Ast.AstFromDual v -> Ast.AstFromDual $ astCastS v-  Ast.AstN1S NegateOp u -> negate (astCastS u)-  Ast.AstN1S AbsOp u -> abs (astCastS u)-  Ast.AstN1S SignumOp u -> signum (astCastS u)-  Ast.AstFromIntegralS v -> astFromIntegralS v-  Ast.AstCastS v -> astCastS v---  Ast.AstIndexS shn v ix -> astIndexS shn (astCastS v) ix-    -- increases work; also index goes into fromIntegral, so we'd loop-  Ast.AstScatterS shn v (vars, ix) -> astScatterS shn (astCastS v) (vars, ix)-  Ast.AstGatherS shn v (vars, ix) -> astGatherS shn (astCastS v) (vars, ix)---  Ast.AstMinIndexS v -> Ast.AstMinIndexS (astCastS v)-  Ast.AstIotaS snat -> Ast.AstIotaS snat---  Ast.AstSliceS i n k v -> astSliceS i n k (astCastS v)-  Ast.AstReverseS v -> astReverseS (astCastS v)-  Ast.AstTransposeS perm v -> astTransposeS perm (astCastS v)-  Ast.AstReshapeS sh v -> astReshapeS sh (astCastS v)-  AstSFromK' a -> astSFromK' (astCastK a)-  _ -> Ast.AstCastS t--astIndexS-  :: forall shm shn s r. AstSpan s-  => ShS shn-  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r) -> AstIxS AstMethodLet shm-  -> AstTensor AstMethodLet s (TKS2 shn r)-astIndexS = astIndexKnobsS defaultKnobs--astIndexKnobsS-  :: forall shm shn s r. AstSpan s-  => SimplifyKnobs-  -> ShS shn-  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)-  -> AstIxS AstMethodLet shm-  -> AstTensor AstMethodLet s (TKS2 shn r)-astIndexKnobsS _ _ v0 ZIS = v0-astIndexKnobsS _ shn v0 (i1 :.$ _)-  | let (lb, ub) = bounds i1--- this doesn't work in GHC 9.10:---      FTKS (snat :$$ _) x = ftkAst v0-  , FTKS (snat :$$ _) x <- ftkAst v0-  , ub < 0 || lb >= fromInteger (fromSNat snat) =-    let ftk = FTKS shn x-    in fromPrimal $ astConcrete ftk (tdefTarget ftk)-astIndexKnobsS knobs shn v0 (Ast.AstCond b i1 i2 :.$ rest0)-  | knobPhase knobs `notElem` [PhaseUnspecified, PhaseVectorization] =-      -- don't undo vectorization tweaks-    astLetFun v0 $ \v ->-    shareIx rest0 $ \rest ->-      astCond b (astIndexKnobsS knobs shn v (i1 :.$ rest))-                (astIndexKnobsS knobs shn v (i2 :.$ rest))-astIndexKnobsS knobs shn v0 ix@((:.$) @in1 @shm1 i1 rest1) =- let FTKS _ x = ftkAst v0-     astIndex-       :: forall shm' shn' s' r'. AstSpan s'-       => ShS shn'-       -> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') r')-       -> AstIxS AstMethodLet shm'-       -> AstTensor AstMethodLet s' (TKS2 shn' r')-     astIndex shn' v2 ix2 = astIndexKnobsS (deVect knobs) shn' v2 ix2-     astGather-       :: forall shm' shn' shp'.-          ShS shn'-       -> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') r)-       -> (AstVarListS shm', AstIxS AstMethodLet shp')-       -> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') r)-     astGather shn' v2 (vars2, ix2) =-       astGatherKnobsS @shm' @shn' @shp' (deVect knobs) shn' v2 (vars2, ix2)- in case v0 of-  Ast.AstProject1{} -> Ast.AstIndexS shn v0 ix-  Ast.AstProject2{} -> Ast.AstIndexS shn v0 ix-  Ast.AstFromVector _ STKS{} l | AstConcreteK it <- i1 ->-    let i = fromIntegral it-    in astIndex shn (l V.! i) rest1-  Ast.AstFromVector _ STKScalar l | AstConcreteK it <- i1, ZIS <- rest1 ->-    let i = fromIntegral it-    in astSFromK' (l V.! i)-  Ast.AstFromVector{} | ZIS <- rest1 ->  -- normal form-    Ast.AstIndexS shn v0 ix-  Ast.AstFromVector snat STKS{} l ->-    shareIx rest1 $ \ !rest2 ->-      Ast.AstIndexS @'[in1] @shn shn (astFromVector snat (STKS shn (ftkToSTK x))-                                      $ V.map (\a -> astIndex shn a rest2) l)-                    (i1 :.$ ZIS)-  Ast.AstSum snat@(SNat @n1) STKS{} v ->-    let perm3 = backpermCycle $ shsLength (shsFromIxS ix) + 1-    in Permutation.permFromList perm3 $ \(perm :: Permutation.Perm perm3P) ->-         gcastWith (unsafeCoerceRefl-                    :: Compare (Rank perm3P) (Rank (n1 : shm ++ shn))-                       :~: LT) $-         gcastWith (unsafeCoerceRefl-                    :: Permutation.PermutePrefix perm3P (n1 : (shm ++ shn))-                       :~: shm ++ (n1 : shn)) $-         fromMaybe (error "astIndexKnobsS: impossible non-permutation")-         $ Permutation.permCheckPermutation perm-         $ astSum snat (STKS shn (ftkToSTK x))-         $ astIndex @shm @(n1 : shn) (snat :$$ shn)-                    (astTransposeS @perm3P @(n1 : shm ++ shn) perm v)-                    ix-  Ast.AstReplicate _ STKS{} v ->-    let ftk = FTKS shn x-        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-    in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of-      AstBoolConst b -> if b then astIndex shn v rest1 else defArr-      _ -> Ast.AstIndexS shn v0 ix-  Ast.AstReplicate _ STKScalar v | ZIS <- rest1 ->-    let ftk = FTKS shn x-        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-    in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of-      AstBoolConst b -> if b then astSFromK' v else defArr-      _ -> Ast.AstIndexS shn v0 ix-  Ast.AstApply{} -> Ast.AstIndexS shn v0 ix-  Ast.AstVar{} -> Ast.AstIndexS shn v0 ix-  Ast.AstCond b v w ->-    shareIx ix $ \ !ix2 ->-      astCond b (astIndexKnobsS knobs shn v ix2)-                (astIndexKnobsS knobs shn w ix2)-{- This is wrong: in a counterfactual case, astLet assigns OOB i to var2,-   violating the invariant about variables bounds:-  Ast.AstBuild1 (SNat @k) STKS{} (var2, v) ->-    let ftk = FTKS shn x-        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-    in astLetFunB i1 $ \i ->-    astCond (0 <=. i &&* i <=. valueOf @k - 1)-            (astIndex shn (astLet var2 i v) rest1)-            defArr -}-  Ast.AstBuild1 (SNat @k) STKS{} (var2, v) ->-    let ftk = FTKS shn x-        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-    in case 0 <=. i1 &&* i1 <=. valueOf @k - 1 of-      AstBoolConst b ->-        if b then astIndex shn (astLet var2 i1 v) rest1 else defArr-      _ -> Ast.AstIndexS shn v0 ix-  Ast.AstBuild1 (SNat @k) STKScalar (var2, v) | ZIS <- rest1 ->-    let ftk = FTKS shn x-        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-    in case 0 <=. i1 &&* i1 <=. valueOf @k - 1 of-      AstBoolConst b ->-        if b then astSFromK' $ astLet var2 i1 v else defArr-      _ -> Ast.AstIndexS shn v0 ix--  Ast.AstLet var u v -> astLet var u (astIndexKnobsS knobs shn v ix)--  Ast.AstPrimalPart{} -> Ast.AstIndexS shn v0 ix  -- must be a NF-  Ast.AstDualPart{} -> Ast.AstIndexS shn v0 ix-  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astIndexKnobsS knobs shn v ix-  Ast.AstFromDual v -> Ast.AstFromDual $ astIndexKnobsS knobs shn v ix--  AstPlusS u v ->-    shareIx ix $ \ !ix2 ->-    astIndexKnobsS knobs shn u ix2 + astIndexKnobsS knobs shn v ix2-  AstTimesS u v ->-    shareIx ix $ \ !ix2 ->-    astIndexKnobsS knobs shn u ix2 * astIndexKnobsS knobs shn v ix2-  Ast.AstN1S NegateOp u -> negate (astIndexKnobsS knobs shn u ix)-  Ast.AstN1S AbsOp u -> abs (astIndexKnobsS knobs shn u ix)-  Ast.AstN1S SignumOp u -> signum (astIndexKnobsS knobs shn u ix)-  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astIndexKnobsS knobs shn u ix)-  Ast.AstR2S opCode u v ->-    shareIx ix-    $ \ !ix2 -> Ast.AstR2S opCode (astIndexKnobsS knobs shn u ix2)-                                  (astIndexKnobsS knobs shn v ix2)-  Ast.AstI2S opCode u v ->-    shareIx ix-    $ \ !ix2 -> Ast.AstI2S opCode (astIndexKnobsS knobs shn u ix2)-                                  (astIndexKnobsS knobs shn v ix2)-  AstConcreteS a | AstConcreteK i <- i1 ->-    let u = withKnownShS (shsFromIxS rest1 `shsAppend` shn) $-            tsindex (Concrete a) (Concrete i :.$ ZIS)-    in astIndex shn (astConcreteS u) rest1-  AstConcreteS{} -> case unRepl1 v0 of-    Just u ->-      let ftk = FTKS shn x-          defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-      in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of-        AstBoolConst b -> if b then astIndex shn u rest1 else defArr-        _ -> Ast.AstIndexS shn v0 ix-    _ -> Ast.AstIndexS shn v0 ix-  Ast.AstFloorS v -> astFloorS $ astIndexKnobsS knobs shn v ix-  Ast.AstFromIntegralS v -> astFromIntegralS $ astIndexKnobsS knobs shn v ix-  Ast.AstCastS t -> astCastS $ astIndexKnobsS knobs shn t ix--  Ast.AstIndexS _ v (ix2 :: AstIxS AstMethodLet sh4)-    | Refl <- lemAppAssoc (Proxy @sh4) (Proxy @shm) (Proxy @shn) ->-      astIndexKnobsS knobs shn v (ix2 `ixsAppend` ix)-  Ast.AstScatterS @shm7 @shn7 @shp7 shn7 v (vars, AstIntVar var5 :.$ ix2)-    | AstIntVar var6 <- i1, var6 == var5 ->-        astIndex shn (astScatterS @shm7 @shn7 @(Tail shp7)-                                  shn7 v (vars, ix2)) rest1-  Ast.AstScatterS @shm7 @shn7 @shp7 shn7-                  v (vars, AstConcreteK i5 :.$ ix2)-    | AstConcreteK i6 <- i1 ->-        if i6 == i5-        then astIndex shn (astScatterS @shm7 @shn7 @(Tail shp7)-                                       shn7 v (vars, ix2)) rest1-        else let ftk = FTKS shn x-             in fromPrimal $ astConcrete ftk (tdefTarget ftk)-  -- AstScatter sh v (vars2, ZIR) ->-  --   AstScatter sh (astIndex (astTranspose perm3 v) ix) (vars2, ZIR)-  Ast.AstScatterS{} ->  -- normal form-    Ast.AstIndexS shn v0 ix-  -- This is not a possible normal form, but pattern needs to be exhaustive.-  Ast.AstGatherS @_ @_ @shp' _shn' v (ZS, ix2) ->-    gcastWith (unsafeCoerceRefl-              :: shp' ++ (in1 ': shm1) ++ shn-                 :~: shp' ++ (in1 ': shm1 ++ shn)) $-    astIndex @(shp' ++ shm) @shn shn v (ix2 `ixsAppend` ix)-  Ast.AstGatherS @_ @shn' @shp' shn'-                 v ((::$) @m71 @shm71 (Const var2) vars, ix2) ->-    gcastWith (unsafeCoerceRefl :: shm71 ++ shn' :~: shm1 ++ shn) $-    let ftk = FTKS shn x-        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-        w :: AstTensor AstMethodLet s (TKS2 (shm1 ++ shn) r)-        w = astGather @shm71 @shn' @shp' shn' v (vars, ix2)-        u = astLet var2 i1 $ astIndex @shm1 @shn shn w rest1-          -- this let makes it impossible to use astCond when i1 is OOB-    in case 0 <=. i1 &&* i1 <=. valueOf @m71 - 1 of-      AstBoolConst b -> if b then u else defArr-      _ -> Ast.AstIndexS shn v0 ix-  Ast.AstMinIndexS @n1 @shz v -> case ftkAst v of-    FTKS nsh _ -> case shsLast nsh of-     nl@(SNat @nl) ->-      let shnl = shn `shsAppend` (nl :$$ ZSS)-      in gcastWith (unsafeCoerceRefl-                    :: Permutation.Index 0 (shn ++ '[nl])-                       ': Drop 1 (shn ++ '[nl])-                       :~: shn ++ '[nl]) $-         gcastWith (unsafeCoerceRefl-                    :: Init (shn ++ '[nl]) :~: shn) $-         gcastWith (unsafeCoerceRefl-                    :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) $-         Ast.AstMinIndexS @(Head (shn ++ '[nl]))-                          @(Tail (shn ++ '[nl]))-         $ astIndexKnobsS @shm @(shn ++ '[nl]) knobs shnl v ix-  Ast.AstMaxIndexS @n1 @shz v -> case ftkAst v of-    FTKS nsh _ -> case shsLast nsh of-     nl@(SNat @nl) ->-      let shnl = shn `shsAppend` (nl :$$ ZSS)-      in gcastWith (unsafeCoerceRefl-                    :: Permutation.Index 0 (shn ++ '[nl])-                       ': Drop 1 (shn ++ '[nl])-                       :~: shn ++ '[nl]) $-         gcastWith (unsafeCoerceRefl-                    :: Init (shn ++ '[nl]) :~: shn) $-         gcastWith (unsafeCoerceRefl-                    :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) $-         Ast.AstMaxIndexS @(Head (shn ++ '[nl]))-                          @(Tail (shn ++ '[nl]))-         $ astIndexKnobsS @shm @(shn ++ '[nl]) knobs shnl v ix-  Ast.AstIotaS (SNat @k) -> case testEquality shn ZSS of-    Just Refl ->-      let ftk = FTKS ZSS x-          defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-      in astLetFunB i1 $ \i ->-        astCond (0 <=. i &&* i <=. valueOf @k - 1)-                (astFromIntegralS $ astSFromK' i)-                defArr-    _ -> error "astIndexKnobsS: shape not []"-  Ast.AstAppendS u v | FTKS (SNat @m :$$ _) _ <- ftkAst u ->-    astLetFunB i1 $ \i ->-    shareIx rest1 $ \ !rest2 ->-    let ulen = valueOf @m-        ix1 = i :.$ rest2-        ix2 = i - ulen :.$ rest2-    in case ulen <=. i of-      AstBoolConst b -> if b then astIndex shn v ix2 else astIndex shn u ix1-      bExpr ->-        -- This results in a larger term, so we consider this late.-        if knobPhase knobs == PhaseExpansion-        then astCond bExpr (astIndex shn v ix2) (astIndex shn u ix1)-        else Ast.AstIndexS shn v0 ix-  Ast.AstSliceS i@(SNat @i) (SNat @n) k@SNat v ->-    astLetFunB i1 $ \iShared ->-    let ftk = FTKS shn x-        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-        b = (if sNatValue i == 0 then true else 0 <=. iShared )-            &&* (if sNatValue k == 0 then true else iShared <=. valueOf @n - 1)-        ii = valueOf @i + iShared-    in astCond b (astIndex shn v (ii :.$ rest1)) defArr-  Ast.AstReverseS v ->-    let iRev = valueOf @in1 - 1 - i1-    in astIndex shn v (iRev :.$ rest1)-  Ast.AstTransposeS @_ @sh2 perm v-    | gcompare (shsRank (shsFromIxS ix)) (Permutation.permRank perm) /= GLT ->-      -- TODO: remake once there's an S version of permInverse:-      permInverse perm $ \(permR :: Permutation.Perm permR) _ ->-      let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix permR shm)-          ix2 = ixsPermutePrefix permR ix-      in gcastWith (unsafeCoerceRefl-                    :: sh2 :~: Permutation.PermutePrefix permR shm ++ shn) $-         astIndex @(Permutation.PermutePrefix permR shm) shn v ix2-  Ast.AstTransposeS @perm perm v-    | knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion] ->-      astIndex shn (astTransposeAsGatherS @perm (deVect knobs) perm v) ix-  Ast.AstTransposeS{} -> Ast.AstIndexS shn v0 ix-  -- This results in a larger term, so we consider this late.-  Ast.AstReshapeS sh v | knobPhase knobs == PhaseExpansion-                         || shsLength sh <= 1 ->-    astIndex shn (astReshapeAsGatherS (deVect knobs) sh v) ix-  Ast.AstReshapeS{} -> Ast.AstIndexS shn v0 ix--  AstFromS' ftkz v -> case matchingFTK (ftkAst v) ftkz of-    Just Refl -> astIndexKnobsS knobs shn v ix-      -- rare, usually simplifies away earlier-    Nothing -> error "astIndexKnobsS: wrong tensor kinds in AstFromS"-  -- These conversions need to stay down, so this is NF, see vectorization.-  Ast.AstConvert{} -> Ast.AstIndexS shn v0 ix--  -- These should not appear here unless via wacky tests.-  Ast.AstDot1InS{} -> Ast.AstIndexS shn v0 ix-  Ast.AstMatmul2S{} -> Ast.AstIndexS shn v0 ix---- TODO: compared to tletIx, it adds many lets, not one, but does not--- create other (and non-simplified!) big terms and also uses astIsSmall,--- so it's probably more efficient. Use this instead of tletIx--- or design something even better.------ See https://github.com/Mikolaj/horde-ad/issues/119.-shareIx :: AstSpan s-        => AstIxS AstMethodLet shm-        -> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)-        -> AstTensor AstMethodLet s y-{-# NOINLINE shareIx #-}-shareIx ix f = unsafePerformIO $ do-  let shareI :: AstInt AstMethodLet-             -> IO ( Maybe (IntVarName, AstInt AstMethodLet)-                   , AstInt AstMethodLet )-      shareI i | astIsSmall True i = return (Nothing, i)-      shareI i =-        -- i can be OOB, so we can't use shape to determine its bounds-        let bds = bounds i-        in funToAstIntVarIO (Just bds) $ \ (!varFresh, !astVarFresh) ->-                                           (Just (varFresh, i), astVarFresh)-  (bindings, ix2) <- mapAndUnzipM shareI (Foldable.toList ix)-  return $! foldr (uncurry astLet)-                  (withKnownShS (shsFromIxS ix) $ f $ fromList ix2)-                  (catMaybes bindings)---- TODO: fuse scatters, scatter and sum, and perhaps more (fromList?)-astScatterS :: forall shm shn shp r s. AstSpan s-            => ShS shn-            -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)-            -> (AstVarListS shm, AstIxS AstMethodLet shp)-            -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)-astScatterS _shn v (ZS, ZIS) = v-astScatterS shn v0 (_vars, ix@((:.$) @k i1 _))-  | let (lb, ub) = bounds i1-        FTKS _ x = ftkAst v0-  , ub < 0 || lb >= valueOf @k =-    let ftk = FTKS (shsFromIxS ix `shsAppend` shn) x-    in fromPrimal $ astConcrete ftk (tdefTarget ftk)-astScatterS shn v (vars, (:.$) @k (AstConcreteK _) rest)-  | Just Refl <- sameNat (SNat @k) (SNat @1)-  , FTKS _ x <- ftkAst v =-    astReplicate (SNat @1) (STKS (shsFromIxS rest `shsAppend` shn) (ftkToSTK x))-    $ astScatterS shn v (vars, rest)-astScatterS shn v (Const var ::$ (vars :: AstVarListS sh3), ix)-  | not $ var `varNameInIxS` ix-  , FTKS _ x <- ftkAst v =-      astScatterS @sh3 @shn @shp shn-        (astSum SNat (STKS (shsFromListS vars `shsAppend` shn) (ftkToSTK x)) v)-        (vars, ix)--- TODO? astScatterS v (ZR, ix) = update (rzero sh 0) ix v-astScatterS shn (Ast.AstLet var u v) (vars, ix) =-      astLet var u (astScatterS @shm @shn @shp shn v (vars, ix))-astScatterS shn (Ast.AstFromPrimal v) (vars, ix) =-  Ast.AstFromPrimal $ astScatterS @shm @shn @shp shn v (vars, ix)-astScatterS shn (Ast.AstFromDual v) (vars, ix) =-  Ast.AstFromDual $ astScatterS @shm @shn @shp shn v (vars, ix)-astScatterS shn v (vars, ix) = Ast.AstScatterS @shm @shn @shp shn v (vars, ix)--flipCompare :: forall (a :: Nat) b. Compare a b ~ GT => Compare b a :~: LT-flipCompare = unsafeCoerceRefl--astGatherS-  :: forall shm shn shp r s. AstSpan s-  => ShS shn-  -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)-  -> (AstVarListS shm, AstIxS AstMethodLet shp)-  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)-astGatherS = astGatherKnobsS @shm @shn @shp defaultKnobs---- Assumption: vars0 don't not occur in v0. The assumption only holds--- when newly generated variables are fresh, which is the case as long--- as resetVarCounter is not used. The assumption makes it easier to spot--- bugs or corruption, hence we assert it in the code below.-astGatherKnobsS-  :: forall shm shn shp r s. AstSpan s-  => SimplifyKnobs-  -> ShS shn-  -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)-  -> (AstVarListS shm, AstIxS AstMethodLet shp)-  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)-astGatherKnobsS _ _ v0 (!vars0, !_ix0)-  | any (`varNameInAst` v0) $ listsToList vars0 =-    error $ "astGatherKnobsS: gather vars in v0: " ++ show (vars0, v0)-astGatherKnobsS knobs shn v0 (ZS, ix0) = astIndexKnobsS knobs shn v0 ix0-astGatherKnobsS _ _ v0 (vars0, ZIS) =-  astReplicateNS @shm @shn (shsFromListS vars0) v0-astGatherKnobsS _ shn v0 (vars, i1 :.$ _)-  | let (lb, ub) = bounds i1--- this doesn't work in GHC 9.10:---      FTKS (snat :$$ _) x = ftkAst v0-  , FTKS (snat :$$ _) x <- ftkAst v0-  , ub < 0 || lb >= fromInteger (fromSNat snat) =-    let ftk = FTKS (shsFromListS vars `shsAppend` shn) x-    in fromPrimal $ astConcrete ftk (tdefTarget ftk)-astGatherKnobsS knobs shn v0 (vars0@(var1 ::$ vars1), ix0)-  | not (getConst var1 `varNameInIxS` ix0) =-    let k :$$ sh' = shsFromListS vars0-        FTKS _ x = ftkAst v0-    in astReplicate k (STKS (sh' `shsAppend` shn) (ftkToSTK x))-                    (astGatherKnobsS @(Tail shm) @shn @shp-                                     knobs shn v0 (vars1, ix0))-astGatherKnobsS knobs shn v0 (vars0@(_ ::$ _), ix0@(_ :.$ _))-  | let ixInit = ixsInit ix0-        varInit = listsInit vars0-        varLast = getConst $ listsLast vars0-  , AstIntVar ixvarLast <- ixsLast ix0-  , ixvarLast == varLast-  , not (varLast `varNameInIxS` ixInit)-  , kLast@SNat <- shsLast (shsFromListS vars0)-  , Just Refl <- testEquality kLast (shsLast (shsFromIxS ix0)) =-    gcastWith (unsafeCoerceRefl-               :: Init shp ++ (Last shm ': shn) :~: shp ++ shn) $-    gcastWith (unsafeCoerceRefl-               :: Init shm ++ (Last shm ': shn) :~: shm ++ shn) $-    astGatherKnobsS @(Init shm) @(Last shm ': shn) @(Init shp)-                    knobs (kLast :$$ shn) v0 (varInit, ixInit)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstCond-           (Ast.AstBoolAnd a@(AstLeqInt (AstConcreteK j) AstIntVar{}) b)-           v w :.$ prest )-    | j <= 0 || j >= valueOf @m || ixIsSmall prest =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstLet varN uN-           (Ast.AstCond-              (Ast.AstBoolAnd a@(AstLeqInt (AstConcreteK j) AstIntVar{}) b)-              v w) :.$ prest )-    | j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstCond-           (Ast.AstBoolAnd a@(AstLeqInt (AstConcreteK j)-                                        (Ast.AstN1K NegateOp AstIntVar{})) b)-           v w :.$ prest )-    | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstLet varN uN-           (Ast.AstCond-              (Ast.AstBoolAnd a@(AstLeqInt (AstConcreteK j)-                                           (Ast.AstN1K NegateOp AstIntVar{})) b)-              v w) :.$ prest )-    | - j + 1 <= 0 || - j + 1 >= valueOf @m-      || ixIsSmall prest && astIsSmall True uN =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)--astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstCond-           (Ast.AstBoolAnd-              a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK j) AstIntVar{})) b)-           v w :.$ prest )-    | j <= 0 || j >= valueOf @m || ixIsSmall prest =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstLet varN uN-           (Ast.AstCond-              (Ast.AstBoolAnd-                 a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK j) AstIntVar{})) b)-              v w) :.$ prest )-    | j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstCond-           (Ast.AstBoolAnd-              a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK j)-                                           (Ast.AstN1K NegateOp-                                                       AstIntVar{}))) b)-           v w :.$ prest )-    | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstLet varN uN-           (Ast.AstCond-              (Ast.AstBoolAnd-                 a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK j)-                                              (Ast.AstN1K NegateOp-                                                          AstIntVar{}))) b)-              v w) :.$ prest )-    | - j + 1 <= 0 || - j + 1 >= valueOf @m-      || ixIsSmall prest && astIsSmall True uN =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)--astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstCond-      (Ast.AstBoolAnd-         a@(Ast.AstBoolNot-              (Ast.AstBoolAnd (AstLeqInt (AstConcreteK j) AstIntVar{}) _)) b)-      v w :.$ prest )-    | j <= 0 || j >= valueOf @m || ixIsSmall prest =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstLet varN uN-           (Ast.AstCond-              (Ast.AstBoolAnd-                 a@(Ast.AstBoolNot-                      (Ast.AstBoolAnd (AstLeqInt (AstConcreteK j)-                                                 AstIntVar{}) _)) b)-              v w) :.$ prest )-    | j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstCond-           (Ast.AstBoolAnd-              a@(Ast.AstBoolNot-                   (Ast.AstBoolAnd (AstLeqInt (AstConcreteK j)-                                              (Ast.AstN1K NegateOp-                                                          AstIntVar{})) _)) b)-           v w :.$ prest )-    | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstLet varN uN-           (Ast.AstCond-              (Ast.AstBoolAnd-                 a@(Ast.AstBoolNot-                      (Ast.AstBoolAnd-                         (AstLeqInt (AstConcreteK j)-                                    (Ast.AstN1K NegateOp-                                                AstIntVar{})) _)) b)-              v w) :.$ prest )-    | - j + 1 <= 0 || - j + 1 >= valueOf @m-      || ixIsSmall prest && astIsSmall True uN =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)--astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstCond-           (Ast.AstBoolAnd-              a@(Ast.AstBoolNot-                   (Ast.AstBoolAnd-                      (Ast.AstBoolNot-                         (AstLeqInt (AstConcreteK j) AstIntVar{})) _)) b)-           v w :.$ prest )-    | j <= 0 || j >= valueOf @m || ixIsSmall prest =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstLet varN uN-           (Ast.AstCond-              (Ast.AstBoolAnd-                 a@(Ast.AstBoolNot-                      (Ast.AstBoolAnd-                         (Ast.AstBoolNot (AstLeqInt (AstConcreteK j)-                                                    AstIntVar{})) _)) b)-              v w) :.$ prest )-    | j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstCond-           (Ast.AstBoolAnd-              a@(Ast.AstBoolNot-                   (Ast.AstBoolAnd-                      (Ast.AstBoolNot-                         (AstLeqInt (AstConcreteK j)-                                   (Ast.AstN1K NegateOp AstIntVar{}))) _)) b)-           v w :.$ prest )-    | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)-astGatherKnobsS-  knobs shn v0-  ( vars@((::$) @m _ _)-  , Ast.AstLet varN uN-           (Ast.AstCond-              (Ast.AstBoolAnd-                 a@(Ast.AstBoolNot-                      (Ast.AstBoolAnd-                         (Ast.AstBoolNot-                            (AstLeqInt (AstConcreteK j)-                                       (Ast.AstN1K NegateOp-                                                   AstIntVar{}))) _)) b)-              v w) :.$ prest )-    | - j + 1 <= 0 || - j + 1 >= valueOf @m-      || ixIsSmall prest && astIsSmall True uN =-  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared-  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)--- Rules with AstConcreteK on the right hand side of AstPlusK are--- not needed, thanks to the normal form of AstPlusK rewriting.-astGatherKnobsS knobs shn v0-  (vars, AstPlusK (AstConcreteK i64) i1 :.$ prest)-  | let (lb, ub) = bounds i1-  , lb >= 0  -- if not, we may need to apply astReverse first-  , FTKS (SNat @p :$$ _) x <- ftkAst v0 =-    if i64 >= 0 then-      withSNat (fromIntegral i64) $ \(SNat @i) ->-      withSNat (fromIntegral $ min (valueOf @p - i64) (ub + 1)) $ \(SNat @k) ->-        gcastWith (unsafeCoerceRefl :: (i + k <=? p) :~: True) $-        let v2 = astSliceS (SNat @i) (SNat @k) (SNat @(p - (i + k))) v0-        in astGatherKnobsS knobs shn v2 (vars, i1 :.$ prest)-             -- this gather may still index out of bounds, which is fine-    else-      withSNat (- fromIntegral i64) $ \(SNat @i) ->-        let ftk = FTKS (SNat @i :$$ shsFromIxS prest `shsAppend` shn) x-            v2 = fromPrimal (astConcrete ftk (tdefTarget ftk))-                 `astAppendS`-                 v0-        in astGatherKnobsS knobs shn v2 (vars, i1 :.$ prest)-             -- this gather may still index out of bounds, which is fine-astGatherKnobsS knobs shn v0-  (vars, Ast.AstLet varN uN (AstPlusK (AstConcreteK i64) i1) :.$ prest)-  | let (lb, ub) = bounds i1-  , lb >= 0  -- if not, we may need to apply astReverse first-  , FTKS (SNat @p :$$ _) x <- ftkAst v0 =-    if i64 >= 0 then-      withSNat (fromIntegral i64) $ \(SNat @i) ->-      withSNat (fromIntegral $ min (valueOf @p - i64) (ub + 1)) $ \(SNat @k) ->-        gcastWith (unsafeCoerceRefl :: (i + k <=? p) :~: True) $-        let v2 = astSliceS (SNat @i) (SNat @k) (SNat @(p - (i + k))) v0-        in astGatherKnobsS knobs shn v2 (vars, Ast.AstLet varN uN i1 :.$ prest)-             -- this gather may still index out of bounds, which is fine-    else-      withSNat (- fromIntegral i64) $ \(SNat @i) ->-        let ftk = FTKS (SNat @i :$$ shsFromIxS prest `shsAppend` shn) x-            v2 = fromPrimal (astConcrete ftk (tdefTarget ftk))-                 `astAppendS`-                 v0-        in astGatherKnobsS knobs shn v2 (vars, Ast.AstLet varN uN i1 :.$ prest)-             -- this gather may still index out of bounds, which is fine-astGatherKnobsS knobs shn v0-  ( vars@((::$) @m (Const varm) mrest)-  , Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar varp)) i1 i2-    :.$ prest )-  | varNameToAstVarId varm == varNameToAstVarId varp-  , j <= 0 || j >= valueOf @m || ixIsSmall prest =-    if | j <= 0 ->-         astGatherKnobsS knobs shn v0 (vars, i1 :.$ prest)-       | j >= valueOf @m ->-         astGatherKnobsS knobs shn v0 (vars, i2 :.$ prest)-       | otherwise ->-         withSNat (fromIntegral j) $ \(SNat @j) ->-         gcastWith (unsafeCoerceRefl :: (j <=? m) :~: True) $-         astLetFun v0 $ \v ->-         let varm2 = mkAstVarName (varNameToFTK varm)-                                  (Just (0, j - 1))-                                  (varNameToAstVarId varm)-             varm3 = mkAstVarName (varNameToFTK varm)-                                  (Just (0, valueOf @m - j - 1))-                                  (varNameToAstVarId varm)-         in astGatherKnobsS knobs shn v-              ( (::$) @j (Const varm2) mrest-              , substituteAstIxS (astVar varm2)-                                 varm (i2 :.$ prest) )-            `astAppendS`-            astGatherKnobsS knobs shn v-              ( (::$) @(m - j) (Const varm3) mrest-              , substituteAstIxS (AstConcreteK j + astVar varm3)-                                 varm (i1 :.$ prest) )-astGatherKnobsS knobs shn v0-  ( vars@((::$) @m (Const varm) mrest)-  , Ast.AstCond (AstLeqInt (AstConcreteK j)-                           (Ast.AstN1K NegateOp (AstIntVar varp))) i1 i2-    :.$ prest )-  | varNameToAstVarId varm == varNameToAstVarId varp-  , - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =-    if | - j + 1 <= 0 ->-         astGatherKnobsS knobs shn v0 (vars, i2 :.$ prest)-       | - j + 1 >= valueOf @m ->-         astGatherKnobsS knobs shn v0 (vars, i1 :.$ prest)-       | otherwise ->-         withSNat (- fromIntegral j + 1) $ \(SNat @mj) ->-         gcastWith (unsafeCoerceRefl :: (mj <=? m) :~: True) $-         astLetFun v0 $ \v ->-         let varm2 = mkAstVarName (varNameToFTK varm)-                                  (Just (0, valueOf @mj - 1))-                                  (varNameToAstVarId varm)-             varm3 = mkAstVarName (varNameToFTK varm)-                                  (Just (0, valueOf @m - valueOf @mj - 1))-                                  (varNameToAstVarId varm)-         in astGatherKnobsS knobs shn v-              ( (::$) @mj (Const varm2) mrest-              , substituteAstIxS (astVar varm2)--- TODO: when I use AstIntVar here, which is wrong, I get phantom errors;--- make sure this vanished after the fixes in HEAD-                                 varm (i1 :.$ prest) )-            `astAppendS`-            astGatherKnobsS knobs shn v-              ( (::$) @(m - mj) (Const varm3) mrest-              , substituteAstIxS (AstConcreteK (- j + 1) + astVar varm3)-                                 varm (i2 :.$ prest))-astGatherKnobsS knobs shn v0-  ( vars@((::$) @m (Const varm) mrest)-  , Ast.AstLet varN uN-      (Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar varp)) i1 i2)-      :.$ prest )-  | varNameToAstVarId varm == varNameToAstVarId varp-  , j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =-    if | j <= 0 ->-         astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i1 :.$ prest)-       | j >= valueOf @m ->-         astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i2 :.$ prest)-       | otherwise ->-         withSNat (fromIntegral j) $ \(SNat @j) ->-         gcastWith (unsafeCoerceRefl :: (j <=? m) :~: True) $-         astLetFun v0 $ \v ->-         let varm2 = mkAstVarName (varNameToFTK varm)-                                  (Just (0, j - 1))-                                  (varNameToAstVarId varm)-             varm3 = mkAstVarName (varNameToFTK varm)-                                  (Just (0, valueOf @m - j - 1))-                                  (varNameToAstVarId varm)-         in astGatherKnobsS knobs shn v-              ( (::$) @j (Const varm2) mrest-              , substituteAstIxS (astVar varm2) varm-                                 (Ast.AstLet varN uN i2 :.$ prest) )-            `astAppendS`-            astGatherKnobsS knobs shn v-              ( (::$) @(m - j) (Const varm3) mrest-              , substituteAstIxS (AstConcreteK j + astVar varm3)-                                 varm (Ast.AstLet varN uN i1 :.$ prest) )-astGatherKnobsS knobs shn v0-  ( vars@((::$) @m (Const varm) mrest)-  , Ast.AstLet varN uN-      (Ast.AstCond (AstLeqInt (AstConcreteK j)-                              (Ast.AstN1K NegateOp (AstIntVar varp))) i1 i2)-    :.$ prest )-  | varNameToAstVarId varm == varNameToAstVarId varp-  , - j + 1 <= 0 || - j + 1 >= valueOf @m-    || ixIsSmall prest && astIsSmall True uN =-    if | - j + 1 <= 0 ->-         astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i2 :.$ prest)-       | - j + 1 >= valueOf @m ->-         astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i1 :.$ prest)-       | otherwise ->-         withSNat (- fromIntegral j + 1) $ \(SNat @mj) ->-         gcastWith (unsafeCoerceRefl :: (mj <=? m) :~: True) $-         astLetFun v0 $ \v ->-         let varm2 = mkAstVarName (varNameToFTK varm)-                                  (Just (0, valueOf @mj - 1))-                                  (varNameToAstVarId varm)-             varm3 = mkAstVarName (varNameToFTK varm)-                                  (Just (0, valueOf @m - valueOf @mj - 1))-                                  (varNameToAstVarId varm)-         in astGatherKnobsS knobs shn v-              ( (::$) @mj (Const varm2) mrest-              , substituteAstIxS (astVar varm2)-                                 varm (Ast.AstLet varN uN i1 :.$ prest) )-            `astAppendS`-            astGatherKnobsS knobs shn v-              ( (::$) @(m - mj) (Const varm3) mrest-              , substituteAstIxS (AstConcreteK (- j + 1) + astVar varm3)-                                 varm (Ast.AstLet varN uN i2 :.$ prest))-astGatherKnobsS knobs shn v0-  ( (::$) @m @shmTail (Const varm) mrest-  , (:.$) @p @shpTail (AstIntVar varp) prest )-  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]-      -- prevent a loop-  , varm == varp-  , not (varm `varNameInIxS` prest)-  , FTKS _ x <- ftkAst v0 =-    withSNat (min (valueOf @p) (valueOf @m)) $ \(SNat @m2) ->-    gcastWith (unsafeCoerceRefl :: (m2 <=? p) :~: True) $-    gcastWith (unsafeCoerceRefl :: (m2 <=? m) :~: True) $-    Permutation.permFromList (permCycle $ shsLength (shsFromListS mrest) + 1)-    $ \(permVars :: Permutation.Perm permVars) ->-    Permutation.permFromList (backpermCycle $ shsLength (shsFromIxS prest) + 1)-    $ \(permIx :: Permutation.Perm permIx) ->-       gcastWith (unsafeCoerceRefl-                  :: m2 ': shmTail ++ shn-                     :~: Permutation.PermutePrefix-                           permVars (shmTail ++ (m2 ': shn))) $-       gcastWith (unsafeCoerceRefl-                  :: shpTail ++ (m2 ': shn)-                     :~: Permutation.PermutePrefix-                           permIx (m2 ': shpTail ++ shn)) $-       gcastWith (unsafeCoerceRefl-                  :: (Rank permVars <=? Rank (shmTail ++ (m2 ': shn)))-                     :~: True) $-       gcastWith (unsafeCoerceRefl-                  :: (Rank permIx <=? Rank (m2 ': shpTail ++ shn)) :~: True) $-       fromMaybe (error "astGatherKnobsS: impossible non-permutation")-       $ Permutation.permCheckPermutation permVars-       $ fromMaybe (error "astGatherKnobsS: impossible non-permutation")-       $ Permutation.permCheckPermutation permIx-       $ let v2 = astTransposeS permIx-                  $ astSliceS (SNat @0) (SNat @m2) (SNat @(p - m2)) v0-             u = astGatherKnobsS knobs (SNat @m2 :$$ shn) v2 (mrest, prest)-             ftk =-               FTKS (SNat @(m - m2) :$$ shsFromListS mrest `shsAppend` shn) x-         in astTransposeS permVars u-            `astAppendS`-            fromPrimal (astConcrete ftk (tdefTarget ftk))-astGatherKnobsS knobs shn v7@(Ast.AstFromVector _ (STKS _ x2) l)-                ( ((::$) @m1' @shm4 (Const var4) vrest4)-                , ((:.$) @_ @shp1' i4 rest4) )-  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]-  , let g = case i4 of-          AstIntVar var | var == var4 -> Just id-          AstTimesK (AstConcreteK n) (AstIntVar var)-            | var == var4 -> Just (n *)-          -- TODO: add more or define evaluation-          _ -> Nothing-  , Just h <- g-  , ixIsSmall rest4 =-    let f i =-          let subRest4 = substituteAstIxS (AstConcreteK i) var4 rest4-              j = fromIntegral $ h i-          in if j >= V.length l-             then let FTKS _ x = ftkAst v7-                      ftk = FTKS (shsFromListS vrest4 `shsAppend` shn) x-                  in fromPrimal $ astConcrete ftk (tdefTarget ftk)-             else astGatherKnobsS @shm4 @shn @shp1' knobs shn-                                  (l V.! j) (vrest4, subRest4)-    in astFromVector (SNat @m1')-                     (STKS (shsFromListS vrest4 `shsAppend` shn) x2)-       $ V.fromList $ map f [0 .. valueOf @m1' - 1]-astGatherKnobsS knobs shn v0 (vars0, i1 :.$ rest1)-  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]-      -- prevent a loop-  , not (any (`varNameInAst` i1) $ listsToList vars0) =-    astGatherKnobsS @shm @shn-      knobs shn-      (astIndexKnobsS knobs (shsFromIxS rest1 `shsAppend` shn) v0 (i1 :.$ ZIS))-      (vars0, rest1)-astGatherKnobsS knobs shn v0-  (vars@((::$) @m _ _), ix@(i1 :.$ prest))-  | let intInteresting = \case-          AstPlusK (AstConcreteK _) i2-            | fst (bounds i2) >= 0 -> True-          Ast.AstLet _ _ (AstPlusK (AstConcreteK _) i2)-            | fst (bounds i2) >= 0 -> True-          Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar var)) _ _-            | j <= 0 || j >= valueOf @m || ixIsSmall prest-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          Ast.AstCond (AstLeqInt (AstConcreteK j)-                                 (Ast.AstN1K NegateOp (AstIntVar var))) _ _-            | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          Ast.AstLet _ uN-            (Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar var)) _ _)-            | j <= 0 || j >= valueOf @m-              || ixIsSmall prest && astIsSmall True uN-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          Ast.AstLet _ uN-            (Ast.AstCond (AstLeqInt (AstConcreteK j)-                                    (Ast.AstN1K NegateOp (AstIntVar var))) _ _)-            | - j + 1 <= 0 || - j + 1 >= valueOf @m-              || ixIsSmall prest && astIsSmall True uN-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          Ast.AstCond-            (Ast.AstBoolAnd-               (AstLeqInt (AstConcreteK j) (AstIntVar var)) _) _ _-            | j <= 0 || j >= valueOf @m || ixIsSmall prest-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          Ast.AstCond-            (Ast.AstBoolAnd-               (AstLeqInt (AstConcreteK j)-                          (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _-            | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          Ast.AstLet _ uN-            (Ast.AstCond-               (Ast.AstBoolAnd-                  (AstLeqInt (AstConcreteK j) (AstIntVar var)) _) _ _)-            | j <= 0 || j >= valueOf @m-              || ixIsSmall prest && astIsSmall True uN-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          Ast.AstLet _ uN-            (Ast.AstCond-               (Ast.AstBoolAnd-                  (AstLeqInt (AstConcreteK j)-                             (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _)-            | - j + 1 <= 0 || - j + 1 >= valueOf @m-              || ixIsSmall prest && astIsSmall True uN-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          AstIntVar var-            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]-            , null $ drop 1 $ filter (var `varNameInAst`) (Foldable.toList ix)-            , any ((== varNameToAstVarId var) . varNameToAstVarId)-                  (listsToList vars) -> True-          ik | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]-             , not (any (`varNameInAst` ik) $ listsToList vars) -> True-          -- We can't reorder ix for the gather(fromVector) rule above,-          -- because it becomes gather(transpose); we can only reorder vars.-          _ -> False-  , not (intInteresting i1)  -- now vars may need to be reordered, too-  , Just i <- findIndex intInteresting-                        (Foldable.toList ix) = assert (i > 0) $-    Permutation.permFromList (backpermCycle $ i + 1)-    $ \(perm :: Permutation.Perm perm) ->-    gcastWith (unsafeCoerceRefl-               :: Permutation.PermutePrefix perm (shp ++ shn)-                  :~: Permutation.PermutePrefix perm shp ++ shn) $-    gcastWith (unsafeCoerceRefl :: (Rank perm <=? Rank (shp ++ shn)) :~: True) $-    fromMaybe (error "astGatherKnobsS: impossible non-permutation")-    $ Permutation.permCheckPermutation perm-    $ let v2 = astTransposeS perm v0-      in astGatherKnobsS knobs shn v2 (vars, ixsPermutePrefix perm ix)-        -- this call is guaranteed to simplify as above, so the tranpose-        -- won't reduce it back to the original and cause a loop-astGatherKnobsS knobs shn v0-  (vars@((::$) @m _ _), ix@(i1 :.$ prest))-  | let varInteresting = \case-          Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar var)) _ _-            | j <= 0 || j >= valueOf @m || ixIsSmall prest ->-              Just var-          Ast.AstCond (AstLeqInt (AstConcreteK j)-                                 (Ast.AstN1K NegateOp (AstIntVar var))) _ _-            | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest ->-              Just var-          Ast.AstLet _ uN-            (Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar var)) _ _)-            | j <= 0 || j >= valueOf @m-              || ixIsSmall prest && astIsSmall True uN ->-              Just var-          Ast.AstLet _ uN-            (Ast.AstCond (AstLeqInt (AstConcreteK j)-                                    (Ast.AstN1K NegateOp-                                                (AstIntVar var))) _ _)-            | - j + 1 <= 0 || - j + 1 >= valueOf @m-              || ixIsSmall prest && astIsSmall True uN ->-              Just var-          Ast.AstCond-            (Ast.AstBoolAnd-               (AstLeqInt (AstConcreteK j) (AstIntVar var)) _) _ _-            | j <= 0 || j >= valueOf @m || ixIsSmall prest ->-              Just var-          Ast.AstCond-            (Ast.AstBoolAnd-               (AstLeqInt (AstConcreteK j)-                          (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _-            | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest ->-              Just var-          Ast.AstLet _ uN-            (Ast.AstCond-               (Ast.AstBoolAnd-                  (AstLeqInt (AstConcreteK j) (AstIntVar var)) _) _ _)-            | j <= 0 || j >= valueOf @m-              || ixIsSmall prest && astIsSmall True uN ->-              Just var-          Ast.AstLet _ uN-            (Ast.AstCond-               (Ast.AstBoolAnd-                  (AstLeqInt (AstConcreteK j)-                             (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _)-            | - j + 1 <= 0 || - j + 1 >= valueOf @m-              || ixIsSmall prest && astIsSmall True uN ->-              Just var-          AstIntVar var-            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]-            , not (var `varNameInIxS` prest) -> Just var-          i4  -- has to be last, because ix can't be reordered-            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]-            , Ast.AstFromVector{} <- v0-            , ixIsSmall prest-            , mvar <- case i4 of-                AstIntVar var -> Just var-                AstTimesK AstConcreteK{} (AstIntVar var) -> Just var-                _ -> Nothing-            , Just{} <- mvar -> mvar-          _ -> Nothing-  , Just varp <- varInteresting i1-  , Just i <- findIndex ((== varNameToAstVarId varp) . varNameToAstVarId)-                        (listsToList vars) = assert (i > 0) $-    Permutation.permFromList (backpermCycle $ i + 1)-    $ \(permWhole :: Permutation.Perm permWhole) ->-    permInverse permWhole $ \(invperm :: Nested.Perm invperm) _ ->-    gcastWith (unsafeCoerceRefl-               :: shm ++ shn-                  :~: Permutation.PermutePrefix permWhole-                        (Permutation.PermutePrefix invperm shm ++ shn)) $-    gcastWith (unsafeCoerceRefl-               :: (Rank permWhole-                   <=? Rank (Permutation.PermutePrefix invperm shm ++ shn))-                  :~: True) $-    fromMaybe (error "astGatherKnobsS: impossible non-permutation")-    $ Permutation.permCheckPermutation permWhole-    $ astTransposeS permWhole-    $ astGatherKnobsS knobs shn v0 (listsPermutePrefix invperm vars, ix)-        -- this call is guaranteed to simplify as above, so the tranpose-        -- won't reduce it back to the original and cause a loop-astGatherKnobsS knobs shn v4 (vars4, ix4@((:.$) @in1 @shp1' i4 rest4))-  | FTKS _ x <- ftkAst v4 = case v4 of-    Ast.AstProject1{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstProject2{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    {- Ast.AstFromVector{} | gatherFromNF (shsFromIxS rest4) vars4 ix4 ->-        -- normal form-      Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4) -}-    {- this rule seems counterproductive in many cases, so disabled until-       we can detect cases where it helps:-    Ast.AstFromVector snat STKS{} l ->-      -- Term rest4 is duplicated without sharing and we can't help it,-      -- because it needs to be in scope of vars4, so we can't use tlet.-      funToVarsIxS @shm (shsFromListS vars4) $ \ (!varsFresh, IxS !ixFresh) ->-        let f v = astGather @shm @shn @shp1' shn v (vars4, rest4)-            -- This subst doesn't currently break sharing because it's a rename.-            subst i =-              foldr (\(i2, var2) v2 -> substituteAst i2 var2 v2)-                    i-                    (listsToList $ zipSizedS ixFresh vars4)-            i5 = subst i4-       in astGather @shm @shn @(p1' ': shm)-                    shn (astFromVector snat (STKS (shsFromListS varsFresh-                                                    `shsAppend` shn)-                                                   (ftkToSTK x))-                          $ V.map f l)-                    (varsFresh, i5 :.$ IxS ixFresh) -}-    Ast.AstFromVector{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    -- This accomplishes fusion if v is a gather or anything-    -- that gather can fuse with, but at the cost of an extra transpose-    -- that doesn't fuse here unless astTransposeAsGatherS is used.-    -- Since the transpose is O(1), let's leave this as is.-    Ast.AstSum snat@(SNat @n1) STKS{} v ->-      let perm3 = backpermCycle $ shsLength (shsFromIxS ix4) + 1-          perm4 = permCycle $ shsLength (shsFromListS vars4) + 1-      in Permutation.permFromList perm3-         $ \(perm3S :: Permutation.Perm perm3P) ->-         gcastWith (unsafeCoerceRefl-                    :: Compare (Rank perm3P) (Rank (n1 : shp ++ shn))-                       :~: LT) $-         gcastWith (unsafeCoerceRefl-                    :: Permutation.PermutePrefix perm3P (n1 : (shp ++ shn))-                       :~: shp ++ (n1 : shn)) $-         fromMaybe (error "astGatherKnobsS: impossible non-permutation")-         $ Permutation.permCheckPermutation perm3S-         $ Permutation.permFromList perm4-         $ \(perm4S :: Permutation.Perm perm4P) ->-         gcastWith (unsafeCoerceRefl-                    :: Compare (Rank perm4P) (Rank (shm ++ (n1 : shn)))-                       :~: LT) $-         gcastWith (unsafeCoerceRefl-                    :: Permutation.PermutePrefix perm4P (shm ++ (n1 : shn))-                       :~: n1 : (shm ++ shn)) $-         fromMaybe (error "astGatherKnobsS: impossible non-permutation")-         $ Permutation.permCheckPermutation perm4S-         $ let innerGather =-                 astGather @shm @(n1 : shn) @shp-                           (snat :$$ shn) (astTransposeS perm3S v) (vars4, ix4)-           in astSum snat (STKS (shsFromListS vars4 `shsAppend` shn)-                                (ftkToSTK x))-              $ astTransposeS perm4S innerGather-                {- TODO: disabled until we can reliably fuse back to transpose-                if not (knobExpand knobs)-                then astTransposeS perm4S innerGather-                else astTransposeAsGatherS knobs perm4S innerGather -}-    Ast.AstReplicate _ STKS{} v ->-      let ftk = FTKS (shsFromListS vars4 `shsAppend` shn) x-          defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-      -- This boolean term may have free variables that act as universally-      -- quantified.-      in case 0 <=. i4 &&* i4 <=. valueOf @in1 - 1 of-        AstBoolConst b ->-          if b then astGather @shm @shn @shp1' shn v (vars4, rest4) else defArr-        _ -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstReplicate (SNat @k) STKScalar v | ZIS <- rest4 ->-      let ftk = FTKS (shsFromListS vars4 `shsAppend` shn) x-          defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-      in case 0 <=. i4 &&* i4 <=. valueOf @k - 1 of-        AstBoolConst b ->-          if b then astGather @shm @shn @shp1'-                              shn (astSFromK' v) (vars4, rest4) else defArr-        _ -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstApply{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstVar{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstCond b v w | ixIsSmall ix4 ->-      astCond b (astGather @shm @shn @shp shn v (vars4, ix4))-                (astGather @shm @shn @shp shn w (vars4, ix4))-    Ast.AstCond{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstBuild1{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    AstConcreteS{} -> case unRepl1 v4 of-      Just u ->-        let ftk = FTKS (shsFromListS vars4 `shsAppend` shn) x-            defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)-        in case 0 <=. i4 &&* i4 <=. valueOf @in1 - 1 of-          AstBoolConst b ->-            if b-            then astGather @shm @shn @shp1' shn u (vars4, rest4)-            else defArr-          _ -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-      _ -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-             -- free variables possible in the index, so can't compute the array--    Ast.AstLet var u v ->-      astLet var u (astGather @shm @shn @shp shn v (vars4, ix4))--    Ast.AstPrimalPart{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstDualPart{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstFromPrimal v ->-      Ast.AstFromPrimal $ astGather @shm @shn @shp shn v (vars4, ix4)-    Ast.AstFromDual v ->-      Ast.AstFromDual $ astGather @shm @shn @shp shn v (vars4, ix4)--    -- Going inside a binary ops usually makes a term more expensive-    -- to interpret and inverting that requires comparing two arguments,-    -- so it's not practical.-    AstPlusS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    AstTimesS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstN1S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstR1S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstR2S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstI2S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstFloorS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstFromIntegralS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstCastS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)--    {- is reverted in astGatherKnobsS immediatedly; only do in expansion phase?-    Ast.AstIndexS @shm2 _shn2 v2 (i2 :.$ ZIS) ->-        astGather @shm @shn @(shm2 ++ shp) shn v2 (vars4, i2 :.$ ix4) -}-    Ast.AstIndexS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstScatterS @shm7 @shn7 @shp7 shn7 v (vars, AstIntVar var5 :.$ ix2)-      | AstIntVar var6 <- i4, var6 == var5 ->-          astGather @shm @shn @shp1' shn-                    (astScatterS @shm7 @shn7 @(Tail shp7) shn7-                                 v (vars, ix2))-                    (vars4, rest4)-    Ast.AstScatterS{} ->  -- normal form-      Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstGatherS @shm2 @shn2 @shp2 shn2 v2 (vars2, ix2)-                   | SNat @rank4 <- ixsRank ix4-                   , SNat @rank2 <- listsRank vars2 ->-      let subst :: AstIxS AstMethodLet shm7 -> AstVarListS shm7-                -> AstInt AstMethodLet-                -> AstInt AstMethodLet-          subst (IxS ix) vars t0 =-            foldr (\ (v, i) -> substituteAst i v)-                  t0 (listsFold (\(Fun.Pair (Const v) (Const i)) -> [(v, i)])-                      $ listsZip vars ix)-          inBounds :: AstIxS AstMethodLet shm7 -> AstVarListS shm7 -> Bool-          inBounds (IxS ix) vars =-            let inb (v, i) =-                  let (lbi, ubi) = bounds i-                  in case varNameToBounds v of-                    Nothing -> True-                    Just (lbv, ubv) -> lbv <= lbi && ubi <= ubv-            in all inb (listsFold (\(Fun.Pair (Const v) (Const i)) -> [(v, i)])-                        $ listsZip vars ix)-          IxS list4 = ix4-          composedGather ::  -- rank4 <= rank2-            Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) r))-          composedGather =-            -- we have: shm2 ++ shn2 == shp ++ shn-            -- so from ranks:-            gcastWith (unsafeCoerceRefl :: TakeLen shp shm2 :~: shp) $-            -- and from congruence:---            gcastWith (unsafeCoerceRefl---                       :: DropLen shp shm2 ++ shn2 :~: shn) $-            -- from congruence:-            gcastWith (unsafeCoerceRefl-                       :: (shm ++ DropLen shp shm2) ++ shn2-                          :~: shm ++ shn) $-            let vars2p = listsTakeLen list4 vars2-                vars22 = listsDropLen list4 vars2-                ix22 = fmap (subst ix4 vars2p) ix2-                list422 = vars4 `listsAppend` vars22-            in if ixIsSmall ix4 && inBounds ix4 vars2p-               then Just $ astGather shn2 v2 (list422, ix22)-               else Nothing-          assimilatedGather ::  -- rank2 <= rank4-            Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) r))-          assimilatedGather =-            -- we have: shm2 ++ shn2 == shp ++ shn-            -- so from ranks:-            gcastWith (unsafeCoerceRefl :: TakeLen shm2 shp :~: shm2) $-            -- and from congruence:---            gcastWith (unsafeCoerceRefl---                       :: DropLen shm2 shp ++ shn :~: shn2) $-            -- from congruence:-            gcastWith (unsafeCoerceRefl-                       :: (shp2 ++ DropLen shm2 shp) ++ shn-                          :~: shp2 ++ shn2) $-            let ix42 = IxS $ listsTakeLen vars2 list4-                ix44 = IxS $ listsDropLen vars2 list4-                ix22 = fmap (subst ix42 vars2) ix2-                ix2244 = ix22 `ixsAppend` ix44-            in if ixIsSmall ix42 && inBounds ix42 vars2-               then Just $ astGather shn v2 (vars4, ix2244)-               else Nothing-      in fromMaybe (Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4))-         $ case cmpNat (Proxy @rank4) (Proxy @rank2) of-        LTI -> composedGather-        EQI -> assimilatedGather-        GTI -> gcastWith (flipCompare @rank4 @rank2) assimilatedGather-    Ast.AstMinIndexS @n @sh v -> case ftkAst v of-     FTKS nsh _ -> case shsLast nsh of-      nl@(SNat @nl) ->-        let shnl = shn `shsAppend` (nl :$$ ZSS)-        in gcastWith (unsafeCoerceRefl-                     :: shp ++ (shn ++ '[nl]) :~: n ': sh) $-           gcastWith (unsafeCoerceRefl-                      :: Head (shm ++ (shn ++ '[nl]))-                         ': Tail (shm ++ (shn ++ '[nl]))-                         :~: shm ++ (shn ++ '[nl])) $-           gcastWith (unsafeCoerceRefl-                      :: Init (shm ++ (shn ++ '[nl]))-                      :~: shm ++ shn) $-           Ast.AstMinIndexS @(Head (shm ++ (shn ++ '[nl])))-                            @(Tail (shm ++ (shn ++ '[nl])))-           $ astGather shnl v (vars4, ix4)-    Ast.AstMaxIndexS @n @sh v -> case ftkAst v of-     FTKS nsh _ -> case shsLast nsh of-      nl@(SNat @nl) ->-        let shnl = shn `shsAppend` (nl :$$ ZSS)-        in gcastWith (unsafeCoerceRefl-                     :: shp ++ (shn ++ '[nl]) :~: n ': sh) $-           gcastWith (unsafeCoerceRefl-                      :: Head (shm ++ (shn ++ '[nl]))-                         ': Tail (shm ++ (shn ++ '[nl]))-                         :~: shm ++ (shn ++ '[nl])) $-           gcastWith (unsafeCoerceRefl-                      :: Init (shm ++ (shn ++ '[nl]))-                      :~: shm ++ shn) $-           Ast.AstMaxIndexS @(Head (shm ++ (shn ++ '[nl])))-                            @(Tail (shm ++ (shn ++ '[nl])))-           $ astGather shnl v (vars4, ix4)-    Ast.AstIotaS{} ->  -- probably nothing can be simplified; a normal form-      Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstAppendS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-      -- fusing would result in gather([gather, gather]), so no gain-    Ast.AstSliceS{}-> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-      -- slicing is O(1) so no point fusing and complicating the expression;-      -- if it did not simplify further with slice, it wouldn't with gather-    Ast.AstReverseS{}-> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-      -- reversing is O(1)-    Ast.AstTransposeS @perm @sh perm v | FTKS sh _ <- ftkAst v ->-      let rankPerm = Permutation.permRank perm-      in case gcompare (ixsRank ix4) rankPerm of-        GLT ->  -- TODO: this does not provide any proof, so use cmpNat :(-          if knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion]-          then astGather @shm @shn @shp-                         shn (astTransposeAsGatherS knobs perm v) (vars4, ix4)-          else Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-        _ ->-          gcastWith (lemRankMapJust $ shsTakeLen perm sh) $-          gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $-          permInverse perm-          $ \(invperm :: Nested.Perm invperm) proof ->-            case proof (ssxFromShX $ shxFromShS $ shsTakeLen perm sh) of-              Refl ->-                -- from PermutePrefix and ranks:-                gcastWith-                  (unsafeCoerceRefl-                   :: Permutation.PermutePrefix invperm shp ++ shn-                      :~: Permutation.PermutePrefix invperm (shp ++ shn)) $-                -- from AstTransposeS:---                gcastWith---                  (unsafeCoerceRefl---                   :: Permutation.PermutePrefix invperm (shp ++ shn)---                      :~: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)) $-                -- from PermutePrefix:---                gcastWith---                  (unsafeCoerceRefl---                   :: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)---                      :~: {-1-} Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))---                          ++ {-2-} DropLen invperm (Permutation.PermutePrefix perm sh)) $-                -- 1. from PermutePrefix:---                gcastWith---                  (unsafeCoerceRefl---                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))---                      :~: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))) $-                -- ranks-                gcastWith-                  (unsafeCoerceRefl-                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))-                      :~: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))) $-                -- from permInverse but MapJust-unwrapped:-                gcastWith-                  (unsafeCoerceRefl-                   :: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))-                      :~: TakeLen perm sh) $-                -- end of 1.-                -- 2. from PermutePrefix---                gcastWith---                  (unsafeCoerceRefl---                   :: DropLen invperm (Permutation.PermutePrefix perm sh)---                      :~: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)) $-                -- ranks-                gcastWith-                  (unsafeCoerceRefl-                   :: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)-                      :~: DropLen perm sh) $-                -- end of 2.-                -- from TakeLen:-                gcastWith-                  (unsafeCoerceRefl-                   :: TakeLen perm sh ++ DropLen perm sh :~: sh) $-                let invix4 = ixsPermutePrefix invperm ix4-                in astGather shn v (vars4, invix4)-    Ast.AstReshapeS sh v ->-      if shsLength sh <= 1-      then astGather @shm @shn @shp shn-                     (astReshapeAsGatherS knobs sh v) (vars4, ix4)-      else Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)--    AstFromS' ftkz v -> case matchingFTK (ftkAst v) ftkz of-      Just Refl -> astGather @shm @shn @shp shn v (vars4, ix4)-        -- rare, usually simplifies away earlier-      Nothing -> error "astGatherKnobsS: wrong tensor kinds in AstFromS"-    -- These conversions need to stay down.-    Ast.AstConvert{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)--    -- These should not appear here unless via wacky tests.-    Ast.AstDot1InS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)-    Ast.AstMatmul2S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)- where-  astGather-    :: forall shm' shn' shp' s' r'. AstSpan s'-    => ShS shn'-    -> AstTensor AstMethodLet s' (TKS2 (shp' ++ shn') r')-    -> (AstVarListS shm', AstIxS AstMethodLet shp')-    -> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') r')-  astGather shn' v2 (vars2, ix2) =-    astGatherKnobsS @shm' @shn' @shp' knobs shn' v2 (vars2, ix2)---- Normal form of chains of appends has the append constructor on the right.-astAppendS :: AstSpan s-           => AstTensor AstMethodLet s (TKS2 (m ': sh) r)-           -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)-           -> AstTensor AstMethodLet s (TKS2 ((m + n) ': sh) r)-astAppendS u v | FTKS (SNat' @0 :$$ _) _ <- ftkAst u = v-astAppendS u v | FTKS (SNat' @0 :$$ _) _ <- ftkAst v = u-astAppendS (Ast.AstFromVector (SNat @k1) stk2@STKS{} l1)-           (Ast.AstFromVector (SNat @k2) STKS{} l2) =-  astFromVector (SNat @(k1 + k2)) stk2 $ l1 V.++ l2-astAppendS (Ast.AstFromVector (SNat @k1) stk2@STKScalar l1)-           (Ast.AstFromVector (SNat @k2) STKScalar l2) =-  astFromVector (SNat @(k1 + k2)) stk2 $ l1 V.++ l2-astAppendS (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) =-  Ast.AstFromPrimal $ astAppendS u v-astAppendS (Ast.AstFromDual u) (Ast.AstFromDual v) =-  Ast.AstFromDual $ astAppendS u v-astAppendS (AstConcreteS u) (AstConcreteS v) =-  astConcreteS (tsappend (Concrete u) (Concrete v))-astAppendS (AstConcreteS u) (Ast.AstAppendS (AstConcreteS v) w) =-  astAppendS (astConcreteS (tsappend (Concrete u) (Concrete v))) w-astAppendS (Ast.AstAppendS v u) w = astAppendS v (astAppendS u w)-astAppendS u v = Ast.AstAppendS u v--astSliceS :: forall i n k sh s r. AstSpan s-          => SNat i -> SNat n -> SNat k-          -> AstTensor AstMethodLet s (TKS2 (i + n + k ': sh) r)-          -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)-astSliceS SNat SNat SNat (Ast.AstFromVector _ stk@STKS{} l) =-  astFromVector (SNat @n) stk $ V.take (valueOf @n) $ V.drop (valueOf @i) l-astSliceS SNat SNat SNat (Ast.AstFromVector _ stk@STKScalar l) =-  astFromVector (SNat @n) stk $ V.take (valueOf @n) $ V.drop (valueOf @i) l-astSliceS SNat SNat SNat (Ast.AstReplicate _ stk@STKS{} v) =-  astReplicate (SNat @n) stk v-astSliceS SNat SNat SNat (Ast.AstReplicate _ stk@STKScalar v) =-  astReplicate (SNat @n) stk v-astSliceS (SNat' @0) SNat (SNat' @0) v = v-astSliceS SNat (SNat' @1) SNat v | FTKS (_ :$$ sh) x <- ftkAst v =-  astReplicate (SNat @1) (STKS sh (ftkToSTK x))-               (astIndexS sh v (valueOf @i :.$ ZIS))-astSliceS SNat SNat SNat (Ast.AstGatherS shn v (Const var ::$ vars, ix)) =-  let varn = mkAstVarName (varNameToFTK var)-                          (Just (0, valueOf @n - 1))-                          (varNameToAstVarId var)-      ivar = valueOf @i + astVar varn-      ix2 = substituteAstIxS ivar var ix  -- cheap subst, because ivar is tiny-  in astGatherS shn v (Const varn ::$ vars, ix2)-astSliceS i n@(SNat @n0) _k (Ast.AstAppendS v1 v2)-  | FTKS (m1@(SNat @m1) :$$ _) _ <- ftkAst v1-  , FTKS (m2@(SNat @m2) :$$ _) _ <- ftkAst v2 =-    let i1 = sNatValue i `min` sNatValue m1-        n1 = sNatValue n `min` (sNatValue m1 - i1)-        k1 = sNatValue m1 - i1 - n1-        i2' = sNatValue i `max` sNatValue m1-        i2 = i2' - sNatValue m1-        n2 = sNatValue n - n1-        k2 = sNatValue m2 - i2 - n2-    in withSNat i1 $ \si1@(SNat @i1) ->-       withSNat n1 $ \sn1@(SNat @n1) ->-       withSNat k1 $ \sk1@(SNat @k1) ->-       withSNat i2 $ \si2@(SNat @i2) ->-       withSNat n2 $ \sn2@(SNat @n2) ->-       withSNat k2 $ \sk2@(SNat @k2) ->-         gcastWith (unsafeCoerceRefl :: n1 + n2 :~: n0) $-         gcastWith (unsafeCoerceRefl :: i1 + n1 + k1 :~: m1) $-         gcastWith (unsafeCoerceRefl :: i2 + n2 + k2 :~: m2) $-         astAppendS (astSliceS si1 sn1 sk1 v1) (astSliceS si2 sn2 sk2 v2)-astSliceS i n k (Ast.AstSliceS i2 _n2 k2 v) =-  astSliceS (snatPlus i i2) n (snatPlus k k2) v-astSliceS i n k (Ast.AstReverseS v) = astReverseS (astSliceS k n i v)--- This enlarges the term and increases computation, but sometimes--- it permits eliminating the AstFromVector node altogether, so we risk it--- for cases that commonly emerge from conditionals.-astSliceS i n@SNat k (Ast.AstTransposeS-                        perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil)-                        (Ast.AstFromVector (SNat' @2) (STKS (_ :$$ sh) x) l)) =-    Ast.AstTransposeS perm-    $ astFromVector (SNat @2) (STKS (n :$$ sh) x) (V.map (astSliceS i n k) l)--- TODO: generalize (maybe the above, too) using unReplN, but it's hard.--- TODO: does it really work only for replicate-like things in-between?-astSliceS i n@SNat k-          (Ast.AstTransposeS perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil) t)-  | STKS (snat :$$ _ :$$ sh3) x <- ftkToSTK $ ftkAst t-  , Just u <- unRepl1 t =-    astTransposeS perm-    $ astReplicate snat (STKS (n :$$ sh3) x)-    $ astSliceS i n k u-astSliceS i n k v1 = case v1 of-  Ast.AstCond b a2 a3 ->-    astCond b (astSliceS i n k a2) (astSliceS i n k a3)-  Ast.AstLet var u v -> astLet var u (astSliceS i n k v)-  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astSliceS i n k v-  Ast.AstFromDual v -> Ast.AstFromDual $ astSliceS i n k v-  AstPlusS u v -> astSliceS i n k u + astSliceS i n k v-  AstTimesS u v -> astSliceS i n k u * astSliceS i n k v-  Ast.AstN1S NegateOp u -> negate (astSliceS i n k u)-  Ast.AstN1S AbsOp u -> abs (astSliceS i n k u)-  Ast.AstN1S SignumOp u -> signum (astSliceS i n k u)-  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astSliceS i n k u)-  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astSliceS i n k u)-                                             (astSliceS i n k v)-  Ast.AstI2S opCode u v -> Ast.AstI2S opCode (astSliceS i n k u)-                                             (astSliceS i n k v)-  Ast.AstFloorS a -> astFloorS $ astSliceS i n k a-  Ast.AstFromIntegralS v -> astFromIntegralS $ astSliceS i n k v-  Ast.AstCastS v -> astCastS $ astSliceS i n k v-  AstConcreteS v -> astConcreteS (tsslice i n k $ Concrete v)-  _ -> Ast.AstSliceS i n k v1--astReverseS :: forall n sh s r. AstSpan s-            => AstTensor AstMethodLet s (TKS2 (n ': sh) r)-            -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)-astReverseS (Ast.AstFromVector snat stk l) =-  astFromVector snat stk $ V.reverse l-astReverseS (Ast.AstReplicate snat stk v) = astReplicate snat stk v-astReverseS (Ast.AstCond b a2 a3) = astCond b (astReverseS a2) (astReverseS a3)-astReverseS (Ast.AstLet var u v) = astLet var u (astReverseS v)-astReverseS (Ast.AstFromPrimal v) = Ast.AstFromPrimal $ astReverseS v-astReverseS (Ast.AstFromDual v) = Ast.AstFromDual $ astReverseS v-astReverseS (Ast.AstGatherS @shm @shn @shp-                            shn v ((::$) @k (Const var) vars, ix)) =-  let ivar = valueOf @k - astVar var-      ix2 = substituteAstIxS ivar var ix  -- cheap subst, because ivar is tiny-  in astGatherS @shm @shn @shp shn v (Const var ::$ vars, ix2)-astReverseS (Ast.AstReverseS v) = v-astReverseS (Ast.AstTransposeS perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil) t)-  | STKS (snat :$$ sh2@(_ :$$ _)) x <- ftkToSTK $ ftkAst t-  , Just u <- unRepl1 t =-    astTransposeS perm-    $ astReplicate snat (STKS sh2 x)-    $ astReverseS u-astReverseS (AstConcreteS v) = astConcreteS (tsreverse $ Concrete v)-astReverseS v = Ast.AstReverseS v---- TODO: try to completely cover the AstGatherS case here, which would permit--- not expanding to astTransposeAsGatherS in astGatherCase--- | Beware, this does not do full simplification, which often requires--- the gather form, so astTransposeAsGather needs to be called in addition--- if full simplification is required.-astTransposeS-  :: forall perm sh s r.-     (Permutation.IsPermutation perm, Rank perm <= Rank sh, AstSpan s)-  => Permutation.Perm perm -> AstTensor AstMethodLet s (TKS2 sh r)-  -> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)-astTransposeS perm t =-  gcastWith (unsafeCoerceRefl-             :: TakeLen perm sh ++ DropLen perm sh :~: sh) $ case perm of- PNil -> t- PCons (SNat' @0) PNil ->-   gcastWith (unsafeCoerceRefl :: Permutation.PermutePrefix '[0] sh :~: sh) $-   t- _ | FTKS sh _ <- ftkAst t-   , Just u2 <- unReplN @_ @_ @(DropLen perm sh) (shsTakeLen perm sh) t ->-   astReplicateNS (shsPermute perm (shsTakeLen perm sh)) u2- _ -> case t of-  Ast.AstFromVector snat@(SNat @n) (STKS @sh2 sh2 x) l-    | SNat' @0 `PCons` _ <- perm -> case permUnShift1 perm of-      (perm2 :: Permutation.Perm perm2) ->-        fromMaybe (error "astTransposeS: impossible non-permutation")-        $ Permutation.permCheckPermutation perm2-        $ gcastWith (unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)-        -- for GHC 9.10 only:-        $ gcastWith (unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)-        $ gcastWith (unsafeCoerceRefl-                     :: Permutation.PermutePrefix perm (n : sh2)-                        :~: n : Permutation.PermutePrefix perm2 sh2)-        $ astFromVector snat (STKS (shsPermutePrefix perm2 sh2) x)-                        (V.map (astTransposeS perm2) l)-  Ast.AstSum snat@(SNat @n) (STKS sh x) v ->-    let zsuccP :: Permutation.Perm (0 : Permutation.MapSucc perm)-        zsuccP = Permutation.permShift1 perm-    in-      gcastWith (unsafeCoerceRefl :: Rank (0 : Permutation.MapSucc perm)-                                     :~: 1 + Rank perm) $-      gcastWith (unsafeCoerceRefl-                 :: Permutation.PermutePrefix-                      (0 : Permutation.MapSucc perm) (n : sh)-                    :~: n : Permutation.PermutePrefix perm sh) $-      fromMaybe (error "astTransposeS: impossible non-permutation")-      $ Permutation.permCheckPermutation zsuccP-      $ astSum snat (STKS (shsPermutePrefix perm sh) x) $ astTransposeS zsuccP v-  _ | SNat' @0 `PCons` _ <- perm-    , STKS ((:$$) @n @sh2 snat sh2) x <- ftkToSTK $ ftkAst t-    , Just u <- unRepl1 t -> case permUnShift1 perm of-      (perm2 :: Permutation.Perm perm2) ->-        fromMaybe (error "astTransposeS: impossible non-permutation")-        $ Permutation.permCheckPermutation perm2-        $ gcastWith (unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)-        -- for GHC 9.10 only:-        $ gcastWith (unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)-        $ gcastWith (unsafeCoerceRefl-                     :: Permutation.PermutePrefix perm (n : sh2)-                        :~: n : Permutation.PermutePrefix perm2 sh2)-        $ astReplicate snat (STKS (shsPermutePrefix perm2 sh2) x)-                       (astTransposeS perm2 u)-  -- This increases term size and work, so limited to size 2.-  Ast.AstReplicate snat1@SNat _-                   (Ast.AstFromVector snat2@(SNat' @2) stk2@STKScalar l)-    | SNat' @1 `PCons` SNat' @0 `PCons` PNil <- perm ->-      astFromVector snat2 (STKS (snat1 :$$ ZSS) stk2)-                    (V.map (astReplicate snat1 stk2) l)-  -- TODO: generalize-  Ast.AstReplicate snat1@SNat _-                   (Ast.AstFromVector snat2@(SNat' @2) stk2@(STKS sh x) l)-    | SNat' @1 `PCons` SNat' @0 `PCons` PNil <- perm ->-      astFromVector snat2 (STKS (snat1 :$$ sh) x)-                    (V.map (astReplicate snat1 stk2) l)-  Ast.AstCond b u v -> astCond b (astTransposeS perm u) (astTransposeS perm v)--  Ast.AstLet var u v -> astLet var u (astTransposeS perm v)--  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astTransposeS perm v-  Ast.AstFromDual v -> Ast.AstFromDual $ astTransposeS perm v--  AstPlusS u v -> astTransposeS perm u + astTransposeS perm v-  AstTimesS u v -> astTransposeS perm u * astTransposeS perm v-  Ast.AstN1S NegateOp u -> negate (astTransposeS perm u)-  Ast.AstN1S AbsOp u -> abs (astTransposeS perm u)-  Ast.AstN1S SignumOp u -> signum (astTransposeS perm u)-  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astTransposeS perm u)-  Ast.AstR2S opCode u v ->-    Ast.AstR2S opCode (astTransposeS perm u) (astTransposeS perm v)-  Ast.AstI2S opCode u v ->-    Ast.AstI2S opCode (astTransposeS perm u) (astTransposeS perm v)-  AstConcreteS v -> astConcreteS (tstranspose perm $ Concrete v)---  Ast.AstFloorS v -> astFloorS $ astTransposeS perm v---  Ast.AstFromIntegralS v -> astFromIntegralS $ astTransposeS perm v---  Ast.AstCastS v -> astCastS $ astTransposeS perm v--  Ast.AstIndexS @shm shn v ix | SNat @n <- ixsRank ix ->-    Permutation.permFromList-      (Permutation.permToList'-       $ iterate (unsafeCoerce Permutation.permShift1) perm-         !! (valueOf @n))  -- this has a fake type, but that's fine-      $ \ (permn :: Perm permn) ->-        fromMaybe (error "astTransposeS: impossible non-permutation")-        $ Permutation.permCheckPermutation permn-        $ gcastWith (unsafeCoerceRefl-                     :: Permutation.PermutePrefix permn (shm ++ sh)-                        :~: shm ++ Permutation.PermutePrefix perm sh)--- should suffice, but it doesn't---      $ gcastWith (unsafeCoerceRefl :: Rank permn :~: n + Rank perm)-        $ gcastWith (unsafeCoerceRefl-                     :: (Rank permn <=? Rank (shm ++ sh)) :~: True)-        $ astIndexS (shsPermutePrefix perm shn) (astTransposeS permn v) ix-  Ast.AstScatterS @shm @shn @shp shn v (vars, ix)-    -- TODO: should the below be backpermute or permute?-    | gcompare (Permutation.permRank perm) (ixsRank ix) /= GGT ->-        let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix perm shp)-            ix2 = ixsPermutePrefix perm ix-        in gcastWith (unsafeCoerceRefl-                      :: Permutation.PermutePrefix perm shp ++ shn-                         :~: Permutation.PermutePrefix perm (shp ++ shn)) $-           astScatterS @shm @shn @(Permutation.PermutePrefix perm shp)-                       shn v (vars, ix2)-  Ast.AstGatherS @shm @shn @shp shn v (vars, ix)-    -- TODO: should the below be backpermute or permute?-    | gcompare (Permutation.permRank perm) (listsRank vars) /= GGT ->-        let vars2 :: AstVarListS (Permutation.PermutePrefix perm shm)-            vars2 = listsPermutePrefix perm vars-        in gcastWith (unsafeCoerceRefl-                      :: Permutation.PermutePrefix perm shm ++ shn-                         :~: Permutation.PermutePrefix perm (shm ++ shn)) $-           astGatherS @(Permutation.PermutePrefix perm shm) @shn @shp-                      shn v (vars2, ix)-  Ast.AstAppendS u v | SNat' @0 `PCons` _ <- perm-                     , FTKS ((:$$) @m @sh2 _ _) _ <- ftkAst u-                     , FTKS ((:$$) @n _ _) _ <- ftkAst v ->-    gcastWith (unsafeCoerceRefl-               :: Permutation.PermutePrefix perm ((m + n) : sh2)-                  :~: m + n : Tail (Permutation.PermutePrefix-                                      perm (m : sh2))) $-    gcastWith (unsafeCoerceRefl-               :: Permutation.PermutePrefix perm ((m + n) : sh2)-                  :~: m + n : Tail (Permutation.PermutePrefix-                                      perm (n : sh2))) $-    astAppendS (astTransposeS perm u) (astTransposeS perm v)-  Ast.AstSliceS i n@(SNat @n) k u | SNat' @0 `PCons` _ <- perm-                                  , FTKS ((:$$) @ink @sh2 _ _) _ <- ftkAst u ->-    gcastWith (unsafeCoerceRefl-               :: Permutation.PermutePrefix perm (n : sh2)-                  :~: n : Tail (Permutation.PermutePrefix perm (ink : sh2))) $-    astSliceS i n k (astTransposeS perm u)-  Ast.AstReverseS u | SNat' @0 `PCons` _ <- perm ->-    astReverseS (astTransposeS perm u)-  Ast.AstTransposeS @_ @sh2 perm2 u | FTKS sh2 _ <- ftkAst u ->-    -- TODO: try to perform at type level-    let permV = Permutation.permToList' perm-        perm2V = Permutation.permToList' perm2-        perm2Matched =-          perm2V-          ++ take (length permV - length perm2V) (drop (length perm2V) [0 ..])-        perm3V = normalizePermutationHack-                 $ backpermutePrefixList permV perm2Matched-    in Permutation.permFromList perm3V $ \(perm3 :: Permutation.Perm perm3) ->-      fromMaybe (error "astTransposeS: impossible non-permutation")-      $ Permutation.permCheckPermutation perm3-      $ gcastWith (unsafeCoerceRefl-                   :: Permutation.PermutePrefix perm3 sh2-                      :~: Permutation.PermutePrefix perm sh) $-        case compare (length perm3V) (shsLength sh2) of-          LT -> gcastWith (unsafeCoerceRefl-                           :: Compare (Rank perm3) (Rank sh2) :~: LT) $-                astTransposeS perm3 u-          EQ -> gcastWith (unsafeCoerceRefl-                           :: Compare (Rank perm3) (Rank sh2) :~: EQ) $-                astTransposeS perm3 u-          GT -> error "astTransposeS: GT"-  u -> Ast.AstTransposeS perm u---- TODO: try to cover the AstGatherS case here, which would permit--- not expanding to astReshapeAsGatherS in astGatherCase--- | Beware, this does not do full simplification, which often requires--- the gather form, so astReshapeAsGather needs to be called in addition--- if full simplification is required.-astReshapeS :: forall sh sh2 x s. (Product sh ~ Product sh2, AstSpan s)-            => ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh x)-            -> AstTensor AstMethodLet s (TKS2 sh2 x)-astReshapeS sh2 t = case t of-  Ast.AstReplicate (SNat' @1) STKS{} x -> astReshapeS sh2 x-  _ | Just u <- unRepl t-    , Refl <- lemAppNil @sh2 -> astReplicateNS sh2 u-  _ | STKS ((:$$) @_ @sh1 k _) x <- ftkToSTK $ ftkAst t-    , Just u <- unRepl1 t-    , (:$$) @_ @rest2 k2 rest2 <- sh2-    , Just Refl <- testEquality k k2 ->-      gcastWith (unsafeCoerceRefl :: Product rest2 :~: Product sh1) $-      astReplicate k (STKS rest2 x) $ astReshapeS rest2 u-  Ast.AstLet var u v -> astLet var u (astReshapeS @_ @sh2 sh2 v)-  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astReshapeS sh2 v-  Ast.AstFromDual v -> Ast.AstFromDual $ astReshapeS sh2 v-  -- Reshaping can be costly, so we don't touch AstTimesS, etc.-  Ast.AstN1S NegateOp u -> negate (astReshapeS @_ @sh2 sh2 u)-  Ast.AstN1S AbsOp u -> abs (astReshapeS @_ @sh2 sh2 u)-  Ast.AstN1S SignumOp u -> signum (astReshapeS @_ @sh2 sh2 u)-  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astReshapeS @_ @sh2 sh2 u)-  Ast.AstReshapeS _ v -> astReshapeS @_ @sh2 sh2 v-  _ | FTKS sh _ <- ftkAst t -> case testEquality sh sh2 of-    Just Refl -> t-    _ -> Ast.AstReshapeS sh2 t--astConvert-  :: AstSpan s-  => TKConversion y z -> AstTensor AstMethodLet s y-  -> AstTensor AstMethodLet s z-astConvert c a = case (ftkAst a, convertFTK c (ftkAst a)) of-  (yftk, zftk) | Just Refl <- matchingFTK yftk zftk -> a-    -- this covers the ConvId case and more, so not simplifying c at worst-    -- causes c to take more memory but doesn't inhibit rewriting-  _ | Ast.AstConvert c2 t2 <- a -> astConvert (ConvCmp c c2) t2-  (FTKScalar @ry, zftk@(FTKS ZSS (FTKScalar @rz)))-    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->-      astConvertSFromK c zftk a-  (FTKR shr xy, zftk@(FTKS sh xz))-    | Just Refl <- matchingFTK xy xz-    , Just Refl <- testEquality (shrRank shr) (shsRank sh) ->-      astConvertSFromR c zftk a-  (FTKX shx xy, zftk@(FTKS sh xz))-    | Just Refl <- matchingFTK xy xz-    , Just Refl <- testEquality (shxRank shx) (shsRank sh) ->-      astConvertSFromX c zftk a-  (FTKS{}, zftk) -> astConvertFromS c zftk a-  (yftk@FTKProduct{}, zftk) | checkFtkAstFromS yftk zftk ->-    astConvertFromS c zftk a-  (yftk@FTKProduct{}, zftk) | checkFtkAstSFrom yftk zftk ->-    astConvertSFrom c zftk a-  _ -> Ast.AstConvert c a---- We are pulling conversions from shaped tensors up, generally.--- z is shaped or a product (presumably with some shaped components,--- but it's not checked; the other components are supposed to be--- converted identically, which is not checked in c, either).-astConvertFromS-  :: AstSpan s-  => TKConversion y z -> FullShapeTK z -> AstTensor AstMethodLet s y-  -> AstTensor AstMethodLet s z-astConvertFromS c zftk a = case (zftk, a) of-  (_, Ast.AstConvert c2 a2) -> astConvert (ConvCmp c c2) a2-  (FTKScalar @r1, AstConcreteS @r2 v)-    | ZSS <- Nested.sshape v-    , Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->-      AstConcreteK (Nested.sunScalar v)-  (_, Ast.AstFromPrimal v) ->-    Ast.AstFromPrimal $ astConvertFromS c zftk v-      -- a rare case where we don't pull up but push down so that conversions-      -- don't end up interspersed with AstFromPrimal and similar-  (_, Ast.AstFromDual v) ->-    Ast.AstFromDual $ astConvertFromS c zftk v-  (FTKScalar, Ast.AstCond b v1 v2) ->-    astCond b (astConvertFromS c FTKScalar v1)-              (astConvertFromS c FTKScalar v2)-      -- for scalars, we don't pull up but push down-  (FTKScalar, Ast.AstLet var u v) ->-    astLet var u (astConvertFromS c FTKScalar v)-  (FTKScalar, AstPlusS u v) ->-    astConvertFromS c FTKScalar u + astConvertFromS c FTKScalar v-  (FTKScalar, AstTimesS u v) ->-    astConvertFromS c FTKScalar u * astConvertFromS c FTKScalar v-  (FTKScalar, Ast.AstN1S NegateOp u) ->-    negate (astConvertFromS c FTKScalar u)-  (FTKScalar, Ast.AstN1S AbsOp u) ->-    abs (astConvertFromS c FTKScalar u)-  (FTKScalar, Ast.AstN1S SignumOp u) ->-    signum (astConvertFromS c FTKScalar u)-  (FTKScalar @r1, Ast.AstI2S @r2 QuotOp u v)-    | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->-      astConvertFromS c FTKScalar u-      `quotH` astConvertFromS c FTKScalar v-  (FTKScalar @r1, Ast.AstI2S @r2 RemOp u v)-    | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->-      astConvertFromS c FTKScalar u-      `remH` astConvertFromS c FTKScalar v-  _ -> Ast.AstConvert c a  -- by default we pull up---- We are pushing conversions to shaped tensors down, into concrete values--- and towards variables, so that the conversions often cancel out.-astConvertSFromK :: forall r s. AstSpan s-                 => TKConversion (TKScalar r) (TKS '[] r)-                 -> FullShapeTK (TKS '[] r)-                 -> AstTensor AstMethodLet s (TKScalar r)-                 -> AstTensor AstMethodLet s (TKS '[] r)-astConvertSFromK c zftk@(FTKS ZSS FTKScalar) a0 = case a0 of-  Ast.AstConvert c2 a2 -> astConvert (ConvCmp c c2) a2-  {- TODO: this is the right thing to do, but it results in unreadable-     terms with big tconvert wrappers over product type variables-  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->-    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t-  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->-    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}-  Ast.AstProject1{} -> Ast.AstConvert c a0-  Ast.AstProject2{} -> Ast.AstConvert c a0-  Ast.AstSum snat@SNat STKScalar a -> astSum snat (STKS ZSS STKScalar) a-  AstConcreteK k -> AstConcreteS $ Nested.sscalar k-  Ast.AstFloorK{} -> Ast.AstConvert c a0-  Ast.AstFromIntegralK{} -> Ast.AstConvert c a0-  Ast.AstCastK{} -> Ast.AstConvert c a0-  AstPlusK u v ->-    astConvertSFromK c zftk u + astConvertSFromK c zftk v-  AstTimesK u v ->-    astConvertSFromK c zftk u * astConvertSFromK c zftk v-  Ast.AstN1K NegateOp u ->-    negate (astConvertSFromK c zftk u)-  Ast.AstN1K AbsOp u ->-    abs (astConvertSFromK c zftk u)-  Ast.AstN1K SignumOp u ->-    signum (astConvertSFromK c zftk u)-  Ast.AstR1K opCode u -> Ast.AstR1S opCode (astConvertSFromK c zftk u)-  Ast.AstR2K opCode u v ->-    Ast.AstR2S opCode (astConvertSFromK c zftk u)-                      (astConvertSFromK c zftk v)-  Ast.AstI2K QuotOp u v ->-    astConvertSFromK c zftk u `quotH` astConvertSFromK c zftk v-  Ast.AstI2K RemOp u v ->-    astConvertSFromK c zftk u `remH` astConvertSFromK c zftk v-  Ast.AstApply{} -> Ast.AstConvert c a0-  Ast.AstVar{} -> Ast.AstConvert c a0-  Ast.AstCond b v w -> astCond b (astConvertSFromK c zftk v)-                                 (astConvertSFromK c zftk w)-  Ast.AstLet var u v -> astLet var u (astConvertSFromK c zftk v)-  Ast.AstPrimalPart a -> astPrimalPart $ astConvertSFromK c zftk a-  Ast.AstDualPart a -> astDualPart $ astConvertSFromK c zftk a-  Ast.AstFromPrimal a -> Ast.AstFromPrimal $ astConvertSFromK c zftk a-  Ast.AstFromDual a -> Ast.AstFromDual $ astConvertSFromK c zftk a--astConvertSFromR :: forall sh x s. AstSpan s-                 => TKConversion (TKR2 (Rank sh) x) (TKS2 sh x)-                 -> FullShapeTK (TKS2 sh x)-                 -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)-                 -> AstTensor AstMethodLet s (TKS2 sh x)-astConvertSFromR c zftk@(FTKS sh x) a0 = case a0 of-  Ast.AstConvert c2 a2 -> astConvert (ConvCmp c c2) a2-  {- TODO: this is the right thing to do, but it results in unreadable-     terms with big tconvert wrappers over product type variables-  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->-    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t-  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->-    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}-  Ast.AstProject1{} -> Ast.AstConvert c a0-  Ast.AstProject2{} -> Ast.AstConvert c a0-  -- TODO: here and elsewhere, make sure the generated c2 is unique/correct-  Ast.AstFromVector snat@SNat (STKR @n _ xstk) l -> case sh of-    snat2 :$$ rest | Just Refl <- sameNat snat snat2-                   , Refl <- lemRankReplicate (Proxy @n) ->-      let c2 = ConvCmp (ConvXS' (FTKS rest x)) ConvRX-      in astFromVector snat (STKS rest xstk)-                       (V.map (astConvert c2) l)-    _ -> error "astConvertSFromR: impossible shape"-  Ast.AstSum snat@SNat (STKR @n _ xstk) a-    | Refl <- lemRankReplicate (Proxy @(1 + n)) ->-      let c2 = ConvCmp (ConvXS' (FTKS (snat :$$ sh) x)) ConvRX-          !a2 = astConvert c2 a-      in astSum snat (STKS sh xstk) a2-  Ast.AstReplicate snat@SNat (STKR @n _ xstk) a -> case sh of-    snat2 :$$ rest | Just Refl <- sameNat snat snat2-                   , Refl <- lemRankReplicate (Proxy @n) ->-      let c2 = ConvCmp (ConvXS' (FTKS rest x)) ConvRX-          !a2 = astConvert c2 a-      in astReplicate snat (STKS rest xstk) a2-    _ -> error "astConvertSFromR: impossible shape"-  Ast.AstApply{} -> Ast.AstConvert c a0-  Ast.AstVar{} -> Ast.AstConvert c a0-  Ast.AstCond b v w -> astCond b (astConvertSFromR c zftk v)-                                 (astConvertSFromR c zftk w)-  Ast.AstBuild1 snat@SNat (STKR @n _ xstk) (var, a) -> case sh of-    snat2 :$$ rest | Just Refl <- sameNat snat snat2-                   , Refl <- lemRankReplicate (Proxy @n) ->-      let c2 = ConvCmp (ConvXS' (FTKS rest x)) ConvRX-          !a2 = astConvert c2 a-      in Ast.AstBuild1 snat (STKS rest xstk) (var, a2)-    _ -> error "astConvertSFromR: impossible shape"-  Ast.AstLet var u v -> astLet var u (astConvertSFromR c zftk v)-  Ast.AstPrimalPart a -> astPrimalPart $ astConvertSFromR c zftk a-  Ast.AstDualPart a -> astDualPart $ astConvertSFromR c zftk a-  Ast.AstFromPrimal a ->-    Ast.AstFromPrimal $ astConvertSFromR c zftk a-  Ast.AstFromDual a ->-    Ast.AstFromDual $ astConvertSFromR c zftk a--astConvertSFromX :: forall sh shx x s. (AstSpan s, Rank shx ~ Rank sh)-                 => TKConversion (TKX2 shx x) (TKS2 sh x)-                 -> FullShapeTK (TKS2 sh x)-                 -> AstTensor AstMethodLet s (TKX2 shx x)-                 -> AstTensor AstMethodLet s (TKS2 sh x)-astConvertSFromX c zftk@(FTKS sh x) a0 = case a0 of-  Ast.AstConvert c2 a2 -> astConvert (ConvCmp c c2) a2-  {- TODO: this is the right thing to do, but it results in unreadable-     terms with big tconvert wrappers over product type variables-  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->-    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t-  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->-    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}-  Ast.AstProject1{} -> Ast.AstConvert c a0-  Ast.AstProject2{} -> Ast.AstConvert c a0-  Ast.AstFromVector snat@SNat (STKX @shx2 _ xstk) l -> case sh of-    (:$$) @_ @rest snat2 rest | Just Refl <- sameNat snat snat2 ->-      -- This is needed only for GHC 9.10.-      gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank shx2) $-      let c2 = ConvXS' (FTKS rest x)-      in astFromVector snat (STKS rest xstk)-                       (V.map (astConvert c2) l)-    _ -> error "astConvertSFromX: impossible shape"-  Ast.AstSum snat@SNat (STKX _ xstk) a ->-      let c2 = ConvXS' (FTKS (snat :$$ sh) x)-          !a2 = astConvert c2 a-      in astSum snat (STKS sh xstk) a2-  Ast.AstReplicate snat@SNat (STKX @shx2 _ xstk) a -> case sh of-    (:$$) @_ @rest snat2 rest | Just Refl <- sameNat snat snat2 ->-      -- This is needed only for GHC 9.10.-      gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank shx2) $-      let c2 = ConvXS' (FTKS rest x)-          !a2 = astConvert c2 a-      in astReplicate snat (STKS rest xstk) a2-    _ -> error "astConvertSFromX: impossible shape"-  Ast.AstApply{} -> Ast.AstConvert c a0-  Ast.AstVar{} -> Ast.AstConvert c a0-  Ast.AstCond b v w -> astCond b (astConvertSFromX c zftk v)-                                 (astConvertSFromX c zftk w)-  Ast.AstBuild1 snat@SNat (STKX @shx2 _ xstk) (var, a) -> case sh of-    (:$$) @_ @rest snat2 rest | Just Refl <- sameNat snat snat2 ->-      -- This is needed only for GHC 9.10.-      gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank shx2) $-      let c2 = ConvXS' (FTKS rest x)-          !a2 = astConvert c2 a-      in Ast.AstBuild1 snat (STKS rest xstk) (var, a2)-    _ -> error "astConvertSFromX: impossible shape"-  Ast.AstLet var u v -> astLet var u (astConvertSFromX c zftk v)-  Ast.AstPrimalPart a -> astPrimalPart $ astConvertSFromX c zftk a-  Ast.AstDualPart a -> astDualPart $ astConvertSFromX c zftk a-  Ast.AstFromPrimal a -> Ast.AstFromPrimal $ astConvertSFromX c zftk a-  Ast.AstFromDual a -> Ast.AstFromDual $ astConvertSFromX c zftk a---- We are pushing conversions to shaped tensors down, into concrete values--- and towards variables, so that the conversions often cancel out.-astConvertSFrom :: forall y z s. AstSpan s-                => TKConversion y z -> FullShapeTK z-                -> AstTensor AstMethodLet s y-                -> AstTensor AstMethodLet s z-astConvertSFrom c zftk t = case (zftk, ftkAst t) of-  (_, yftk) | Just Refl <- matchingFTK yftk zftk -> t-  _ | Ast.AstConvert c2 t2 <- t -> astConvert (ConvCmp c c2) t2-  (FTKS ZSS (FTKScalar @rz), FTKScalar @ry) ->-    case testEquality (typeRep @ry) (typeRep @rz) of-      Just Refl -> astConvertSFromK c zftk t-      Nothing -> error "astConvertSFrom: tensor kinds don't match"-  (FTKS shz zx, FTKR shr yx)->-    case (matchingFTK yx zx, testEquality (shsRank shz) (shrRank shr)) of-      (Just Refl, Just Refl) -> astConvertSFromR c zftk t-      _ -> error "astConvertSFrom: tensor kinds don't match"-  (FTKS shz zx, FTKX shy yx) ->-    case (matchingFTK yx zx, testEquality (shsRank shz) (shxRank shy)) of-      (Just Refl, Just Refl) -> astConvertSFromX c zftk t-      _ -> error "astConvertSFrom: tensor kinds don't match"-  (FTKProduct zftk1 zftk2, FTKProduct yftk1 yftk2) -> case t of-    Ast.AstPair a1 a2 ->-      -- Here we can't always use the c the user presumably wrote,-      -- so we always create a canonical one.-      astPair (astConvertSFrom (convSFrom yftk1 (ftkToSTK zftk1)) zftk1 a1)-              (astConvertSFrom (convSFrom yftk2 (ftkToSTK zftk2)) zftk2 a2)-    _ -> Ast.AstConvert c t  -- don't introduce let just to push a conversion-  (_, yftk) ->-    error $ "astConvertSFrom: wrong tensor kinds: " ++ show (yftk, zftk, t)--pattern AstSFromK' :: () => sh ~ '[]-                   => AstTensor AstMethodLet s (TKScalar r)-                   -> AstTensor AstMethodLet s (TKS sh r)-pattern AstSFromK' t <--  Ast.AstConvert c (checkPatternAstSFromK' c -> Just (Refl, t))--checkPatternAstSFromK' :: TKConversion y (TKS2 sh (TKScalar r))-                       -> AstTensor AstMethodLet s y-                       -> Maybe ( sh :~: '[]-                                , AstTensor AstMethodLet s (TKScalar r) )-checkPatternAstSFromK' c t-  | FTKScalar @ry <- ftkAst t-  , FTKS ZSS (FTKScalar @r) <- convertFTK c (ftkAst t)-  , Just Refl <- testEquality (typeRep @ry) (typeRep @r) = Just (Refl, t)-checkPatternAstSFromK' _ _ = Nothing---- The smart constructor with the prime suffix create canonical conversions--- instead of taking a conversion as an argument.-astFromS' :: forall y z s. AstSpan s-          => FullShapeTK z -> AstTensor AstMethodLet s y-          -> AstTensor AstMethodLet s z-astFromS' zftk t = astConvertFromS (convFromS (ftkAst t) zftk) zftk t--astSFrom' :: forall y z s. AstSpan s-          => SingletonTK z -> AstTensor AstMethodLet s y-          -> AstTensor AstMethodLet s z-astSFrom' zstk t = astConvert (convSFrom (ftkAst t) zstk) t--astKFromS' :: forall r s. (AstSpan s, GoodScalar r)-           => AstTensor AstMethodLet s (TKS2 '[] (TKScalar r))-           -> AstTensor AstMethodLet s (TKScalar r)-astKFromS' t = astConvertFromS (ConvCmp ConvX0 ConvSX) FTKScalar t---- Or should we take SNat (Rank sh) to help proving n ~ Rank sh?-astRFromS' :: forall sh x s. AstSpan s-           => AstTensor AstMethodLet s (TKS2 sh x)-           -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)-astRFromS' t | FTKS sh x <- ftkAst t-             , Refl <- lemRankMapJust sh =-  let zftk = FTKR (shrFromShS sh) x-  in astConvertFromS (ConvCmp (ConvXR (ftkToSTK x)) ConvSX) zftk t--astXFromS' :: forall sh shx x s. (AstSpan s, Rank sh ~ Rank shx)-           => StaticShX shx -> AstTensor AstMethodLet s (TKS2 sh x)-           -> AstTensor AstMethodLet s (TKX2 shx x)-astXFromS' ssx t | FTKS sh x <- ftkAst t-                 , Refl <- lemRankMapJust sh =-  let zftk = FTKX (shCastSX ssx sh) x-  in astConvertFromS (ConvCmp (ConvXX' zftk) ConvSX) zftk t--astSFromK' :: forall r s. (GoodScalar r, AstSpan s)-           => AstTensor AstMethodLet s (TKScalar r)-           -> AstTensor AstMethodLet s (TKS '[] r)-astSFromK' a =-  let c2 = ConvCmp ConvXS (Conv0X STKScalar)-  in astConvertSFromK c2 (FTKS ZSS FTKScalar) a--astSFromR' :: forall sh s r. AstSpan s-           => ShS sh -> AstTensor AstMethodLet s (TKR2 (Rank sh) r)-           -> AstTensor AstMethodLet s (TKS2 sh r)-astSFromR' sh t = case ftkAst t of-  FTKR _ x | Refl <- lemRankReplicate (Proxy @(Rank sh)) ->-    let zftk = FTKS sh x-    in astConvertSFromR (ConvCmp (ConvXS' zftk) ConvRX) zftk t--astSFromX' :: forall sh sh' s x. (AstSpan s, Rank sh ~ Rank sh')-           => ShS sh -> AstTensor AstMethodLet s (TKX2 sh' x)-           -> AstTensor AstMethodLet s (TKS2 sh x)-astSFromX' sh t = case ftkAst t of-  FTKX _ x ->-    let zftk = FTKS sh x-    in astConvertSFromX (ConvXS' zftk) zftk t--astSum0S :: AstSpan s-         => AstTensor AstMethodLet s (TKS2 sh x)-         -> AstTensor AstMethodLet s (TKS2 '[] x)-astSum0S t = case t of-  Ast.AstSum SNat _ u -> astSum0S u-  Ast.AstReplicate snat (STKS _ STKScalar) u ->-    astSum0S u * (fromPrimal $ AstConcreteS-                  $ Nested.sscalar $ fromInteger $ fromSNat snat)-  Ast.AstReplicate snat STKScalar u ->-    astSFromK' u * (fromPrimal $ AstConcreteS-                    $ Nested.sscalar $ fromInteger $ fromSNat snat)-  Ast.AstLet var u v -> astLet var u (astSum0S v)-  AstTimesS t1 t2 -> astDot0S t1 t2-  AstConcreteS v ->-    withKnownShS (Nested.sshape v) $-    astConcreteS $ tssum0 (Concrete v)-  Ast.AstFromPrimal u -> Ast.AstFromPrimal $ astSum0S u-  Ast.AstFromDual u -> Ast.AstFromDual $ astSum0S u-  Ast.AstReverseS u -> astSum0S u-  Ast.AstTransposeS _ u -> astSum0S u-  Ast.AstReshapeS _ u -> astSum0S u-  Ast.AstN1S NegateOp u -> negate $ astSum0S u-  Ast.AstSum0S u -> astSum0S u-  Ast.AstDot0S{} -> t-  Ast.AstDot1InS _ _ t1 t2 -> astDot0S t1 t2-  Ast.AstMatmul2S m@SNat SNat p@SNat m1 m2 ->-    astDot0S (astTransposeS (Permutation.makePerm @'[1, 0])-                            (astReplicate p knownSTK m1))-             (astTransposeS (Permutation.makePerm @'[0, 2, 1])-                            (astReplicate m knownSTK m2))-  _ -> Ast.AstSum0S t--astDot0S :: (GoodScalar r, AstSpan s)-         => AstTensor AstMethodLet s (TKS sh r)-         -> AstTensor AstMethodLet s (TKS sh r)-         -> AstTensor AstMethodLet s (TKS '[] r)-astDot0S t1 t2 = case (t1, t2) of-  (AstConcreteS v1, AstConcreteS v2) ->-    withKnownShS (Nested.sshape v1) $-    astConcreteS $ tsdot0 (Concrete v1) (Concrete v2)-  _ | FTKS (snat :$$ _) _ <- ftkAst t1-    , Just u1 <- unRepl1 t1-    , Just u2 <- unRepl1 t2 ->-      astDot0S u1 u2 * (fromPrimal $ AstConcreteS-                        $ Nested.sscalar $ fromInteger $ fromSNat snat)-  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->-    Ast.AstFromPrimal $ astDot0S u1 u2-  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->-    Ast.AstFromDual $ astDot0S u1 u2-  (Ast.AstTransposeS @_ @sh1 perm1 u1, Ast.AstTransposeS @_ @sh2 perm2 u2)-    | Just Refl <- testEquality perm1 perm2 ->-      gcastWith (unsafeCoerceRefl :: sh1 :~: sh2) $-      astDot0S u1 u2-  (Ast.AstReverseS u1, Ast.AstReverseS u2) -> astDot0S u1 u2-  (Ast.AstN1S NegateOp u1, Ast.AstN1S NegateOp u2) -> astDot0S u1 u2-  _ -> Ast.AstDot0S t1 t2--astDot1InS :: forall sh n r s. GoodScalar r-           => ShS sh -> SNat n-           -> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)-           -> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)-           -> AstTensor AstMethodLet s (TKS sh r)-astDot1InS sh n@SNat t1 t2 = case (t1, t2) of-  (AstConcreteS v1, AstConcreteS v2) ->-    withKnownShS sh $-    astConcreteS $ tsdot1In @_ @sh (SNat @n) (Concrete v1) (Concrete v2)-  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->-    Ast.AstFromPrimal $ astDot1InS sh n u1 u2-  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->-    Ast.AstFromDual $ astDot1InS sh n u1 u2-  _ -> Ast.AstDot1InS sh n t1 t2--astMatmul2S :: GoodScalar r-            => SNat m -> SNat n -> SNat p-            -> AstTensor AstMethodLet s (TKS '[m, n] r)-            -> AstTensor AstMethodLet s (TKS '[n, p] r)-            -> AstTensor AstMethodLet s (TKS '[m, p] r)-astMatmul2S m@SNat n@SNat p@SNat t1 t2 = case (t1, t2) of-  (AstConcreteS v1, AstConcreteS v2) ->-    astConcreteS $ tsmatmul2 (Concrete v1) (Concrete v2)-  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->-    Ast.AstFromPrimal $ astMatmul2S m n p u1 u2-  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->-    Ast.AstFromDual $ astMatmul2S m n p u1 u2-  _ -> Ast.AstMatmul2S m n p t1 t2----- * ConvertTensor instances needed for unwinding in astConcrete--instance AstSpan s => ConvertTensor (AstTensor AstMethodLet s) where-  tconvert c _astk v = astConvert c v--  rfromX a = case ftkAst a of-    FTKX sh' _ ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        withKnownShS sh $-        rfromS $ sfromX @_ @sh a-  xfromR a = case ftkAst a of-    FTKR shr _ ->-      withShsFromShR shr $ \(sh :: ShS sh) ->-        withKnownShS sh $-        xfromS @_ @sh $ sfromR a--  sfromR = astSFromR' knownShS-  sfromX = astSFromX' knownShS-  xfromS = astXFromS' knownShX--  rzip @_ @_ @n a-   | Refl <- lemRankReplicate (Proxy @n) = case ftkAst a of-    FTKProduct (FTKR _sh y) (FTKR _ z) ->-      let c = ConvCmp-                (ConvXR (ftkToSTK (FTKProduct y z)))-                (ConvCmp-                   (ConvZip (ftkToSTK y) (ftkToSTK z))-                   (ConvT2 ConvRX ConvRX))-      in astConvert c a-  runzip @_ @_ @n a-   | Refl <- lemRankReplicate (Proxy @n) = case ftkAst a of-    FTKR _sh (FTKProduct y z) ->-      let c = ConvCmp-                (ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))-                (ConvCmp-                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))-                   ConvRX)-      in astConvert c a-  szip a = case ftkAst a of-    FTKProduct (FTKS _sh y) (FTKS _ z) ->-      let c = ConvCmp-                ConvXS-                (ConvCmp-                   (ConvZip (ftkToSTK y) (ftkToSTK z))-                   (ConvT2 ConvSX ConvSX))-      in astConvert c a-  sunzip a = case ftkAst a of-    FTKS _sh (FTKProduct y z) ->-      let c = ConvCmp-                (ConvT2 ConvXS ConvXS)-                (ConvCmp-                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))-                   ConvSX)-      in astConvert c a-  xzip a = case ftkAst a of-    FTKProduct (FTKX _sh y) (FTKX _ z) ->-      let c = ConvZip (ftkToSTK y) (ftkToSTK z)-      in astConvert c a-  xunzip a = case ftkAst a of-    FTKX _sh (FTKProduct y z) ->-      let c = ConvUnzip (ftkToSTK y) (ftkToSTK z)-      in astConvert c a--  xnestR @sh1 @m @x sh1 a-    | Refl <- lemRankReplicate (Proxy @m) =-      let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)-                            (TKX2 sh1 (TKR2 m x))-          c = ConvCmp-                (ConvXX (ConvXR (knownSTK @x)))-                (ConvNest @_ @_ @(Replicate m Nothing)-                          (STKX sh1 (knownSTK @x)))-      in astConvert c a-  xnestS @_ @_ @x sh1 a =-    let c = ConvCmp (ConvXX ConvXS)-                    (ConvNest (STKX sh1 (knownSTK @x)))-    in astConvert c a-  xnest @_ @_ @x sh1 a =-    let c = ConvNest (STKX sh1 (knownSTK @x))-    in astConvert c a-  xunNestR a =-    let c = ConvCmp ConvUnnest-                    (ConvXX ConvRX)-    in astConvert c a-  xunNestS a =-    let c = ConvCmp ConvUnnest-                    (ConvXX ConvSX)-    in astConvert c a-  xunNest a =-    let c = ConvUnnest-    in astConvert c a--  tpairConv = astPair-  tunpairConv t = (astProject1 t, astProject2 t)----- * Helper combinators---- All but the last case are shortcuts for common forms.-astConcrete :: FullShapeTK y -> Concrete y-            -> AstTensor AstMethodLet PrimalSpan y-astConcrete ftk v = case ftk of-  FTKScalar -> astConcreteK v-  FTKR sh' FTKScalar ->-    withShsFromShR sh' $ \(sh :: ShS sh) ->-      withKnownShS sh $-      astFromS' ftk $ astConcreteS (sfromR @_ @sh v)-  FTKS _ FTKScalar -> astConcreteS v-  FTKX sh' FTKScalar ->-    withShsFromShX sh' $ \(sh :: ShS sh) ->-      withKnownShS sh $-      astFromS' ftk $ astConcreteS (sfromX @_ @sh v)-  FTKProduct ftk1 ftk2 ->-    astPair (astConcrete ftk1 (tproject1 v)) (astConcrete ftk2 (tproject2 v))-  _ -> concreteTarget astConcreteK astConcreteS astFromS' (ftkToSTK ftk) v--astLetFun :: forall y z s s2. (AstSpan s, AstSpan s2)-          => AstTensor AstMethodLet s y-          -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)-          -> AstTensor AstMethodLet s2 z-astLetFun = astLetFunBounds Nothing--astLetFunB :: forall z s s2. (AstSpan s, AstSpan s2)-           => AstTensor AstMethodLet s (TKScalar Int64)-           -> (AstTensor AstMethodLet s (TKScalar Int64)-               -> AstTensor AstMethodLet s2 z)-           -> AstTensor AstMethodLet s2 z-astLetFunB w = astLetFunBounds (Just $ bounds w) w--astLetFunBounds :: forall y z s s2. (AstSpan s, AstSpan s2)-                => Maybe (Int64, Int64) -> AstTensor AstMethodLet s y-                -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)-                -> AstTensor AstMethodLet s2 z-astLetFunBounds _ a f | astIsSmall True a = f a-astLetFunBounds mbs a f = case a of-  AstFromS' FTKScalar _ ->-    let (var, ast) = funToAst2 (ftkAst a) mbs f-    in astLet var a ast-  AstFromS' @y2 ftkz v ->-    let (var, ast) = funToAst2 (ftkAst v) mbs (f . astFromS' @y2 ftkz)-    in astLet var v ast-  _ -> case ftkAst a of-    ftk@(FTKR @_ @x sh' x) ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        let (var, ast) =-              funToAst2 (FTKS sh x) mbs-                        (f . astFromS' @(TKS2 sh x) ftk)-        in astLet var (astSFromR' sh a) ast-             -- safe, because subsitution ruled out above-    ftk@(FTKX @_ @x sh' x) ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        let (var, ast) =-              funToAst2 (FTKS sh x) mbs-                        (f . astFromS' @(TKS2 sh x) ftk)-        in astLet var (astSFromX' sh a) ast-    -- calling recursively for product may be not worth it-    ftk -> let (var, ast) = funToAst2 ftk mbs f-           in astLet var a ast--astReplicateNS :: forall shn shp s x. AstSpan s-               => ShS shn -> AstTensor AstMethodLet s (TKS2 shp x)-               -> AstTensor AstMethodLet s (TKS2 (shn ++ shp) x)-astReplicateNS shn v | STKS shp x <- ftkToSTK (ftkAst v) =-  let go :: ShS shn' -> AstTensor AstMethodLet s (TKS2 (shn' ++ shp) x)-      go ZSS = v-      go (snat :$$ shn2) =-        astReplicate snat (STKS (shn2 `shsAppend` shp) x) $ go shn2-  in go shn--unRepl :: AstSpan s-       => AstTensor AstMethodLet s (TKS2 sh x)-       -> Maybe (AstTensor AstMethodLet s (TKS2 '[] x))-unRepl (Ast.AstReplicate _ (STKS ZSS _) u) = Just u-unRepl (Ast.AstReplicate _ STKScalar u) = Just $ astSFromK' u-unRepl (Ast.AstReplicate _ STKS{} u) = unRepl u-unRepl (AstConcreteS a) = AstConcreteS . Nested.sscalar <$> sunReplicateScal a-unRepl (Ast.AstCond b v1 v2) = do-  u1 <- unRepl v1-  u2 <- unRepl v2-  return $! astCond b u1 u2-unRepl (Ast.AstLet var u t) = Ast.AstLet var u <$> unRepl t-unRepl (Ast.AstPrimalPart t) = astPrimalPart <$> unRepl t-unRepl (Ast.AstDualPart t) = astDualPart <$> unRepl t-unRepl _ = Nothing--unRepl1 :: AstSpan s-        => AstTensor AstMethodLet s (TKS2 (n ': sh) x)-        -> Maybe (AstTensor AstMethodLet s (TKS2 sh x))-unRepl1 (Ast.AstReplicate _ STKS{} u) = Just u-unRepl1 (Ast.AstReplicate _ STKScalar u) = Just $ astSFromK' u-unRepl1 (AstConcreteS a) = AstConcreteS <$> sunReplicate1 a-unRepl1 (Ast.AstCond b v1 v2) = do-  u1 <- unRepl1 v1-  u2 <- unRepl1 v2-  return $! astCond b u1 u2-unRepl1 (Ast.AstLet var u t) = Ast.AstLet var u <$> unRepl1 t-unRepl1 (Ast.AstPrimalPart t) = astPrimalPart <$> unRepl1 t-unRepl1 (Ast.AstDualPart t) = astDualPart <$> unRepl1 t-unRepl1 _ = Nothing--unReplN :: AstSpan s-        => ShS shm -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)-        -> Maybe (AstTensor AstMethodLet s (TKS2 shn x))-unReplN ZSS a = Just a-unReplN (_ :$$ ZSS) (Ast.AstReplicate _ STKScalar u) = Just $ astSFromK' u-unReplN (_ :$$ sh) (Ast.AstReplicate _ STKS{} u) = unReplN sh u-unReplN shm (AstConcreteS a) = AstConcreteS <$> sunReplicateN shm a-unReplN shm (Ast.AstCond b v1 v2) = do-  u1 <- unReplN shm v1-  u2 <- unReplN shm v2-  return $! astCond b u1 u2-unReplN shm (Ast.AstLet var u t) = Ast.AstLet var u <$> unReplN shm t-unReplN shm (Ast.AstPrimalPart t) = astPrimalPart <$> unReplN shm t-unReplN shm (Ast.AstDualPart t) = astDualPart <$> unReplN shm t-unReplN _ _ = Nothing----- * Substitution wrappers---- | We assume no variable is shared between a binding and its nested binding--- and nobody substitutes into variables that are bound.--- This keeps the substitution code simple, because we never need to compare--- variables to any variable in the bindings.-substituteAst :: forall s s2 y z. (AstSpan s, AstSpan s2)-              => AstTensor AstMethodLet s2 z -> AstVarName s2 z-              -> AstTensor AstMethodLet s y-              -> AstTensor AstMethodLet s y-substituteAst i var v1 =-  fromMaybe v1 $ substitute1Ast i var v1--substituteAstIxS-  :: AstSpan s-  => AstTensor AstMethodLet s y -> AstVarName s y -> AstIxS AstMethodLet sh-  -> AstIxS AstMethodLet sh-substituteAstIxS i var ix =-  fromMaybe ix $ substitute1AstIxS i var ix--substituteAstBool-  :: AstSpan s-  => AstTensor AstMethodLet s y -> AstVarName s y -> AstBool AstMethodLet-  -> AstBool AstMethodLet-substituteAstBool i var v1 =-  fromMaybe v1 $ substitute1AstBool i var v1----- * Substitution workers---- Invariant: if the variable has bounds, the expression can only have--- values within the bounds (regardless of what the `bounds` call would say).-substitute1Ast :: forall s s2 y z. (AstSpan s, AstSpan s2)-               => AstTensor AstMethodLet s2 z -> AstVarName s2 z-               -> AstTensor AstMethodLet s y-               -> Maybe (AstTensor AstMethodLet s y)-substitute1Ast i var = subst where- subst :: forall s3 y2. AstSpan s3-       => AstTensor AstMethodLet s3 y2-       -> Maybe (AstTensor AstMethodLet s3 y2)- subst = \case-  Ast.AstPair u v ->-    case (subst u, subst v) of-      (Nothing, Nothing) -> Nothing-      (mu, mv) -> Just $ astPair (fromMaybe u mu) (fromMaybe v mv)-  Ast.AstProject1 a -> astProject1 <$> subst a-  Ast.AstProject2 a -> astProject2 <$> subst a-  Ast.AstFromVector snat stk args ->-    let margs = V.map subst args-    in if V.any isJust margs-       then Just $ astFromVector snat stk $ V.zipWith fromMaybe args margs-       else Nothing-  Ast.AstSum snat stk v -> astSum snat stk <$> subst v-  Ast.AstReplicate snat stk v ->-    astReplicate snat stk <$> subst v-  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->-      case ( substitute1AstHFun i var f, substitute1AstHFun i var df-           , substitute1AstHFun i var rf, subst acc0-           , subst es ) of-        (Nothing, Nothing, Nothing, Nothing, Nothing) -> Nothing-        (mf, mdf, mrf, macc0, mes) ->-          Just $ astMapAccumRDer k bftk eftk-                                 (fromMaybe f mf)-                                 (fromMaybe df mdf)-                                 (fromMaybe rf mrf)-                                 (fromMaybe acc0 macc0)-                                 (fromMaybe es mes)-  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->-      case ( substitute1AstHFun i var f, substitute1AstHFun i var df-           , substitute1AstHFun i var rf, subst acc0-           , subst es ) of-        (Nothing, Nothing, Nothing, Nothing, Nothing) -> Nothing-        (mf, mdf, mrf, macc0, mes) ->-          Just $ astMapAccumLDer k bftk eftk-                                 (fromMaybe f mf)-                                 (fromMaybe df mdf)-                                 (fromMaybe rf mrf)-                                 (fromMaybe acc0 macc0)-                                 (fromMaybe es mes)-  Ast.AstApply t ll ->-    case ( substitute1AstHFun i var t-         , subst ll ) of-      (Nothing, Nothing) -> Nothing-      (mt, mll) -> Just $ astApply (fromMaybe t mt) (fromMaybe ll mll)-  Ast.AstVar var2 ->-    -- We can't assert anything here, because only all runtime values need-    -- to be in bounds and bounds approximations don't have to agree.-    if varNameToAstVarId var == varNameToAstVarId var2-    then case sameAstSpan @s3 @s2 of-      Just Refl -> case testEquality var var2 of-        Just Refl -> case i of-          Ast.AstVar var3 | FTKScalar <- varNameToFTK var3 ->-            let (lb, ub) = fromMaybe (-1000000000, 1000000000)-                           $ varNameToBounds var-                (lb2, ub2) = fromMaybe (-1000000000, 1000000000)-                             $ varNameToBounds var2-                (lb3, ub3) = fromMaybe (-1000000000, 1000000000)-                             $ varNameToBounds var3-                bs = (max (max lb lb2) lb3, min (min ub ub2) ub3)-                  -- We know all bounds approximations have to be correct-                  -- so we can intersect them.-            in Just $ astVar $ mkAstVarName (varNameToFTK var3)-                                            (Just bs)-                                            (varNameToAstVarId var3)-          _ -> Just i-        _ -> error $ "substitute1Ast: kind of the variable "-                     ++ show var2 ++ ": " ++ show (varNameToFTK var)-                     ++ ", payload kind: " ++ show (varNameToFTK var2)-                     ++ ", payload: " ++ show i-      _ -> error "substitute1Ast: span"-    else Nothing-  Ast.AstCond b v w ->-    case ( substitute1AstBool i var b-         , subst v-         , subst w ) of-      (Nothing, Nothing, Nothing) -> Nothing-      (mb, mv, mw) ->-        Just $ astCond (fromMaybe b mb) (fromMaybe v mv) (fromMaybe w mw)-  Ast.AstBuild1 k stk (var2, v) ->-    assert (varNameToAstVarId var2 /= varNameToAstVarId var) $-    Ast.AstBuild1 k stk . (var2,) <$> subst v--  Ast.AstLet var2 u v ->-    case (subst u, subst v) of-      (Nothing, Nothing) -> Nothing-      (mu, mv) -> Just $ astLet var2 (fromMaybe u mu) (fromMaybe v mv)--  Ast.AstPrimalPart a -> astPrimalPart <$> subst a-  Ast.AstDualPart a -> astDualPart <$> subst a-  Ast.AstFromPrimal a -> Ast.AstFromPrimal <$> subst a-  Ast.AstFromDual a -> Ast.AstFromDual <$> subst a--  AstPlusK u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ fromMaybe u mu + fromMaybe v mv-       else Nothing-  AstTimesK u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ fromMaybe u mu * fromMaybe v mv-       else Nothing-  Ast.AstN1K NegateOp u -> negate <$> subst u-  Ast.AstN1K AbsOp u -> abs <$> subst u-  Ast.AstN1K SignumOp u -> signum <$> subst u-  Ast.AstR1K opCode u -> Ast.AstR1K opCode <$> subst u-  Ast.AstR2K opCode u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ Ast.AstR2K opCode (fromMaybe u mu) (fromMaybe v mv)-       else Nothing-  Ast.AstI2K QuotOp u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ quotH (fromMaybe u mu) (fromMaybe v mv)-       else Nothing-  Ast.AstI2K RemOp u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ remH (fromMaybe u mu) (fromMaybe v mv)-       else Nothing-  Ast.AstConcreteK{} -> Nothing-  Ast.AstFloorK a -> astFloorK <$> subst a-  Ast.AstFromIntegralK v -> astFromIntegralK <$> subst v-  Ast.AstCastK v -> astCastK <$> subst v--  AstPlusS u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ fromMaybe u mu + fromMaybe v mv-       else Nothing-  AstTimesS u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ fromMaybe u mu * fromMaybe v mv-       else Nothing-  Ast.AstN1S NegateOp u -> negate <$> subst u-  Ast.AstN1S AbsOp u -> abs <$> subst u-  Ast.AstN1S SignumOp u -> signum <$> subst u-  Ast.AstR1S opCode u -> Ast.AstR1S opCode <$> subst u-  Ast.AstR2S opCode u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ Ast.AstR2S opCode (fromMaybe u mu) (fromMaybe v mv)-       else Nothing-  Ast.AstI2S QuotOp u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ quotH (fromMaybe u mu) (fromMaybe v mv)-       else Nothing-  Ast.AstI2S RemOp u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ remH (fromMaybe u mu) (fromMaybe v mv)-       else Nothing-  Ast.AstConcreteS{} -> Nothing-  Ast.AstFloorS a -> astFloorS <$> subst a-  Ast.AstFromIntegralS a -> astFromIntegralS <$> subst a-  Ast.AstCastS v -> astCastS <$> subst v--  Ast.AstIndexS shn v ix ->-    case (subst v, substitute1AstIxS i var ix) of-      (Nothing, Nothing) -> Nothing-      (mv, mix) -> Just $ astIndexS shn (fromMaybe v mv) (fromMaybe ix mix)-  Ast.AstScatterS shn v (vars, ix) ->-    case (subst v, substitute1AstIxS i var ix) of-      (Nothing, Nothing) -> Nothing-      (mv, mix) -> Just $ astScatterS shn-                                      (fromMaybe v mv)-                                      (vars, fromMaybe ix mix)-  Ast.AstGatherS shn v (vars, ix) ->-    case (subst v, substitute1AstIxS i var ix) of-      (Nothing, Nothing) -> Nothing-      (mv, mix) -> Just $ astGatherS shn-                                     (fromMaybe v mv)-                                     (vars, fromMaybe ix mix)-  Ast.AstMinIndexS a -> Ast.AstMinIndexS <$> subst a-  Ast.AstMaxIndexS a -> Ast.AstMaxIndexS <$> subst a-  Ast.AstIotaS{} -> Nothing-  Ast.AstAppendS x y ->-    case (subst x, subst y) of-      (Nothing, Nothing) -> Nothing-      (mx, my) -> Just $ astAppendS (fromMaybe x mx) (fromMaybe y my)-  Ast.AstSliceS i2 n k v -> astSliceS i2 n k <$> subst v-  Ast.AstReverseS v -> astReverseS <$> subst v-  Ast.AstTransposeS perm v -> astTransposeS perm <$> subst v-  Ast.AstReshapeS sh v -> astReshapeS sh <$> subst v--  Ast.AstConvert c v -> astConvert c <$> subst v--  Ast.AstSum0S v -> astSum0S <$> subst v-  Ast.AstDot0S u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ astDot0S (fromMaybe u mu) (fromMaybe v mv)-       else Nothing-  Ast.AstDot1InS sh n u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ astDot1InS sh n (fromMaybe u mu) (fromMaybe v mv)-       else Nothing-  Ast.AstMatmul2S m n p u v ->-    let mu = subst u-        mv = subst v-    in if isJust mu || isJust mv-       then Just $ astMatmul2S m n p (fromMaybe u mu) (fromMaybe v mv)-       else Nothing--substitute1AstIxS-  :: AstSpan s2-  => AstTensor AstMethodLet s2 y -> AstVarName s2 y -> AstIxS AstMethodLet sh-  -> Maybe (AstIxS AstMethodLet sh)-substitute1AstIxS i var ix =-  let mix = fmap (substitute1Ast i var) ix-  in if any isJust mix-     then Just $ ixsZipWith fromMaybe ix mix-     else Nothing--substitute1AstHFun-  :: forall s s2 s3 x y z.-     AstTensor AstMethodLet s3 z -> AstVarName s3 z -> AstHFun s s2 x y-  -> Maybe (AstHFun s s2 x y)-substitute1AstHFun _i _var AstLambda{} =-  Nothing  -- no outside free variables--substitute1AstBool :: AstSpan s2-                   => AstTensor AstMethodLet s2 y -> AstVarName s2 y-                   -> AstBool AstMethodLet-                   -> Maybe (AstBool AstMethodLet)-substitute1AstBool i var = subst where- subst :: AstBool AstMethodLet-       -> Maybe (AstBool AstMethodLet)- subst = \case-  Ast.AstBoolConst{} -> Nothing-  Ast.AstBoolNot arg -> notB <$> subst arg-  Ast.AstBoolAnd arg1 arg2 ->-    let mb1 = subst arg1-        mb2 = subst arg2-    in if isJust mb1 || isJust mb2-       then Just $ fromMaybe arg1 mb1 &&* fromMaybe arg2 mb2-       else Nothing-  Ast.AstLeqK arg1 arg2 ->-    let mr1 = substitute1Ast i var arg1-        mr2 = substitute1Ast i var arg2-    in if isJust mr1 || isJust mr2-       then Just $ fromMaybe arg1 mr1 <=. fromMaybe arg2 mr2-       else Nothing-  Ast.AstLeqS arg1 arg2 ->-    let mr1 = substitute1Ast i var arg1-        mr2 = substitute1Ast i var arg2-    in if isJust mr1 || isJust mr2-       then Just $ fromMaybe arg1 mr1 <=. fromMaybe arg2 mr2-       else Nothing+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif+{-# LANGUAGE UnboxedTuples #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-orphans #-}+{-# OPTIONS_GHC -fmax-pmcheck-models=50 #-}+-- {-# OPTIONS_GHC -freduction-depth=10000 #-}+-- {-# OPTIONS_GHC -fconstraint-solver-iterations=10000 #-}+-- | This module holds smart constructors for AST, that is,+-- term-simplifying combinators corresponding to the Ast constructors.+-- The combinators simplify only on the basis of inspecting the roots of their+-- argument term trees. If the arguments get modified,+-- the modified forms are again inspected and potentially simplified again.+--+-- The limited simplification via combinators is enough to uncover redexes+-- for the factorization rules to fire and to undo some of the complication+-- introduced by vectorization. The intention is to leave intact as much+-- as possible of the original terms provided by the user while making+-- sure subterms introduced by vectorization are maximally simplified.+module HordeAd.Core.AstSimplify+  ( RewritePhase(..), SimplifyKnobs (..), defaultKnobs+  , -- * The simplifying combinators, one for almost each AST constructor+    astPair, astProject1, astProject2, astMapAccumLDer, astApply, astLet+  , astPrimalPart, astDualPart, astPlainPart++  , astPlusK, astTimesK, astN1K, astR1K, astR2K, astI2K, astConcreteK+  , astFloorK, astFromIntegralK, astCastK+  , astArgMinK, astArgMaxK, astIndexK, astIndexKnobsK++  , astPlusS, astTimesS, astN1S, astR1S, astR2S, astI2S, astConcreteS+  , astFloorS, astFromIntegralS, astCastS, astIndexS, astIndexKnobsS++  , astCondKInitial, astCondK, astCondSInitial, astCondS+  , astFromVectorK, astFromVectorS+  , astSumK, astSumS, astScatterS, astScatterKnobsS+  , astReplicateK, astReplicateS, astGatherS, astGatherKnobsS+  , astAppendS, astSliceS, astReverseS, astTransposeS, astReshapeS++  , astConvert+  , astConvDownSFromR, astConvDownSFromX, astConvUpRFromS, astConvUpXFromS++  , astSumKContract, astDot0, astDot1InS, astMatmul2S++  , astBoolNotK, astBoolNotS, astBoolAndK, astBoolAndS, astLeqK, astLeq, astLeqS++    -- * Helper combinators+  , astConcrete, astLetFun, astLetRefresh, unRepl1++    -- * Substitution operations+  , substituteAst, substituteAstIxS+  ) where++import Prelude++import Control.Exception.Assert.Sugar+import Control.Monad (mapAndUnzipM, mplus)+import Data.Default+import Data.Foldable qualified as Foldable+import Data.List (elemIndex, findIndex)+import Data.Maybe (catMaybes, fromMaybe, isJust)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Type.Ord (Compare)+import Data.Vector.Generic qualified as V+import Data.Vector.Strict qualified as Data.Vector+import GHC.TypeLits+  ( Nat+  , OrderingI (..)+  , cmpNat+  , sameNat+  , type (*)+  , type (+)+  , type (-)+  , type (<=)+  , type (<=?)+  )+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)+import Unsafe.Coerce (unsafeCoerce)++import Data.Array.Nested (MapJust, Replicate, type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert+  (shrFromShS, shxFromShS, withShsFromShR, withShsFromShX)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation (DropLen, Perm (..), TakeLen, permInverse)+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped qualified as Shaped+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types+  (Head, Init, Last, Tail, fromSNat', pattern SZ, snatMul, snatPlus, unsafeCoerceRefl)++import HordeAd.Core.Ast+  ( AstTensor (AstConcreteK, AstConcreteS, AstPlusK, AstPlusS, AstTimesK, AstTimesS)+  )+import HordeAd.Core.Ast hiding (AstTensor (..))+import HordeAd.Core.Ast qualified as Ast (AstTensor(..))+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersAst+  (eqY, eqUnknownShapes, sunReplicate1, sunReplicateN, sunReplicate, unReplC, unAstK, unAstS)+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Conversion+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind++data RewritePhase =+    PhaseUnspecified+  | PhaseVectorization+  | PhaseSimplification+  | PhaseExpansion+  | PhaseContraction+ deriving Eq++newtype SimplifyKnobs = SimplifyKnobs+  { knobPhase :: RewritePhase+  }++defaultKnobs :: SimplifyKnobs+defaultKnobs = SimplifyKnobs PhaseUnspecified++-- @PhaseVectorization@ should only affect the topmost redex.+deVect :: SimplifyKnobs -> SimplifyKnobs+deVect (SimplifyKnobs PhaseVectorization) = SimplifyKnobs PhaseUnspecified+deVect knobs = knobs+++-- * Expressing operations as Gather; introduces new variable names++-- | We keep AstTranspose terms for as long as possible, because+-- they are small and fuse nicely in many cases. For some forms of indexing+-- and nesting with reshape and gather they don't fuse, which is when+-- this function is invoked.+astTransposeAsGatherS+  :: forall perm sh s r. KnownSpan s+  => SimplifyKnobs -> Permutation.Perm perm+  -> AstTensor AstMethodLet s (TKS2 sh r)+  -> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)+{-# NOINLINE astTransposeAsGatherS #-}+astTransposeAsGatherS knobs perm v =+  let FTKS sh _ = ftkAst v+      shn = shsDropLenPerm perm sh+      shp = shsTakeLenPerm perm sh+      shpPermuted = shsPermute perm shp+  in funToVarsIxS @_ @AstMethodLet shpPermuted $ \vars ix ->+    -- See astGatherCase.AstTransposeS for similar code with more comments.+    gcastWith (lemRankMapJust shp) $+    gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $+    permInverse perm $ \(invperm :: Nested.Perm invperm) proof ->+      case proof (ssxFromShX $ shxFromShS shp) of+        Refl ->+          gcastWith (unsafeCoerceRefl+                     :: DropLen invperm+                          (Permutation.Permute perm (TakeLen perm sh))+                        :~: '[]) $+          gcastWith (lemAppNil+                       @(Permutation.Permute invperm+                           (Permutation.Permute perm (TakeLen perm sh)))) $+          -- Seriously? This should be deduced from the above:+          gcastWith (unsafeCoerceRefl+                     :: Permutation.PermutePrefix invperm+                          (Permutation.Permute perm (TakeLen perm sh))+                        :~: Permutation.Permute invperm+                          (Permutation.Permute perm (TakeLen perm sh))) $+          -- This should follow from @proof@, if not for MapJust:+          gcastWith (unsafeCoerceRefl+                     :: Permutation.Permute invperm+                          (Permutation.Permute perm (TakeLen perm sh))+                        :~: TakeLen perm sh) $+          let asts :: AstIxS AstMethodLet (TakeLen perm sh)+              !asts = ixsPermutePrefix invperm ix+          in gcastWith (unsafeCoerceRefl+                        :: TakeLen perm sh ++ DropLen perm sh :~: sh) $+             astGatherKnobsS knobs shpPermuted shn shp v (vars, asts)++-- TODO: To make this less disastrous, we need to add an extra constructor+-- to AstIndex with the semantics "this index reshaped from shIn to shOut"+-- that fuses perfectly with itself and absorbs normal indexes+-- by substitution. Or perhaps make this the only constructor, with normal+-- indexes represented as "this index reshaped from sh to sh".+-- Or only extend AstGather and possibly also AstIndex with the extra+-- shIn and shOut arguments. This complicates any code related to+-- AstGather and AstIndex, but often prevents nested reshapes from affecting+-- term size in any way. But we'd need to be careful to avoid breaking such+-- an index into components, because that forces index normalization,+-- e.g., index(gather) can no longer simplify recursively by one index+-- component at a time (probably possible only if the index is shorter+-- that the list of variables fo the gather). There are probably bad cases+-- where term size blowup can't be avoided, because the index has to be+-- normalized between each reshape.+--+-- | This generates big terms that don't simplify well,+-- so we keep the AstReshape form until simplification gets stuck.+-- In fact, to simplify the terms we'd need advanced solving of equations+-- in integer arithmetic modulo.+astReshapeAsGatherS+  :: forall sh sh2 x s. KnownSpan s+  => SimplifyKnobs -> ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh x)+  -> AstTensor AstMethodLet s (TKS2 sh2 x)+{-# NOINLINE astReshapeAsGatherS #-}+astReshapeAsGatherS knobs shOut v | Refl <- lemAppNil @sh2+                                  , Refl <- lemAppNil @sh+                                  , FTKS shIn _ <- ftkAst v =+  funToVarsIxS shOut $ \vars ix ->+    let iUnshared :: AstInt AstMethodLet+        iUnshared = ixsToLinear shOut ix+-- This can't be done, because i depends on vars:+--     astLetFun iUnshared $ \i ->+        i = iUnshared  -- sharing broken+        asts :: AstIxS AstMethodLet sh+        !asts = fromLinearIdxS shIn i+    in astGatherKnobsS knobs shOut ZSS shIn v (vars, asts)++-- I can't switch to ixxFromLinear from ox-arrays+-- even just because IntegralH is not available in ox-arrays.+--+-- The inlines are not needed due to function arguments,+-- because there are none, but due to observed significant speedup+-- they provide in tests that intensively use these operations.+-- Maybe specialization doesn't quite work for them? Not verified.+--+-- | Given a linear index into the buffer, get the corresponding+-- multidimensional index.+--+-- If any of the dimensions is 0, the linear index has to be 0+-- (which we can't assert, because j may be a term and so == lies),+-- which is fine, because that's pointing at the start of the empty buffer.+fromLinearIdxS :: forall sh j. IntegralH j+               => ShS sh -> j -> IxS sh j+fromLinearIdxS = \sh lin -> case go sh lin of (# _, ix #) -> ix+  where+    go :: ShS sh1 -> j -> (# j, IxS sh1 j #)+    go ZSS !n = (# n, ZIS #)+    go ((:$$) n sh) lin =+      let (# tensLin, idxInTens #) = go sh lin+          tensLin' = tensLin `quotH` fromIntegral (fromSNat' n)+          i = tensLin `remH` fromIntegral (fromSNat' n)+      in (# tensLin', i :.$ idxInTens #)+++-- * The simplifying combinators, one for almost each AST constructor++astPair :: KnownSpan s+        => AstTensor AstMethodLet s x -> AstTensor AstMethodLet s y+        -> AstTensor AstMethodLet s (TKProduct x y)+astPair (Ast.AstProject1 t1) (Ast.AstProject2 t2)+  | Just Refl <- eqUnknownShapes t1 t2 = t1+astPair (Ast.AstFromPrimal v1) (Ast.AstFromPrimal v2) =+  fromPrimal $ astPair v1 v2+astPair (Ast.AstFromDual v1) (Ast.AstFromDual v2) =+  fromDual $ astPair v1 v2+astPair (Ast.AstFromPlain v1) (Ast.AstFromPlain v2) =+  fromPlain $ astPair v1 v2+astPair (AstConvUp c1 ftk1 v1) (AstConvUp c2 ftk2 v2) =+  astConvertUp (ConvT2 c1 c2) (FTKProduct ftk1 ftk2) (astPair v1 v2)+astPair (AstConvUp c1 ftk1 v1) v2 =+  astConvertUp (ConvT2 c1 ConvId) (FTKProduct ftk1 (ftkAst v2)) (astPair v1 v2)+astPair v1 (AstConvUp c2 ftk2 v2) =+  astConvertUp (ConvT2 ConvId c2) (FTKProduct (ftkAst v1) ftk2) (astPair v1 v2)+astPair v1 v2 = Ast.AstPair v1 v2++astProject1+  :: forall x z s. KnownSpan s+  => AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s x+astProject1 u = case u of+  Ast.AstPair x _z -> x+  Ast.AstMapAccumLDer k bftk eftk+                      (AstLambda varf vf)+                      (AstLambda vard vd)+                      (AstLambda varr vr)+                      acc0 es | Nothing <- matchingFTK bftk ftkUnit ->+    let varf2 = varf+        vf2 = astPair (astProject1 vf) (fromPlain $ AstConcreteK Z1)+        vard2 = vard+        vd2 = astPair (astProject1 vd) (fromPlain $ AstConcreteK Z1)+        accftk = ftkAst acc0+        ftkr2 = FTKProduct+                  (adFTK $ FTKProduct accftk ftkUnit)+                  (FTKProduct accftk eftk)+        varr2 = reshapeVarName ftkr2 varr+        astr2 = astVar varr2+        zero = fromPlain $ astConcrete (adFTK bftk) (tdefTarget (adFTK bftk))+        vr2 = substituteAst+                (astPair (astPair (astProject1 (astProject1 astr2))+                                  zero)+                         (astProject2 astr2))+                varr vr+    in Ast.AstProject1+       $ astMapAccumLDer k ftkUnit eftk (AstLambda varf2 vf2)+                                        (AstLambda vard2 vd2)+                                        (AstLambda varr2 vr2) acc0 es+  Ast.AstLet var t v -> astLet var t (astProject1 v)+  Ast.AstPrimalPart v -> astPrimalPart $ astProject1 v+  Ast.AstDualPart v -> astDualPart $ astProject1 v+  Ast.AstPlainPart v -> astPlainPart $ astProject1 v+  Ast.AstFromPrimal u1 -> fromPrimal $ astProject1 u1+  Ast.AstFromDual u1 -> fromDual $ astProject1 u1+  Ast.AstFromPlain u1 -> fromPlain $ astProject1 u1+  Ast.AstConvert (ConvT2 c1 _c2) t -> astConvert c1 $ astProject1 t+    -- Even Down conversions are not pushed down projections, so AstConvUp+    -- is not necessary here for consistency of rewriting. The AstConvUp case+    -- commented out below probably doesn't add anything, because ConvCmp is not+    -- a normal form of conversions of products (and ConvZip is ruled out).+    -- AstConvUp (FTKProduct zftk _) t | FTKProduct yftk _ <- ftkAst t ->+    --   astConvert (convUp yftk zftk) (astProject1 t)+  _ -> Ast.AstProject1 u++astProject2+  :: forall x z s. KnownSpan s+  => AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s z+astProject2 u = case u of+  Ast.AstPair _x z -> z+  Ast.AstLet var t v -> astLet var t (astProject2 v)+  Ast.AstPrimalPart v -> astPrimalPart $ astProject2 v+  Ast.AstDualPart v -> astDualPart $ astProject2 v+  Ast.AstPlainPart v -> astPlainPart $ astProject2 v+  Ast.AstFromPrimal u1 -> fromPrimal $ astProject2 u1+  Ast.AstFromDual u1 -> fromDual $ astProject2 u1+  Ast.AstFromPlain u1 -> fromPlain $ astProject2 u1+  Ast.AstConvert (ConvT2 _c1 c2) t -> astConvert c2 $ astProject2 t+  _ -> Ast.AstProject2 u++astMapAccumLDer+  :: forall accy by ey k s. KnownSpan s+  => SNat k+  -> FullShapeTK by+  -> FullShapeTK ey+  -> AstHFun s+             (TKProduct accy ey) (TKProduct accy by)+  -> AstHFun s+             (TKProduct (ADTensorKind (TKProduct accy ey))+                        (TKProduct accy ey))+             (ADTensorKind (TKProduct accy by))+  -> AstHFun s+             (TKProduct (ADTensorKind (TKProduct accy by))+                        (TKProduct accy ey))+             (ADTensorKind (TKProduct accy ey))+  -> AstTensor AstMethodLet s accy+  -> AstTensor AstMethodLet s (BuildTensorKind k ey)+  -> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))+astMapAccumLDer k bftk eftk (AstLambda varf vf)+                            (AstLambda vard vd)+                            (AstLambda varr vr)+                (AstConvUp @accyFrom c accftk acc0From) es =+  let accftkFrom = ftkAst acc0From+      ftkf2 = FTKProduct accftkFrom eftk+      varf2 = reshapeVarName ftkf2 varf+      astf2 = astVar varf2+      vf2 =+        let subbed =+              substituteAst+                (astPair (astConvertUp c accftk (astProject1 astf2))+                         (astProject2 astf2))+                varf vf+        in astConvDown @(TKProduct accy by)+                       (FTKProduct accftkFrom bftk)+                       subbed+      ftkd2 = FTKProduct+                (adFTK $ FTKProduct accftkFrom eftk)+                (FTKProduct accftkFrom eftk)+      vard2 = reshapeVarName ftkd2 vard+      astd2 = astVar vard2+      vd2 =+        let subbed =+              substituteAst+                (astPair (astPair (astConvUp @(ADTensorKind accyFrom)+                                     (adFTK accftk)+                                     (astProject1 (astProject1 astd2)))+                                     -- TODO: adTKConversion c?+                                  (astProject2 (astProject1 astd2)))+                         (astPair (astConvertUp c accftk+                                     (astProject1 (astProject2 astd2)))+                                  (astProject2 (astProject2 astd2))))+                vard vd+        in astConvDown @(ADTensorKind (TKProduct accy by))+                       (adFTK $ FTKProduct accftkFrom bftk)+                       subbed+      ftkr2 = FTKProduct+                (adFTK $ FTKProduct accftkFrom bftk)+                (FTKProduct accftkFrom eftk)+      varr2 = reshapeVarName ftkr2 varr+      astr2 = astVar varr2+      vr2 =+        let subbed =+              substituteAst+                (astPair (astPair (astConvUp @(ADTensorKind accyFrom)+                                     (adFTK accftk)+                                     (astProject1 (astProject1 astr2)))+                                  (astProject2 (astProject1 astr2)))+                         (astPair (astConvertUp c accftk+                                     (astProject1 (astProject2 astr2)))+                                  (astProject2 (astProject2 astr2))))+                varr vr+        in astConvDown @(ADTensorKind (TKProduct accy ey))+                       (adFTK $ FTKProduct accftkFrom eftk)+                       subbed+  in astConvertUp (ConvT2 c ConvId)+                  (FTKProduct accftk (buildFTK k bftk))+     $ astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)+                                   (AstLambda vard2 vd2)+                                   (AstLambda varr2 vr2)+                                   acc0From es+astMapAccumLDer k bftk eftk (AstLambda varf vf)+                            (AstLambda vard vd)+                            (AstLambda varr vr)+                acc0 (AstConvUp @esShsFrom _c _esShsFTK esFrom) =+  let accftk = ftkAst acc0+      esShsFrom = ftkAst esFrom+      esShsFromSTK = ftkToSTK esShsFrom+  in case razeSTK esShsFromSTK of+    (eftkFromSTK :: SingletonTK eyFrom) ->+      gcastWith (unsafeCoerceRefl :: BuildTensorKind k eyFrom :~: esShsFrom) $+      let eftkFrom = razeFTK k eftkFromSTK esShsFrom+          ftkf2 = FTKProduct accftk eftkFrom+          varf2 = reshapeVarName ftkf2 varf+          astf2 = astVar varf2+          vf2 =+            let subbed =+                  substituteAst+                    (astPair (astProject1 astf2)+                             (astConvUp @eyFrom eftk (astProject2 astf2)))+                    varf vf+            in subbed+          ftkd2 = FTKProduct+                    (adFTK $ FTKProduct accftk eftkFrom)+                    (FTKProduct accftk eftkFrom)+          vard2 = reshapeVarName ftkd2 vard+          astd2 = astVar vard2+          vd2 =+            let subbed =+                  substituteAst+                    (astPair (astPair (astProject1 (astProject1 astd2))+                                      (astConvUp @(ADTensorKind eyFrom)+                                         (adFTK eftk)+                                         (astProject2 (astProject1 astd2))))+                             (astPair (astProject1 (astProject2 astd2))+                                      (astConvUp @eyFrom eftk+                                         (astProject2 (astProject2 astd2)))))+                    vard vd+            in subbed+          ftkr2 = FTKProduct+                    (adFTK $ FTKProduct accftk bftk)+                    (FTKProduct accftk eftkFrom)+          varr2 = reshapeVarName ftkr2 varr+          astr2 = astVar varr2+          vr2 =+            let subbed =+                  substituteAst+                    (astPair (astProject1 astr2)+                             (astPair (astProject1 (astProject2 astr2))+                                      (astConvUp @eyFrom eftk+                                         (astProject2 (astProject2 astr2)))))+                    varr vr+            in astConvDown @(ADTensorKind (TKProduct accy ey))+                           (adFTK $ FTKProduct accftk eftkFrom)+                           subbed+      in astMapAccumLDer k bftk eftkFrom (AstLambda varf2 vf2)+                                         (AstLambda vard2 vd2)+                                         (AstLambda varr2 vr2)+                                         acc0 esFrom+astMapAccumLDer k bftk eftk f df rf acc0 es =+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es++astApply :: forall x z s. KnownSpan s+         => AstHFun s x z -> AstTensor AstMethodLet s x+         -> AstTensor AstMethodLet s z+astApply (AstLambda !var !v) u = astLet var u v++-- Invariant: if the variable has bounds, the expression can only have+-- values within the bounds (regardless of what the `bounds` call would say).+astLet :: forall y z s s2. KnownSpan s2+       => AstVarName '(s, y) -> AstTensor AstMethodLet s y+       -> AstTensor AstMethodLet s2 z+       -> AstTensor AstMethodLet s2 z+astLet _var _u v@AstConcreteK{} = v+astLet _var _u v@AstConcreteS{} = v+astLet _var _u v@Ast.AstIotaS{} = v+astLet var u v@(Ast.AstVar var2) =+  if varNameToAstVarId var2 == varNameToAstVarId var+  then case testEquality var var2 of+    Just Refl -> u+    _ -> error "astLet: wrong variable types at AstVar"+  else v+astLet var u v@(Ast.AstPrimalPart (Ast.AstVar var2)) =  -- a common noop+  if varNameToAstVarId var2 == varNameToAstVarId var+  then case testEquality var var2 of+    Just Refl -> astPrimalPart u+    _ -> error "astLet: wrong variable types at AstPrimalPart"+  else v+astLet var u v@(Ast.AstDualPart (Ast.AstVar var2)) =  -- a noop+  if varNameToAstVarId var2 == varNameToAstVarId var+  then case testEquality var var2 of+    Just Refl -> astDualPart u+    _ -> error "astLet: wrong variable types at AstDualPart"+  else v+astLet var u v@(Ast.AstPlainPart (Ast.AstVar var2)) =  -- a noop+  if varNameToAstVarId var2 == varNameToAstVarId var+  then case testEquality var var2 of+    Just Refl -> astPlainPart u+    _ -> error "astLet: wrong variable types at AstPlainPart"+  else v+astLet var u v | astIsSmall True u =+  substituteAst u var v+astLet var u (Ast.AstFromPrimal v0) = fromPrimal $ astLet var u v0+astLet var u (Ast.AstFromDual v0) = fromDual $ astLet var u v0+astLet var u (Ast.AstFromPlain v0) = fromPlain $ astLet var u v0+astLet var (Ast.AstFromPrimal u) v =+  withKnownSpan (varNameToSpan var)+  $ astLetFun u $ \ !ast1 -> substituteAst (fromPrimal ast1) var v+astLet var (Ast.AstFromDual u) v =+  astLetFun u $ \ !ast1 -> substituteAst (fromDual ast1) var v+astLet var (Ast.AstFromPlain u) v =+  withKnownSpan (varNameToSpan var)+  $ astLetFun u $ \ !ast1 -> substituteAst (fromPlain ast1) var v+astLet var (Ast.AstPair u1 u2) v =+  withKnownSpan (varNameToSpan var) $+  astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->+    substituteAst (Ast.AstPair ast1 ast2) var v+astLet var (Ast.AstLet varN uN u12@Ast.AstPair{}) v =+  withKnownSpan (varNameToSpan var) $+  astLetRefresh varN uN u12 $ \u3 -> case u3 of+    Ast.AstPair u1' u2' ->+      astLetFun u1' $ \ !ast1 -> astLetFun u2' $ \ !ast2 ->+        substituteAst (Ast.AstPair ast1 ast2) var v+    _ -> astLet var (astLet varN uN u3) v+-- This is a common case, e.g., from representing conditionals.+astLet var (Ast.AstFromVectorK shm u) v | V.length u == 2 =+  withKnownSpan (varNameToSpan var) $+  astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->+    substituteAst (Ast.AstFromVectorK shm+                   $ V.fromListN 2 [ast1, ast2]) var v+astLet var (Ast.AstLet varN uN+              u12@(Ast.AstFromVectorK _ u)) v | V.length u == 2 =+  withKnownSpan (varNameToSpan var) $+  astLetRefresh varN uN u12 $ \u3 -> case u3 of+    Ast.AstFromVectorK shm u' ->+      astLetFun (u' V.! 0) $ \ !ast1 -> astLetFun (u' V.! 1) $ \ !ast2 ->+        substituteAst (Ast.AstFromVectorK shm+                       $ V.fromListN 2 [ast1, ast2]) var v+    _ -> astLet var (astLet varN uN u3) v+astLet var (Ast.AstFromVectorS shm u) v | V.length u == 2 =+  withKnownSpan (varNameToSpan var) $+  astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->+    substituteAst (Ast.AstFromVectorS shm+                   $ V.fromListN 2 [ast1, ast2]) var v+astLet var (Ast.AstLet varN uN+              u12@(Ast.AstFromVectorS _ u)) v | V.length u == 2 =+  withKnownSpan (varNameToSpan var) $+  astLetRefresh varN uN u12 $ \u3 -> case u3 of+    Ast.AstFromVectorS shm u' ->+      astLetFun (u' V.! 0) $ \ !ast1 -> astLetFun (u' V.! 1) $ \ !ast2 ->+        substituteAst (Ast.AstFromVectorS shm+                       $ V.fromListN 2 [ast1, ast2]) var v+    _ -> astLet var (astLet varN uN u3) v+astLet var (Ast.AstReplicateK shm a) v =+  let var2 = reshapeVarName (ftkAst a) var+      ast = Ast.AstReplicateK shm $ astVar var2+  in astLet var2 a (substituteAst ast var v)+astLet var (Ast.AstReplicateS shm a) v =+  let var2 = reshapeVarName (ftkAst a) var+      ast = Ast.AstReplicateS shm $ astVar var2+  in astLet var2 a (substituteAst ast var v)+astLet var (Ast.AstTransposeS perm a) v =+  let var2 = reshapeVarName (ftkAst a) var+      ast = Ast.AstTransposeS perm $ astVar var2+  in astLet var2 a (substituteAst ast var v)+astLet var (AstConvUp c zftk a) v =+  let var2 = reshapeVarName (ftkAst a) var+      ast = withKnownSpan (varNameToSpan var)+            $ astConvertUp c zftk $ astVar var2+  in astLet var2 a (substituteAst ast var v)+astLet var u (AstConvUp c zftk a) =+  astConvertUp c zftk $ astLet var u a+astLet var u v = Ast.AstLet var u v++astPrimalPart :: KnownSpan s+              => AstTensor AstMethodLet s y+              -> AstTensor AstMethodLet (PrimalStepSpan s) y+astPrimalPart t = case t of+  Ast.AstPair t1 t2 -> astPair (astPrimalPart t1) (astPrimalPart t2)+  -- We really need equality saturation...+  Ast.AstProject1 u@Ast.AstMapAccumLDer{} -> astProject1 $ astPrimalPart u+  Ast.AstProject2 u@Ast.AstMapAccumLDer{} -> astProject2 $ astPrimalPart u+  Ast.AstProject1 u@(Ast.AstProject1 Ast.AstMapAccumLDer{}) ->+    astProject1 $ astPrimalPart u+  Ast.AstProject1 u@(Ast.AstProject2 Ast.AstMapAccumLDer{}) ->+    astProject1 $ astPrimalPart u+  Ast.AstProject2 u@(Ast.AstProject2 Ast.AstMapAccumLDer{}) ->+    astProject2 $ astPrimalPart u+  Ast.AstProject2 u@(Ast.AstProject1 Ast.AstMapAccumLDer{}) ->+    astProject2 $ astPrimalPart u+  Ast.AstProject1{} -> Ast.AstPrimalPart t+  Ast.AstProject2{} -> Ast.AstPrimalPart t+  Ast.AstMapAccumLDer k bftk eftk (AstLambda varf vf)+                                  (AstLambda vard vd)+                                  (AstLambda varr vr) acc0 es ->+    let varf2 = respanVarName varf+        astf2 = fromPrimal $ astVar varf2+        vf2 = astPrimalPart $ substituteAst astf2 varf vf+        vard2 = respanVarName vard+        astd2 = fromPrimal $ astVar vard2+        vd2 = astPrimalPart $ substituteAst astd2 vard vd+        varr2 = respanVarName varr+        astr2 = fromPrimal $ astVar varr2+        vr2 = astPrimalPart $ substituteAst astr2 varr vr+    in astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)+                                   (AstLambda vard2 vd2)+                                   (AstLambda varr2 vr2)+                       (astPrimalPart acc0) (astPrimalPart es)+  Ast.AstApply{} -> Ast.AstPrimalPart t+  Ast.AstVar{} -> Ast.AstPrimalPart t  -- the only normal form+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = astPrimalPart v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v -> astLet var u (astPrimalPart v)++  Ast.AstPrimalPart{} -> Ast.AstPrimalPart t+  Ast.AstDualPart{} -> Ast.AstPrimalPart t+  Ast.AstPlainPart{} -> Ast.AstPrimalPart t  -- impossible, but fine+  Ast.AstFromPrimal v -> v+  Ast.AstFromDual v ->+    let ftk = ftkAst v+    in fromPlain $ astConcrete ftk (tdefTarget ftk)+  Ast.AstFromPlain v -> fromPlain v++  AstPlusK u v -> astPrimalPart u `astPlusK` astPrimalPart v+  AstTimesK u v -> astPrimalPart u `astTimesK` astPrimalPart v+  Ast.AstN1K opCode u -> astN1K opCode (astPrimalPart u)+  Ast.AstR1K opCode u -> astR1K opCode (astPrimalPart u)+  Ast.AstR2K opCode u v -> astR2K opCode (astPrimalPart u) (astPrimalPart v)+  Ast.AstI2K opCode u v -> astI2K opCode (astPrimalPart u) (astPrimalPart v)+  AstConcreteK{} -> Ast.AstPrimalPart t+  Ast.AstFloorK{} -> Ast.AstPrimalPart t+  Ast.AstFromIntegralK{} -> Ast.AstPrimalPart t+  Ast.AstCastK v -> astCastK $ astPrimalPart v+  Ast.AstArgMinK{} -> Ast.AstPrimalPart t+  Ast.AstArgMaxK{} -> Ast.AstPrimalPart t+  Ast.AstIndexK v ix -> astIndexK (astPrimalPart v) ix++  AstPlusS u v -> astPrimalPart u `astPlusS` astPrimalPart v+  AstTimesS u v -> astPrimalPart u `astTimesS` astPrimalPart v+  Ast.AstN1S opCode u -> astN1S opCode (astPrimalPart u)+  Ast.AstR1S opCode u -> astR1S opCode (astPrimalPart u)+  Ast.AstR2S opCode u v -> astR2S opCode (astPrimalPart u) (astPrimalPart v)+  Ast.AstI2S opCode u v -> astI2S opCode (astPrimalPart u) (astPrimalPart v)+  AstConcreteS{} -> Ast.AstPrimalPart t+  Ast.AstFloorS{} -> Ast.AstPrimalPart t+  Ast.AstFromIntegralS{} -> Ast.AstPrimalPart t+  Ast.AstCastS v -> astCastS $ astPrimalPart v+  Ast.AstArgMinS{} -> Ast.AstPrimalPart t+  Ast.AstArgMaxS{} -> Ast.AstPrimalPart t+  Ast.AstIndexS shn v ix -> astIndexS shn (astPrimalPart v) ix++  Ast.AstCondK b a2 a3 -> astCondK b (astPrimalPart a2) (astPrimalPart a3)+  Ast.AstCondS b a2 a3 -> astCondS b (astPrimalPart a2) (astPrimalPart a3)+  Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map astPrimalPart l)+  Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map astPrimalPart l)+  Ast.AstSumK v -> astSumK $ astPrimalPart v+  Ast.AstSumS shm v -> astSumS shm $ astPrimalPart v+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterS shm shn shp (astPrimalPart v) (vars, ix)+  Ast.AstReplicateK shm v -> astReplicateK shm (astPrimalPart v)+  Ast.AstReplicateS shm v -> astReplicateS shm (astPrimalPart v)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherS shm shn shp (astPrimalPart v) (vars, ix)+  Ast.AstIotaS{} -> Ast.AstPrimalPart t+  Ast.AstAppendS x y -> astAppendS (astPrimalPart x) (astPrimalPart y)+  Ast.AstSliceS i n k v -> astSliceS i n k (astPrimalPart v)+  Ast.AstReverseS v -> astReverseS (astPrimalPart v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astPrimalPart v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astPrimalPart v)++  -- Most conversions need to stay down here to cancel out.+  AstConvUp c zftk a -> astConvertUp c zftk $ astPrimalPart a+  Ast.AstConvert{} -> Ast.AstPrimalPart t++  Ast.AstDot0 u v -> astDot0 (astPrimalPart u) (astPrimalPart v)+  Ast.AstDot1InS sh n u v -> astDot1InS sh n (astPrimalPart u) (astPrimalPart v)+  Ast.AstMatmul2S m n p u v ->+    astMatmul2S m n p (astPrimalPart u) (astPrimalPart v)++  Ast.AstBoolNotK{} -> Ast.AstPrimalPart t+  Ast.AstBoolNotS{} -> Ast.AstPrimalPart t+  Ast.AstBoolAndK{} -> Ast.AstPrimalPart t+  Ast.AstBoolAndS{} -> Ast.AstPrimalPart t+  Ast.AstLeqK{} -> Ast.AstPrimalPart t+  Ast.AstLeq{} -> Ast.AstPrimalPart t+  Ast.AstLeqS{} -> Ast.AstPrimalPart t++-- Note how this can't be pushed down into, say, multiplication, because it+-- multiplies the dual part by the primal part. Addition is fine, though.+astDualPart :: AstTensor AstMethodLet FullSpan y+            -> AstTensor AstMethodLet DualSpan y+astDualPart t = case t of+  Ast.AstPair t1 t2 -> astPair (astDualPart t1) (astDualPart t2)+  Ast.AstProject1{} -> Ast.AstDualPart t+  Ast.AstProject2{} -> Ast.AstDualPart t+  Ast.AstMapAccumLDer k bftk eftk (AstLambda varf vf)+                                  (AstLambda vard vd)+                                  (AstLambda varr vr) acc0 es ->+    let varf2 = respanVarName varf+        astf2 = fromDual $ astVar varf2+        vf2 = astDualPart $ substituteAst astf2 varf vf+        vard2 = respanVarName vard+        astd2 = fromDual $ astVar vard2+        vd2 = astDualPart $ substituteAst astd2 vard vd+        varr2 = respanVarName varr+        astr2 = fromDual $ astVar varr2+        vr2 = astDualPart $ substituteAst astr2 varr vr+    in astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)+                                   (AstLambda vard2 vd2)+                                   (AstLambda varr2 vr2)+                       (astDualPart acc0) (astDualPart es)+  Ast.AstApply{} -> Ast.AstDualPart t+  Ast.AstVar{} -> Ast.AstDualPart t+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = astDualPart v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v -> astLet var u (astDualPart v)++  Ast.AstFromPrimal v ->+    let ftk = ftkAst v+    in Ast.AstDualPart $ fromPlain+       $ astConcrete ftk (tdefTarget ftk)+           -- let's hope this is smaller than v+  Ast.AstFromDual v -> v+  Ast.AstFromPlain v ->  -- TODO: are deltas zeroed enough here?+    let ftk = ftkAst v+    in Ast.AstDualPart $ fromPlain+       $ astConcrete ftk (tdefTarget ftk)+           -- let's hope this is smaller than v++  AstPlusK u v -> astDualPart u `astPlusK` astDualPart v+  -- This one is mathematically wrong, dual numbers don't mult like that:+  -- AstTimesK u v -> astDualPart u * astDualPart v+  Ast.AstN1K NegateOp u -> negate (astDualPart u)+  {- Some of these are wrong, so let's be conservative:+  Ast.AstN1K AbsOp u -> abs (astDualPart u)+  Ast.AstN1K SignumOp u -> signum (astDualPart u)+  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astDualPart u)+  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (astDualPart u) (astDualPart v)+  Ast.AstI2K QuotOp u v -> quotH (astDualPart u) (astDualPart v)+  Ast.AstI2K RemOp u v -> remH (astDualPart u) (astDualPart v)+  -}+  Ast.AstCastK v -> astCastK $ astDualPart v+  Ast.AstIndexK v ix -> astIndexK (astDualPart v) ix++  AstPlusS u v -> astDualPart u `astPlusS` astDualPart v+  -- This one is mathematically wrong, dual numbers don't mult like that:+  -- AstTimesS u v -> astDualPart u * astDualPart v+  Ast.AstN1S NegateOp u -> negate (astDualPart u)+  {- Some of these are wrong, so let's be conservative:+  Ast.AstN1S AbsOp u -> abs (astDualPart u)+  Ast.AstN1S SignumOp u -> signum (astDualPart u)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astDualPart u)+  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astDualPart u)+                                             (astDualPart v)+  Ast.AstI2S opCode u v -> Ast.AstI2S opCode (astDualPart u)+                                             (astDualPart v)+  -}+  Ast.AstCastS v -> astCastS $ astDualPart v+  Ast.AstIndexS shn v ix -> astIndexS shn (astDualPart v) ix++  Ast.AstCondK b a2 a3 -> astCondK b (astDualPart a2) (astDualPart a3)+  Ast.AstCondS b a2 a3 -> astCondS b (astDualPart a2) (astDualPart a3)+  Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map astDualPart l)+  Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map astDualPart l)+  Ast.AstSumK v -> astSumK $ astDualPart v+  Ast.AstSumS shm v -> astSumS shm $ astDualPart v+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterS shm shn shp (astDualPart v) (vars, ix)+  Ast.AstReplicateK shm v -> astReplicateK shm (astDualPart v)+  Ast.AstReplicateS shm v -> astReplicateS shm (astDualPart v)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherS shm shn shp (astDualPart v) (vars, ix)+  Ast.AstAppendS x y -> astAppendS (astDualPart x) (astDualPart y)+  Ast.AstSliceS i n k v -> astSliceS i n k (astDualPart v)+  Ast.AstReverseS v -> astReverseS (astDualPart v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astDualPart v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astDualPart v)++  -- Most conversions need to stay down here to cancel out.+  AstConvUp c zftk a -> astConvertUp c zftk $ astDualPart a+  Ast.AstConvert{} -> Ast.AstDualPart t++  Ast.AstDot0 u v -> astDot0 (astDualPart u) (astDualPart v)+  Ast.AstDot1InS sh n u v -> astDot1InS sh n (astDualPart u) (astDualPart v)+  Ast.AstMatmul2S m n p u v ->+    astMatmul2S m n p (astDualPart u) (astDualPart v)++  _ -> Ast.AstDualPart t++astPlainPart :: forall y s. KnownSpan s+             => AstTensor AstMethodLet s y+             -> AstTensor AstMethodLet PlainSpan y+astPlainPart t = case t of+  _ | SPlainSpan <- knownSpan @s -> t+  Ast.AstPair t1 t2 -> astPair (astPlainPart t1) (astPlainPart t2)+  -- We really need equality saturation...+  Ast.AstProject1 u@Ast.AstMapAccumLDer{} -> astProject1 $ astPlainPart u+  Ast.AstProject2 u@Ast.AstMapAccumLDer{} -> astProject2 $ astPlainPart u+  Ast.AstProject1 u@(Ast.AstProject1 Ast.AstMapAccumLDer{}) ->+    astProject1 $ astPlainPart u+  Ast.AstProject1 u@(Ast.AstProject2 Ast.AstMapAccumLDer{}) ->+    astProject1 $ astPlainPart u+  Ast.AstProject2 u@(Ast.AstProject2 Ast.AstMapAccumLDer{}) ->+    astProject2 $ astPlainPart u+  Ast.AstProject2 u@(Ast.AstProject1 Ast.AstMapAccumLDer{}) ->+    astProject2 $ astPlainPart u+  Ast.AstProject1{} -> Ast.AstPlainPart t+  Ast.AstProject2{} -> Ast.AstPlainPart t+  Ast.AstMapAccumLDer k bftk eftk (AstLambda varf vf)+                                  (AstLambda vard vd)+                                  (AstLambda varr vr) acc0 es ->+    let varf2 = respanVarName varf+        astf2 = fromPlain $ astVar varf2+        vf2 = astPlainPart $ substituteAst astf2 varf vf+        vard2 = respanVarName vard+        astd2 = fromPlain $ astVar vard2+        vd2 = astPlainPart $ substituteAst astd2 vard vd+        varr2 = respanVarName varr+        astr2 = fromPlain $ astVar varr2+        vr2 = astPlainPart $ substituteAst astr2 varr vr+    in astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)+                                   (AstLambda vard2 vd2)+                                   (AstLambda varr2 vr2)+                       (astPlainPart acc0) (astPlainPart es)+  Ast.AstApply{} -> Ast.AstPlainPart t+  Ast.AstVar{} -> Ast.AstPlainPart t  -- the only normal form+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = astPlainPart v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v -> astLet var u (astPlainPart v)++  Ast.AstPrimalPart v -> astPlainPart v+  Ast.AstDualPart{} -> Ast.AstPlainPart t+  Ast.AstFromPrimal v -> astPlainPart v+  Ast.AstFromDual v ->+    let ftk = ftkAst v+    in astConcrete ftk (tdefTarget ftk)+  Ast.AstFromPlain v -> v++  AstPlusK u v -> astPlainPart u `astPlusK` astPlainPart v+  AstTimesK u v -> astPlainPart u `astTimesK` astPlainPart v+  Ast.AstN1K opCode u -> astN1K opCode (astPlainPart u)+  Ast.AstR1K opCode u -> astR1K opCode (astPlainPart u)+  Ast.AstR2K opCode u v -> astR2K opCode (astPlainPart u) (astPlainPart v)+  Ast.AstI2K opCode u v -> astI2K opCode (astPlainPart u) (astPlainPart v)+  Ast.AstCastK v -> astCastK $ astPlainPart v+  Ast.AstIndexK v ix -> astIndexK (astPlainPart v) ix++  AstPlusS u v -> astPlainPart u `astPlusS` astPlainPart v+  AstTimesS u v -> astPlainPart u `astTimesS` astPlainPart v+  Ast.AstN1S opCode u -> astN1S opCode (astPlainPart u)+  Ast.AstR1S opCode u -> astR1S opCode (astPlainPart u)+  Ast.AstR2S opCode u v -> astR2S opCode (astPlainPart u) (astPlainPart v)+  Ast.AstI2S opCode u v -> astI2S opCode (astPlainPart u) (astPlainPart v)+  Ast.AstCastS v -> astCastS $ astPlainPart v+  Ast.AstIndexS shn v ix -> astIndexS shn (astPlainPart v) ix++  Ast.AstCondK b a2 a3 -> astCondK b (astPlainPart a2) (astPlainPart a3)+  Ast.AstCondS b a2 a3 -> astCondS b (astPlainPart a2) (astPlainPart a3)+  Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map astPlainPart l)+  Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map astPlainPart l)+  Ast.AstSumK v -> astSumK $ astPlainPart v+  Ast.AstSumS shm v -> astSumS shm $ astPlainPart v+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterS shm shn shp (astPlainPart v) (vars, ix)+  Ast.AstReplicateK shm v -> astReplicateK shm (astPlainPart v)+  Ast.AstReplicateS shm v -> astReplicateS shm (astPlainPart v)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherS shm shn shp (astPlainPart v) (vars, ix)+  Ast.AstAppendS x y -> astAppendS (astPlainPart x) (astPlainPart y)+  Ast.AstSliceS i n k v -> astSliceS i n k (astPlainPart v)+  Ast.AstReverseS v -> astReverseS (astPlainPart v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astPlainPart v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astPlainPart v)++  -- Most conversions need to stay down here to cancel out.+  AstConvUp c zftk a -> astConvertUp c zftk $ astPlainPart a+  Ast.AstConvert{} -> Ast.AstPlainPart t++  Ast.AstDot0 u v -> astDot0 (astPlainPart u) (astPlainPart v)+  Ast.AstDot1InS sh n u v -> astDot1InS sh n (astPlainPart u) (astPlainPart v)+  Ast.AstMatmul2S m n p u v ->+    astMatmul2S m n p (astPlainPart u) (astPlainPart v)++-- TODO: perhaps aim for a polynomial normal form? but that requires global+-- inspection of the whole expression+-- TODO: let's aim at SOP (Sum-of-Products) form, just as+-- ghc-typelits-natnormalise does. Also, let's associate to the right+-- and let's push negation down.+--+-- Not considered are rules that would require comparing non-constant terms+-- or that would duplicate a non-constant term, as well as most rules+-- informed by inequalities, expressed via max or min, such as+-- max n (signum (abs x)) | n <= 0 --> signum (abs x).+--+-- | Integer terms need to be simplified, because large ones are often created+-- due to vectorization and simplification, e.g., via astTransposeAsGather+-- or astReshapeAsGather and can be a deciding factor in whether+-- the other tensor terms can be simplified in turn.+--+-- Additions are flattened into a list-like form.+-- AstConcreteK or AstFromPlain (AstConcreteK), if any, is always the first+-- element of the list of summands.+astPlusK  :: (NumScalar r, KnownSpan s)+          => AstTensor AstMethodLet s (TKScalar r)+          -> AstTensor AstMethodLet s (TKScalar r)+          -> AstTensor AstMethodLet s (TKScalar r)+{-# INLINEABLE astPlusK #-}+astPlusK = \cases+  (Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astPlusK u' v+  u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astPlusK u v'+  (Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astPlusK u v+  (Ast.AstDualPart u) (Ast.AstDualPart v) -> dualPart $ astPlusK u v+  (Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)+    | Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->+      plainPart $ astPlusK u v+  -- This is, sadly, unsound, because @fromPrimal . primalPart@ is not identity.+  -- u (Ast.AstFromPrimal v) -> fromPrimal $ astPlusK (astPrimalPart u) v+  (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astPlusK u v+  (Ast.AstFromDual u) (Ast.AstFromDual v) -> fromDual $ astPlusK u v+  (Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astPlusK u v+  (AstPlusK u w) v -> astPlusK u (astPlusK w v)  -- flattening+  u v | Just 0 <- unAstK u -> v+  u v | Just 0 <- unAstK v -> u+  u v | Just u0 <- unAstK u+      , Just v0 <- unAstK v -> fromPlain $ AstConcreteK (u0 + v0)+  u v | Just v0 <- unAstK v -> AstPlusK (fromPlain $ AstConcreteK v0) u+  u (AstPlusK v w) | Just v0 <- unAstK v ->+    astPlusK (fromPlain $ AstConcreteK v0) u + w++  (Ast.AstN1K NegateOp t1) t2 | eqY t1 t2 -> fromPlain $ AstConcreteK 0+  (Ast.AstN1K NegateOp t1) (AstPlusK t2 u) | eqY t1 t2 -> u+  t2 (Ast.AstN1K NegateOp t1) | eqY t1 t2 -> fromPlain $ AstConcreteK 0+  t2 (AstPlusK (Ast.AstN1K NegateOp t1) u) | eqY t1 t2 -> u++  (Ast.AstI2K RemOp (Ast.AstN1K NegateOp t1) n)+    (Ast.AstI2K RemOp t2 n')+      | eqY t1 t2 && eqY n n' -> fromPlain $ AstConcreteK 0+  (Ast.AstI2K RemOp (Ast.AstN1K NegateOp t1) n)+    (AstPlusK (Ast.AstI2K RemOp t2 n') u)+      | eqY t1 t2 && eqY n n' -> u+  (Ast.AstI2K RemOp t1 n')+    (Ast.AstI2K RemOp (Ast.AstN1K NegateOp t2) n)+      | eqY t1 t2 && eqY n n' -> fromPlain $ AstConcreteK 0+  (Ast.AstI2K RemOp t1 n')+    (AstPlusK (Ast.AstI2K RemOp (Ast.AstN1K NegateOp t2) n) u)+      | eqY t1 t2 && eqY n n' -> u++  t1 t2 | eqY t1 t2 -> fromPlain (AstConcreteK 2) `astTimesK` t1+  t1 (AstPlusK t2 w)+    | eqY t1 t2 -> fromPlain (AstConcreteK 2) `astTimesK` t1 + w+  t1 (AstTimesK n t2)+    | Just n0 <- unAstK n+    , eqY t1 t2 -> fromPlain (AstConcreteK (n0 + 1)) `astTimesK` t1+  t1 (AstPlusK (AstTimesK n t2) w)+    | Just n0 <- unAstK n+    , eqY t1 t2 -> fromPlain (AstConcreteK (n0 + 1)) `astTimesK` t1 + w+  (AstTimesK n t2) t1+    | Just n0 <- unAstK n+    , eqY t1 t2 -> fromPlain (AstConcreteK (n0 + 1)) `astTimesK` t1+  (AstTimesK n t2) (AstPlusK t1 w)+    | Just n0 <- unAstK n+    , eqY t1 t2 -> fromPlain (AstConcreteK (n0 + 1)) `astTimesK` t1 + w++  (AstTimesK n1 t1) (AstTimesK n2 t2)+    | Just n10 <- unAstK n1+    , Just n20 <- unAstK n2+    , eqY t1 t2 -> fromPlain (AstConcreteK (n10 + n20)) `astTimesK` t1+  (AstTimesK n1 t1) (AstPlusK (AstTimesK n2 t2) w)+    | Just n10 <- unAstK n1+    , Just n20 <- unAstK n2+    , eqY t1 t2 -> fromPlain (AstConcreteK (n10 + n20)) `astTimesK` t1 + w+  u v -> AstPlusK u v++-- Just as with AstPlusK, factors are flattened and a constant comes first.+-- Also, negation comes first.+astTimesK :: (NumScalar r, KnownSpan s)+          => AstTensor AstMethodLet s (TKScalar r)+          -> AstTensor AstMethodLet s (TKScalar r)+          -> AstTensor AstMethodLet s (TKScalar r)+{-# INLINEABLE astTimesK #-}+astTimesK = \cases+  (Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astTimesK u' v+  u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astTimesK u v'+  (Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astTimesK u v+  (Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)+    | Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->+      plainPart $ astTimesK u v+  (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astTimesK u v+  Ast.AstFromDual{} Ast.AstFromDual{} -> fromPlain $ AstConcreteK 0+  (Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astTimesK u v+  (AstTimesK u w) v -> astTimesK u (astTimesK w v)  -- flattening+  u _ | Just 0 <- unAstK u -> u+  _ v | Just 0 <- unAstK v -> v+  u v | Just 1 <- unAstK u -> v+  u v | Just 1 <- unAstK v -> u+  u v | Just (-1) <- unAstK u -> negate v+  u v | Just (-1) <- unAstK v -> negate u+  u v | Just u0 <- unAstK u+      , Just v0 <- unAstK v -> fromPlain $ AstConcreteK (u0 * v0)+  u v | Just v0 <- unAstK v -> AstTimesK (fromPlain $ AstConcreteK v0) u+  u (AstTimesK v w) | Just v0 <- unAstK v ->+    astTimesK (fromPlain $ AstConcreteK v0) u * w++  (Ast.AstN1K NegateOp u) (Ast.AstN1K NegateOp v) -> astTimesK u v+  (Ast.AstN1K NegateOp u) (AstTimesK (Ast.AstN1K NegateOp v) w) ->+    astTimesK u v * w+  u (Ast.AstN1K NegateOp v) -> astTimesK (negate u) v+  u (AstTimesK (Ast.AstN1K NegateOp v) w) -> astTimesK (negate u) v * w++  (Ast.AstR2K DivideOp u1 u2) v -> astR2K DivideOp (u1 * v) u2+  u (Ast.AstR2K DivideOp v1 v2) -> astR2K DivideOp (u * v1) v2+  (Ast.AstR1K RecipOp u) v -> astR2K DivideOp v u+  u (Ast.AstR1K RecipOp v) -> astR2K DivideOp u v++  -- This breaks sharing, because although u is concrete and so doesn't+  -- have to be shared, the multiplication is not shared --- we end up+  -- with one addition and two multiplications, not one. Similarly below.+  -- However, three instead of two cheap scalar operations is benign+  -- and there's no risk of exponential blowup via duplicating+  -- variables, so the ability to simplify better is worth the overhead.+  -- OTOH, we don't declare scalar operations small in astIsSmall,+  -- because that could lead to such operations being inlined an unbounded+  -- number of times into expressions.+  -- Note that due to non-scalar versions of these rules being banned,+  -- we get different terms depending on the form of conversions+  -- and rank 0 arrays.+  u (AstPlusK v w) | Just u0 <- unAstK u ->+    astPlusK (astTimesK (fromPlain $ AstConcreteK u0) v)+             (astTimesK (fromPlain $ AstConcreteK u0) w)+  u (AstTimesK (AstPlusK v w) x) | Just u0 <- unAstK u ->+    astPlusK (astTimesK (fromPlain $ AstConcreteK u0) v)+             (astTimesK (fromPlain $ AstConcreteK u0) w) * x++  {- TODO: such rules increase the number of occurrences of a variable+     and trade multiplication and quotient for an equally problematic remnant,+     so they are disabled until we find a way to profit from them.+  -- With static shapes, the second argument to QuotOp and RemOp+  -- is often a constant, which makes such rules worth including,+  -- since they are likely to fire. To help them fire, we avoid changing+  -- that constant, if possible, e.g., in rules for NegateOp.+  (AstConcreteK n) (Ast.AstI2K QuotOp (Ast.AstVar var) (AstConcreteK n'))+    | n == n' ->+      AstPlusK+        (astVar var)+        (negate (Ast.AstI2K RemOp (astVar var) (AstConcreteK n))) -}++  u v -> AstTimesK u v++astN1K :: (NumScalar r, KnownSpan s)+       => OpCodeNum1 -> AstTensor AstMethodLet s (TKScalar r)+       -> AstTensor AstMethodLet s (TKScalar r)+astN1K opCode t = case t of+  Ast.AstLet var n k -> astLet var n (astN1K opCode k)+  Ast.AstPrimalPart n -> primalPart (astN1K opCode n)+  Ast.AstDualPart n -> dualPart (astN1K opCode n)+  Ast.AstPlainPart n -> plainPart (astN1K opCode n)+  Ast.AstFromPrimal n -> fromPrimal (astN1K opCode n)+  Ast.AstFromDual n -> fromDual (astN1K opCode n)+  Ast.AstFromPlain n -> fromPlain (astN1K opCode n)+  Ast.AstCondK b n k -> astCondK b (astN1K opCode n) (astN1K opCode k)+  _ -> case (opCode, t) of+    (NegateOp, AstConcreteK n) -> AstConcreteK (negate n)+    (NegateOp, AstPlusK u v) -> AstPlusK (negate u) (negate v)+    (NegateOp, AstTimesK u v) -> astTimesK (negate u) v+    (NegateOp, Ast.AstN1K NegateOp u) -> u+    (NegateOp, Ast.AstN1K SignumOp u) -> astN1K SignumOp (negate u)+    (NegateOp, Ast.AstI2K QuotOp u v) -> astI2K QuotOp (negate u) v+      -- v is positive and let's keep it so+    (NegateOp, Ast.AstI2K RemOp u v) -> astI2K RemOp (negate u) v+    (AbsOp, AstConcreteK n) -> AstConcreteK (abs n)+    (AbsOp, Ast.AstN1K AbsOp u) -> astN1K AbsOp u+    (AbsOp, Ast.AstN1K NegateOp u) -> astN1K AbsOp u+    (SignumOp, AstConcreteK n) -> AstConcreteK (signum n)+    (SignumOp, Ast.AstN1K SignumOp u) -> astN1K SignumOp u+    _ -> Ast.AstN1K opCode t++astR1K :: (NumScalar r, Differentiable r, KnownSpan s)+       => OpCode1 -> AstTensor AstMethodLet s (TKScalar r)+       -> AstTensor AstMethodLet s (TKScalar r)+astR1K opCode = \case+  Ast.AstLet var n k -> astLet var n (astR1K opCode k)+  Ast.AstPrimalPart u -> primalPart $ astR1K opCode u+  Ast.AstPlainPart u -> plainPart $ astR1K opCode u+  Ast.AstFromPrimal u -> fromPrimal $ astR1K opCode u+  Ast.AstFromPlain u -> fromPlain $ astR1K opCode u+  Ast.AstCondK b n k -> astCondK b (astR1K opCode n) (astR1K opCode k)+  AstConcreteK u -> case opCode of+    RecipOp -> AstConcreteK $ recip u+    ExpOp -> AstConcreteK $ exp u+    LogOp -> AstConcreteK $ log u+    SqrtOp -> AstConcreteK $ sqrt u+    SinOp -> AstConcreteK $ sin u+    CosOp -> AstConcreteK $ cos u+    TanOp -> AstConcreteK $ tan u+    AsinOp -> AstConcreteK $ asin u+    AcosOp -> AstConcreteK $ acos u+    AtanOp -> AstConcreteK $ atan u+    SinhOp -> AstConcreteK $ sinh u+    CoshOp -> AstConcreteK $ cosh u+    TanhOp -> AstConcreteK $ tanh u+    AsinhOp -> AstConcreteK $ asinh u+    AcoshOp -> AstConcreteK $ acosh u+    AtanhOp -> AstConcreteK $ atanh u+  u -> Ast.AstR1K opCode u++astR2K :: (NumScalar r, Differentiable r, KnownSpan s)+       => OpCode2+       -> AstTensor AstMethodLet s (TKScalar r)+       -> AstTensor AstMethodLet s (TKScalar r)+       -> AstTensor AstMethodLet s (TKScalar r)+astR2K opCode = \cases+  (Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astR2K opCode u' v+  u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astR2K opCode u v'+  (Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astR2K opCode u v+  (Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)+    | Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->+      plainPart $ astR2K opCode u v+  (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astR2K opCode u v+  (Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astR2K opCode u v+  (AstConcreteK u) (AstConcreteK v) -> case opCode of+    DivideOp -> AstConcreteK $ u / v+    PowerOp -> AstConcreteK $ u ** v+    LogBaseOp -> AstConcreteK $ logBase u v+    Atan2Op -> AstConcreteK $ atan2H u v+  u v -> case opCode of+    DivideOp -> case u of+      _ | Just 0 <- unAstK u -> u+      _ | Just 1 <- unAstK u -> recip v+      _ | Just 1 <- unAstK v -> u+      _ | Just 0 <- unAstK v -> u  -- the partiality-removal hack+      AstTimesK n t+        | eqY n v -> t+      _ -> Ast.AstR2K DivideOp u v+             -- TODO: add other rules that are relatively numerically stable+    _ -> Ast.AstR2K opCode u v++astI2K :: (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)+       => OpCodeIntegral2+       -> AstTensor AstMethodLet s (TKScalar r)+       -> AstTensor AstMethodLet s (TKScalar r)+       -> AstTensor AstMethodLet s (TKScalar r)+astI2K opCode = \cases+  (Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astI2K opCode u' v+  u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astI2K opCode u v'+  (Ast.AstPrimalPart n) (Ast.AstPrimalPart k) ->+    primalPart (astI2K opCode n k)+  (Ast.AstPlainPart @_ @s1 n) (Ast.AstPlainPart @_ @s2 k)+    | Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->+      plainPart (astI2K opCode n k)+  (Ast.AstFromPrimal n) (Ast.AstFromPrimal k) ->+    fromPrimal (astI2K opCode n k)+  (Ast.AstFromPlain n) (Ast.AstFromPlain k) ->+    fromPlain (astI2K opCode n k)+  u v | let t = Ast.AstI2K opCode u v+      , Just (u1, u2) <- bounds t+      , u1 == u2 ->+        fromPlain $ AstConcreteK u1+  u v -> case opCode of+    QuotOp -> case (u, v) of+      _ | Just u0 <- unAstK u+        , Just v0 <- unAstK v -> fromPlain $ AstConcreteK (quotH u0 v0)+      _ | Just 0 <- unAstK u -> u+      _ | Just 1 <- unAstK v -> u+      _ | Just 0 <- unAstK v -> u  -- the partiality-removal hack+      _ | Just v0 <- unAstK v, v0 < 0 ->+          astI2K QuotOp (negate u) (fromPlain $ AstConcreteK $ negate v0)+      (Ast.AstI2K RemOp _ k, _)+        | Just k0 <- unAstK k+        , Just v0 <- unAstK v+        , v0 >= k0 && k0 >= 0 -> fromPlain $ AstConcreteK 0+      (Ast.AstI2K QuotOp u0 v0, _) ->+        astI2K QuotOp u0 (astTimesK v0 v)+      (AstTimesK n t, _)+        | eqY n v -> t+      _ -> Ast.AstI2K QuotOp u v+    RemOp -> case (u, v) of+      _ | Just u0 <- unAstK u+        , Just v0 <- unAstK v -> fromPlain $ AstConcreteK (remH u0 v0)+      _ | Just 0 <- unAstK u -> u+      _ | Just 1 <- unAstK v -> fromPlain $ AstConcreteK 0+      _ | Just 0 <- unAstK v -> fromPlain $ AstConcreteK 0  -- the hack+      _ | Just w <- unAstK v, w < 0 ->+          astI2K RemOp u (fromPlain $ AstConcreteK $ negate w)+      (Ast.AstI2K RemOp t k, _)+        | Just k0 <- unAstK k+        , Just v0 <- unAstK v+        , v0 >= k0 && k0 >= 0 -> astI2K RemOp t k+      (Ast.AstI2K RemOp t k, _)+        | Just k0 <- unAstK k+        , Just v0 <- unAstK v+        , remH k0 v0 == 0 && k0 > 0 -> astI2K RemOp t v+      (AstTimesK n _, _)+        | Just n0 <- unAstK n+        , Just v0 <- unAstK v+        , remH n0 v0 == 0 -> fromPlain $ AstConcreteK 0+      _ -> Ast.AstI2K RemOp u v++astConcreteK :: GoodScalar r+             => Concrete (TKScalar r)+             -> AstTensor AstMethodLet PlainSpan (TKScalar r)+astConcreteK = AstConcreteK . unConcrete++astFloorK :: (NumScalar r1, Differentiable r1, NumScalar r2, Integral r2)+          => AstTensor AstMethodLet PlainSpan (TKScalar r1)+          -> AstTensor AstMethodLet PlainSpan (TKScalar r2)+astFloorK t = case t of+  Ast.AstLet var u v -> astLet var u (astFloorK v)+  -- This increases work and term size, because conditional is eager.+  -- Ast.AstCond b a2 a3 -> Ast.AstCond b (astFloorK a2) (astFloorK a3)+  -- These values are small, so we can simplify them ASAP.+  AstConcreteK k -> astConcreteK (tkfloor $ Concrete k)+  Ast.AstFloorK v -> astFloorK v+  Ast.AstFromIntegralK v -> astFromIntegralK v+  Ast.AstCastK v -> astFloorK v+  _ -> Ast.AstFloorK t++-- Beware that increasing the number of calls to this constructor+-- sometimes increases runtime, because not enough copies cancel out.+astFromIntegralK :: forall r1 r2. (NumScalar r1, Integral r1, NumScalar r2)+                 => AstTensor AstMethodLet PlainSpan (TKScalar r1)+                 -> AstTensor AstMethodLet PlainSpan (TKScalar r2)+astFromIntegralK t = case t of+  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t+  Ast.AstLet var u v -> astLet var u (astFromIntegralK v)+  AstConcreteK k -> astConcreteK (tkfromIntegral $ Concrete k)+  Ast.AstN1K opCode u -> astN1K opCode (astFromIntegralK u)+  Ast.AstFromIntegralK v -> astFromIntegralK v+  _ -> Ast.AstFromIntegralK t++astCastK :: forall r1 r2 s.+            ( NumScalar r1, Differentiable r1, NumScalar r2, Differentiable r2+            , KnownSpan s )+         => AstTensor AstMethodLet s (TKScalar r1)+         -> AstTensor AstMethodLet s (TKScalar r2)+astCastK t = case t of+  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t+  Ast.AstLet var u v -> astLet var u (astCastK v)+  AstConcreteK k -> astConcreteK (tkcast $ Concrete k)+  -- TODO: which should go deeper, casts or fromPrimal? Or maybe alternate+  -- in different phases to make sure both can cancel out?+  -- Rethink. For now, astFromPrimal is not called, to avoid loops.+  -- The same with many others.+  Ast.AstFromPrimal v -> fromPrimal $ astCastK v+  Ast.AstFromDual v -> fromDual $ astCastK v+  Ast.AstFromPlain v -> fromPlain $ astCastK v+  Ast.AstN1K opCode u -> astN1K opCode (astCastK u)+  Ast.AstR1K opCode u -> astR1K opCode (astCastK u)+  Ast.AstFromIntegralK v -> astFromIntegralK v+  Ast.AstCastK v -> astCastK v+  _ -> Ast.AstCastK t++astArgMinK :: forall n r. NumScalar r+           => AstTensor AstMethodLet PlainSpan (TKS '[n] r)+           -> AstTensor AstMethodLet PlainSpan (TKScalar Int)+astArgMinK t = case t of+  Ast.AstLet var u v -> astLet var u (astArgMinK v)+  AstConcreteS v -> astConcreteK (tkargMin $ Concrete v)+  Ast.AstN1S NegateOp u -> astArgMaxK u+  Ast.AstFromIntegralS v -> astArgMinK v+  Ast.AstCastS v -> astArgMinK v+  Ast.AstReplicateK{} -> 0  -- we mask undefined if n is zero; tough luck+  Ast.AstReplicateS{} -> 0+  Ast.AstIotaS{} -> 0+  _ -> Ast.AstArgMinK t++astArgMaxK :: forall n r. NumScalar r+           => AstTensor AstMethodLet PlainSpan (TKS '[n] r)+           -> AstTensor AstMethodLet PlainSpan (TKScalar Int)+astArgMaxK t = case t of+  Ast.AstLet var u v -> astLet var u (astArgMaxK v)+  AstConcreteS v -> astConcreteK (tkargMax $ Concrete v)+  Ast.AstN1S NegateOp u -> astArgMinK u+  Ast.AstFromIntegralS v -> astArgMaxK v+  Ast.AstCastS v -> astArgMaxK v+  Ast.AstReplicateK{} -> 0  -- we mask undefined if n is zero; tough luck+  Ast.AstReplicateS{} -> 0+  Ast.AstIotaS{} -> 0+  _ -> Ast.AstArgMaxK t++astIndexK+  :: forall shm s r. (GoodScalar r, KnownSpan s)+  => AstTensor AstMethodLet s (TKS shm r) -> AstIxS AstMethodLet shm+  -> AstTensor AstMethodLet s (TKScalar r)+astIndexK = astIndexKnobsK defaultKnobs++astIndexKnobsK+  :: forall shm s r. (GoodScalar r, KnownSpan s)+  => SimplifyKnobs+  -> AstTensor AstMethodLet s (TKS shm r) -> AstIxS AstMethodLet shm+  -> AstTensor AstMethodLet s (TKScalar r)+astIndexKnobsK _ v0 ZIS = kfromS v0+astIndexKnobsK _ v0 (i1 :.$ _)+  | Just (lb, ub) <- intBounds i1+  , FTKS (snat :$$ _) _ <- ftkAst v0+  , ub < 0 || lb >= fromSNat' snat = fromPlain $ AstConcreteK def+astIndexKnobsK knobs u0 (Ast.AstCondK b v w :.$ rest0)+  | knobPhase knobs `notElem` [PhaseUnspecified, PhaseVectorization] =+      -- don't undo vectorization tweaks+    astLetFun u0 $ \u ->+    shareIx rest0 $ \rest ->+    astCondK b (astIndexKnobsK knobs u (v :.$ rest))+               (astIndexKnobsK knobs u (w :.$ rest))+astIndexKnobsK knobs v0 ix@(i1 :.$ rest1)+  | FTKS ((:$$) @in1 @shm1 SNat shm1) _ <- ftkAst v0+  , let tryRecursing :: AstTensor AstMethodLet s (TKS shm1 r)+                     -> AstTensor AstMethodLet s (TKScalar r)+        tryRecursing u =+          case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of+            AstConcreteK b ->+              if b then astIndex u rest1 else fromPlain $ AstConcreteK def+            _ -> Ast.AstIndexK v0 ix = case v0 of+  Ast.AstProject1{} -> Ast.AstIndexK v0 ix+  Ast.AstProject2{} -> Ast.AstIndexK v0 ix+  Ast.AstApply{} -> Ast.AstIndexK v0 ix+  Ast.AstVar{} -> Ast.AstIndexK v0 ix+  Ast.AstBuild1 _ STKScalar (var2, v) | ZIS <- rest1 ->+    tryRecursing $ sfromK $ astLet var2 i1 v+  Ast.AstBuild1 _ STKS{} (var2, v) ->+    tryRecursing $ astLet var2 i1 v++  Ast.AstLet var u v -> astLetRefresh var u v $ \v' -> astIndex v' ix++  Ast.AstPrimalPart{} -> Ast.AstIndexK v0 ix  -- must be a NF+  Ast.AstDualPart{} -> Ast.AstIndexK v0 ix+  Ast.AstPlainPart{} -> Ast.AstIndexK v0 ix+  Ast.AstFromPrimal v -> fromPrimal $ astIndex v ix+  Ast.AstFromDual v -> fromDual $ astIndex v ix+  Ast.AstFromPlain v -> fromPlain $ astIndex v ix++  AstPlusS u v ->+    shareIx ix $ \ !ix2 -> astPlusK (astIndex u ix2)+                                    (astIndex v ix2)+  AstTimesS u v ->+    shareIx ix $ \ !ix2 -> astTimesK (astIndex u ix2)+                                     (astIndex v ix2)+  Ast.AstN1S opCode u -> astN1K opCode (astIndex u ix)+  Ast.AstR1S opCode u -> astR1K opCode (astIndex u ix)+  Ast.AstR2S opCode u v ->+    shareIx ix $ \ !ix2 -> astR2K opCode (astIndex u ix2)+                                         (astIndex v ix2)+  Ast.AstI2S opCode u v ->+    shareIx ix $ \ !ix2 -> astI2K opCode (astIndex u ix2)+                                         (astIndex v ix2)+  AstConcreteS a | AstConcreteK i <- i1 ->+    let u = withKnownShS shm1 $+            tsindex (Concrete a) (Concrete i :.$ ZIS)+    in astIndex (astConcreteS u) rest1+  AstConcreteS{} -> case unRepl1 v0 of+    Just u -> tryRecursing u+    _ -> Ast.AstIndexK v0 ix+  Ast.AstFloorS v -> astFloorK $ astIndex v ix+  Ast.AstFromIntegralS v -> astFromIntegralK $ astIndex v ix+  Ast.AstCastS t -> astCastK $ astIndex t ix+  Ast.AstArgMinS @k @shz v | FTKS nsh _ <- ftkAst v+                           , nl@(SNat @nl) <- shsLast nsh ->+    gcastWith (unsafeCoerceRefl :: shm ++ '[nl] :~: k ': shz) $+    astArgMinK @nl+    $ astIndexKnobsS (deVect knobs) (nl :$$ ZSS) v ix+  Ast.AstArgMaxS @k @shz v | FTKS nsh _ <- ftkAst v+                           , nl@(SNat @nl) <- shsLast nsh ->+    gcastWith (unsafeCoerceRefl :: shm ++ '[nl] :~: k ': shz) $+    astArgMaxK @nl+    $ astIndexKnobsS (deVect knobs) (nl :$$ ZSS) v ix+  Ast.AstIndexS _ v ix2 -> astIndex v (ix2 `ixsAppend` ix)++  Ast.AstCondS b v w ->+    shareIx ix $ \ !ix2 -> astCondK b (astIndex v ix2)+                                      (astIndex w ix2)+  Ast.AstFromVectorK (_ :$$ ZSS) l | AstConcreteK i <- i1, ZIS <- rest1 ->+    tryRecursing $ sfromK (l V.! i)+  Ast.AstFromVectorK{} ->  -- normal form+    Ast.AstIndexK v0 ix+  Ast.AstFromVectorS (_ :$$ ZSS) l | AstConcreteK i <- i1 ->+    tryRecursing (l V.! i)+  Ast.AstFromVectorS @shm2 @shn2 shm2 l | shsLength shm2 < length ix+                                        , SNat <- shsRank shm2 ->+    gcastWith (unsafeCoerceRefl :: shn2 :~: Drop (Rank shm2) shm) $+    gcastWith (unsafeCoerceRefl :: Take (Rank shm2) shm :~: shm2) $+    shareIx ix $ \ !ix2 ->+      let ixRest2 = ixsDrop @(Rank shm2) ix2+      in Ast.AstIndexK (astFromVectorK @shm2 shm2+                        $ V.map (\a -> astIndex a ixRest2) l)+                       (ixsTake @(Rank shm2) ix2)+  Ast.AstFromVectorS{} ->  -- normal form (shsLength shm2 >= length ix)+    Ast.AstIndexK v0 ix+  Ast.AstSumS @shm2 shm2 v | Dict0 <- numFromTKAllNum (Proxy @r) ->+    let perm1 = backpermCycleN (shsLength shm2) (shsLength shm2 + length ix)+    in Permutation.permFromListCont perm1 $ \(perm :: Permutation.Perm perm) ->+         gcastWith (unsafeCoerceRefl+                    :: (Rank perm <=? Rank (shm2 ++ shm)) :~: True) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix perm (shm2 ++ shm)+                       :~: shm ++ shm2) $+         fromMaybe (error "astIndexKnobsK: impossible non-permutation")+         $ Permutation.permCheckPermutation perm+         $ astSumK+         $ astIndexKnobsS @shm (deVect knobs) shm2 (astTransposeS perm v) ix+  Ast.AstScatterS shm7 shn7 shp7 v (vars, i5 :.$ ix2) | eqY i5 i1 ->+    astIndex (astScatter shm7 shn7 (shsTail shp7) v (vars, ix2)) rest1+  Ast.AstScatterS _ _ _ _ (_, AstConcreteK{} :.$ _)+    | AstConcreteK{} <- i1 ->+      fromPlain $ AstConcreteK def  -- from above we know i5 /= i1+  Ast.AstScatterS{} ->  -- normal form+    Ast.AstIndexK v0 ix+  Ast.AstReplicateK (_ :$$ shmRest) v ->+    tryRecursing $ astReplicateK shmRest v+  Ast.AstReplicateS ZSS v -> astIndex v ix+  Ast.AstReplicateS (_ :$$ shmRest) v ->+    tryRecursing $ astReplicateS shmRest v+  -- This is not a possible normal form, but pattern needs to be exhaustive.+  Ast.AstGatherS @_ @_ @shp' _ _ _ v (ZS, ix2) ->+    astIndex @(shp' ++ shm) v (ix2 `ixsAppend` ix)+  Ast.AstGatherS (_ :$$ (shm71 :: ShS shm71)) shn' shp'+                 v (var2 ::$ vars, ix2) ->+    let w :: AstTensor AstMethodLet s (TKS shm1 r)+        w = astGather shm71 shn' shp' v (vars, ix2)+        u = astLet var2 i1 $ astIndex @shm1 w rest1+          -- This let makes it impossible to use astCond when i1 is OOB.+          -- No astLetRefresh is needed, because the smart constructor+          -- for AstGatherS ensures var2 doesn't occur in v, even bound.+    in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of+      AstConcreteK b -> if b then u else fromPlain $ AstConcreteK def+      _ -> Ast.AstIndexK v0 ix+  Ast.AstIotaS _ ->+    astLetFun i1 $ \i ->+      astCondKInitial (0 <=. i &&* i <=. valueOf @in1 - 1)+                      (astFromIntegralK i)+                      0+  Ast.AstAppendS u v | FTKS (SNat @m :$$ _) _ <- ftkAst u ->+    if knobPhase knobs == PhaseExpansion then+      astLetFun i1 $ \i ->+      shareIx rest1 $ \ !rest2 ->+      let ulen = valueOf @m+          ix1 = i :.$ rest2+          ix2 = i - ulen :.$ rest2+      in case ulen <=. i of+        AstConcreteK b -> if b then astIndex v ix2 else astIndex u ix1+        bExpr ->+          -- This results in a larger term, so we consider this late.+          astCondKInitial bExpr (astIndex v ix2) (astIndex u ix1)+    else+      -- We don't need to share i1 and rest1 here and, in fact, we can't+      -- because we'd loop via building AstLet around the resulting AstIndexS.+      let ulen = valueOf @m+          ix1 = i1 :.$ rest1+          ix2 = i1 - ulen :.$ rest1+      in case ulen <=. i1 of+        AstConcreteK b -> if b then astIndex v ix2 else astIndex u ix1+        _ -> Ast.AstIndexK v0 ix+  Ast.AstSliceS i@(SNat @i) (SNat @n) k@SNat v ->+    astLetFun i1 $ \iShared ->+    let b = (if fromSNat' i == 0 then true else 0 <=. iShared)+            &&* (if fromSNat' k == 0 then true else iShared <=. valueOf @n - 1)+        ii = valueOf @i + iShared+    in astCondKInitial b (astIndex v (ii :.$ rest1))+       $ fromPlain $ AstConcreteK def+  Ast.AstReverseS v ->+    let iRev = valueOf @in1 - 1 - i1+    in astIndex v (iRev :.$ rest1)+  Ast.AstTransposeS @_ @sh2 perm v+    | fromSNat' (Permutation.permRank perm) <= length ix ->+      -- TODO: remake once there's an S version of permInverse:+      permInverse perm $ \(permR :: Permutation.Perm permR) _ ->+      let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix permR shm)+          ix2 = ixsPermutePrefix permR ix+      in gcastWith (unsafeCoerceRefl+                    :: sh2 :~: Permutation.PermutePrefix permR shm) $+         astIndex @(Permutation.PermutePrefix permR shm) v ix2+  Ast.AstTransposeS @perm perm v ->+    astIndex (astTransposeAsGatherS @perm (deVect knobs) perm v) ix+  Ast.AstReshapeS sh v ->+    astIndex (astReshapeAsGatherS (deVect knobs) sh v) ix++  -- Conversions to shaped need to stay down, so this is NF.+  Ast.AstConvert{} -> Ast.AstIndexK v0 ix++  Ast.AstDot1InS _ n u v ->+    astDot0 (astIndexKnobsS (deVect knobs) (n :$$ ZSS) u ix)+            (astIndexKnobsS (deVect knobs) (n :$$ ZSS) v ix)+  Ast.AstMatmul2S{} -> Ast.AstIndexK v0 ix  -- TODO++  Ast.AstBoolNotS a -> Ast.AstBoolNotK $ astIndex a ix+  Ast.AstBoolAndS b1 b2 -> Ast.AstBoolAndK (astIndex b1 ix)+                                           (astIndex b2 ix)+  Ast.AstLeqS _shb sh arg1 arg2 ->+    Ast.AstLeq (astIndexKnobsS (deVect knobs) sh arg1 ix)+               (astIndexKnobsS (deVect knobs) sh arg2 ix)+ where+  astIndex+    :: forall shm' s' r'. (GoodScalar r', KnownSpan s')+    => AstTensor AstMethodLet s' (TKS shm' r')+    -> AstIxS AstMethodLet shm'+    -> AstTensor AstMethodLet s' (TKScalar r')+  astIndex = astIndexKnobsK (deVect knobs)+  astGather+    :: forall shm' shn' shp' x'.+       ShS shm' -> ShS shn' -> ShS shp'+    -> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') x')+    -> (AstVarListS shm', AstIxS AstMethodLet shp')+    -> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') x')+  astGather = astGatherKnobsS (deVect knobs)+  astScatter+    :: forall shm' shn' shp' x'. TKAllNum x'+    => ShS shm' -> ShS shn' -> ShS shp'+    -> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') x')+    -> (AstVarListS shm', AstIxS AstMethodLet shp')+    -> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') x')+  astScatter = astScatterKnobsS (deVect knobs)++-- Just as with AstPlusK, summands are flattened and a constant comes first.+-- We don't sum fromVectors, because it's costly and even if everything+-- nicely simplifies, we end up with an equaly big vector of the results.+-- OTOH, AstSum* and AstDot* eliminate at least one dimension, so they+-- are worth it, at least for small enough arguments and if the result+-- simplifies decently.+astPlusS :: (NumScalar r, KnownSpan s)+         => AstTensor AstMethodLet s (TKS sh r)+         -> AstTensor AstMethodLet s (TKS sh r)+         -> AstTensor AstMethodLet s (TKS sh r)+astPlusS = \cases+  u v | FTKS (snat :$$ _) _ <- ftkAst u+      , Just u0 <- unRepl1 u+      , Just v0 <- unRepl1 v ->+        astReplicateS (snat :$$ ZSS) $ astPlusS u0 v0+  (Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astPlusS u v+  (Ast.AstDualPart u) (Ast.AstDualPart v) -> dualPart $ astPlusS u v+  (Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)+    | Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->+      plainPart $ astPlusS u v+  (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astPlusS u v+  (Ast.AstFromDual u) (Ast.AstFromDual v) -> fromDual $ astPlusS u v+  (Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astPlusS u v+  (AstPlusS u w) v -> astPlusS u (astPlusS w v)  -- flattening+  u v | Just 0 <- unReplC u -> v+  u v | Just 0 <- unReplC v -> u+  u v | Just u0 <- unAstS u+      , Just v0 <- unAstS v -> fromPlain $ AstConcreteS (u0 + v0)+  u v | Just v0 <- unAstS v -> AstPlusS (fromPlain $ AstConcreteS v0) u+  u (AstPlusS v w) | Just v0 <- unAstS v ->+    astPlusS (fromPlain $ AstConcreteS v0) u + w++  (Ast.AstN1S NegateOp t1) t2 | eqY t1 t2 ->+    fromPlain $ AstConcreteS $ defTargetRep $ ftkAst t1+  (Ast.AstN1S NegateOp t1) (AstPlusS t2 u) | eqY t1 t2 -> u+  t2 (Ast.AstN1S NegateOp t1) | eqY t1 t2 ->+    fromPlain $ AstConcreteS $ defTargetRep $ ftkAst t1+  t2 (AstPlusS (Ast.AstN1S NegateOp t1) u) | eqY t1 t2 -> u++  (AstTimesS n1 t1) (AstTimesS n2 t2)+    | Just n10 <- unAstS n1+    , Just n20 <- unAstS n2+    , eqY t1 t2 -> fromPlain (AstConcreteS (n10 + n20)) `astTimesS` t1+  (AstTimesS n1 t1) (AstPlusS (AstTimesS n2 t2) w)+    | Just n10 <- unAstS n1+    , Just n20 <- unAstS n2+    , eqY t1 t2 -> fromPlain (AstConcreteS (n10 + n20)) `astTimesS` t1 + w++  (AstConvUpSFromK u) (AstConvUpSFromK v) -> sfromK $ astPlusK u v+  u (AstConvUpSFromK v) -> sfromK $ astPlusK (kfromS u) v+  (AstConvUpSFromK u) v -> sfromK $ astPlusK u (kfromS v)++  u@Ast.AstScatterS{} v@Ast.AstScatterS{} ->+    fromMaybe (AstPlusS u v) $ astPlusSMaybe u v+  u@Ast.AstScatterS{} v@(AstPlusS u2@Ast.AstScatterS{} w) ->+    maybe (AstPlusS u v) (+ w) $ astPlusSMaybe u u2+  u@(Ast.AstTransposeS perm1 a1@(Ast.AstScatterS @_ @shn1 @shp1 _ _ _ _ _))+    v@(Ast.AstTransposeS perm2 a2@(Ast.AstScatterS @_ @shn2 @shp2 _ _ _ _ _))+      | Just Refl <- testEquality perm1 perm2 ->+        gcastWith (unsafeCoerceRefl :: shp1 ++ shn1 :~: shp2 ++ shn2) $+        maybe (AstPlusS u v) (astTransposeS perm1) $ astPlusSMaybe a1 a2+  u@(Ast.AstTransposeS perm1 a1@(Ast.AstScatterS @_ @shn1 @shp1 _ _ _ _ _))+    v@(AstPlusS (Ast.AstTransposeS perm2 a2@(Ast.AstScatterS+                                               @_ @shn2 @shp2 _ _ _ _ _)) w)+      | Just Refl <- testEquality perm1 perm2 ->+        gcastWith (unsafeCoerceRefl :: shp1 ++ shn1 :~: shp2 ++ shn2) $+        maybe (AstPlusS u v) (\res -> astTransposeS perm1 res + w)+        $ astPlusSMaybe a1 a2+  u v -> AstPlusS u v++astPlusSMaybe :: (NumScalar r, KnownSpan s)+              => AstTensor AstMethodLet s (TKS sh r)+              -> AstTensor AstMethodLet s (TKS sh r)+              -> Maybe (AstTensor AstMethodLet s (TKS sh r))+astPlusSMaybe = \cases+  (Ast.AstScatterS @_ @_ @shp (m :$$ shmRest) shn shp+                              u (var ::$ vars, ix))+    (Ast.AstScatterS @_ @_ @shp2 (m2 :$$ shmRest2) shn2 shp2+                                 u2 (var2 ::$ vars2, ix2))+      | Just Refl <- testEquality shmRest shmRest2+      , Just Refl <- testEquality shn shn2 ->+        gcastWith (unsafeCoerceRefl :: shp :~: shp2) $+        fuseScatters2 (m :$$ shmRest) shn shp u (var ::$ vars, ix)+                      (m2 :$$ shmRest2) shn2 shp2 u2 (var2 ::$ vars2, ix2)+        -- Same as above, but with arguments reversed, in case it works better:+        `mplus`+        fuseScatters2 (m2 :$$ shmRest2) shn2 shp2 u2 (var2 ::$ vars2, ix2)+                      (m :$$ shmRest) shn shp u (var ::$ vars, ix)+{- TODO: not needed ATM, though this could simplify astScatterS+  (Ast.AstScatterS @_ @_ @shp ZSS shn shp+                              u (_, ix))+    (Ast.AstScatterS @_ @_ @shp2 ZSS shn2 shp2+                                 u2 (_, ix2))+      | Just Refl <- testEquality shn shn2 ->+        gcastWith (unsafeCoerceRefl :: shp :~: shp2) $+        fuseScatters0 shn shp u ix+                      shn2 shp2 u2 ix2+  (Ast.AstScatterS @_ @_ @shp (m :$$ shmRest) shn shp+                              u (var ::$ vars, ix))+    (Ast.AstScatterS @_ @_ @shp2 shm2 shn2 shp2+                                 u2 (vars2, ix2))+      | Just Refl <- testEquality shmRest shm2+      , Just Refl <- testEquality shn shn2 ->+        gcastWith (unsafeCoerceRefl :: shp :~: shp2) $+        fuseScatters1 (m :$$ shmRest) shn shp u (var ::$ vars, ix)+                      shm2 shn2 shp2 u2 (vars2, ix2)+  (Ast.AstScatterS @_ @_ @shp2 shm2 shn2 shp2+                               u2 (vars2, ix2))+    (Ast.AstScatterS @_ @_ @shp (m :$$ shmRest) shn shp+                                u (var ::$ vars, ix))+      | Just Refl <- testEquality shmRest shm2+      , Just Refl <- testEquality shn shn2 ->+        gcastWith (unsafeCoerceRefl :: shp :~: shp2) $+        fuseScatters1 (m :$$ shmRest) shn shp u (var ::$ vars, ix)+                      shm2 shn2 shp2 u2 (vars2, ix2) -}+  _ _ -> Nothing++{- TODO: not needed ATM+fuseScatters0 :: forall shn shp x s. (TKAllNum x, KnownSpan s)+              => ShS shn -> ShS shp+              -> AstTensor AstMethodLet s (TKS2 shn x)+              -> AstIxS AstMethodLet shp+              -> ShS shn -> ShS shp+              -> AstTensor AstMethodLet s (TKS2 shn x)+              -> AstIxS AstMethodLet shp+              -> Maybe (AstTensor AstMethodLet s (TKS2 (shp ++ shn) x))+fuseScatters0 shn shp u ix+              _shn2 _shp2 u2 ix2 =+  funToAstIntMaybe (Just (0, 1)) $ \(var', ast') ->+    let g i i2 = astCond (AstConcreteK 1 <=. ast')+                         i2 i+        ix3 = ixsZipWith g ix ix2+        a = astFromVector (SNat @2) (STKS shn (stkAstX u)) (V.fromList [u, u2])+    in if any isCond ix3+       then Nothing  -- a risk of nested conditional build-up+       else Just $ astScatterS (SNat @2 :$$ ZSS) shn shp+                               a (var' ::$ ZS, ix3)++fuseScatters1 :: forall m shm shn shp x s. (TKAllNum x, KnownSpan s)+              => ShS (m : shm) -> ShS shn -> ShS shp+              -> AstTensor AstMethodLet s (TKS2 ((m : shm) ++ shn) x)+              -> (AstVarListS (m : shm), AstIxS AstMethodLet shp)+              -> ShS shm -> ShS shn -> ShS shp+              -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+              -> (AstVarListS shm, AstIxS AstMethodLet shp)+              -> Maybe (AstTensor AstMethodLet s (TKS2 (shp ++ shn) x))+fuseScatters1 (m :$$ shmRest) shn shp u (var ::$ vars, ix)+              _shm2 _shn2 _shp2 u2 (vars2, ix2) =+  let ub = fromSNat' m+      var' = reboundsVarName (0, ub) var+      f i2 = foldr (\(v, v2) -> substituteAst (astVar v) v2) i2+                   (ixsZip vars vars2)+      g i i2 = astCond (AstConcreteK (fromSNat' m) <=. astVar var')+                       (f i2) i  -- i has, effectively, varFalse in it+      ix3 = ixsZipWith g ix ix2+      a = astReplicateS (SNat @1 :$$ ZSS) u2+  in if any isCond ix3+     then Nothing  -- a risk of nested conditional build-up+     else Just $ astScatterS (SS m :$$ shmRest) shn shp+                             (astAppendS u a) (var' ::$ vars, ix3) -}++fuseScatters2 :: forall m m2 shm shn shp x s. (TKAllNum x, KnownSpan s)+              => ShS (m : shm) -> ShS shn -> ShS shp+              -> AstTensor AstMethodLet s (TKS2 ((m : shm) ++ shn) x)+              -> (AstVarListS (m : shm), AstIxS AstMethodLet shp)+              -> ShS (m2 : shm) -> ShS shn -> ShS shp+              -> AstTensor AstMethodLet s (TKS2 ((m2 : shm) ++ shn) x)+              -> (AstVarListS (m2 : shm), AstIxS AstMethodLet shp)+              -> Maybe (AstTensor AstMethodLet s (TKS2 (shp ++ shn) x))+fuseScatters2 (m :$$ shmRest) shn shp u (var ::$ vars, ix)+              (m2 :$$ _shmRest2) _shn2 _shp2 u2 (var2 ::$ vars2, ix2) =+  let ub = fromSNat' m + fromSNat' m2 - 1+      var' = reboundsVarName (0, ub) var+      astVarTrue = astVar $ reboundsVarName (fromSNat' m, ub) var+      f i2 = substituteAst (astVarTrue - AstConcreteK (fromSNat' m)) var2+             $ foldr (\(v, v2) -> substituteAst (astVar v) v2)+                     i2+                     (ixsZip (unAstVarListS vars) (unAstVarListS vars2))+      g i i2 = astCondK (AstConcreteK (fromSNat' m) <=. astVar var')+                        (f i2) i  -- i has, effectively, varFalse in it+      ix3 = ixsZipWith g ix ix2+  in if any isCond ix3+     then Nothing  -- a risk of nested conditional build-up+     else Just $ astScatterS (snatPlus m m2 :$$ shmRest) shn shp+                             (astAppendS u u2) (var' ::$ vars, ix3)++-- Checks whether it's a conditional or something easily obtained+-- from conditionals via rewrite rules.+isCond :: AstInt AstMethodLet -> Bool+isCond (Ast.AstLet _var _u v) = isCond v+isCond Ast.AstCondK{} = True+isCond (AstPlusK _ Ast.AstCondK{}) = True+isCond (AstPlusK Ast.AstCondK{} _) = True+isCond (AstTimesK _ Ast.AstCondK{}) = True+isCond (AstTimesK Ast.AstCondK{} _) = True+isCond _ = False++-- Just as with astTimesK, factors are flattened and a constant comes first.+-- Also, negation comes first.+astTimesS :: (NumScalar r, KnownSpan s)+          => AstTensor AstMethodLet s (TKS sh r)+          -> AstTensor AstMethodLet s (TKS sh r)+          -> AstTensor AstMethodLet s (TKS sh r)+astTimesS = \cases+  u v | FTKS (snat :$$ _) _ <- ftkAst u+      , Just u0 <- unRepl1 u+      , Just v0 <- unRepl1 v ->+        astReplicateS (snat :$$ ZSS) $ astTimesS u0 v0+  (Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astTimesS u v+  (Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)+    | Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->+      plainPart $ astTimesS u v+  (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astTimesS u v+  u@Ast.AstFromDual{} Ast.AstFromDual{} ->+    fromPlain $ AstConcreteS $ defTargetRep $ ftkAst u+  (Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astTimesS u v+  (AstTimesS u w) v -> astTimesS u (astTimesS w v)  -- flattening+  u _ | Just 0 <- unReplC u -> u+  _ v | Just 0 <- unReplC v -> v+  u v | Just 1 <- unReplC u -> v+  u v | Just 1 <- unReplC v -> u+  u v | Just (-1) <- unReplC u -> negate v+  u v | Just (-1) <- unReplC v -> negate u+  u v | Just u0 <- unAstS u+      , Just v0 <- unAstS v -> fromPlain $ AstConcreteS (u0 * v0)+  u v | Just v0 <- unAstS v -> AstTimesS (fromPlain $ AstConcreteS v0) u+  u (AstTimesS v w) | Just v0 <- unAstS v ->+    astTimesS (fromPlain $ AstConcreteS v0) u * w++  (Ast.AstN1S NegateOp u) (Ast.AstN1S NegateOp v) -> astTimesS u v+  (Ast.AstN1S NegateOp u) (AstTimesS (Ast.AstN1S NegateOp v) w) ->+    astTimesS u v * w+  u (Ast.AstN1S NegateOp v) -> astTimesS (negate u) v+  u (AstTimesS (Ast.AstN1S NegateOp v) w) -> astTimesS (negate u) v * w++  (Ast.AstR2S DivideOp u1 u2) v -> astR2S DivideOp (u1 * v) u2+  u (Ast.AstR2S DivideOp v1 v2) -> astR2S DivideOp (u * v1) v2+  (Ast.AstR1S RecipOp u) v -> astR2S DivideOp v u+  u (Ast.AstR1S RecipOp v) -> astR2S DivideOp u v++  (Ast.AstScatterS shm shn shp v (vars, ix)) u+    | Just w <- unRepl u, FTKS shv _ <- ftkAst v ->+      Ast.AstScatterS shm shn shp (v `astTimesS` astReplicateK shv (kfromS w))+                                  (vars, ix)+  (Ast.AstScatterS shm shn shp v (vars, ix)) (AstTimesS u a)+    | Just w <- unRepl u, FTKS shv _ <- ftkAst v ->+      astTimesS+        (Ast.AstScatterS shm shn shp (v `astTimesS` astReplicateK+                                                      shv (kfromS w))+                                     (vars, ix)) a+  u (Ast.AstScatterS shm shn shp v (vars, ix))+    | Just w <- unRepl u, FTKS shv _ <- ftkAst v ->+      Ast.AstScatterS shm shn shp (astReplicateK shv (kfromS w) `astTimesS` v)+                                  (vars, ix)+  u  (AstTimesS (Ast.AstScatterS shm shn shp v (vars, ix)) a)+    | Just w <- unRepl u, FTKS shv _ <- ftkAst v ->+      astTimesS+        (Ast.AstScatterS shm shn shp (astReplicateK shv (kfromS w)+                                      `astTimesS` v)+                                     (vars, ix)) a+  (Ast.AstGatherS shm shn shp v (vars, ix)) u+    | Just w <- unRepl u, FTKS shv _ <- ftkAst v ->+      Ast.AstGatherS shm shn shp (v `astTimesS` astReplicateK shv (kfromS w))+                                 (vars, ix)+  (Ast.AstGatherS shm shn shp v (vars, ix)) (AstTimesS u a)+    | Just w <- unRepl u, FTKS shv _ <- ftkAst v ->+      astTimesS+        (Ast.AstGatherS shm shn shp (v `astTimesS` astReplicateK shv (kfromS w))+                                    (vars, ix)) a+  u (Ast.AstGatherS shm shn shp v (vars, ix))+    | Just w <- unRepl u, FTKS shv _ <- ftkAst v ->+      Ast.AstGatherS shm shn shp (astReplicateK shv (kfromS w) `astTimesS` v)+                                 (vars, ix)+  u (AstTimesS (Ast.AstGatherS shm shn shp v (vars, ix)) a)+    | Just w <- unRepl u, FTKS shv _ <- ftkAst v ->+      astTimesS+        (Ast.AstGatherS shm shn shp (astReplicateK shv (kfromS w) `astTimesS` v)+                                    (vars, ix)) a++  (AstConvUpSFromK u) (AstConvUpSFromK v) -> sfromK $ astTimesK u v+  u (AstConvUpSFromK v) -> sfromK $ astTimesK (kfromS u) v+  (AstConvUpSFromK u) v -> sfromK $ astTimesK u (kfromS v)++  -- This breaks sharing, because although u is concrete and so doesn't+  -- have to be shared, the multiplication is not shared --- we end up+  -- with one addition and two multiplications, not one.+  -- u@AstConcreteS{} * AstPlusS v w -> AstPlusS (astTimesS u v) (u * w)++  u v -> AstTimesS u v++astN1S :: (NumScalar r, KnownSpan s)+       => OpCodeNum1 -> AstTensor AstMethodLet s (TKS sh r)+       -> AstTensor AstMethodLet s (TKS sh r)+astN1S opCode t = case t of+  _ | FTKS (snat :$$ _) _ <- ftkAst t+    , Just t0 <- unRepl1 t ->+      astReplicateS (snat :$$ ZSS) (astN1S opCode t0)+-- This increases work:+-- Ast.AstCond b n k -> astCond b (astN1S opCode n) (astN1S opCode k)+  Ast.AstLet var n k -> astLet var n (astN1S opCode k)+  Ast.AstPrimalPart n -> primalPart (astN1S opCode n)+  Ast.AstDualPart n -> dualPart (astN1S opCode n)+  Ast.AstPlainPart n -> plainPart (astN1S opCode n)+  Ast.AstFromPrimal n -> fromPrimal (astN1S opCode n)+  Ast.AstFromDual n -> fromDual (astN1S opCode n)+  Ast.AstFromPlain n -> fromPlain (astN1S opCode n)+  AstConvUpSFromK n -> sfromK $ astN1K opCode n+  _ -> case (opCode, t) of+    (NegateOp, AstConcreteS n) -> AstConcreteS (negate n)+    (NegateOp, AstPlusS u v) -> AstPlusS (negate u) (negate v)+    (NegateOp, AstTimesS u v) -> astTimesS (negate u) v+    (NegateOp, Ast.AstN1S NegateOp u) -> u+    (NegateOp, Ast.AstN1S SignumOp u) -> astN1S SignumOp (negate u)+    (NegateOp, Ast.AstI2S QuotOp u v) -> astI2S QuotOp (negate u) v+        -- v is likely positive and let's keep it so+    (NegateOp, Ast.AstI2S RemOp u v) -> astI2S RemOp (negate u) v+    (NegateOp, Ast.AstScatterS shm shn shp v (vars, ix)) ->+      astScatterS shm shn shp (negate v) (vars, ix)+    (NegateOp, Ast.AstGatherS shm shn shp v (vars, ix)) ->+      astGatherS shm shn shp (negate v) (vars, ix)+    (AbsOp, AstConcreteS n) -> AstConcreteS (abs n)+    (AbsOp, Ast.AstN1S AbsOp u) -> astN1S AbsOp u+    (AbsOp, Ast.AstN1S NegateOp u) -> astN1S AbsOp u+    (SignumOp, AstConcreteS n) -> AstConcreteS (signum n)+    (SignumOp, Ast.AstN1S SignumOp u) -> astN1S SignumOp u+    _ -> Ast.AstN1S opCode t++astR1S :: (NumScalar r, Differentiable r, KnownSpan s)+       => OpCode1 -> AstTensor AstMethodLet s (TKS sh r)+       -> AstTensor AstMethodLet s (TKS sh r)+astR1S opCode = \case+  t | FTKS (snat :$$ _) _ <- ftkAst t+    , Just t0 <- unRepl1 t ->+      astReplicateS (snat :$$ ZSS) (astR1S opCode t0)+-- This increases work:+--  Ast.AstCond b n k -> astCond b (astR1S opCode n) (astR1S opCode k)+  Ast.AstLet var n k -> astLet var n (astR1S opCode k)+  Ast.AstPrimalPart u -> primalPart $ astR1S opCode u+  Ast.AstPlainPart u -> plainPart $ astR1S opCode u+  Ast.AstFromPrimal u -> fromPrimal $ astR1S opCode u+  Ast.AstFromPlain u -> fromPlain $ astR1S opCode u+  AstConvUpSFromK u -> sfromK $ astR1K opCode u+  AstConcreteS u -> case opCode of+    RecipOp -> AstConcreteS $ recip u+    ExpOp -> AstConcreteS $ exp u+    LogOp -> AstConcreteS $ log u+    SqrtOp -> AstConcreteS $ sqrt u+    SinOp -> AstConcreteS $ sin u+    CosOp -> AstConcreteS $ cos u+    TanOp -> AstConcreteS $ tan u+    AsinOp -> AstConcreteS $ asin u+    AcosOp -> AstConcreteS $ acos u+    AtanOp -> AstConcreteS $ atan u+    SinhOp -> AstConcreteS $ sinh u+    CoshOp -> AstConcreteS $ cosh u+    TanhOp -> AstConcreteS $ tanh u+    AsinhOp -> AstConcreteS $ asinh u+    AcoshOp -> AstConcreteS $ acosh u+    AtanhOp -> AstConcreteS $ atanh u+  u -> Ast.AstR1S opCode u++astR2S :: (NumScalar r, Differentiable r, KnownSpan s)+       => OpCode2+       -> AstTensor AstMethodLet s (TKS sh r)+       -> AstTensor AstMethodLet s (TKS sh r)+       -> AstTensor AstMethodLet s (TKS sh r)+astR2S opCode = \cases+  u v | FTKS (snat :$$ _) _ <- ftkAst u+      , Just u0 <- unRepl1 u+      , Just v0 <- unRepl1 v ->+        astReplicateS (snat :$$ ZSS) (astR2S opCode u0 v0)+  (Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astR2S opCode u v+  (Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)+    | Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->+      plainPart $ astR2S opCode u v+  (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astR2S opCode u v+  (Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astR2S opCode u v+  (AstConvUpSFromK n) (AstConvUpSFromK k) ->+    sfromK $ astR2K opCode n k+  n (AstConvUpSFromK k) ->+    sfromK $ astR2K opCode (kfromS n) k+  (AstConvUpSFromK n) k ->+    sfromK $ astR2K opCode n (kfromS k)+  (AstConcreteS u) (AstConcreteS v) -> case opCode of+    DivideOp -> AstConcreteS $ u / v+    PowerOp -> AstConcreteS $ u ** v+    LogBaseOp -> AstConcreteS $ logBase u v+    Atan2Op -> AstConcreteS $ atan2H u v+  u v -> case opCode of+    DivideOp -> case u of+      _ | Just 0 <- unReplC u -> u+      _ | Just 1 <- unReplC u -> recip v+      _ | Just 1 <- unReplC v -> u+      _ | Just 0 <- unReplC v -> u  -- the partiality-removal hack+      AstTimesS n t+        | eqY n v -> t+      _ -> Ast.AstR2S DivideOp u v+             -- TODO: add other rules that are relatively numerically stable+    _ -> Ast.AstR2S opCode u v++astI2S :: (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)+       => OpCodeIntegral2+       -> AstTensor AstMethodLet s (TKS sh r)+       -> AstTensor AstMethodLet s (TKS sh r)+       -> AstTensor AstMethodLet s (TKS sh r)+astI2S opCode = \cases+  u v | FTKS (snat :$$ _) _ <- ftkAst u+      , Just u0 <- unRepl1 u+      , Just v0 <- unRepl1 v ->+        astReplicateS (snat :$$ ZSS) (astI2S opCode u0 v0)+  (Ast.AstPrimalPart n) (Ast.AstPrimalPart k) ->+    primalPart (astI2S opCode n k)+  (Ast.AstPlainPart @_ @s1 n) (Ast.AstPlainPart @_ @s2 k)+    | Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->+      plainPart (astI2S opCode n k)+  (Ast.AstFromPrimal n) (Ast.AstFromPrimal k) ->+    fromPrimal (astI2S opCode n k)+  (Ast.AstFromPlain n) (Ast.AstFromPlain k) ->+    fromPlain (astI2S opCode n k)+  (AstConvUpSFromK n) (AstConvUpSFromK k) ->+    sfromK $ astI2K opCode n k+  n (AstConvUpSFromK k) ->+    sfromK $ astI2K opCode (kfromS n) k+  (AstConvUpSFromK n) k ->+    sfromK $ astI2K opCode n (kfromS k)+  u v -> case opCode of+    QuotOp -> case (u, v) of+      _ | Just u0 <- unAstS u+        , Just v0 <- unAstS v -> fromPlain $ AstConcreteS (quotH u0 v0)+      _ | Just 0 <- unReplC u -> u+      _ | Just 1 <- unReplC v -> u+      _ | Just 0 <- unReplC v -> u  -- the partiality-removal hack+      (Ast.AstI2S QuotOp u0 v0, _) ->+        astI2S QuotOp u0 (astTimesS v0 v)+      (AstTimesS n t, _)+        | eqY n v -> t+      _ -> Ast.AstI2S QuotOp u v+    RemOp -> case (u, v) of+      _ | Just u0 <- unAstS u+        , Just v0 <- unAstS v -> fromPlain $ AstConcreteS (remH u0 v0)+      _ | Just 0 <- unReplC u -> u+      _ | Just 1 <- unReplC v ->+          fromPlain $ AstConcreteS $ defTargetRep $ ftkAst u+      _ | Just 0 <- unReplC v ->  -- the partiality-removal hack+          fromPlain $ AstConcreteS $ defTargetRep $ ftkAst u+      _ -> Ast.AstI2S RemOp u v++astConcreteS :: GoodScalar r+             => Concrete (TKS sh r)+             -> AstTensor AstMethodLet PlainSpan (TKS sh r)+astConcreteS = AstConcreteS . unConcrete++-- We don't floor concrete tensors ASAP but wait until AstTraverse,+-- because the flooring may fuse or simplify away yet.+astFloorS :: forall r1 r2 sh.+             (NumScalar r1, Differentiable r1, NumScalar r2, Integral r2)+          => AstTensor AstMethodLet PlainSpan (TKS sh r1)+          -> AstTensor AstMethodLet PlainSpan (TKS sh r2)+astFloorS t = case t of+  _ | FTKS (snat :$$ _) _ <- ftkAst t+    , Just u <- unRepl1 t ->+      astReplicateS (snat :$$ ZSS) (astFloorS u)+  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->+    let !v2 = astFloorS v+    in Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v2)+  Ast.AstBuild1 snat STKScalar (var, v) ->+    let !v2 = astFloorK v+    in Ast.AstBuild1 snat STKScalar (var, v2)+  Ast.AstLet var u v -> astLet var u (astFloorS v)+  Ast.AstFloorS v -> astFloorS v+  Ast.AstFromIntegralS v -> astFromIntegralS v+  Ast.AstCastS v -> astFloorS v+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterS shm shn shp (astFloorS v) (vars, ix)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherS shm shn shp (astFloorS v) (vars, ix)+  Ast.AstIotaS snat -> Ast.AstIotaS snat+  Ast.AstReverseS v -> astReverseS (astFloorS v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astFloorS v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astFloorS v)+  AstConvUpSFromK a -> sfromK (astFloorK a)+  _ -> Ast.AstFloorS t++astFromIntegralS :: forall r1 r2 sh. (NumScalar r1, NumScalar r2, Integral r1)+                 => AstTensor AstMethodLet PlainSpan (TKS sh r1)+                 -> AstTensor AstMethodLet PlainSpan (TKS sh r2)+astFromIntegralS t = case t of+  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t+  _ | FTKS (snat :$$ _) _ <- ftkAst t+    , Just u <- unRepl1 t ->+      astReplicateS (snat :$$ ZSS) (astFromIntegralS u)+--  Ast.AstFromVector snat (STKS sh STKScalar) l ->+--   astFromVector snat (STKS sh STKScalar) (V.map astFromIntegralS l)+--  Ast.AstFromVector snat STKScalar l ->+--   astFromVector snat STKScalar (V.map astFromIntegralK l)+  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->+    let !v2 = astFromIntegralS v+    in Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v2)+  Ast.AstBuild1 snat STKScalar (var, v) ->+    let !v2 = astFromIntegralK v+    in Ast.AstBuild1 snat STKScalar (var, v2)+  Ast.AstLet var u v -> astLet var u (astFromIntegralS v)+  Ast.AstN1S opCode u -> astN1S opCode (astFromIntegralS u)+  Ast.AstFromIntegralS v -> astFromIntegralS v+--  Ast.AstIndexS shn v ix -> astIndexS shn (astFromIntegralS v) ix+    -- increases work; also index goes into fromIntegral, so we'd loop+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterS shm shn shp (astFromIntegralS v) (vars, ix)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherS shm shn shp (astFromIntegralS v) (vars, ix)+  Ast.AstIotaS snat -> Ast.AstIotaS snat+--  Ast.AstSliceS i n k v -> astSliceS i n k (astFromIntegralS v)+  Ast.AstReverseS v -> astReverseS (astFromIntegralS v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astFromIntegralS v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astFromIntegralS v)+  AstConvUpSFromK a -> sfromK (astFromIntegralK a)+  _ -> Ast.AstFromIntegralS t++astCastS :: forall r1 r2 s sh.+            ( NumScalar r1, Differentiable r1+            , NumScalar r2, Differentiable r2, KnownSpan s )+         => AstTensor AstMethodLet s (TKS sh r1)+         -> AstTensor AstMethodLet s (TKS sh r2)+astCastS t = case t of+  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t+  _ | FTKS (snat :$$ _) _ <- ftkAst t+    , Just u <- unRepl1 t ->+      astReplicateS (snat :$$ ZSS) (astCastS u)+--  Ast.AstFromVector snat (STKS sh STKScalar) l ->+--   astFromVector snat (STKS sh STKScalar) (V.map astCastS l)+--  Ast.AstFromVector snat STKScalar l ->+--   astFromVector snat STKScalar (V.map astCastK l)+  {- This (and other similar rules) is bad, because it changes semantics+     and also impacts performance negatively (a is larger than sum a):+  Ast.AstSumS shm a -> astSumS shm (astCastS a) -}+  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->+    let !v2 = astCastS v+    in Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v2)+  Ast.AstBuild1 snat STKScalar (var, v) ->+    let !v2 = astCastK v+    in Ast.AstBuild1 snat STKScalar (var, v2)+  Ast.AstLet var u v -> astLet var u (astCastS v)+  -- These (rarely) loop if ast*Part is used instead of these versions.+  Ast.AstPrimalPart a -> primalPart $ astCastS a+  Ast.AstDualPart a -> dualPart $ astCastS a+  Ast.AstPlainPart a -> plainPart $ astCastS a+  Ast.AstFromPrimal v -> fromPrimal $ astCastS v+  Ast.AstFromDual v -> fromDual $ astCastS v+  Ast.AstFromPlain v -> fromPlain $ astCastS v+  Ast.AstN1S opCode u -> astN1S opCode (astCastS u)+  Ast.AstR1S opCode u -> astR1S opCode (astCastS u)+  Ast.AstFromIntegralS v -> astFromIntegralS v+  Ast.AstCastS v -> astCastS v+--  Ast.AstIndexS shn v ix -> astIndexS shn (astCastS v) ix+    -- increases work; also index goes into astCastS, so we'd loop+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterS shm shn shp (astCastS v) (vars, ix)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherS shm shn shp (astCastS v) (vars, ix)+  Ast.AstIotaS snat -> Ast.AstIotaS snat+--  Ast.AstSliceS i n k v -> astSliceS i n k (astCastS v)+  Ast.AstReverseS v -> astReverseS (astCastS v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astCastS v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astCastS v)+  AstConvUpSFromK a -> sfromK (astCastK a)+  _ -> Ast.AstCastS t++astIndexS+  :: forall shm shn s x. KnownSpan s+  => ShS shn+  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x) -> AstIxS AstMethodLet shm+  -> AstTensor AstMethodLet s (TKS2 shn x)+astIndexS = astIndexKnobsS defaultKnobs++astIndexKnobsS+  :: forall shm shn s x. KnownSpan s+  => SimplifyKnobs+  -> ShS shn+  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+  -> AstIxS AstMethodLet shm+  -> AstTensor AstMethodLet s (TKS2 shn x)+astIndexKnobsS _ _ v0 ZIS = v0+astIndexKnobsS _ shn v0 (i1 :.$ _)+  | Just (lb, ub) <- intBounds i1+  , FTKS (snat :$$ _) x <- ftkAst v0+  , ub < 0 || lb >= fromSNat' snat =+    let ftk = FTKS shn x+    in fromPlain $ astConcrete ftk (tdefTarget ftk)+astIndexKnobsS knobs shn u0 (Ast.AstCondK b v w :.$ rest0)+  | knobPhase knobs `notElem` [PhaseUnspecified, PhaseVectorization] =+      -- don't undo vectorization tweaks+    astLetFun u0 $ \u ->+    shareIx rest0 $ \rest ->+    astCondS b (astIndexKnobsS knobs shn u (v :.$ rest))+               (astIndexKnobsS knobs shn u (w :.$ rest))+astIndexKnobsS knobs shn v0 ix@(i1 :.$ rest1)+  | FTKS shmshn x <- ftkAst v0+  , (:$$) @in1 @shm1 SNat _ <- shsTakeIx @shm @shn Proxy ix shmshn+  , let tryRecursing :: AstTensor AstMethodLet s (TKS2 (shm1 ++ shn) x)+                     -> AstTensor AstMethodLet s (TKS2 shn x)+        tryRecursing u =+          let ftk = FTKS shn x+              defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)+          in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of+            AstConcreteK b -> if b then astIndex shn u rest1 else defArr+            _ -> Ast.AstIndexS shn v0 ix = case v0 of+  Ast.AstProject1{} -> Ast.AstIndexS shn v0 ix+  Ast.AstProject2{} -> Ast.AstIndexS shn v0 ix+  Ast.AstApply{} -> Ast.AstIndexS shn v0 ix+  Ast.AstVar{} -> Ast.AstIndexS shn v0 ix+  {- This is wrong: in a counterfactual case, astLet assigns OOB i to var2,+     violating the invariant about variables bounds:+  Ast.AstBuild1 (SNat @k) STKS{} (var2, v) ->+    let ftk = FTKS shn x+        defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)+    in astLetFun i1 $ \i ->+    astCondInitial (0 <=. i &&* i <=. valueOf @k - 1)+            (astIndex shn (astLet var2 i v) rest1)+            defArr -}+  Ast.AstBuild1 _ STKScalar (var2, v) | ZIS <- rest1 ->+    tryRecursing $ sfromK $ astLet var2 i1 v+  Ast.AstBuild1 _ STKS{} (var2, v) ->+    tryRecursing $ astLet var2 i1 v++  Ast.AstLet var u v -> astLetRefresh var u v $ \v' -> astIndex shn v' ix++  Ast.AstPrimalPart{} -> Ast.AstIndexS shn v0 ix  -- must be a NF+  Ast.AstDualPart{} -> Ast.AstIndexS shn v0 ix+  Ast.AstPlainPart{} -> Ast.AstIndexS shn v0 ix+  Ast.AstFromPrimal v -> fromPrimal $ astIndex shn v ix+  Ast.AstFromDual v -> fromDual $ astIndex shn v ix+  Ast.AstFromPlain v -> fromPlain $ astIndex shn v ix++  AstPlusS u v ->+    shareIx ix $ \ !ix2 -> astPlusS (astIndex shn u ix2)+                                    (astIndex shn v ix2)+  AstTimesS u v ->+    shareIx ix $ \ !ix2 -> astTimesS (astIndex shn u ix2)+                                     (astIndex shn v ix2)+  Ast.AstN1S opCode u -> astN1S opCode (astIndex shn u ix)+  Ast.AstR1S opCode u -> astR1S opCode (astIndex shn u ix)+  Ast.AstR2S opCode u v ->+    shareIx ix $ \ !ix2 -> astR2S opCode (astIndex shn u ix2)+                                         (astIndex shn v ix2)+  Ast.AstI2S opCode u v ->+    shareIx ix $ \ !ix2 -> astI2S opCode (astIndex shn u ix2)+                                         (astIndex shn v ix2)+  AstConcreteS a | AstConcreteK i <- i1 ->+    let u = withKnownShS (shsTail shmshn) $+            tsindex (Concrete a) (Concrete i :.$ ZIS)+    in astIndex shn (astConcreteS u) rest1+  AstConcreteS{} -> case unRepl1 v0 of+    Just u -> tryRecursing u+    _ -> Ast.AstIndexS shn v0 ix+  Ast.AstFloorS v -> astFloorS $ astIndex shn v ix+  Ast.AstFromIntegralS v -> astFromIntegralS $ astIndex shn v ix+  Ast.AstCastS t -> astCastS $ astIndex shn t ix+  Ast.AstArgMinS @n1 @shz v | FTKS nsh _ <- ftkAst v+                            , nl@(SNat @nl) <- shsLast nsh ->+    let shnl = shn `shsAppend` (nl :$$ ZSS)+    in gcastWith (unsafeCoerceRefl+                  :: Head (shn ++ '[nl]) ': Tail (shn ++ '[nl])+                     :~: shn ++ '[nl]) $+       gcastWith (unsafeCoerceRefl :: Init (shn ++ '[nl]) :~: shn) $+       gcastWith (unsafeCoerceRefl :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) $+       Ast.AstArgMinS @(Head (shn ++ '[nl])) @(Tail (shn ++ '[nl]))+       $ astIndex shnl v ix+  Ast.AstArgMaxS @n1 @shz v | FTKS nsh _ <- ftkAst v+                            , nl@(SNat @nl) <- shsLast nsh ->+    let shnl = shn `shsAppend` (nl :$$ ZSS)+    in gcastWith (unsafeCoerceRefl+                  :: Head (shn ++ '[nl]) ': Tail (shn ++ '[nl])+                     :~: shn ++ '[nl]) $+       gcastWith (unsafeCoerceRefl :: Init (shn ++ '[nl]) :~: shn) $+       gcastWith (unsafeCoerceRefl :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) $+       Ast.AstArgMaxS @(Head (shn ++ '[nl])) @(Tail (shn ++ '[nl]))+       $ astIndex shnl v ix+  Ast.AstIndexS @shm4 _ v ix2+    | Refl <- lemAppAssoc (Proxy @shm4) (Proxy @shm) (Proxy @shn) ->+      astIndex shn v (ix2 `ixsAppend` ix)++  Ast.AstCondS b v w ->+    shareIx ix $ \ !ix2 -> astCondS b (astIndex shn v ix2)+                                      (astIndex shn w ix2)+  Ast.AstFromVectorK (_ :$$ ZSS) l | AstConcreteK i <- i1, ZIS <- rest1 ->+    tryRecursing $ sfromK (l V.! i)+  Ast.AstFromVectorK{} ->  -- normal form+    Ast.AstIndexS shn v0 ix+  Ast.AstFromVectorS (_ :$$ ZSS) l | AstConcreteK i <- i1 ->+    tryRecursing (l V.! i)+  Ast.AstFromVectorS @shm2 @shn2 shm2 l | shsLength shm2 < length ix+                                        , SNat <- shsRank shm2 ->+    gcastWith (unsafeCoerceRefl :: shn2 :~: Drop (Rank shm2) shm ++ shn) $+    gcastWith (unsafeCoerceRefl :: Take (Rank shm2) shm :~: shm2) $+    shareIx ix $ \ !ix2 ->+      let ixRest2 = ixsDrop @(Rank shm2) ix2+      in Ast.AstIndexS shn+                       (astFromVectorS @shm2 @shn shm2+                        $ V.map (\a -> astIndex shn a ixRest2) l)+                       (ixsTake @(Rank shm2) ix2)+  Ast.AstFromVectorS{} ->  -- normal form (shsLength shm2 >= length ix)+    Ast.AstIndexS shn v0 ix+  Ast.AstSumS @shm2 shm2 v ->+    let perm1 = backpermCycleN (shsLength shm2) (shsLength shm2 + length ix)+    in Permutation.permFromListCont perm1 $ \(perm :: Permutation.Perm perm) ->+         gcastWith (unsafeCoerceRefl+                    :: (Rank perm <=? Rank (shm2 ++ (shm ++ shn))) :~: True) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix perm (shm2 ++ (shm ++ shn))+                       :~: shm ++ (shm2 ++ shn)) $+         fromMaybe (error "astIndexKnobsS: impossible non-permutation")+         $ Permutation.permCheckPermutation perm+         $ astSumS shm2+         $ astIndex (shm2 `shsAppend` shn) (astTransposeS perm v) ix+  Ast.AstScatterS shm7 shn7 shp7 v (vars, i5 :.$ ix2) | eqY i5 i1 ->+    astIndex shn (astScatter shm7 shn7 (shsTail shp7) v (vars, ix2)) rest1+  Ast.AstScatterS _ _ _ _ (_, AstConcreteK{} :.$ _)+    | AstConcreteK{} <- i1 ->  -- from above we know i5 /= i1+      let ftk = FTKS shn x+      in fromPlain $ astConcrete ftk (tdefTarget ftk)+  Ast.AstScatterS{} ->  -- normal form+    Ast.AstIndexS shn v0 ix+  Ast.AstReplicateK (_ :$$ shmRest) v ->+    tryRecursing $ astReplicateK shmRest v+  Ast.AstReplicateS ZSS v -> astIndex shn v ix+  Ast.AstReplicateS (_ :$$ shmRest) v ->+    tryRecursing $ astReplicateS shmRest v+  -- This is not a possible normal form, but pattern needs to be exhaustive.+  Ast.AstGatherS @_ @_ @shp' _ _ _ v (ZS, ix2)+    | Refl <- lemAppAssoc (Proxy @shp') (Proxy @(in1 ': shm1)) (Proxy @shn) ->+      astIndex @(shp' ++ shm) @shn shn v (ix2 `ixsAppend` ix)+  Ast.AstGatherS @_ @shn' (_ :$$ (shm71 :: ShS shm71)) shn' shp'+                 v (var2 ::$ vars, ix2) ->+    gcastWith (unsafeCoerceRefl :: shm71 ++ shn' :~: shm1 ++ shn) $+    let ftk = FTKS shn x+        defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)+        w :: AstTensor AstMethodLet s (TKS2 (shm1 ++ shn) x)+        w = astGather shm71 shn' shp' v (vars, ix2)+        u = astLet var2 i1 $ astIndex @shm1 @shn shn w rest1+          -- This let makes it impossible to use astCond when i1 is OOB.+          -- No astLetRefresh is needed, because the smart constructor+          -- for AstGatherS ensures var2 doesn't occur in v, even bound.+    in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of+      AstConcreteK b -> if b then u else defArr+      _ -> Ast.AstIndexS shn v0 ix+  Ast.AstIotaS _ -> case shn of+    ZSS ->+      let ftk = FTKS ZSS x+          defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)+      in astLetFun i1 $ \i ->+           astCondSInitial (0 <=. i &&* i <=. valueOf @in1 - 1)+                           (astFromIntegralS $ sfromK i)+                           defArr+    _ -> error "astIndexKnobsS: shn not []"+  Ast.AstAppendS u v | FTKS (SNat @m :$$ _) _ <- ftkAst u ->+    if knobPhase knobs == PhaseExpansion then+      astLetFun i1 $ \i ->+      shareIx rest1 $ \ !rest2 ->+      let ulen = valueOf @m+          ix1 = i :.$ rest2+          ix2 = i - ulen :.$ rest2+      in case ulen <=. i of+        AstConcreteK b -> if b then astIndex shn v ix2 else astIndex shn u ix1+        bExpr ->+          -- This results in a larger term, so we consider this late.+          astCondSInitial bExpr (astIndex shn v ix2) (astIndex shn u ix1)+    else+      -- We don't need to share i1 and rest1 here and, in fact, we can't+      -- because we'd loop via building AstLet around the resulting AstIndexS.+      let ulen = valueOf @m+          ix1 = i1 :.$ rest1+          ix2 = i1 - ulen :.$ rest1+      in case ulen <=. i1 of+        AstConcreteK b -> if b then astIndex shn v ix2 else astIndex shn u ix1+        _ -> Ast.AstIndexS shn v0 ix+  Ast.AstSliceS i@(SNat @i) (SNat @n) k@SNat v ->+    astLetFun i1 $ \iShared ->+    let ftk = FTKS shn x+        defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)+        b = (if fromSNat' i == 0 then true else 0 <=. iShared)+            &&* (if fromSNat' k == 0 then true else iShared <=. valueOf @n - 1)+        ii = valueOf @i + iShared+    in astCondSInitial b (astIndex shn v (ii :.$ rest1)) defArr+  Ast.AstReverseS v ->+    let iRev = valueOf @in1 - 1 - i1+    in astIndex shn v (iRev :.$ rest1)+  Ast.AstTransposeS @_ @sh2 perm v+    | fromSNat' (Permutation.permRank perm) <= length ix ->+      -- TODO: remake once there's an S version of permInverse:+      permInverse perm $ \(invperm :: Permutation.Perm invperm) _ ->+      let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix invperm shm)+          ix2 = ixsPermutePrefix invperm ix+      in gcastWith (unsafeCoerceRefl+                    :: sh2 :~: Permutation.PermutePrefix invperm shm ++ shn) $+         astIndex shn v ix2+  -- This results in a larger term, so we consider this late.+  Ast.AstTransposeS @perm perm v+    | knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion] ->+      astIndex shn (astTransposeAsGatherS @perm (deVect knobs) perm v) ix+  Ast.AstTransposeS{} -> Ast.AstIndexS shn v0 ix+  -- This results in a larger term, so we only permit it for the flat case.+  Ast.AstReshapeS sh v | shsLength sh <= 1 ->  -- this is flatten+    astIndex shn (astReshapeAsGatherS (deVect knobs) sh v) ix+  Ast.AstReshapeS{} -> Ast.AstIndexS shn v0 ix++  -- Conversions to shaped need to stay down, so this is NF.+  Ast.AstConvert{} -> Ast.AstIndexS shn v0 ix++  Ast.AstDot1InS @_ @n _ n u v+    | Refl <- lemAppAssoc (Proxy @shm) (Proxy @shn) (Proxy @'[n]) ->+      astDot1InS shn n (astIndex (shn `shsAppend` (n :$$ ZSS)) u ix)+                       (astIndex (shn `shsAppend` (n :$$ ZSS)) v ix)+  Ast.AstMatmul2S{} -> Ast.AstIndexS shn v0 ix  -- TODO++  Ast.AstBoolNotS a -> Ast.AstBoolNotS $ astIndex shn a ix+  Ast.AstBoolAndS b1 b2 -> Ast.AstBoolAndS (astIndex shn b1 ix)+                                           (astIndex shn b2 ix)+  Ast.AstLeqS @_ @sh _shb sh arg1 arg2+    | Refl <- lemAppAssoc (Proxy @shm) (Proxy @shn) (Proxy @sh) ->+      Ast.AstLeqS shn sh (astIndex (shn `shsAppend` sh) arg1 ix)+                         (astIndex (shn `shsAppend` sh) arg2 ix)+ where+  astIndex+    :: forall shm' shn' s' x'. KnownSpan s'+    => ShS shn'+    -> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') x')+    -> AstIxS AstMethodLet shm'+    -> AstTensor AstMethodLet s' (TKS2 shn' x')+  astIndex = astIndexKnobsS (deVect knobs)+  astGather+    :: forall shm' shn' shp' x'.+       ShS shm' -> ShS shn' -> ShS shp'+    -> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') x')+    -> (AstVarListS shm', AstIxS AstMethodLet shp')+    -> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') x')+  astGather = astGatherKnobsS (deVect knobs)+  astScatter+    :: forall shm' shn' shp' x'. TKAllNum x'+    => ShS shm' -> ShS shn' -> ShS shp'+    -> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') x')+    -> (AstVarListS shm', AstIxS AstMethodLet shp')+    -> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') x')+  astScatter = astScatterKnobsS (deVect knobs)++-- TODO: compared to tletIx, it adds many lets, not one, but does not+-- create other (and non-simplified!) big terms and also uses astIsSmall,+-- so it's probably more efficient. Use this instead of tletIx+-- or design something even better.+--+-- See https://github.com/Mikolaj/horde-ad/issues/119.+shareIx :: forall s shm y. KnownSpan s+        => AstIxS AstMethodLet shm+        -> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)+        -> AstTensor AstMethodLet s y+{-# NOINLINE shareIx #-}+shareIx ix f = unsafePerformIO $ do+  let shareI :: AstInt AstMethodLet+             -> IO ( Maybe (IntVarName, AstInt AstMethodLet)+                   , AstInt AstMethodLet )+      shareI i | astIsSmall True i = return (Nothing, i)+      shareI i =+        -- i can be OOB, so we can't use shape to determine its bounds+        funToAstIntMaybeIO (intBounds i) $ \ (!varFresh, !astVarFresh) ->+          (Just (varFresh, i), astVarFresh)+  (bindings, ix2) <- mapAndUnzipM shareI (Foldable.toList ix)+  return $! foldl' (\v (var, u) -> astLet var u v)+                   (f $ ixsFromIxS ix ix2)+                   (catMaybes bindings)++astCondKInitial :: forall r s. (GoodScalar r, KnownSpan s)+                => AstBool AstMethodLet+                -> AstTensor AstMethodLet s (TKScalar r)+                -> AstTensor AstMethodLet s (TKScalar r)+                -> AstTensor AstMethodLet s (TKScalar r)+astCondKInitial b v w = case b of+  AstLeqInt (AstConcreteK j) (AstIntVar var)+    | Just (lb, ub) <- varNameToBounds var ->+      let varTrue = reboundsVarName (max lb j, ub) var+          varFalse = reboundsVarName (lb, min ub (j - 1)) var+      in astCondK b (substituteAst (astVar varTrue) var v)+                    (substituteAst (astVar varFalse) var w)+  _ -> astCondK b v w++astCondK :: forall r s. (GoodScalar r, KnownSpan s)+         => AstBool AstMethodLet+         -> AstTensor AstMethodLet s (TKScalar r)+         -> AstTensor AstMethodLet s (TKScalar r)+         -> AstTensor AstMethodLet s (TKScalar r)+astCondK b v w | Just b0 <- unAstK b = if b0 then v else w+astCondK _b u v | eqY u v = u+astCondK (Ast.AstBoolNotK b) v w = astCondK b w v+astCondK (Ast.AstLet var n b) u v =+  astLetRefresh var n b $ \b' -> astCondKInitial b' u v+astCondK b (Ast.AstLet var n u) v =+  astLetRefresh var n u $ \u' -> astCondK b u' v+astCondK b u (Ast.AstLet var n v) =+  astLetRefresh var n v $ \v' -> astCondK b u v'+astCondK b (Ast.AstFromPrimal v) (Ast.AstFromPrimal w) =+  fromPrimal $ astCondK b v w+astCondK b (Ast.AstFromDual v) (Ast.AstFromDual w) =+  fromDual $ astCondK b v w+astCondK b (Ast.AstFromPlain v) (Ast.AstFromPlain w) =+  fromPlain $ astCondK b v w+astCondK b (AstPlusK u1 u2) (AstPlusK v1 v2) | eqY u1 v1 =+  astPlusK u1 (astCondK b u2 v2)+astCondK b (AstPlusK u1 u2) v1 | eqY u1 v1 =+  astPlusK u1 (astCondK b u2 0)+astCondK b u1 (AstPlusK v1 v2) | eqY u1 v1 =+  astPlusK u1 (astCondK b 0 v2)+astCondK b (AstPlusK u1 u2) (AstPlusK v1 v2) | eqY u2 v2 =+  astPlusK (astCondK b u1 v1) u2+astCondK b (AstPlusK u1 u2) v2 | eqY u2 v2 =+  astPlusK (astCondK b u1 0) u2+astCondK b u2 (AstPlusK v1 v2) | eqY u2 v2 =+  astPlusK (astCondK b 0 v1) u2+astCondK b (AstTimesK u1 u2) (AstTimesK v1 v2) | eqY u1 v1 =+  astTimesK u1 (astCondK b u2 v2)+astCondK b (AstTimesK u1 u2) v1 | eqY u1 v1 =+  astTimesK u1 (astCondK b u2 0)+astCondK b u1 (AstTimesK v1 v2) | eqY u1 v1 =+  astTimesK u1 (astCondK b 0 v2)+astCondK b (AstTimesK u1 u2) (AstTimesK v1 v2) | eqY u2 v2 =+  astTimesK (astCondK b u1 v1) u2+astCondK b u2 (AstTimesK v1 v2) | eqY u2 v2 =+  astTimesK (astCondK b 0 v1) u2+astCondK b (AstTimesK u1 u2) v2 | eqY u2 v2 =+  astTimesK (astCondK b u1 0) u2+-- These two normal forms prevent simple cases of an exponential blow-up+-- when substituting into nested conditionals in the rule that follows them.+-- For now they don't seem to be needed, though.+-- astCondK b v@Ast.AstCond{} w = Ast.AstCond b v w+-- astCondK b v w@Ast.AstCond{} = Ast.AstCond b v w+astCondK (AstLeqInt (AstConcreteK j) var0@(AstIntVar var)) v w+  | Just Refl <- testEquality (typeRep @r) (typeRep @Int)+  , Just (lb, ub) <- varNameToBounds var+  , let varTrue = reboundsVarName (max lb j, ub) var+        varFalse = reboundsVarName (lb, min ub (j - 1)) var+  , AstConcreteK dv <- v - fromPlain (astVar varTrue)+  , AstConcreteK dw <- w - fromPlain (astVar varFalse)+  , dv == dw =+    fromPlain $ AstConcreteK dv + var0+astCondK b@(AstLeqInt AstConcreteK{}+                     (Ast.AstN1K NegateOp AstIntVar{})) v w =+  astCondKInitial (astBoolNotK b) w v  -- eliminate the NegateOp+{- TODO: Investigate how much stronger this rule is than the above (it's slower)+astCondK (AstLeqInt (AstConcreteK k) var0@(AstIntVar var)) v w+  | FTKScalar @r <- ftkAst v+  , Just Refl <- testEquality (typeRep @r) (typeRep @Int)+  , Just (lb, ub) <- varNameToBounds var+  , let varTrue = reboundsVarName (max lb k, ub) var+        varFalse = reboundsVarName (lb, min ub (k - 1)) var+        d = v - fromPlain (astVar varTrue)+  , eqY w (substituteAst (astVar varFalse) varTrue d+           + fromPlain (astVar varFalse)) =+    substituteAst var0 varTrue d + fromPlain var0 -}+{- Disabled, because scatters and gathers simplify better with conditionals+   than with quotH, so the rules need to be restricted at least.+-- TODO: prove nothing wrong comes out of dividing by zero here:+astCondK (AstLeqInt (AstConcreteK k) var0@(AstIntVar var)) v w+  | FTKScalar @r <- ftkAst v+  , Just Refl <- testEquality (typeRep @r) (typeRep @Int)+  , Just (lb, ub) <- varNameToBounds var+  , let varTrue = reboundsVarName (max lb k, ub) var+        varFalse = reboundsVarName (lb, min ub (k - 1)) var+        d = fromPlain (AstConcreteK k)  -- a simplified guess+            `quotH` substituteAst (AstConcreteK k) var v+  , eqY (fromPlain (astVar varTrue) `quotH` d)+        (substituteAst (astVar varTrue) var v)+  , eqY (fromPlain (astVar varFalse) `quotH` d)+        (substituteAst (astVar varFalse) var w) =+    fromPlain var0 `quotH` d+astCondK (AstLeqInt (AstConcreteK k) var0@(AstIntVar var)) v w+  | FTKScalar @r <- ftkAst v+  , Just Refl <- testEquality (typeRep @r) (typeRep @Int)+  , Just (lb, ub) <- varNameToBounds var+  , let varTrue = reboundsVarName (max lb k, ub) var+        varFalse = reboundsVarName (lb, min ub (k - 1)) var+        d = fromPlain (AstConcreteK (k - 1))  -- the only change vs above+            `quotH` substituteAst (AstConcreteK (k - 1)) var w+  , eqY (fromPlain (astVar varTrue) `quotH` d)+        (substituteAst (astVar varTrue) var v)+  , eqY (fromPlain (astVar varFalse) `quotH` d)+        (substituteAst (astVar varFalse) var w) =+    fromPlain var0 `quotH` d -}+astCondK b v w = Ast.AstCondK b v w++astCondSInitial :: forall sh x s. KnownSpan s+                => AstBool AstMethodLet+                -> AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKS2 sh x)+astCondSInitial b v w = case b of+  AstLeqInt (AstConcreteK j) (AstIntVar var)+    | Just (lb, ub) <- varNameToBounds var ->+      let varTrue = reboundsVarName (max lb j, ub) var+          varFalse = reboundsVarName (lb, min ub (j - 1)) var+      in astCondS b (substituteAst (astVar varTrue) var v)+                    (substituteAst (astVar varFalse) var w)+  _ -> astCondS b v w++astCondS :: forall sh x s. KnownSpan s+         => AstBool AstMethodLet+         -> AstTensor AstMethodLet s (TKS2 sh x)+         -> AstTensor AstMethodLet s (TKS2 sh x)+         -> AstTensor AstMethodLet s (TKS2 sh x)+astCondS b v w | Just b0 <- unAstK b = if b0 then v else w+astCondS _b u v | eqY u v = u+astCondS b v w | FTKS (snat :$$ _) _ <- ftkAst v+               , Just v1 <- unRepl1 v+               , Just w1 <- unRepl1 w =+  astReplicateS (snat :$$ ZSS) (astCondS b v1 w1)+astCondS (Ast.AstBoolNotK b) v w = astCondS b w v+astCondS (Ast.AstLet var n b) u v =+  astLetRefresh var n b $ \b' -> astCondS b' u v+astCondS b (Ast.AstLet var n u) v =+  astLetRefresh var n u $ \u' -> astCondS b u' v+astCondS b u (Ast.AstLet var n v) =+  astLetRefresh var n v $ \v' -> astCondS b u v'+astCondS b (Ast.AstFromPrimal v) (Ast.AstFromPrimal w) =+  fromPrimal $ astCondS b v w+astCondS b (Ast.AstFromDual v) (Ast.AstFromDual w) =+  fromDual $ astCondS b v w+astCondS b (Ast.AstFromPlain v) (Ast.AstFromPlain w) =+  fromPlain $ astCondS b v w+astCondS _b (Ast.AstIotaS n) Ast.AstIotaS{} = Ast.AstIotaS n+astCondS b v w | FTKS ZSS FTKScalar <- ftkAst v =+  sfromK (astCondK b (kfromS v) (kfromS w))+astCondS b v w = Ast.AstCondS b v w++astFromVectorK :: forall shm r s. (GoodScalar r, KnownSpan s)+               => ShS shm+               -> Data.Vector.Vector (AstTensor AstMethodLet s (TKScalar r))+               -> AstTensor AstMethodLet s (TKS shm r)+astFromVectorK ZSS l = assert (V.length l == 1) $ sfromK $ l V.! 0+astFromVectorK _ l | V.null l = error "astFromVectorK: empty vector"+astFromVectorK shm l | SNat' @1 <- shsProduct shm =+  astReplicateK shm (l V.! 0)+astFromVectorK shm l = fromMaybe (Ast.AstFromVectorK shm l) $+  -- This disables some rules, e.g., indexing or summing of fromVector+  -- of concrete arrays, but allocating an extra array of the same size+  -- as the fromVector is not a big deal and early rules are better+  -- then the same rules in contraction phase.+  (let unConc :: AstTensor AstMethodLet s (TKScalar r)+              -> Maybe (Concrete (TKScalar r))+       unConc t = Concrete <$> unAstK t+    in case V.mapM unConc l of+         Just l4 -> Just $ fromPlain $ astConcreteS (tsfromVectorLinear shm l4)+         Nothing -> Nothing)+  `mplus`+  (let unFromPrimal :: AstTensor AstMethodLet s2 y+                    -> Maybe (AstTensor AstMethodLet (PrimalStepSpan s2) y)+       unFromPrimal (Ast.AstFromPrimal t) = Just t+       unFromPrimal _ = Nothing+   in case V.mapM unFromPrimal l of+     Just l2 -> Just $ fromPrimal $ astFromVectorK shm l2+     Nothing -> Nothing)+  `mplus`+  (case knownSpan @s of+     SFullSpan ->+       let unFromDual :: AstTensor AstMethodLet FullSpan y+                      -> Maybe (AstTensor AstMethodLet DualSpan y)+           unFromDual (Ast.AstFromDual t) = Just t+           unFromDual _ = Nothing+       in case V.mapM unFromDual l of+         Just l2 -> Just $ fromDual $ astFromVectorK shm l2+         Nothing -> Nothing+     _ -> Nothing)+  `mplus`+  (let unFromPlain :: AstTensor AstMethodLet s2 y+                   -> Maybe (AstTensor AstMethodLet PlainSpan y)+       unFromPlain (Ast.AstFromPlain t) = Just t+       unFromPlain _ = Nothing+   in case V.mapM unFromPlain l of+     Just l2 -> Just $ fromPlain $ astFromVectorK shm l2+     Nothing -> Nothing)++astFromVectorS :: forall shm shn x s. KnownSpan s+               => ShS shm+               -> Data.Vector.Vector (AstTensor AstMethodLet s (TKS2 shn x))+               -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+astFromVectorS ZSS l = assert (V.length l == 1) $ l V.! 0+astFromVectorS _ l | V.null l = error "astFromVectorS: empty vector"+astFromVectorS shm l | SNat' @1 <- shsProduct shm =+  astReplicateS shm (l V.! 0)+astFromVectorS shm l = fromMaybe (Ast.AstFromVectorS shm l) $+  (let unFromPrimal :: AstTensor AstMethodLet s2 y+                    -> Maybe (AstTensor AstMethodLet (PrimalStepSpan s2) y)+       unFromPrimal (Ast.AstFromPrimal t) = Just t+       unFromPrimal _ = Nothing+   in case V.mapM unFromPrimal l of+     Just l2 -> Just $ fromPrimal $ astFromVectorS shm l2+     Nothing -> Nothing)+  `mplus`+  (case knownSpan @s of+     SFullSpan ->+       let unFromDual :: AstTensor AstMethodLet FullSpan y+                      -> Maybe (AstTensor AstMethodLet DualSpan y)+           unFromDual (Ast.AstFromDual t) = Just t+           unFromDual _ = Nothing+       in case V.mapM unFromDual l of+         Just l2 -> Just $ fromDual $ astFromVectorS shm l2+         Nothing -> Nothing+     _ -> Nothing)+  `mplus`+  (let unFromPlain :: AstTensor AstMethodLet s2 y+                   -> Maybe (AstTensor AstMethodLet PlainSpan y)+       unFromPlain (Ast.AstFromPlain t) = Just t+       unFromPlain _ = Nothing+   in case V.mapM unFromPlain l of+     Just l2 -> Just $ fromPlain $ astFromVectorS shm l2+     Nothing -> Nothing)+  `mplus`+  (case ftkAst (l V.! 0) of+     FTKS ZSS FTKScalar | Refl <- lemAppNil @shm ->+       Just $ astFromVectorK shm $ V.map kfromS l+     FTKS shn FTKScalar ->+       -- This disables some rules, e.g., indexing or summing of fromVector+       -- of concrete arrays, but allocating an extra array of the same size+       -- as the fromVector is not a big deal and early rules are better+       -- then the same rules in contraction phase.+       let unConc :: AstTensor AstMethodLet s (TKS shn r)+                  -> Maybe (Concrete (TKS shn r))+           unConc t = Concrete <$> unAstS t+       in case V.mapM unConc l of+         Just l4 ->+           withKnownShS shn $+           Just $ fromPlain $ astConcreteS (tsfromVectorN shm l4)+         Nothing -> Nothing+     FTKS{} -> Nothing)++astSumK :: forall shm r s. (NumScalar r, KnownSpan s)+        => AstTensor AstMethodLet s (TKS shm r)+        -> AstTensor AstMethodLet s (TKScalar r)+astSumK t0 | FTKS shm FTKScalar <- ftkAst t0 = case t0 of+  _ | SZ <- shsProduct shm -> 0+  _ | SNat' @1 <- shsProduct shm ->+      astIndexK t0 (AstConcreteK <$> ixsZero shm)+  Ast.AstLet var u v -> astLet var u (astSumK v)+  Ast.AstFromPrimal u -> fromPrimal $ astSumK u+  Ast.AstFromDual u -> fromDual $ astSumK u+  Ast.AstFromPlain u -> fromPlain $ astSumK u+  Ast.AstN1S NegateOp u -> negate $ astSumK u+  AstTimesS t1 t2+    | Just u <- unRepl t1 ->+      astTimesK (kfromS u) (astSumK t2)+  AstConcreteS v ->+    withKnownShS (Nested.sshape v) $+    astConcreteK $ tssum0 (Concrete v)+  Ast.AstFromVectorK _ l | V.length l < 100 ->+    foldr1 astPlusK l+  Ast.AstFromVectorS _ l | V.length l < 100 ->+    astSumK $ foldr1 astPlusS l+  Ast.AstSumS _ u -> astSumK u+  Ast.AstScatterS shm1@(k2 :$$ _) shn1 shp1 v (vars, i1 :.$ rest)+    | -- This boolean term may have free variables that act as universally+      -- quantified.+      AstConcreteK True <- 0 <=. i1+                           &&* i1 <=. AstConcreteK (fromSNat' k2 - 1) ->+      astSumK $ astScatterS shm1 shn1 (shsTail shp1) v (vars, rest)+  Ast.AstReplicateK _ u ->+    astTimesK u+              (fromPlain $ AstConcreteK $ fromIntegral $ shsSize shm)+  Ast.AstReplicateS shm2 u ->+    astTimesK (astSumK u)+              (fromPlain $ AstConcreteK $ fromIntegral $ shsSize shm2)+  _ | ZSS <- shm -> kfromS t0+  _ | snat :$$ _ <- shm+    , Just u <- unRepl1 t0 ->+      astTimesK (astSumK u)+                (fromPlain $ AstConcreteK $ fromIntegral $ fromSNat' snat)+  -- TODO: if ever needed, we may also add case+  -- Ast.AstCond AstReplicate AstReplicate+  -- or use unRepl and unRepl1 to handle it.+  Ast.AstIotaS n ->+    AstConcreteK $ fromIntegral $ fromSNat' n * (fromSNat' n - 1) `div` 2+  Ast.AstAppendS u v -> astSumK u `astPlusK` astSumK v+  Ast.AstReverseS u -> astSumK u+  Ast.AstTransposeS _ u -> astSumK u+  Ast.AstReshapeS _ u -> astSumK u+  _ -> Ast.AstSumK t0++astSumS :: forall shm shn x s. (KnownSpan s, TKAllNum x)+        => ShS shm -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+        -> AstTensor AstMethodLet s (TKS2 shn x)+astSumS ZSS t0 = t0+astSumS shm@(snat :$$ _) t0 | FTKS shmshn x <- ftkAst t0+                            , SNat <- shsRank shm+                            , let shn = shsDrop @(Rank shm) shmshn =+                                gcastWith (unsafeCoerceRefl+                                           :: Drop (Rank shm) (shm ++ shn)+                                              :~: shn) $+                                case t0 of+  _ | SZ <- shsProduct shmshn ->+      let ftk = FTKS shn x+      in fromPlain $ astConcrete ftk (tdefTarget ftk)+  _ | SNat' @1 <- shsProduct shm ->+      astIndexS shn t0 (AstConcreteK <$> ixsZero shm)+        -- astReshape slows down the CNNO test+  _ | ZSS <- shn+    , FTKScalar @r <- x+    , Dict0 <- numFromTKAllNum (Proxy @r) ->+      sfromK $ astSumK t0+  -- This keeps tensors alive for longer, but it enables new simplifications,+  -- while hiding a sum inside let not often prevents other simplifications,+  -- because there are few redexes with sum but not at the top.+  Ast.AstLet var u v -> astLet var u (astSumS shm v)+  Ast.AstFromPrimal v -> fromPrimal $ astSumS shm v+  Ast.AstFromDual v -> fromDual $ astSumS shm v+  Ast.AstFromPlain v -> fromPlain $ astSumS shm v+  Ast.AstN1S NegateOp u -> negate $ astSumS shm u+  -- This exchanges a multiplication at rank n+m for one+  -- at rank n, which should be faster. We choose t1, since it's likely to be+  -- concrete and so it's easier to see if it's replicated and also t2 is+  -- likely to be AstTimes, to be processed recusively. This rule may distrupt+  -- recognizing matmul, sdot1In, etc., but it sometimes leads to a smaller+  -- and better sdot1In.+  AstTimesS t1 t2+    | Just u <- unReplN shm t1+    , FTKScalar @r <- x+    , Dict0 <- numFromTKAllNum (Proxy @r) ->+      astTimesS u (astSumS shm t2)+  AstConcreteS t ->+    withKnownShS shm $+    withKnownShS shn $+    astConcreteS $ tssumN @_ @shm $ Concrete t+  Ast.AstFromVectorS (_ :$$ ZSS) l+    | V.length l < 100+    , FTKScalar @r <- x+    , Dict0 <- numFromTKAllNum (Proxy @r) ->+      astSumS (shsTail shm) $ foldr1 astPlusS l+  Ast.AstSumS @shm2 shm2 u+    | Refl <- lemAppAssoc (Proxy @shm2) (Proxy @shm) (Proxy @shn)  ->+      astSumS (shm2 `shsAppend` shm) u+  Ast.AstScatterS shm1@(k2 :$$ _) shn1 shp1 v (vars, i1 :.$ rest)+    | -- This boolean term may have free variables that act as universally+      -- quantified.+      AstConcreteK True <- 0 <=. i1+                           &&* i1 <=. AstConcreteK (fromSNat' k2 - 1) ->+      astSumS (shsTail shm)+      $ astScatterS shm1 shn1 (shsTail shp1) v (vars, rest)+  Ast.AstReplicateK _ u | FTKScalar @r <- x+                        , Dict0 <- numFromTKAllNum (Proxy @r) ->+    astReplicateK shn+    $ astTimesK u (fromPlain $ AstConcreteK $ fromIntegral $ shsSize shm)+  _ | Just u <- unRepl1 t0+    , FTKScalar @r <- x+    , Dict0 <- numFromTKAllNum (Proxy @r) ->+      astTimesS (astSumS (shsTail shm) u)+                (astReplicateK shn $ fromPlain $ AstConcreteK $ fromIntegral+                 $ fromSNat' snat)+  Ast.AstAppendS u v+    | FTKScalar @r <- x+    , Dict0 <- numFromTKAllNum (Proxy @r)+    , FTKS (m :$$ _) _ <- ftkAst u+    , FTKS (n :$$ _) _ <- ftkAst v ->+      astSumS (m :$$ shsTail shm) u `astPlusS` astSumS (n :$$ shsTail shm) v+  Ast.AstReverseS v -> astSumS shm v+  -- See the analogous astSliceS rule.+  Ast.AstTransposeS perm t+    | SNat <- Permutation.permRank perm+    , LTI <- cmpNat (Permutation.permRank perm) (shsRank shm)+    , FTKS @shkshn shkshn _ <- ftkAst t ->+      gcastWith (unsafeCoerceRefl :: Drop (Rank shm) shkshn :~: shn) $+      gcastWith (unsafeCoerceRefl ::+                 Take (Rank shm) shkshn ++ Drop (Rank shm) shkshn :~: shkshn) $+      astSumS (shsTake @(Rank shm) shkshn) t+  Ast.AstTransposeS perm t+    | SNat <- Permutation.permRank perm+    , EQI <- cmpNat (Permutation.permRank perm) (shsRank shm)+    , FTKS @shkshn shkshn _ <- ftkAst t ->+      gcastWith (unsafeCoerceRefl :: Drop (Rank shm) shkshn :~: shn) $+      gcastWith (unsafeCoerceRefl ::+                 Take (Rank shm) shkshn ++ Drop (Rank shm) shkshn :~: shkshn) $+      astSumS (shsTake @(Rank shm) shkshn) t+  _ -> Ast.AstSumS shm t0++-- TODO: fuse scatters, scatter and sum, and perhaps more (fromList?)+astScatterS :: forall shm shn shp x s. (KnownSpan s, TKAllNum x)+            => ShS shm -> ShS shn -> ShS shp+            -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+            -> (AstVarListS shm, AstIxS AstMethodLet shp)+            -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) x)+astScatterS = astScatterKnobsS defaultKnobs++astScatterKnobsS :: forall shm shn shp x s. (KnownSpan s, TKAllNum x)+                 => SimplifyKnobs+                 -> ShS shm -> ShS shn -> ShS shp+                 -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+                 -> (AstVarListS shm, AstIxS AstMethodLet shp)+                 -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) x)+astScatterKnobsS _ _ _ _ v0 (!vars0, !_ix0)+  | Foldable.any (`varNameInAst` v0) vars0 =+    error $ "astScatterKnobsS: scatter vars in v0: " ++ show (vars0, v0)+astScatterKnobsS _ shm _ _ v0 (_, ZIS) = astSumS shm v0+astScatterKnobsS _ _ shn shp@(SZ :$$ _) v0 _+  | FTKS _ x <- ftkAst v0 =+    let ftk = FTKS (shp `shsAppend` shn) x+    in fromPlain $ astConcrete ftk (tdefTarget ftk)+astScatterKnobsS _ _ shn shp@(k :$$ _) v0 (_,  i1 :.$ _)+  | Just (lb, ub) <- intBounds i1+  , let FTKS _ x = ftkAst v0+  , ub < 0 || lb >= fromSNat' k =+    let ftk = FTKS (shp `shsAppend` shn) x+    in fromPlain $ astConcrete ftk (tdefTarget ftk)+astScatterKnobsS knobs shm shn shp v0 (vars0@(_ ::$ _), ix0@(_ :.$ _))+  | let ixInit = ixsInit ix0+        varInit = AstVarListS $ ixsInit $ unAstVarListS vars0+        varLast = ixsLast $ unAstVarListS vars0+  , AstIntVar ixvarLast <- ixsLast ix0+  , ixvarLast == varLast+  , not (varLast `varNameInIxS` ixInit)+  , kLast@SNat <- shsLast shm+  , Just Refl <- testEquality kLast (shsLast shp) =+    gcastWith (unsafeCoerceRefl+               :: Init shp ++ (Last shm ': shn) :~: shp ++ shn) $+    gcastWith (unsafeCoerceRefl+               :: Init shm ++ (Last shm ': shn) :~: shm ++ shn) $+    astScatterKnobsS knobs (shsInit shm) (kLast :$$ shn) (shsInit shp)+                     v0 (varInit, ixInit)+astScatterKnobsS knobs shm@(m :$$ _) shn (p@(SNat @p) :$$ ZSS) v0+  ( varm ::$ mrest+  , Ast.AstI2K QuotOp (AstIntVar varm') (AstConcreteK i0) :.$ ZIS )+  | varm == varm'+  , i0 > 0  -- ensured by other rules; makes this easier to reason about+  , FTKS _ x <- ftkAst v0+  , let k0 = (fromSNat' m + i0 - 1) `quot` i0+  , fromSNat' p >= k0 =  -- most likely, because not OOB+    withSNat i0 $ \i ->+    withSNat k0 $ \k@(SNat @k) ->+    withSNat (fromSNat' p - k0) $ \z2@(SNat @z2) ->+      let ftk = FTKS (z2 :$$ shn) x+      in gcastWith (unsafeCoerceRefl :: k + z2 :~: p) $+         (`astAppendS` fromPlain (astConcrete ftk (tdefTarget ftk)))+         $ astSumS (i :$$ ZSS)+         $ astReshapeS (i :$$ k :$$ shn)+         -- This looks like we are creating a larger outcome tensor than before,+         -- but the following reduces to a nested sum of v0 with some transposes+         -- and appends, so we likely create something much smaller than v0+         -- despite the outermost dimension not shrinking as it did before.+         $ astScatterKnobsS knobs shm shn (snatMul k i :$$ ZSS)+                            v0 (varm ::$ mrest, astVar varm :.$ ZIS)+             -- this gather may still index out of bounds, which is fine+astScatterKnobsS knobs shm shn shp@(SNat' @1 :$$ _)+                 v (vars, AstConcreteK _ :.$ rest) =  -- if OOB, covered above+  astReplicateS (SNat @1 :$$ ZSS)+  $ astScatterKnobsS knobs shm shn (shsTail shp) v (vars, rest)+astScatterKnobsS knobs shm shn shp (Ast.AstLet var u v) (vars, ix) =+  astLetRefresh var u v+  $ \v' -> astScatterKnobsS knobs shm shn shp v' (vars, ix)+astScatterKnobsS knobs shm shn shp (Ast.AstFromPrimal v) (vars, ix) =+  fromPrimal $ astScatterKnobsS knobs shm shn shp v (vars, ix)+astScatterKnobsS knobs shm shn shp (Ast.AstFromDual v) (vars, ix) =+  fromDual $ astScatterKnobsS knobs shm shn shp v (vars, ix)+astScatterKnobsS knobs shm shn shp (Ast.AstFromPlain v) (vars, ix) =+  fromPlain $ astScatterKnobsS knobs shm shn shp v (vars, ix)+astScatterKnobsS knobs ZSS _shn shp+                 (Ast.AstScatterS @_ @shn2 @shp2 shm2 shn2 shp2 v (vars, ix2))+                 (ZS, ix) =  -- oneHot (scatter) fusion+  gcastWith (unsafeCoerceRefl :: shp ++ shn :~: shp ++ shp2 ++ shn2) $+  astScatterKnobsS knobs shm2 shn2 (shp `shsAppend` shp2)+                   v (vars, ix `ixsAppend` ix2)+astScatterKnobsS knobs (snat@(SNat' @1) :$$ ZSS) _shn shp+                 (Ast.AstReplicateS (_ :$$ ZSS)+                   (Ast.AstScatterS @_ @shn2 @shp2 shm2 shn2 shp2+                                    v (vars, ix2)))+                 (var ::$ ZS, ix) =  -- oneHot1 (scatter) fusion+  gcastWith (unsafeCoerceRefl :: shp ++ shn :~: shp ++ shp2 ++ shn2) $+  astScatterKnobsS knobs (snat :$$ shm2) shn2 (shp `shsAppend` shp2)+                   (Ast.AstReplicateS (snat :$$ ZSS) v)+                   (var ::$ vars, ix `ixsAppend` ix2)+astScatterKnobsS knobs (snat@(SNat' @1) :$$ ZSS) _shn shp+                 (Ast.AstReplicateS (_ :$$ ZSS)  -- TODO: generalize+                   (Ast.AstScatterS @_ @shn2 @shp2 shm2 shn2 shp2+                                    v (vars, ix2)))+                 (var ::$ ZS, ix) =  -- oneHot1 (scatter) fusion+  gcastWith (unsafeCoerceRefl :: shp ++ shn :~: shp ++ shp2 ++ shn2) $+  astScatterKnobsS knobs (snat :$$ shm2) shn2 (shp `shsAppend` shp2)+                   (Ast.AstReplicateS (snat :$$ ZSS) v)+                   (var ::$ vars, ix `ixsAppend` ix2)+astScatterKnobsS knobs shm shn shp (Ast.AstTransposeS @perm @sh perm v)+                 (vars, ix)  -- see the same case of astGatherKnobsS+  | FTKS sh _ <- ftkAst v+  , knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion]+  , let rankPerm = Permutation.permRank perm+  , fromSNat' rankPerm <= length vars =+          gcastWith (lemRankMapJust $ shsTakeLenPerm perm sh) $+          gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $+          permInverse perm+          $ \(invperm :: Nested.Perm invperm) proof ->+            case proof (ssxFromShX $ shxFromShS $ shsTakeLenPerm perm sh) of+              Refl ->+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.PermutePrefix invperm shm ++ shn+                      :~: Permutation.PermutePrefix invperm (shm ++ shn)) $+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))+                      :~: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))) $+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))+                      :~: TakeLen perm sh) $+                gcastWith+                  (unsafeCoerceRefl+                   :: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)+                      :~: DropLen perm sh) $+                gcastWith+                  (unsafeCoerceRefl+                   :: TakeLen perm sh ++ DropLen perm sh :~: sh) $+                let invvars = AstVarListS+                              $ ixsPermutePrefix invperm $ unAstVarListS vars+                in astScatterKnobsS knobs (shsPermutePrefix invperm shm) shn shp+                                    v (invvars, ix)+astScatterKnobsS knobs+                 shm@(SNat @m :$$ (_ :: ShS shmTail))+                 shn+                 shp@(SNat @p :$$ (_ :: ShS shpTail))+                 v0+  ( varm ::$ mrest+  , AstIntVar varp :.$ prest )+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+      -- prevent a loop+  , varm == varp+  , not (varm `varNameInIxS` prest)+  , FTKS _ x <- ftkAst v0 =+    withSNat (min (valueOf @p) (valueOf @m)) $ \(SNat @m2) ->+    gcastWith (unsafeCoerceRefl :: (m2 <=? p) :~: True) $+    gcastWith (unsafeCoerceRefl :: (m2 <=? m) :~: True) $+    Permutation.permFromListCont (permCycle+                                  $ shsLength (shsTail shp) + 1)+    $ \(permVars :: Permutation.Perm permVars) ->+    Permutation.permFromListCont (backpermCycle+                                  $ shsLength (shsTail shm) + 1)+    $ \(permIx :: Permutation.Perm permIx) ->+       gcastWith (unsafeCoerceRefl+                  :: m2 ': shpTail ++ shn+                     :~: Permutation.PermutePrefix+                           permVars (shpTail ++ (m2 ': shn))) $+       gcastWith (unsafeCoerceRefl+                  :: shmTail ++ (m2 ': shn)+                     :~: Permutation.PermutePrefix+                           permIx (m2 ': shmTail ++ shn)) $+       gcastWith (unsafeCoerceRefl+                  :: (Rank permVars <=? Rank (shpTail ++ (m2 ': shn)))+                     :~: True) $+       gcastWith (unsafeCoerceRefl+                  :: (Rank permIx <=? Rank (m2 ': shmTail ++ shn)) :~: True) $+       fromMaybe (error "astScatterKnobsS: impossible non-permutation")+       $ Permutation.permCheckPermutation permVars+       $ fromMaybe (error "astScatterKnobsS: impossible non-permutation")+       $ Permutation.permCheckPermutation permIx+       $ let v2 = astTransposeS permIx+                  $ astSliceS (SNat @0) (SNat @m2) (SNat @(m - m2)) v0+             u = astScatterKnobsS+                   knobs (shsTail shm) (SNat @m2 :$$ shn) (shsTail shp)+                   v2 (mrest, prest)+             ftk = FTKS (SNat @(p - m2) :$$ shsTail shp `shsAppend` shn) x+         in astTransposeS permVars u+            `astAppendS`+            fromPlain (astConcrete ftk (tdefTarget ftk))+astScatterKnobsS knobs shm@(SNat' @1 :$$ ZSS) shn shp+                 v@Ast.AstReplicateS{} (vars, ix)+  | knobPhase knobs /= PhaseContraction =+    Ast.AstScatterS shm shn shp v (vars, ix)  -- oneHot1 NF+-- The above normal form that prevents the use of the rule below+-- and the subsequent rule is to keep one-hots and similar in forms+-- easier for fusion of addition of scatters to operate on.+astScatterKnobsS knobs (snat :$$ shm2) shn shp v (var ::$ vars, ix)+  | not $ var `varNameInIxS` ix =  -- eliminates oneHot1+    astScatterKnobsS knobs shm2 shn shp+                     (astSumS (snat :$$ ZSS) v) (vars, ix)+astScatterKnobsS knobs ZSS shn shp u (ZS, ix)  -- oneHot1 intro+  | knobPhase knobs /= PhaseContraction =+    let shm = SNat @1 :$$ ZSS+        a = Ast.AstReplicateS (SNat @1 :$$ ZSS) u+    in funToVarsIxS shm $ \(var ::$ ZS) _ ->+         Ast.AstScatterS shm shn shp a (var ::$ ZS, ix)+astScatterKnobsS knobs shm shn shp@(p@(SNat @p) :$$ _) v0+  (vars, AstPlusK (AstConcreteK i0) i1 :.$ prest)+  | knobPhase knobs `elem` [PhaseExpansion, PhaseContraction]+      -- give time for additions of scatters to fuse; TODO: harden fusion+      -- or keep as normal forms more terms than one-hots, until contraction+  , Just (lb, _) <- intBounds i1+  , lb >= 0  -- ensured by the OOB rule above: fromSNat' p >= i0+    || i0 <= 0+  , FTKS _ x <- ftkAst v0 =+    if i0 >= 0 then+      withSNat i0 $ \i@(SNat @i) ->+      withSNat (fromSNat' p - i0) $ \k@(SNat @k) ->+        gcastWith (unsafeCoerceRefl :: i + k :~: p) $+        let ftk = FTKS (i :$$ shsTail shp `shsAppend` shn) x+        in astAppendS+             (fromPlain (astConcrete ftk (tdefTarget ftk)))+             (astScatterKnobsS knobs shm shn (k :$$ shsTail shp)+                               v0 (vars, i1 :.$ prest))+                -- this gather may still index out of bounds, which is fine+    else+      withSNat (- i0) $ \i ->+        astSliceS i p SZ+        $ astScatterKnobsS knobs shm shn (snatPlus p i :$$ shsTail shp)+                           v0 (vars, i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+astScatterKnobsS knobs shm shn shp@(p@(SNat @p) :$$ _) v0+  (vars, Ast.AstLet varN uN (AstPlusK (AstConcreteK i0) i1) :.$ prest)+  | knobPhase knobs `elem` [PhaseExpansion, PhaseContraction]+      -- give time for additions of scatters to fuse; TODO+  , Just (lb, _) <- intBounds i1+  , lb >= 0  -- ensured by the OOB rule above: fromSNat' p >= i0+    || i0 <= 0+  , FTKS _ x <- ftkAst v0 =+    if i0 >= 0 then+      withSNat i0 $ \i@(SNat @i) ->+      withSNat (fromSNat' p - i0) $ \k@(SNat @k) ->+        gcastWith (unsafeCoerceRefl :: i + k :~: p) $+        let ftk = FTKS (i :$$ shsTail shp `shsAppend` shn) x+        in astAppendS+             (fromPlain (astConcrete ftk (tdefTarget ftk)))+             (astScatterKnobsS knobs shm shn (k :$$ shsTail shp)+                               v0 (vars, astLet varN uN i1 :.$ prest))+                -- this gather may still index out of bounds, which is fine+    else+      withSNat (- i0) $ \i ->+        astSliceS i p SZ+        $ astScatterKnobsS knobs shm shn (snatPlus p i :$$ shsTail shp)+                           v0 (vars, astLet varN uN i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+-- These rules are questionable, because at worst, we allocate twice.+-- However, often the scatters or one of them would simplify away.+-- Note these don't revert the scatter addtion fusion in rules for astPlusS,+-- because the latter don't fire if the conditional can't be simplified away.+astScatterKnobsS knobs shm@(SNat @m :$$ _) shn shp u0+  ( vars@(varm ::$ mrest)+  , Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar varp)) v w+    :.$ prest )+  | varm == varp+  , j0 <= 0 || j0 >= valueOf @m || ixIsSmall prest+  , FTKS _ (FTKScalar @r) <- ftkAst u0+  , Dict0 <- numFromTKAllNum (Proxy @r) =+    if | j0 <= 0 ->+         astScatterKnobsS knobs shm shn shp+                          u0 (vars, v :.$ prest)+       | j0 >= valueOf @m ->+         astScatterKnobsS knobs shm shn shp+                          u0 (vars, w :.$ prest)+       | otherwise ->+         withSNat j0 $ \j@(SNat @j) ->+         withSNat (valueOf @m - j0) $ \msj@(SNat @msj) ->+         gcastWith (unsafeCoerceRefl :: msj + j :~: m) $+         astLetFun u0 $ \u ->+         let u2 = astSliceS SZ j msj u+             u3 = astSliceS j msj SZ u+             varm2 = reboundsVarName (0, j0 - 1) varm+             varm3 = reboundsVarName (0, valueOf @msj - 1) varm+         in astScatterKnobsS knobs (j :$$ shsTail shm) shn shp u2+              ( varm2 ::$ mrest+              , substituteAstIxS (astVar varm2)+                                 varm (w :.$ prest) )+            `astPlusS`+            astScatterKnobsS knobs (msj :$$ shsTail shm) shn shp u3+              ( varm3 ::$ mrest+              , substituteAstIxS (AstConcreteK j0 + astVar varm3)+                                 varm (v :.$ prest) )+astScatterKnobsS knobs shm@(SNat @m :$$ _) shn shp u0+  ( vars@(varm ::$ mrest)+  , Ast.AstLet varN uN+      (Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar varp)) v w)+    :.$ prest )+  | varm == varp+  , let varIn1 = varN `varNameInAst` v+        varIn2 = varN `varNameInAst` w+  , j0 <= 0 || j0 >= valueOf @m+    || ixIsSmall prest && not (varIn1 && varIn2)+  , FTKS _ (FTKScalar @r) <- ftkAst u0+  , Dict0 <- numFromTKAllNum (Proxy @r) =+    if | j0 <= 0 ->+         astScatterKnobsS knobs shm shn shp+                          u0 (vars, astLet varN uN v :.$ prest)+       | j0 >= valueOf @m ->+         astScatterKnobsS knobs shm shn shp+                          u0 (vars, astLet varN uN w :.$ prest)+       | otherwise ->+         withSNat j0 $ \j@(SNat @j) ->+         withSNat (valueOf @m - j0) $ \msj@(SNat @msj) ->+         gcastWith (unsafeCoerceRefl :: msj + j :~: m) $+         astLetFun u0 $ \u ->+         let u2 = astSliceS SZ j msj u+             u3 = astSliceS j msj SZ u+             varm2 = reboundsVarName (0, j0 - 1) varm+             varm3 = reboundsVarName (0, valueOf @msj - 1) varm+         in astScatterKnobsS knobs (j :$$ shsTail shm) shn shp u2+              ( varm2 ::$ mrest+              , substituteAstIxS (astVar varm2)+                                 varm ((if varIn2+                                        then astLet varN uN w+                                        else w) :.$ prest) )+            `astPlusS`+            astScatterKnobsS knobs (msj :$$ shsTail shm) shn shp u3+              ( varm3 ::$ mrest+              , substituteAstIxS (AstConcreteK j0 + astVar varm3)+                                 varm ((if varIn1+                                        then astLet varN uN v+                                        else v) :.$ prest) )+astScatterKnobsS _ shm shn shp v (vars, ix) =+  Ast.AstScatterS shm shn shp v (vars, ix)++astReplicateK :: forall shm r s. (GoodScalar r, KnownSpan s)+              => ShS shm -> AstTensor AstMethodLet s (TKScalar r)+              -> AstTensor AstMethodLet s (TKS shm r)+astReplicateK ZSS t0 = sfromK t0+astReplicateK shm t0 = case t0 of+  Ast.AstLet var t v -> astLet var t (astReplicateK shm v)+  Ast.AstFromPrimal v -> fromPrimal $ astReplicateK shm v+  Ast.AstFromDual v -> fromDual $ astReplicateK shm v+  Ast.AstFromPlain v -> fromPlain $ astReplicateK shm v+  AstConcreteK t | shsSize shm < 100 ->  -- tough trade-offs here+    astConcreteS $ tsreplicate0N shm $ Concrete t+  _ -> Ast.AstReplicateK shm t0++astReplicateS :: forall shm shn x s. KnownSpan s+              => ShS shm -> AstTensor AstMethodLet s (TKS2 shn x)+              -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+astReplicateS ZSS t0 = t0+astReplicateS shm t0 = case t0 of+  -- TODO: This rules is, in principle, very good, because it permits many other+  -- rules to fire. However, one of these other rules is indexing of transpose+  -- that in some cases complicates terms and degrades performance considerably.+  -- Ast.AstLet var t v -> astLet var t (astReplicateS shm v)+  Ast.AstFromPrimal v -> fromPrimal $ astReplicateS shm v+  Ast.AstFromDual v -> fromDual $ astReplicateS shm v+  Ast.AstFromPlain v -> fromPlain $ astReplicateS shm v+  AstConcreteS t | shsSize shm < 100 ->  -- tough trade-offs here+    let shn = Nested.sshape t+    in withKnownShS shn $+       astConcreteS $ tsreplicateN shm $ Concrete t+      -- revisit the trade-offs once we compile instead of interpreting+      -- and so building big blobby concrete arrays is cheap+  Ast.AstReplicateK shm2 v ->+    astReplicateK (shm `shsAppend` shm2) v+  Ast.AstReplicateS @shm2 @shn2 shm2 v+    | Refl <- lemAppAssoc (Proxy @shm) (Proxy @shm2) (Proxy @shn2) ->+      astReplicateS (shm `shsAppend` shm2) v+  _ | FTKS ZSS FTKScalar <- ftkAst t0+    , Refl <- lemAppNil @shm -> Ast.AstReplicateK shm (kfromS t0)+  _ -> Ast.AstReplicateS shm t0+  -- TODO: maybe add a rule and then generalize:+  -- replicate n1 (str (replicate n2 u))+  -- ~> transpose [0, 2, 1] (replicate n1 (replicate n2 u))+  -- but the reverse rule is already in astTransposeS++flipCompare :: forall (a :: Nat) b. Compare a b ~ GT+            => Proxy a -> Proxy b -> Compare b a :~: LT+flipCompare _ _ = unsafeCoerceRefl++astGatherS+  :: forall shm shn shp r s. KnownSpan s+  => ShS shm -> ShS shn -> ShS shp+  -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)+  -> (AstVarListS shm, AstIxS AstMethodLet shp)+  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)+astGatherS = astGatherKnobsS defaultKnobs++-- Assumption: vars0 don't not occur in v0. The assumption only holds+-- when newly generated variables are fresh, which is the case as long+-- as resetVarCounter is not used. The assumption makes it easier to spot+-- bugs or corruption, hence we assert it in the code below.+astGatherKnobsS+  :: forall shm shn shp x s. KnownSpan s+  => SimplifyKnobs+  -> ShS shm -> ShS shn -> ShS shp+  -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) x)+  -> (AstVarListS shm, AstIxS AstMethodLet shp)+  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+astGatherKnobsS _ _ _ _ v0 (!vars0, !_ix0)+  | Foldable.any (`varNameInAst` v0) vars0 =+    error $ "astGatherKnobsS: gather vars in v0: " ++ show (vars0, v0)+astGatherKnobsS knobs _ shn _ v0 (ZS, ix0) = astIndexKnobsS knobs shn v0 ix0+astGatherKnobsS _ shm _ _ v0 (_, ZIS) = astReplicateS shm v0+astGatherKnobsS _ shm@(SZ :$$ _) shn _ v0 _+  | FTKS _ x <- ftkAst v0 =+    let ftk = FTKS (shm `shsAppend` shn) x+    in fromPlain $ astConcrete ftk (tdefTarget ftk)+astGatherKnobsS _ shm shn _shp v0 (_, i1 :.$ _)+  | Just (lb, ub) <- intBounds i1+  , FTKS (snat :$$ _) x <- ftkAst v0+  , ub < 0 || lb >= fromSNat' snat =+    let ftk = FTKS (shm `shsAppend` shn) x+    in fromPlain $ astConcrete ftk (tdefTarget ftk)+astGatherKnobsS knobs shm@(k :$$ _) shn shp v0 (var1 ::$ vars1, ix0)+  | not (var1 `varNameInIxS` ix0) =+    astReplicateS (k :$$ ZSS)+                  (astGatherKnobsS knobs (shsTail shm) shn shp v0 (vars1, ix0))+astGatherKnobsS knobs shm shn shp v0 (vars0@(_ ::$ _), ix0@(_ :.$ _))+  | let ixInit = ixsInit ix0+        varInit = AstVarListS $ ixsInit $ unAstVarListS vars0+        varLast = ixsLast $ unAstVarListS vars0+  , AstIntVar ixvarLast <- ixsLast ix0+  , ixvarLast == varLast+  , not (varLast `varNameInIxS` ixInit)+  , kLast@SNat <- shsLast shm+  , Just Refl <- testEquality kLast (shsLast shp) =+    gcastWith (unsafeCoerceRefl+               :: Init shp ++ (Last shm ': shn) :~: shp ++ shn) $+    gcastWith (unsafeCoerceRefl+               :: Init shm ++ (Last shm ': shn) :~: shm ++ shn) $+    astGatherKnobsS knobs (shsInit shm) (kLast :$$ shn) (shsInit shp)+                    v0 (varInit, ixInit)+astGatherKnobsS+  knobs shm@(m :$$ _) shn shp v0+  ( vars@(varm ::$ _)+  , Ast.AstCondK+      (Ast.AstBoolAndK a@(AstLeqInt (AstConcreteK j0) (AstIntVar varp)) b)+      v w :.$ prest )+  | varm == varp+  , j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest =+    let i = astLetFun w $ \wShared ->+              astCondKInitial a (astCondKInitial b v wShared) wShared+    in astGatherKnobsS knobs shm shn shp v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shm@(m :$$ _) shn shp v0+  ( vars@(varm ::$ _)+  , Ast.AstLet varN uN+      (Ast.AstCondK+         (Ast.AstBoolAndK a@(AstLeqInt (AstConcreteK j0) (AstIntVar varp)) b)+         v w) :.$ prest )+  | varm == varp+  , let varIn1 = varN `varNameInAst` v+        varIn2 = varN `varNameInAst` w+  , j0 <= 0 || j0 >= fromSNat' m+    || ixIsSmall prest && not (varIn1 && varIn2) =+    let i = astLetFun w $ \wShared ->+              astCondKInitial a (astCondKInitial b v wShared) wShared+    in astGatherKnobsS knobs shm shn shp v0 (vars, astLet varN uN i :.$ prest)+astGatherKnobsS+  knobs shm@(m :$$ _) shn shp v0+  ( vars@(varm ::$ _)+  , Ast.AstCondK+      (Ast.AstBoolAndK a@(AstLeqInt (AstConcreteK j0)+                                    (Ast.AstN1K NegateOp (AstIntVar varp))) b)+      v w :.$ prest )+  | varm == varp+  , - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m || ixIsSmall prest =+    let i = astLetFun w $ \wShared ->+              astCondKInitial a (astCondKInitial b v wShared) wShared+    in astGatherKnobsS knobs shm shn shp v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shm@(m :$$ _) shn shp v0+  ( vars@(varm ::$ _)+  , Ast.AstLet varN uN+      (Ast.AstCondK+         (Ast.AstBoolAndK+            a@(AstLeqInt (AstConcreteK j0)+                         (Ast.AstN1K NegateOp (AstIntVar varp))) b)+      v w) :.$ prest )+  | varm == varp+  , let varIn1 = varN `varNameInAst` v+        varIn2 = varN `varNameInAst` w+  , - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m+    || ixIsSmall prest && not (varIn1 && varIn2) =+    let i = astLetFun w $ \wShared ->+              astCondKInitial a (astCondKInitial b v wShared) wShared+    in astGatherKnobsS knobs shm shn shp v0 (vars, astLet varN uN i :.$ prest)++astGatherKnobsS+  knobs shm@(m :$$ _) shn shp v0+  ( vars@(varm ::$ _)+  , Ast.AstCondK+      (Ast.AstBoolAndK+         a@(Ast.AstBoolNotK+              (Ast.AstBoolAndK+                 (AstLeqInt (AstConcreteK j0) (AstIntVar varp)) _)) b)+      v w :.$ prest )+  | varm == varp+  , j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest =+    let i = astLetFun w $ \wShared ->+              astCondKInitial a (astCondKInitial b v wShared) wShared+    in astGatherKnobsS knobs shm shn shp v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shm@(m :$$ _) shn shp v0+  ( vars@(varm ::$ _)+  , Ast.AstLet varN uN+      (Ast.AstCondK+         (Ast.AstBoolAndK+            a@(Ast.AstBoolNotK+                 (Ast.AstBoolAndK (AstLeqInt (AstConcreteK j0)+                                             (AstIntVar varp)) _)) b)+         v w) :.$ prest )+  | varm == varp+  , let varIn1 = varN `varNameInAst` v+        varIn2 = varN `varNameInAst` w+  , j0 <= 0 || j0 >= fromSNat' m+    || ixIsSmall prest && not (varIn1 && varIn2) =+    let i = astLetFun w $ \wShared ->+              astCondKInitial a (astCondKInitial b v wShared) wShared+    in astGatherKnobsS knobs shm shn shp v0 (vars, astLet varN uN i :.$ prest)+astGatherKnobsS+  knobs shm@(m :$$ _) shn shp v0+  ( vars@(varm ::$ _)+  , Ast.AstCondK+      (Ast.AstBoolAndK+         a@(Ast.AstBoolNotK+              (Ast.AstBoolAndK+                 (AstLeqInt (AstConcreteK j0)+                            (Ast.AstN1K NegateOp (AstIntVar varp))) _)) b)+      v w :.$ prest )+  | varm == varp+  , - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m || ixIsSmall prest =+    let i = astLetFun w $ \wShared ->+              astCondKInitial a (astCondKInitial b v wShared) wShared+    in astGatherKnobsS knobs shm shn shp v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shm@(m :$$ _) shn shp v0+  ( vars@(varm ::$ _)+  , Ast.AstLet varN uN+      (Ast.AstCondK+         (Ast.AstBoolAndK+            a@(Ast.AstBoolNotK+                 (Ast.AstBoolAndK+                    (AstLeqInt (AstConcreteK j0)+                               (Ast.AstN1K NegateOp (AstIntVar varp))) _)) b)+         v w) :.$ prest )+  | varm == varp+  , let varIn1 = varN `varNameInAst` v+        varIn2 = varN `varNameInAst` w+  , - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m+    || ixIsSmall prest && not (varIn1 && varIn2) =+    let i = astLetFun w $ \wShared ->+              astCondKInitial a (astCondKInitial b v wShared) wShared+    in astGatherKnobsS knobs shm shn shp v0 (vars, astLet varN uN i :.$ prest)++-- Rules with AstConcreteK on the right hand side of AstPlusK are+-- not needed, thanks to the normal form of AstPlusK rewriting.+astGatherKnobsS knobs shm shn shp v0+  (vars, AstPlusK (AstConcreteK i) i1 :.$ prest)+  | Just (lb, ub) <- intBounds i1+  , lb >= 0  -- if not, we may need to apply astReverse first+    || i <= 0+  , FTKS (SNat @p :$$ _) x <- ftkAst v0 =+    if i >= 0 then+      withSNat i $ \(SNat @i) ->+      withSNat (min (valueOf @p - i) (ub + 1)) $ \(SNat @k) ->+        gcastWith (unsafeCoerceRefl :: (i + k <=? p) :~: True) $+        let v2 = astSliceS (SNat @i) (SNat @k) (SNat @(p - (i + k))) v0+        in astGatherKnobsS knobs shm shn (SNat @k :$$ shsTail shp)+                           v2 (vars, i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+    else+      withSNat (- i) $ \(SNat @i) ->+        let ftk = FTKS (SNat @i :$$ shsTail shp `shsAppend` shn) x+            v2 = fromPlain (astConcrete ftk (tdefTarget ftk))+                 `astAppendS`+                 v0+        in astGatherKnobsS knobs shm shn (SNat @(p + i) :$$ shsTail shp)+                           v2 (vars, i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+astGatherKnobsS knobs shm shn shp v0+  (vars, Ast.AstLet varN uN (AstPlusK (AstConcreteK i) i1) :.$ prest)+  | Just (lb, ub) <- intBounds i1+  , lb >= 0  -- if not, we may need to apply astReverse first+    || i <= 0+  , FTKS (SNat @p :$$ _) x <- ftkAst v0 =+    if i >= 0 then+      withSNat i $ \(SNat @i) ->+      withSNat (min (valueOf @p - i) (ub + 1)) $ \(SNat @k) ->+        gcastWith (unsafeCoerceRefl :: (i + k <=? p) :~: True) $+        let v2 = astSliceS (SNat @i) (SNat @k) (SNat @(p - (i + k))) v0+        in astGatherKnobsS knobs shm shn (SNat @k :$$ shsTail shp)+                           v2 (vars, astLet varN uN i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+    else+      withSNat (- i) $ \(SNat @i) ->+        let ftk = FTKS (SNat @i :$$ shsTail shp `shsAppend` shn) x+            v2 = fromPlain (astConcrete ftk (tdefTarget ftk))+                 `astAppendS`+                 v0+        in astGatherKnobsS knobs shm shn (SNat @(p + i) :$$ shsTail shp)+                           v2 (vars, astLet varN uN i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+astGatherKnobsS knobs (m@(SNat @m) :$$ (shmRest :: ShS shmRest)) shn shp v0+  ( varm ::$ mrest+  , Ast.AstI2K QuotOp (AstIntVar varm') (AstConcreteK i0) :.$ prest )+  | varm == varm'+  , not (varm `varNameInIxS` prest)+  , i0 > 0  -- ensured by other rules; makes this easier to reason about+  , let k0 = (fromSNat' m + i0 - 1) `quot` i0+  , SNat <- shsRank shmRest+  , SNat <- shsRank shn =+    withSNat i0 $ \i@(SNat @i) ->+    withSNat k0 $ \k@(SNat @k) ->+    withSNat (k0 * i0 - fromSNat' m) $ \z@(SNat @z) ->+      let perm = Permutation.makePerm @'[1, 0]+          varm2 = reboundsVarName (0, k0 - 1) varm+      in gcastWith (unsafeCoerceRefl :: m + z :~: k * i) $+         gcastWith (unsafeCoerceRefl+                    :: Rank ([i, k] ++ shmRest ++ shn)+                       :~: 2 + Rank shmRest + Rank shn) $+         astSliceS SZ m z+         $ astReshapeS (snatMul k i :$$ shmRest `shsAppend` shn)+         $ astTransposeS perm+         $ astReplicateS (i :$$ ZSS)+         $ astGatherKnobsS knobs (k :$$ shmRest) shn shp+                           v0 (varm2 ::$ mrest, astVar varm2 :.$ prest)+             -- this gather may still index out of bounds, which is fine+astGatherKnobsS knobs shm@(SNat @m :$$ _) shn shp u0+  ( vars@(varm ::$ mrest)+  , Ast.AstCondK (AstLeqInt (AstConcreteK j) (AstIntVar varp)) v w+    :.$ prest )+  | varm == varp+  , j <= 0 || j >= valueOf @m || ixIsSmall prest =+    if | j <= 0 ->+         astGatherKnobsS knobs shm shn shp+                         u0 (vars, v :.$ prest)+       | j >= valueOf @m ->+         astGatherKnobsS knobs shm shn shp+                         u0 (vars, w :.$ prest)+       | otherwise ->+         withSNat j $ \(SNat @j) ->+         gcastWith (unsafeCoerceRefl :: (j <=? m) :~: True) $+         astLetFun u0 $ \u ->+         let varm2 = reboundsVarName (0, j - 1) varm+             varm3 = reboundsVarName (0, valueOf @m - j - 1) varm+         in astGatherKnobsS knobs (SNat @j :$$ shsTail shm) shn shp u+              ( varm2 ::$ mrest+              , substituteAstIxS (astVar varm2)+                                 varm (w :.$ prest) )+            `astAppendS`+            astGatherKnobsS knobs (SNat @(m - j) :$$ shsTail shm) shn shp u+              ( varm3 ::$ mrest+              , substituteAstIxS (AstConcreteK j + astVar varm3)+                                 varm (v :.$ prest) )+astGatherKnobsS knobs shm@(SNat @m :$$ _) shn shp u0+  ( vars@(varm ::$ mrest)+  , Ast.AstLet varN uN+      (Ast.AstCondK (AstLeqInt (AstConcreteK j) (AstIntVar varp)) v w)+    :.$ prest )+  | varm == varp+  , let varIn1 = varN `varNameInAst` v+        varIn2 = varN `varNameInAst` w+  , j <= 0 || j >= valueOf @m+    || ixIsSmall prest && not (varIn1 && varIn2) =+    if | j <= 0 ->+         astGatherKnobsS knobs shm shn shp+                         u0 (vars, astLet varN uN v :.$ prest)+       | j >= valueOf @m ->+         astGatherKnobsS knobs shm shn shp+                         u0 (vars, astLet varN uN w :.$ prest)+       | otherwise ->+         withSNat j $ \(SNat @j) ->+         gcastWith (unsafeCoerceRefl :: (j <=? m) :~: True) $+         astLetFun u0 $ \u ->+         let varm2 = reboundsVarName (0, j - 1) varm+             varm3 = reboundsVarName (0, valueOf @m - j - 1) varm+         in astGatherKnobsS knobs (SNat @j :$$ shsTail shm) shn shp u+              ( varm2 ::$ mrest+              , substituteAstIxS (astVar varm2)+                                 varm ((if varIn2+                                        then astLet varN uN w+                                        else w) :.$ prest) )+            `astAppendS`+            astGatherKnobsS knobs (SNat @(m - j) :$$ shsTail shm) shn shp u+              ( varm3 ::$ mrest+              , substituteAstIxS (AstConcreteK j + astVar varm3)+                                 varm ((if varIn1+                                        then astLet varN uN v+                                        else v) :.$ prest) )+astGatherKnobsS knobs+                shm@(SNat @m :$$ (_ :: ShS shmTail))+                shn+                shp@(SNat @p :$$ (_ :: ShS shpTail))+                v0+  ( varm ::$ mrest+  , AstIntVar varp :.$ prest )+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+      -- prevent a loop+  , varm == varp+  , not (varm `varNameInIxS` prest)+  , FTKS _ x <- ftkAst v0 =+    withSNat (min (valueOf @p) (valueOf @m)) $ \(SNat @m2) ->+    gcastWith (unsafeCoerceRefl :: (m2 <=? p) :~: True) $+    gcastWith (unsafeCoerceRefl :: (m2 <=? m) :~: True) $+    Permutation.permFromListCont (permCycle+                                  $ shsLength (shsTail shm) + 1)+    $ \(permVars :: Permutation.Perm permVars) ->+    Permutation.permFromListCont (backpermCycle+                                  $ shsLength (shsTail shp) + 1)+    $ \(permIx :: Permutation.Perm permIx) ->+       gcastWith (unsafeCoerceRefl+                  :: m2 ': shmTail ++ shn+                     :~: Permutation.PermutePrefix+                           permVars (shmTail ++ (m2 ': shn))) $+       gcastWith (unsafeCoerceRefl+                  :: shpTail ++ (m2 ': shn)+                     :~: Permutation.PermutePrefix+                           permIx (m2 ': shpTail ++ shn)) $+       gcastWith (unsafeCoerceRefl+                  :: (Rank permVars <=? Rank (shmTail ++ (m2 ': shn)))+                     :~: True) $+       gcastWith (unsafeCoerceRefl+                  :: (Rank permIx <=? Rank (m2 ': shpTail ++ shn)) :~: True) $+       fromMaybe (error "astGatherKnobsS: impossible non-permutation")+       $ Permutation.permCheckPermutation permVars+       $ fromMaybe (error "astGatherKnobsS: impossible non-permutation")+       $ Permutation.permCheckPermutation permIx+       $ let v2 = astTransposeS permIx+                  $ astSliceS (SNat @0) (SNat @m2) (SNat @(p - m2)) v0+             u = astGatherKnobsS+                   knobs (shsTail shm) (SNat @m2 :$$ shn) (shsTail shp)+                   v2 (mrest, prest)+             ftk = FTKS (SNat @(m - m2) :$$ shsTail shm `shsAppend` shn) x+         in astTransposeS permVars u+            `astAppendS`+            fromPlain (astConcrete ftk (tdefTarget ftk))+astGatherKnobsS knobs (m :$$ shm4) ZSS _shp+                v4@(Ast.AstFromVectorK (_ :$$ ZSS) l)+                ( var4 ::$ _vrest4+                , i4 :.$ rest4 )+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+  , V.length l < 100+  , let g i = case substituteAst (AstConcreteK i) var4 i4 of+          AstConcreteK j -> Just j+          _ -> Nothing+  , Just lj <- mapM g [0 .. fromSNat' m - 1]+  , Refl <- lemAppNil @shm = assert (null rest4) $+    let FTKS _ x = ftkAst v4+        ftk = FTKS shm4 x+        zero = fromPlain $ astConcrete ftk (tdefTarget ftk)+        f j =+          if j >= V.length l+          then zero+          else astReplicateK shm4 (l V.! j)+    in astFromVectorS (m :$$ ZSS)+       $ V.fromListN (fromSNat' m) $ map f lj+astGatherKnobsS knobs shm@(m :$$ (shm4 :: ShS shm4)) shn shp+                v4@(Ast.AstFromVectorS (_ :$$ ZSS) l)+                ( var4 ::$ vrest4+                , i4 :.$ rest4 )+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+  , V.length l < 100+  , let g i = case substituteAst (AstConcreteK i) var4 i4 of+          AstConcreteK j -> Just j+          _ -> Nothing+  , Just lj <- mapM g [0 .. fromSNat' m - 1]+  , ixIsSmall rest4 =+    let FTKS _ x = ftkAst v4+        ftk = FTKS (shsTail shm `shsAppend` shn) x+        zero = fromPlain $ astConcrete ftk (tdefTarget ftk)+        f i j =+          let subRest4 = substituteAstIxS (AstConcreteK i) var4 rest4+          in if j >= V.length l+             then zero+             else astGatherKnobsS knobs shm4 shn (shsTail shp)+                                  (l V.! j) (vrest4, subRest4)+    in astFromVectorS (m :$$ ZSS)+       $ V.fromListN (fromSNat' m) $ zipWith f [0 .. fromSNat' m - 1] lj+astGatherKnobsS knobs shm shn shp v0 (vars0, i1 :.$ rest1)+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+      -- prevent a loop+  , not (Foldable.any (`varNameInAst` i1) vars0) =+    astGatherKnobsS+      knobs shm shn (shsTail shp)+      (astIndexKnobsS knobs (shsTail shp `shsAppend` shn) v0 (i1 :.$ ZIS))+      (vars0, rest1)+astGatherKnobsS knobs shm@(m :$$ _) shn shp v0+  (vars, ix@(i1 :.$ prest))+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+      -- prevent a loop+  , let intInteresting = \case+          AstPlusK (AstConcreteK i) i2+            | Just (lb, _) <- intBounds i2+            , lb >= 0 || i <= 0-> True+          Ast.AstLet _ _ (AstPlusK (AstConcreteK i) i2)+            | Just (lb, _) <- intBounds i2+            , lb >= 0 || i <= 0-> True+          Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar var)) _ _+            | j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest+            , Foldable.any (== var) vars -> True+          Ast.AstLet varN _+            (Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar var)) v w)+            | let varIn1 = varN `varNameInAst` v+                  varIn2 = varN `varNameInAst` w+            , j0 <= 0 || j0 >= fromSNat' m+              || ixIsSmall prest && not (varIn1 && varIn2)+            , Foldable.any (== var) vars -> True+          Ast.AstCondK+            (Ast.AstBoolAndK+               (AstLeqInt (AstConcreteK j0) (AstIntVar var)) _) _ _+            | j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest+            , Foldable.any (== var) vars -> True+          Ast.AstCondK+            (Ast.AstBoolAndK+               (AstLeqInt (AstConcreteK j0)+                          (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _+            | - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m || ixIsSmall prest+            , Foldable.any (== var) vars -> True+          Ast.AstLet varN _+            (Ast.AstCondK+               (Ast.AstBoolAndK+                  (AstLeqInt (AstConcreteK j0) (AstIntVar var)) _) v w)+            | let varIn1 = varN `varNameInAst` v+                  varIn2 = varN `varNameInAst` w+            , j0 <= 0 || j0 >= fromSNat' m+              || ixIsSmall prest && not (varIn1 && varIn2)+            , Foldable.any (== var) vars -> True+          Ast.AstLet varN _+            (Ast.AstCondK+               (Ast.AstBoolAndK+                  (AstLeqInt (AstConcreteK j0)+                             (Ast.AstN1K NegateOp (AstIntVar var))) _) v w)+            | let varIn1 = varN `varNameInAst` v+                  varIn2 = varN `varNameInAst` w+            , - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m+              || ixIsSmall prest && not (varIn1 && varIn2)+            , Foldable.any (== var) vars -> True+          AstIntVar var+            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+            , null $ drop 1 $ filter (var `varNameInAst`) (Foldable.toList ix)+            , Foldable.any (== var) vars -> True+          ik | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+             , not (Foldable.any (`varNameInAst` ik) vars) -> True+          -- We can't reorder ix for the gather(fromVector) rule above,+          -- because it becomes gather(transpose); we can only reorder vars.+          _ -> False+  , not (intInteresting i1)  -- now vars may need to be reordered, too+  , Just i <- findIndex intInteresting+                        (Foldable.toList ix) = assert (i > 0) $+    Permutation.permFromListCont (backpermCycle $ i + 1)+    $ \(perm :: Permutation.Perm perm) ->+    gcastWith (unsafeCoerceRefl+               :: Permutation.PermutePrefix perm (shp ++ shn)+                  :~: Permutation.PermutePrefix perm shp ++ shn) $+    gcastWith (unsafeCoerceRefl :: (Rank perm <=? Rank (shp ++ shn)) :~: True) $+    fromMaybe (error "astGatherKnobsS: impossible non-permutation")+    $ Permutation.permCheckPermutation perm+    $ let v2 = astTransposeS perm v0+      in astGatherKnobsS knobs shm shn (shsPermutePrefix perm shp)+                         v2 (vars, ixsPermutePrefix perm ix)+        -- this call is guaranteed to simplify as above, so the transpose+        -- won't reduce it back to the original and cause a loop+astGatherKnobsS knobs shm@(m :$$ _) shn shp v0+  (vars, ix@(i1 :.$ prest))+  | let varInteresting = \case+          Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar var)) _ _+            | j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest ->+              Just var+          Ast.AstLet varN _+            (Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar var)) v w)+            | let varIn1 = varN `varNameInAst` v+                  varIn2 = varN `varNameInAst` w+            , j0 <= 0 || j0 >= fromSNat' m+              || ixIsSmall prest && not (varIn1 && varIn2) ->+              Just var+          Ast.AstCondK+            (Ast.AstBoolAndK+               (AstLeqInt (AstConcreteK j0) (AstIntVar var)) _) _ _+            | j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest ->+              Just var+          Ast.AstCondK+            (Ast.AstBoolAndK+               (AstLeqInt (AstConcreteK j0)+                          (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _+            | - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m || ixIsSmall prest ->+              Just var+          Ast.AstLet varN _+            (Ast.AstCondK+               (Ast.AstBoolAndK+                  (AstLeqInt (AstConcreteK j0) (AstIntVar var)) _) v w)+            | let varIn1 = varN `varNameInAst` v+                  varIn2 = varN `varNameInAst` w+            , j0 <= 0 || j0 >= fromSNat' m+              || ixIsSmall prest && not (varIn1 && varIn2) ->+              Just var+          Ast.AstLet varN _+            (Ast.AstCondK+               (Ast.AstBoolAndK+                  (AstLeqInt (AstConcreteK j0)+                             (Ast.AstN1K NegateOp (AstIntVar var))) _) v w)+            | let varIn1 = varN `varNameInAst` v+                  varIn2 = varN `varNameInAst` w+            , - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m+              || ixIsSmall prest && not (varIn1 && varIn2) ->+              Just var+          AstIntVar var+            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+            , not (var `varNameInIxS` prest) -> Just var+          i4  -- has to be last, because ix can't be reordered+            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+            , Ast.AstFromVectorK (_ :$$ ZSS) l <- v0+            , V.length l < 100+            , let g var4 i = case substituteAst (AstConcreteK i) var4 i4 of+                    AstConcreteK{} -> True+                    _ -> False+                  mvar = Foldable.find (\var ->+                           all (g var) [0 .. fromSNat' m - 1]) vars+            , Just{} <- mvar -> mvar+            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+            , Ast.AstFromVectorS (_ :$$ ZSS) l <- v0+            , V.length l < 100+            , ixIsSmall prest+            , let g var4 i = case substituteAst (AstConcreteK i) var4 i4 of+                    AstConcreteK{} -> True+                    _ -> False+                  mvar = Foldable.find (\var ->+                           all (g var) [0 .. fromSNat' m - 1]) vars+            , Just{} <- mvar -> mvar+          _ -> Nothing+  , Just varp <- varInteresting i1+  , Just i <- elemIndex varp (Foldable.toList vars) = assert (i > 0) $+    Permutation.permFromListCont (backpermCycle $ i + 1)+    $ \(permWhole :: Permutation.Perm permWhole) ->+    permInverse permWhole $ \(invperm :: Nested.Perm invperm) _ ->+    gcastWith (unsafeCoerceRefl+               :: shm ++ shn+                  :~: Permutation.PermutePrefix permWhole+                        (Permutation.PermutePrefix invperm shm ++ shn)) $+    gcastWith (unsafeCoerceRefl+               :: (Rank permWhole+                   <=? Rank (Permutation.PermutePrefix invperm shm ++ shn))+                  :~: True) $+    fromMaybe (error "astGatherKnobsS: impossible non-permutation")+    $ Permutation.permCheckPermutation permWhole+    $ astTransposeS permWhole+    $ astGatherKnobsS knobs (shsPermutePrefix invperm shm) shn shp+                      v0 (AstVarListS+                          $ ixsPermutePrefix invperm (unAstVarListS vars), ix)+        -- this call is guaranteed to simplify as above, so the transpose+        -- won't reduce it back to the original and cause a loop+astGatherKnobsS knobs shm shn shp@(SNat @in1 :$$ (shp1 :: ShS shp1))+                v4 (vars4, ix4@(i4 :.$ rest4))+  | FTKS _ x <- ftkAst v4+  , let tryRecursing :: AstTensor AstMethodLet s (TKS2 (shp1 ++ shn) x)+                     -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+        tryRecursing u =+          let ftk = FTKS (shm `shsAppend` shn) x+              defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)+          -- This boolean term may have free variables that act as+          -- universally quantified.+          in case 0 <=. i4 &&* i4 <=. valueOf @in1 - 1 of+            AstConcreteK b ->+              if b then astGather shm shn shp1 u (vars4, rest4) else defArr+            _ -> Ast.AstGatherS shm shn shp v4 (vars4, ix4) = case v4 of+    Ast.AstProject1{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstProject2{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    {- Ast.AstFromVector{} | gatherFromNF (shsTail shp) vars4 ix4 ->+        -- normal form+      Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4) -}+    {- this rule seems counterproductive in many cases, so disabled until+       we can detect cases where it helps:+    Ast.AstFromVector snat STKS{} l ->+      -- Term rest4 is duplicated without sharing and we can't help it,+      -- because it needs to be in scope of vars4, so we can't use tlet.+      funToVarsIxS @shm shm $ \varsFresh (IxS !ixFresh) ->+        let f v = astGather @shm @shn @shp1 shn v (vars4, rest4)+            -- This subst doesn't currently break sharing because it's a rename.+            subst i =+              Foldable.foldr (\(i2, var2) v2 -> substituteAst i2 var2 v2)+                    i+                    (zipSizedS ixFresh vars4)+            i5 = subst i4+       in astGather @shm @shn @(p1' ': shm)+                    shn (astFromVector snat (STKS (shsFromListS varsFresh+                                                    `shsAppend` shn)+                                                   (ftkToSTK x))+                          $ V.map f l)+                    (varsFresh, i5 :.$ IxS ixFresh) -}+    Ast.AstApply{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstVar{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstBuild1 _ STKScalar (var2, v) | ZIS <- rest4 ->+      tryRecursing $ sfromK $ astLet var2 i4 v+    Ast.AstBuild1 _ STKS{} (var2, v) ->+      tryRecursing $ astLet var2 i4 v++    Ast.AstLet var u v ->+      astLetRefresh var u v $ \v' -> astGather shm shn shp v' (vars4, ix4)++    Ast.AstPrimalPart{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstDualPart{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstPlainPart{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstFromPrimal v -> fromPrimal $ astGather shm shn shp v (vars4, ix4)+    Ast.AstFromDual v -> fromDual $ astGather shm shn shp v (vars4, ix4)+    Ast.AstFromPlain v -> fromPlain $ astGather shm shn shp v (vars4, ix4)++    -- Going inside a binary ops usually makes a term more expensive+    -- to interpret and inverting that requires comparing two arguments,+    -- so it's not practical.+    AstPlusS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    AstTimesS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstN1S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstR1S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstR2S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstI2S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    AstConcreteS{} -> case unRepl1 v4 of+      Just u -> tryRecursing u+      _ -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+             -- free variables possible in the index, so can't compute the array+    Ast.AstFloorS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstFromIntegralS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstCastS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstArgMinS @n1 @shz v  | FTKS nsh _ <- ftkAst v+                               , nl@(SNat @nl) <- shsLast nsh ->+      let shnl = shn `shsAppend` (nl :$$ ZSS)+      in gcastWith (unsafeCoerceRefl+                    :: Head (shm ++ (shn ++ '[nl]))+                       ': Tail (shm ++ (shn ++ '[nl]))+                       :~: shm ++ (shn ++ '[nl])) $+         gcastWith (unsafeCoerceRefl+                    :: Init (shm ++ (shn ++ '[nl])) :~: shm ++ shn) $+         gcastWith (unsafeCoerceRefl+                   :: shp ++ (shn ++ '[nl]) :~: n1 ': shz) $+         Ast.AstArgMinS @(Head (shm ++ (shn ++ '[nl])))+                        @(Tail (shm ++ (shn ++ '[nl])))+         $ astGather shm shnl shp v (vars4, ix4)+    Ast.AstArgMaxS @n1 @shz v  | FTKS nsh _ <- ftkAst v+                               , nl@(SNat @nl) <- shsLast nsh ->+      let shnl = shn `shsAppend` (nl :$$ ZSS)+      in gcastWith (unsafeCoerceRefl+                    :: Head (shm ++ (shn ++ '[nl]))+                       ': Tail (shm ++ (shn ++ '[nl]))+                       :~: shm ++ (shn ++ '[nl])) $+         gcastWith (unsafeCoerceRefl+                    :: Init (shm ++ (shn ++ '[nl])) :~: shm ++ shn) $+         gcastWith (unsafeCoerceRefl+                   :: shp ++ (shn ++ '[nl]) :~: n1 ': shz) $+         Ast.AstArgMaxS @(Head (shm ++ (shn ++ '[nl])))+                        @(Tail (shm ++ (shn ++ '[nl])))+         $ astGather shm shnl shp v (vars4, ix4)++    {- is reverted in astGatherKnobsS immediatedly; only do in expansion phase?+    Ast.AstIndexS @shm2 _shn2 v2 (i2 :.$ ZIS) ->+        astGather @shm @shn @(shm2 ++ shp) shn v2 (vars4, i2 :.$ ix4) -}+    Ast.AstIndexS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstCondS b v w | ixIsSmall ix4 ->+      astCondS b (astGather shm shn shp v (vars4, ix4))+                 (astGather shm shn shp w (vars4, ix4))+    Ast.AstCondS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstFromVectorK{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstFromVectorS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    -- This accomplishes fusion if v is a gather or anything+    -- that gather can fuse with, but at the cost of an extra transpose+    -- that doesn't fuse here unless astTransposeAsGatherS is used.+    -- Since the transpose is O(1), let's leave this as is.+    Ast.AstSumS @shm2 shm2 v ->+      let perm3 = backpermCycleN (shsLength shm2)+                                 (shsLength shm2 + shsLength shp)+          perm4 = permCycleN (shsLength shm2) (shsLength shm2 + shsLength shm)+      in Permutation.permFromListCont perm3+         $ \(perm3S :: Permutation.Perm perm3P) ->+         gcastWith (unsafeCoerceRefl+                    :: (Rank perm3P <=? Rank (shm2 ++ (shp ++ shn))) :~: True) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix perm3P (shm2 ++ (shp ++ shn))+                       :~: shp ++ (shm2 ++ shn)) $+         fromMaybe (error "astGatherKnobsS: impossible non-permutation")+         $ Permutation.permCheckPermutation perm3S+         $ Permutation.permFromListCont perm4+         $ \(perm4S :: Permutation.Perm perm4P) ->+         gcastWith (unsafeCoerceRefl+                    :: (Rank perm4P <=? Rank (shm ++ (shm2 ++ shn))) :~: True) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix perm4P (shm ++ (shm2 ++ shn))+                       :~: shm2 ++ (shm ++ shn)) $+         fromMaybe (error "astGatherKnobsS: impossible non-permutation")+         $ Permutation.permCheckPermutation perm4S+         $ let innerGather =+                 astGather shm (shm2 `shsAppend` shn) shp+                           (astTransposeS perm3S v) (vars4, ix4)+           in astSumS shm2+              $ astTransposeS perm4S innerGather+                {- TODO: disabled until we can reliably fuse back to transpose+                if not (knobExpand knobs)+                then astTransposeS perm4S innerGather+                else astTransposeAsGatherS knobs perm4S innerGather -}+    Ast.AstScatterS shm7 shn7 shp7 v (vars, i5 :.$ ix2) | eqY i5 i4 ->+      astGather shm shn (shsTail shp)+                (astScatterKnobsS knobs shm7 shn7 (shsTail shp7) v (vars, ix2))+                (vars4, rest4)+    Ast.AstScatterS _ _ _ _ (_, AstConcreteK{} :.$ _)+      | AstConcreteK{} <- i4 ->  -- from above we know i5 /= i4+        let ftk = FTKS (shm `shsAppend` shn) x+        in fromPlain $ astConcrete ftk (tdefTarget ftk)+    Ast.AstScatterS{} ->  -- normal form+      Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstReplicateK (_ :$$ shmRest) v ->+      tryRecursing $ astReplicateK shmRest v+    Ast.AstReplicateS ZSS v -> astGather shm shn shp v (vars4, ix4)+    Ast.AstReplicateS (_ :$$ shmRest) v ->+      tryRecursing $ astReplicateS shmRest v+    Ast.AstGatherS @shm2 @shn2 @shp2 shm2 shn2 shp2 v2 (vars2, ix2)+                   | SNat @rank4 <- ixsRank ix4+                   , SNat @rank2 <- ixsRank (unAstVarListS vars2) ->+      let subst :: AstIxS AstMethodLet shm7 -> AstVarListS shm7+                -> AstInt AstMethodLet+                -> AstInt AstMethodLet+          subst ix vars t0 =+            foldr (\ (v, i) -> substituteAst i v)+                  t0 (ixsZip (unAstVarListS vars) ix)+          inBounds :: AstIxS AstMethodLet shm7 -> AstVarListS shm7 -> Bool+          inBounds ix vars =+            let inb (v, i) | Just (lbv, ubv) <- varNameToBounds v+                           , Just (lbi, ubi) <- intBounds i =+                  lbv <= lbi && ubi <= ubv+                inb _ = True+            in all inb (ixsZip (unAstVarListS vars) ix)+          composedGather ::  -- rank4 <= rank2+            Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) x))+          composedGather | SNat <- shsRank shm =+            -- we have: shm2 ++ shn2 == shp ++ shn+            -- so from ranks:+            gcastWith (unsafeCoerceRefl :: Take (Rank shp) shm2 :~: shp) $+            -- and from congruence:+--            gcastWith (unsafeCoerceRefl+--                       :: Drop (Rank shp) shm2 ++ shn2 :~: shn) $+            -- from congruence:+            gcastWith (unsafeCoerceRefl+                       :: (shm ++ Drop (Rank shp) shm2) ++ shn2+                          :~: shm ++ shn) $+            let vars2p = AstVarListS $ ixsTake @(Rank shp) $ unAstVarListS vars2+                vars22 = AstVarListS $ ixsDrop @(Rank shp) $ unAstVarListS vars2+                ix22 = fmap (subst ix4 vars2p) ix2+                list422 = AstVarListS+                          $ unAstVarListS vars4 `ixsAppend` unAstVarListS vars22+            in if ixIsSmall ix4 && inBounds ix4 vars2p+               then Just $ astGather (shm `shsAppend` shsDrop @(Rank shp) shm2)+                                     shn2+                                     shp2+                                     v2 (list422, ix22)+               else Nothing+          assimilatedGather ::  -- rank2 <= rank4+            Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) x))+          assimilatedGather | SNat <- shsRank shm =+            -- we have: shm2 ++ shn2 == shp ++ shn+            -- so from ranks:+            gcastWith (unsafeCoerceRefl :: Take (Rank shm2) shp :~: shm2) $+            -- and from congruence:+--            gcastWith (unsafeCoerceRefl+--                       :: Drop (Rank shm2) shp ++ shn :~: shn2) $+            -- from congruence:+            gcastWith (unsafeCoerceRefl+                       :: (shp2 ++ Drop (Rank shm2) shp) ++ shn+                          :~: shp2 ++ shn2) $+            let ix42 = ixsTake @(Rank shm2) ix4+                ix44 = ixsDrop @(Rank shm2) ix4+                ix22 = fmap (subst ix42 vars2) ix2+                ix2244 = ix22 `ixsAppend` ix44+            in if ixIsSmall ix42 && inBounds ix42 vars2+               then Just $ astGather shm+                                     shn+                                     (shp2 `shsAppend` shsDrop @(Rank shm2) shp)+                                     v2 (vars4, ix2244)+               else Nothing+      in fromMaybe (Ast.AstGatherS shm shn shp v4 (vars4, ix4))+         $ case cmpNat (Proxy @rank4) (Proxy @rank2) of+        LTI -> composedGather+        EQI -> assimilatedGather+        GTI -> gcastWith (flipCompare (Proxy @rank4) (Proxy @rank2))+                         assimilatedGather+    Ast.AstIotaS{} ->  -- probably nothing can be simplified; a normal form+      Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstAppendS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+      -- fusing would result in gather([gather, gather]), so no gain+    Ast.AstSliceS{}-> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+      -- slicing is O(1) so no point fusing and complicating the expression;+      -- if it did not simplify further with slice, it wouldn't with gather+    Ast.AstReverseS{}-> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+      -- reversing is O(1)+    Ast.AstTransposeS @perm @sh perm v+      | knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion]+      , fromSNat' (Permutation.permRank perm) <= length ix4+      , FTKS sh _ <- ftkAst v+      , Refl <- lemRankMapJust $ shsTakeLenPerm perm sh ->+          gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $+          permInverse perm+          $ \(invperm :: Nested.Perm invperm) proof ->+            case proof (ssxFromShX $ shxFromShS $ shsTakeLenPerm perm sh) of+              Refl ->+                -- from PermutePrefix and ranks:+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.PermutePrefix invperm shp ++ shn+                      :~: Permutation.PermutePrefix invperm (shp ++ shn)) $+                -- from AstTransposeS:+--                gcastWith+--                  (unsafeCoerceRefl+--                   :: Permutation.PermutePrefix invperm (shp ++ shn)+--                      :~: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)) $+                -- from PermutePrefix:+--                gcastWith+--                  (unsafeCoerceRefl+--                   :: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)+--                      :~: {-1-} Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))+--                          ++ {-2-} DropLen invperm (Permutation.PermutePrefix perm sh)) $+                -- 1. from PermutePrefix:+--                gcastWith+--                  (unsafeCoerceRefl+--                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))+--                      :~: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))) $+                -- ranks+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))+                      :~: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))) $+                -- from permInverse but MapJust-unwrapped:+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))+                      :~: TakeLen perm sh) $+                -- end of 1.+                -- 2. from PermutePrefix+--                gcastWith+--                  (unsafeCoerceRefl+--                   :: DropLen invperm (Permutation.PermutePrefix perm sh)+--                      :~: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)) $+                -- ranks+                gcastWith+                  (unsafeCoerceRefl+                   :: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)+                      :~: DropLen perm sh) $+                -- end of 2.+                -- from TakeLen:+                gcastWith+                  (unsafeCoerceRefl+                   :: TakeLen perm sh ++ DropLen perm sh :~: sh) $+                let invix4 = ixsPermutePrefix invperm ix4+                in astGather shm shn (shsPermutePrefix invperm shp)+                             v (vars4, invix4)+    Ast.AstTransposeS perm v+      | knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion] ->+        astGather shm shn shp (astTransposeAsGatherS knobs perm v) (vars4, ix4)+    Ast.AstTransposeS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstReshapeS sh v ->+      if shsLength sh <= 1  -- this is flatten+      then astGather shm shn shp (astReshapeAsGatherS knobs sh v) (vars4, ix4)+      else Ast.AstGatherS shm shn shp v4 (vars4, ix4)++    -- Conversions to shaped need to stay down, so this is NF.+    Ast.AstConvert{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)++    Ast.AstDot1InS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)  -- TODO+    Ast.AstMatmul2S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)  -- TODO++    Ast.AstBoolNotS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstBoolAndS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+    Ast.AstLeqS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)+ where+  astGather+    :: forall shm' shn' shp' s' x'. KnownSpan s'+    => ShS shm' -> ShS shn' -> ShS shp'+    -> AstTensor AstMethodLet s' (TKS2 (shp' ++ shn') x')+    -> (AstVarListS shm', AstIxS AstMethodLet shp')+    -> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') x')+  astGather = astGatherKnobsS knobs++-- Normal form of chains of appends has the append constructor on the right.+astAppendS :: KnownSpan s+           => AstTensor AstMethodLet s (TKS2 (m ': sh) x)+           -> AstTensor AstMethodLet s (TKS2 (n ': sh) x)+           -> AstTensor AstMethodLet s (TKS2 ((m + n) ': sh) x)+astAppendS u v | FTKS (SZ :$$ _) _ <- ftkAst u = v+astAppendS u v | FTKS (SZ :$$ _) _ <- ftkAst v = u+astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)+           (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) =+  astFromVectorS (SNat @(k1 + k2) :$$ ZSS) $ l1 V.++ l2+astAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest) a1)+           (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) =+  astFromVectorS (SNat @(1 + k2) :$$ ZSS) $ astReplicateS shmRest a1 `V.cons` l2+astAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest) a1)+           (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) =+  astFromVectorS (SNat @(1 + k2) :$$ ZSS) $ astReplicateK shmRest a1 `V.cons` l2+astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)+           (Ast.AstReplicateS (SNat' @1 :$$ shmRest) a2) =+  astFromVectorS (SNat @(k1 + 1) :$$ ZSS) $ l1 `V.snoc` astReplicateS shmRest a2+astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)+           (Ast.AstReplicateK (SNat' @1 :$$ shmRest) a2) =+  astFromVectorS (SNat @(k1 + 1) :$$ ZSS) $ l1 `V.snoc` astReplicateK shmRest a2+astAppendS (Ast.AstFromVectorK (SNat @k1 :$$ ZSS) l1)+           (Ast.AstFromVectorK (SNat @k2 :$$ ZSS) l2) =+  astFromVectorK (SNat @(k1 + k2) :$$ ZSS) $ l1 V.++ l2+astAppendS (Ast.AstReplicateK (SNat' @1 :$$ ZSS) a1)+           (Ast.AstFromVectorK (SNat @k2 :$$ ZSS) l2) =+  astFromVectorK (SNat @(1 + k2) :$$ ZSS) $ a1 `V.cons` l2+astAppendS (Ast.AstFromVectorK (SNat @k1 :$$ ZSS) l1)+           (Ast.AstReplicateK (SNat' @1 :$$ ZSS) a2) =+  astFromVectorK (SNat @(k1 + 1) :$$ ZSS) $ l1 `V.snoc` a2+astAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest1) a1)+           (Ast.AstReplicateS (SNat' @1 :$$ shmRest2) a2) =+  astFromVectorS (SNat @2 :$$ ZSS)+  $ V.fromList [astReplicateS shmRest1 a1, astReplicateS shmRest2 a2]+astAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest1) a1)+           (Ast.AstReplicateK (SNat' @1 :$$ shmRest2) a2) =+  astFromVectorS (SNat @2 :$$ ZSS)+  $ V.fromList [astReplicateK shmRest1 a1, astReplicateK shmRest2 a2]+astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)+           (Ast.AstAppendS (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) w) =+  astAppendS (astFromVectorS (SNat @(k1 + k2) :$$ ZSS) $ l1 V.++ l2) w+astAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest) a1)+           (Ast.AstAppendS (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) w) =+  astAppendS (astFromVectorS (SNat @(1 + k2) :$$ ZSS)+              $ astReplicateS shmRest a1 `V.cons` l2) w+astAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest) a1)+           (Ast.AstAppendS (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) w) =+  astAppendS (astFromVectorS (SNat @(1 + k2) :$$ ZSS)+              $ astReplicateK shmRest a1 `V.cons` l2) w+astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)+           (Ast.AstAppendS+              (Ast.AstReplicateS (SNat' @1 :$$ shmRest) a2) w) =+  astAppendS (astFromVectorS (SNat @(k1 + 1) :$$ ZSS)+              $ l1 `V.snoc` astReplicateS shmRest a2) w+astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)+           (Ast.AstAppendS+              (Ast.AstReplicateK (SNat' @1 :$$ shmRest) a2) w) =+  astAppendS (astFromVectorS (SNat @(k1 + 1) :$$ ZSS)+              $ l1 `V.snoc` astReplicateK shmRest a2) w+astAppendS (Ast.AstFromVectorK (SNat @k1 :$$ ZSS) l1)+           (Ast.AstAppendS (Ast.AstFromVectorK (SNat @k2 :$$ ZSS) l2) w) =+  astAppendS (astFromVectorK (SNat @(k1 + k2) :$$ ZSS) $ l1 V.++ l2) w+astAppendS (Ast.AstReplicateK (SNat' @1 :$$ ZSS) a1)+           (Ast.AstAppendS (Ast.AstFromVectorK (SNat @k2 :$$ ZSS) l2) w) =+  astAppendS (astFromVectorK (SNat @(1 + k2) :$$ ZSS) $ a1 `V.cons` l2) w+astAppendS (Ast.AstFromVectorK (SNat @k1 :$$ ZSS) l1)+           (Ast.AstAppendS (Ast.AstReplicateK (SNat' @1 :$$ ZSS) a2) w) =+  astAppendS (astFromVectorK (SNat @(k1 + 1) :$$ ZSS) $ l1 `V.snoc` a2) w+astAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest1) a1)+           (Ast.AstAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest2) a2) w) =+  astAppendS+    (astFromVectorS (SNat @2 :$$ ZSS)+     $ V.fromList [astReplicateS shmRest1 a1, astReplicateS shmRest2 a2]) w+astAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest1) a1)+           (Ast.AstAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest2) a2) w) =+  astAppendS+    (astFromVectorS (SNat @2 :$$ ZSS)+     $ V.fromList [astReplicateK shmRest1 a1, astReplicateK shmRest2 a2]) w+astAppendS (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) =+  fromPrimal $ astAppendS u v+astAppendS (Ast.AstFromDual u) (Ast.AstFromDual v) =+  fromDual $ astAppendS u v+astAppendS (Ast.AstFromPlain u) (Ast.AstFromPlain v) =+  fromPlain $ astAppendS u v+astAppendS u v | Just u0 <- unAstS u+               , Just v0 <- unAstS v+               , FTKS _ FTKScalar <- ftkAst u =+  fromPlain $ astConcreteS (tsappend (Concrete u0) (Concrete v0))+astAppendS u (Ast.AstAppendS v w) | Just u0 <- unAstS u+                                  , Just v0 <- unAstS v+                                  , FTKS _ FTKScalar <- ftkAst u =+  astAppendS (fromPlain $ astConcreteS (tsappend (Concrete u0) (Concrete v0))) w+astAppendS (Ast.AstAppendS v u) w = astAppendS v (astAppendS u w)+astAppendS u v = Ast.AstAppendS u v++astSliceS :: forall i n k sh s x. KnownSpan s+          => SNat i -> SNat n -> SNat k+          -> AstTensor AstMethodLet s (TKS2 (i + n + k ': sh) x)+          -> AstTensor AstMethodLet s (TKS2 (n ': sh) x)+astSliceS SZ SNat SZ v = v+astSliceS _ SZ _ v | FTKS (_ :$$ sh) x <- ftkAst v =+  let ftk = FTKS (SZ :$$ sh) x+  in fromPlain $ astConcrete ftk (tdefTarget ftk)+astSliceS SNat SNat SNat (Ast.AstFromVectorK (_ :$$ ZSS) l) =+  astFromVectorK (SNat @n :$$ ZSS) $ V.take (valueOf @n) $ V.drop (valueOf @i) l+astSliceS SNat SNat SNat (Ast.AstFromVectorS (_ :$$ ZSS) l) =+  astFromVectorS (SNat @n :$$ ZSS) $ V.take (valueOf @n) $ V.drop (valueOf @i) l+astSliceS SNat (SNat' @1) SNat v | FTKS (_ :$$ sh) _ <- ftkAst v =+  astReplicateS (SNat @1 :$$ ZSS) (astIndexS sh v (valueOf @i :.$ ZIS))+astSliceS SNat SNat SNat (Ast.AstReplicateK (_ :$$ shmRest) v) =+  astReplicateK (SNat @n :$$ shmRest) v+astSliceS SNat SNat SNat (Ast.AstReplicateS (_ :$$ shmRest) v) =+  astReplicateS (SNat @n :$$ shmRest) v+astSliceS SNat SNat SNat (Ast.AstGatherS shm shn shp v (var ::$ vars, ix)) =+  let varn = reboundsVarName (0, valueOf @n - 1) var+      ivar = valueOf @i + astVar varn+      ix2 = substituteAstIxS ivar var ix  -- cheap subst, because ivar is tiny+  in astGatherS (SNat @n :$$ shsTail shm) shn shp v (varn ::$ vars, ix2)+astSliceS i n@(SNat @n0) _k (Ast.AstAppendS v1 v2)+  | FTKS (m1@(SNat @m1) :$$ _) _ <- ftkAst v1+  , FTKS (m2@(SNat @m2) :$$ _) _ <- ftkAst v2 =+    let i1 = fromSNat' i `min` fromSNat' m1+        n1 = fromSNat' n `min` (fromSNat' m1 - i1)+        k1 = fromSNat' m1 - i1 - n1+        i2' = fromSNat' i `max` fromSNat' m1+        i2 = i2' - fromSNat' m1+        n2 = fromSNat' n - n1+        k2 = fromSNat' m2 - i2 - n2+    in withSNat i1 $ \si1@(SNat @i1) ->+       withSNat n1 $ \sn1@(SNat @n1) ->+       withSNat k1 $ \sk1@(SNat @k1) ->+       withSNat i2 $ \si2@(SNat @i2) ->+       withSNat n2 $ \sn2@(SNat @n2) ->+       withSNat k2 $ \sk2@(SNat @k2) ->+         gcastWith (unsafeCoerceRefl :: n1 + n2 :~: n0) $+         gcastWith (unsafeCoerceRefl :: i1 + n1 + k1 :~: m1) $+         gcastWith (unsafeCoerceRefl :: i2 + n2 + k2 :~: m2) $+         astAppendS (astSliceS si1 sn1 sk1 v1) (astSliceS si2 sn2 sk2 v2)+astSliceS i n k (Ast.AstSliceS i2 _n2 k2 v) =+  astSliceS (snatPlus i i2) n (snatPlus k k2) v+astSliceS i n k (Ast.AstReverseS v) = astReverseS (astSliceS k n i v)+-- This enlarges the term and increases computation, but sometimes+-- it permits eliminating the AstFromVectorS node altogether, so we risk it+-- for cases that commonly emerge from conditionals.+astSliceS i n@SNat k (Ast.AstTransposeS+                        perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil)+                        t@(Ast.AstFromVectorS (_ :$$ ZSS) l))+  | FTKS (SNat' @2 :$$ _ :$$ (_ :: ShS sh2)) _ <- ftkAst t =+    Ast.AstTransposeS perm+    $ astFromVectorS @_ @(n : sh2) (SNat @2 :$$ ZSS) (V.map (astSliceS i n k) l)+-- TODO: generalize (maybe the above, too) using unReplN, but it's hard.+-- TODO: does it really work only for replicate-like things in-between?+astSliceS i n@SNat k+          (Ast.AstTransposeS perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil) t)+  | STKS (snat :$$ _ :$$ _) _ <- ftkToSTK $ ftkAst t+  , Just u <- unRepl1 t =+    astTransposeS perm $ astReplicateS (snat :$$ ZSS) $ astSliceS i n k u+astSliceS i n k v1 = case v1 of+  Ast.AstLet var u v -> astLet var u (astSliceS i n k v)+  Ast.AstFromPrimal v -> fromPrimal $ astSliceS i n k v+  Ast.AstFromDual v -> fromDual $ astSliceS i n k v+  Ast.AstFromPlain v -> fromPlain $ astSliceS i n k v+  AstPlusS u v -> astSliceS i n k u `astPlusS` astSliceS i n k v+  AstTimesS u v -> astSliceS i n k u `astTimesS` astSliceS i n k v+  Ast.AstN1S opCode u -> astN1S opCode (astSliceS i n k u)+  Ast.AstR1S opCode u -> astR1S opCode (astSliceS i n k u)+  Ast.AstR2S opCode u v -> astR2S opCode (astSliceS i n k u)+                                         (astSliceS i n k v)+  Ast.AstI2S opCode u v -> astI2S opCode (astSliceS i n k u)+                                         (astSliceS i n k v)+  Ast.AstFloorS a -> astFloorS $ astSliceS i n k a+  Ast.AstFromIntegralS v -> astFromIntegralS $ astSliceS i n k v+  Ast.AstCastS v -> astCastS $ astSliceS i n k v+  AstConcreteS v -> astConcreteS (tsslice i n k $ Concrete v)+  Ast.AstCondS b a2 a3 -> astCondS b (astSliceS i n k a2) (astSliceS i n k a3)+  _ -> Ast.AstSliceS i n k v1++astReverseS :: forall n sh s r. KnownSpan s+            => AstTensor AstMethodLet s (TKS2 (n ': sh) r)+            -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)+astReverseS (Ast.AstLet var u v) = astLet var u (astReverseS v)+astReverseS (Ast.AstFromPrimal v) = fromPrimal $ astReverseS v+astReverseS (Ast.AstFromDual v) = fromDual $ astReverseS v+astReverseS (Ast.AstFromPlain v) = fromPlain $ astReverseS v+astReverseS (Ast.AstCondS b a2 a3) =+  astCondS b (astReverseS a2) (astReverseS a3)+astReverseS (Ast.AstFromVectorK shm@(_ :$$ ZSS) l) =+  astFromVectorK shm $ V.reverse l+astReverseS (Ast.AstFromVectorS shm@(_ :$$ ZSS) l) =+  astFromVectorS shm $ V.reverse l+astReverseS (Ast.AstReplicateK shm v) = astReplicateK shm v+astReverseS (Ast.AstReplicateS shm v) = astReplicateS shm v+astReverseS (Ast.AstGatherS shm@(SNat @k :$$ _) shn shp v (var ::$ vars, ix)) =+  let ivar = valueOf @k - 1 - astVar var+      ix2 = substituteAstIxS ivar var ix  -- cheap subst, because ivar is tiny+  in astGatherS shm shn shp v (var ::$ vars, ix2)+astReverseS (Ast.AstReverseS v) = v+astReverseS (Ast.AstTransposeS perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil) t)+  | STKS (snat :$$ (_ :$$ _)) _ <- ftkToSTK $ ftkAst t+  , Just u <- unRepl1 t =+    astTransposeS perm $ astReplicateS (snat :$$ ZSS) $ astReverseS u+astReverseS (AstConcreteS v) = astConcreteS (tsreverse $ Concrete v)+astReverseS v = Ast.AstReverseS v++-- TODO: try to completely cover the AstGatherS case here, which would permit+-- not expanding to astTransposeAsGatherS in astGatherCase+-- | Beware, this does not do full simplification, which often requires+-- the gather form, so astTransposeAsGather needs to be called in addition+-- if full simplification is required.+astTransposeS+  :: forall perm sh s r.+     (Permutation.IsPermutation perm, Rank perm <= Rank sh, KnownSpan s)+  => Permutation.Perm perm -> AstTensor AstMethodLet s (TKS2 sh r)+  -> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)+astTransposeS perm t =+  gcastWith (unsafeCoerceRefl+             :: TakeLen perm sh ++ DropLen perm sh :~: sh) $ case perm of+ PNil -> t+ PCons (SNat' @0) PNil ->+   gcastWith (unsafeCoerceRefl :: Permutation.PermutePrefix '[0] sh :~: sh)+   t+ _ | FTKS sh _ <- ftkAst t+   , Just u2 <- unReplN @_ @(DropLen perm sh) (shsTakeLenPerm perm sh) t ->+   astReplicateS (shsPermute perm (shsTakeLenPerm perm sh)) u2+ _ -> case t of+  _ | SNat' @0 `PCons` _ <- perm+    , STKS ((:$$) @n @sh2 snat _) _ <- ftkToSTK $ ftkAst t+    , Just u <- unRepl1 t -> case permUnShift1 perm of+      (perm2 :: Permutation.Perm perm2) ->+        fromMaybe (error "astTransposeS: impossible non-permutation")+        $ Permutation.permCheckPermutation perm2+        $ gcastWith (unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)+        -- for GHC 9.10 only:+        $ gcastWith (unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)+        $ gcastWith (unsafeCoerceRefl+                     :: Permutation.PermutePrefix perm (n : sh2)+                        :~: n : Permutation.PermutePrefix perm2 sh2)+        $ astReplicateS (snat :$$ ZSS) (astTransposeS perm2 u)++  Ast.AstLet var u v -> astLet var u (astTransposeS perm v)++  Ast.AstFromPrimal v -> fromPrimal $ astTransposeS perm v+  Ast.AstFromDual v -> fromDual $ astTransposeS perm v+  Ast.AstFromPlain v -> fromPlain $ astTransposeS perm v++  AstPlusS u v -> astTransposeS perm u `astPlusS` astTransposeS perm v+  AstTimesS u v -> astTransposeS perm u `astTimesS` astTransposeS perm v+  Ast.AstN1S opCode u -> astN1S opCode (astTransposeS perm u)+  Ast.AstR1S opCode u -> astR1S opCode (astTransposeS perm u)+  Ast.AstR2S opCode u v ->+    astR2S opCode (astTransposeS perm u) (astTransposeS perm v)+  Ast.AstI2S opCode u v ->+    astI2S opCode (astTransposeS perm u) (astTransposeS perm v)+  AstConcreteS v -> astConcreteS (tstranspose perm $ Concrete v)+--  Ast.AstFloorS v -> astFloorS $ astTransposeS perm v+--  Ast.AstFromIntegralS v -> astFromIntegralS $ astTransposeS perm v+--  Ast.AstCastS v -> astCastS $ astTransposeS perm v++  Ast.AstIndexS @shm shn v ix | n <- ixsRank ix ->+    Permutation.permFromListCont+      (Permutation.permToList'+       $ iterate (unsafeCoerce Permutation.permShift1) perm+         !! fromSNat' n)  -- this has a fake type, but that's fine+      $ \ (permn :: Perm permn) ->+        fromMaybe (error "astTransposeS: impossible non-permutation")+        $ Permutation.permCheckPermutation permn+        $ gcastWith (unsafeCoerceRefl+                     :: Permutation.PermutePrefix permn (shm ++ sh)+                        :~: shm ++ Permutation.PermutePrefix perm sh)+-- should suffice, but it doesn't+--      $ gcastWith (unsafeCoerceRefl :: Rank permn :~: n + Rank perm)+        $ gcastWith (unsafeCoerceRefl+                     :: (Rank permn <=? Rank (shm ++ sh)) :~: True)+        $ astIndexS (shsPermutePrefix perm shn) (astTransposeS permn v) ix+  Ast.AstCondS b u v -> astCondS b (astTransposeS perm u) (astTransposeS perm v)+  Ast.AstFromVectorS @_ @sh2 shm@(SNat @n :$$ ZSS) l+    | SNat' @0 `PCons` _ <- perm -> case permUnShift1 perm of+      (perm2 :: Permutation.Perm perm2) ->+        fromMaybe (error "astTransposeS: impossible non-permutation")+        $ Permutation.permCheckPermutation perm2+        $ gcastWith (unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)+        -- for GHC 9.10 only:+        $ gcastWith (unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)+        $ gcastWith (unsafeCoerceRefl+                     :: Permutation.PermutePrefix perm (n : sh2)+                        :~: n : Permutation.PermutePrefix perm2 sh2)+        $ astFromVectorS shm (V.map (astTransposeS perm2) l)+  Ast.AstSumS @shm2 shm2 v ->+    let zsuccP :: Permutation.Perm (Permutation.PermId (Rank shm2)+                                    ++ Permutation.MapPlusN (Rank shm2) perm)+        zsuccP = Permutation.permShiftN (shsRank shm2) perm+    in+      gcastWith (unsafeCoerceRefl+                 :: (Rank (Permutation.PermId (Rank shm2)+                           ++ Permutation.MapPlusN (Rank shm2) perm)+                      <=? Rank (shm2 ++ sh)) :~: True) $+      gcastWith (unsafeCoerceRefl+                 :: Permutation.PermutePrefix+                      (Permutation.PermId (Rank shm2)+                       ++ Permutation.MapPlusN (Rank shm2) perm) (shm2 ++ sh)+                    :~: shm2 ++ Permutation.PermutePrefix perm sh) $+      fromMaybe (error $ "astTransposeS: impossible non-permutation: "+                         ++ show (Permutation.permToList' zsuccP))+      $ Permutation.permCheckPermutation zsuccP+      $ astSumS shm2 $ astTransposeS zsuccP v+  Ast.AstScatterS @_ @shn @shp shm shn shp v (vars, ix)+    -- TODO: should the below be backpermute or permute?+    | fromSNat' (Permutation.permRank perm) <= length ix ->+        let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix perm shp)+            ix2 = ixsPermutePrefix perm ix+        in gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm shp ++ shn+                         :~: Permutation.PermutePrefix perm (shp ++ shn)) $+           astScatterS shm shn (shsPermutePrefix perm shp)+                       v (vars, ix2)+  -- This increases term size and work, so limited to size 2.+  -- TODO: generalize+  Ast.AstReplicateS shm@(_ :$$ ZSS)+                    (Ast.AstFromVectorK shm2@(SNat' @2 :$$ ZSS) l)+    | SNat' @1 `PCons` SNat' @0 `PCons` PNil <- perm ->+      astFromVectorS shm2 (V.map (astReplicateK shm) l)+  Ast.AstReplicateS shm@(_ :$$ ZSS)+                    (Ast.AstFromVectorS shm2@(SNat' @2 :$$ ZSS) l)+    | SNat' @1 `PCons` SNat' @0 `PCons` PNil <- perm ->+      astFromVectorS shm2 (V.map (astReplicateS shm) l)+  Ast.AstGatherS @shm @shn shm shn shp v (vars, ix)+    -- TODO: should the below be backpermute or permute?+    | fromSNat' (Permutation.permRank perm) <= length vars ->+        let vars2 :: AstVarListS (Permutation.PermutePrefix perm shm)+            vars2 = AstVarListS $ ixsPermutePrefix perm $ unAstVarListS vars+        in gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm shm ++ shn+                         :~: Permutation.PermutePrefix perm (shm ++ shn)) $+           astGatherS (shsPermutePrefix perm shm) shn shp+                      v (vars2, ix)+  Ast.AstAppendS u v | SNat' @0 `PCons` _ <- perm+                     , FTKS ((:$$) @m @sh2 _ _) _ <- ftkAst u+                     , FTKS ((:$$) @n _ _) _ <- ftkAst v ->+    gcastWith (unsafeCoerceRefl+               :: Permutation.PermutePrefix perm ((m + n) : sh2)+                  :~: m + n : Tail (Permutation.PermutePrefix+                                      perm (m : sh2))) $+    gcastWith (unsafeCoerceRefl+               :: Permutation.PermutePrefix perm ((m + n) : sh2)+                  :~: m + n : Tail (Permutation.PermutePrefix+                                      perm (n : sh2))) $+    astAppendS (astTransposeS perm u) (astTransposeS perm v)+  Ast.AstSliceS i n@(SNat @n) k u | SNat' @0 `PCons` _ <- perm+                                  , FTKS ((:$$) @ink @sh2 _ _) _ <- ftkAst u ->+    gcastWith (unsafeCoerceRefl+               :: Permutation.PermutePrefix perm (n : sh2)+                  :~: n : Tail (Permutation.PermutePrefix perm (ink : sh2))) $+    astSliceS i n k (astTransposeS perm u)+  Ast.AstReverseS u | SNat' @0 `PCons` _ <- perm ->+    astReverseS (astTransposeS perm u)+  Ast.AstTransposeS @_ @sh2 perm2 u | FTKS sh2 _ <- ftkAst u ->+    -- TODO: try to perform at type level+    let permV = Permutation.permToList' perm+        perm2V = Permutation.permToList' perm2+        perm2Matched =+          perm2V+          ++ take (length permV - length perm2V) (drop (length perm2V) [0 ..])+        perm3V = normalizePermutationHack+                 $ backpermutePrefixList permV perm2Matched+    in Permutation.permFromListCont perm3V $ \(perm3+                                               :: Permutation.Perm perm3) ->+      fromMaybe (error "astTransposeS: impossible non-permutation")+      $ Permutation.permCheckPermutation perm3+      $ gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix perm3 sh2+                      :~: Permutation.PermutePrefix perm sh) $+        case compare (length perm3V) (shsLength sh2) of+          LT -> gcastWith (unsafeCoerceRefl+                           :: Compare (Rank perm3) (Rank sh2) :~: LT) $+                astTransposeS perm3 u+          EQ -> gcastWith (unsafeCoerceRefl+                           :: Compare (Rank perm3) (Rank sh2) :~: EQ) $+                astTransposeS perm3 u+          GT -> error "astTransposeS: GT"+  u -> Ast.AstTransposeS perm u++-- TODO: try to cover the AstGatherS case here, which would permit+-- not expanding to astReshapeAsGatherS in astGatherCase+-- | Beware, this does not do full simplification, which often requires+-- the gather form, so astReshapeAsGather needs to be called in addition+-- if full simplification is required.+astReshapeS :: forall sh sh2 x s. (Product sh ~ Product sh2, KnownSpan s)+            => ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh x)+            -> AstTensor AstMethodLet s (TKS2 sh2 x)+astReshapeS sh2 t = case t of+  _ | Just u <- unRepl t+    , Refl <- lemAppNil @sh2 -> astReplicateS sh2 u+  _ | STKS ((:$$) @_ @sh1 k _) _ <- ftkToSTK $ ftkAst t+    , Just u <- unRepl1 t+    , (:$$) @_ @rest2 k2 rest2 <- sh2+    , Just Refl <- testEquality k k2 ->+      gcastWith (unsafeCoerceRefl :: Product rest2 :~: Product sh1) $+      astReplicateS (k :$$ ZSS) $ astReshapeS rest2 u+  Ast.AstLet var u v -> astLet var u (astReshapeS @_ @sh2 sh2 v)+  Ast.AstFromPrimal v -> fromPrimal $ astReshapeS sh2 v+  Ast.AstFromDual v -> fromDual $ astReshapeS sh2 v+  Ast.AstFromPlain v -> fromPlain $ astReshapeS sh2 v+  -- Reshaping can be costly, so we don't touch AstTimesS, etc.+  Ast.AstN1S opCode u -> astN1S opCode (astReshapeS @_ @sh2 sh2 u)+  Ast.AstR1S opCode u -> astR1S opCode (astReshapeS @_ @sh2 sh2 u)+  -- TODO: enable once we rewrite fromVectors with non-singleton shm:+  -- Ast.AstFromVectorK _shm l -> astFromVectorK sh2 l+  Ast.AstReshapeS _ v -> astReshapeS @_ @sh2 sh2 v+  _ | FTKS sh _ <- ftkAst t -> case testEquality sh sh2 of+    Just Refl -> t+    _ -> Ast.AstReshapeS sh2 t++astConvert+  :: KnownSpan s+  => TKConversion y z -> AstTensor AstMethodLet s y+  -> AstTensor AstMethodLet s z+astConvert c a | yftk <- ftkAst a = case (yftk, convertFTK c yftk) of+  (_, zftk) | Just Refl <- matchingFTK yftk zftk -> a+    -- this covers the ConvId case and more, so not simplifying c at worst+    -- causes c to take more memory but doesn't inhibit rewriting+  -- Below we heavily depend on c being semantically determined+  -- by the domain and codomain. We pick the simplest such c.+  (FTKS ZSS (FTKScalar @ry), FTKScalar @rz)+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      kfromS a+  (FTKR ZSR (FTKScalar @ry), FTKScalar @rz)+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      kfromR a+  (FTKX ZSX (FTKScalar @ry), FTKScalar @rz)+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      kfromX a+  (FTKScalar @ry, FTKS ZSS (FTKScalar @rz))+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      sfromK a+  (FTKR shr xy, FTKS sh@(_ :$$ _) xz)+    | Just Refl <- matchingFTK xy xz+    , Just Refl <- testEquality (shrRank shr) (shsRank sh) ->+      astConvDownSFromR sh xz a+  (FTKX shx xy, FTKS sh@(_ :$$ _) xz)+    | Just Refl <- matchingFTK xy xz+    , Just Refl <- testEquality (shxRank shx) (shsRank sh) ->+      astConvDownSFromX sh xz a+  (FTKS sh@(_ :$$ _) xy, FTKR shr xz)+    | Just Refl <- matchingFTK xy xz+    , Just Refl <- testEquality (shrRank shr) (shsRank sh) ->+      astConvUpRFromS sh xz a+  (FTKS sh@(_ :$$ _) xy, FTKX shx xz)+    | Just Refl <- matchingFTK xy xz+    , Just Refl <- testEquality (shxRank shx) (shsRank sh) ->+      astConvUpXFromS shx xz a+  (_, zftk) | Just c2 <- convDownMaybe yftk (ftkToSTK zftk) ->+    astConvertDown c2 zftk a+  (_, zftk) | Just c2 <- convUpMaybe yftk zftk ->+    astConvertUp c2 zftk a+  {- This has barely any effect:+  (FTKR ZSR (FTKScalar @ry), FTKS ZSS (FTKScalar @rz))+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      case kfromR a of+        Ast.AstConvert{} -> Ast.AstConvert c a  -- would likely loop+        b -> sfromK b+  (FTKX ZSX (FTKScalar @ry), FTKS ZSS (FTKScalar @rz))+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      case kfromX a of+        Ast.AstConvert{} -> Ast.AstConvert c a  -- would likely loop+        b -> sfromK b -}+  _ -> case a of  -- normalize somewhat even for, e.g., product to product+    Ast.AstFromPrimal v -> fromPrimal $ astConvert c v+    Ast.AstFromDual v -> fromDual $ astConvert c v+    Ast.AstFromPlain v -> fromPlain $ astConvert c v+    Ast.AstConvert c2 t2 -> astConvert (c `convCmp` c2) t2+    _ -> Ast.AstConvert c a++-- We are pushing conversions to shaped tensors down, into concrete values+-- and towards variables, so that the conversions often cancel out.+astConvertDown :: forall y z s. KnownSpan s+               => TKConversion y z -> FullShapeTK z+               -> AstTensor AstMethodLet s y+               -> AstTensor AstMethodLet s z+astConvertDown c zftk t = case (ftkAst t, zftk) of+  (yftk, _) | Just Refl <- matchingFTK yftk zftk -> t+  (FTKS ZSS (FTKScalar @ry), FTKScalar @rz) ->+    case testEquality (typeRep @ry) (typeRep @rz) of+      Just Refl -> astConvertDownKFromS c t+      Nothing -> error "astConvertDown: tensor kinds don't match"+  (FTKR ZSR (FTKScalar @ry), FTKScalar @rz) ->+    case testEquality (typeRep @ry) (typeRep @rz) of+      Just Refl -> astConvertDownKFromR c t+      Nothing -> error "astConvertDown: tensor kinds don't match"+  (FTKX ZSX (FTKScalar @ry), FTKScalar @rz) ->+    case testEquality (typeRep @ry) (typeRep @rz) of+      Just Refl -> astConvertDownKFromX c t+      Nothing -> error "astConvertDown: tensor kinds don't match"+  (FTKR shr yx, FTKS shz@(_ :$$ _) zx)->+    case (matchingFTK yx zx, testEquality (shsRank shz) (shrRank shr)) of+      (Just Refl, Just Refl) -> astConvertDownSFromR c shz zx t+      _ -> error "astConvertDown: tensor kinds don't match"+  (FTKX shy yx, FTKS shz@(_ :$$ _) zx) ->+    case (matchingFTK yx zx, testEquality (shsRank shz) (shxRank shy)) of+      (Just Refl, Just Refl) -> astConvertDownSFromX c shz zx t+      _ -> error "astConvertDown: tensor kinds don't match"+  (FTKProduct{}, _) -> case t of+    Ast.AstPair a1 a2 | FTKProduct zftk1 zftk2 <- zftk ->+      -- Here we can't always use the c the user presumably wrote,+      -- so we always create a canonical one.+      astPair (astConvDown zftk1 a1) (astConvDown zftk2 a2)+    Ast.AstLet var u v -> astLet var u (astConvertDown c zftk v)+    Ast.AstPrimalPart v -> astPrimalPart $ astConvertDown c zftk v+    Ast.AstDualPart v -> astDualPart $ astConvertDown c zftk v+    Ast.AstPlainPart v -> astPlainPart $ astConvertDown c zftk v+    Ast.AstFromPrimal v -> fromPrimal $ astConvertDown c zftk v+    Ast.AstFromDual v -> fromDual $ astConvertDown c zftk v+    Ast.AstFromPlain v -> fromPlain $ astConvertDown c zftk v+    Ast.AstConvert c2 t2 -> astConvert (c `convCmp` c2) t2+    _ -> Ast.AstConvert c t  -- don't introduce let just to push a conversion+  (yftk, _) ->+    error $ "astConvertDown: wrong tensor kinds: " ++ show (yftk, zftk, t)++astConvertDownKFromS+  :: forall r s. (GoodScalar r, KnownSpan s)+  => TKConversion (TKS '[] r) (TKScalar r)+  -> AstTensor AstMethodLet s (TKS '[] r)+  -> AstTensor AstMethodLet s (TKScalar r)+astConvertDownKFromS c a = case a of+  Ast.AstProject1{} -> Ast.AstConvert c a+  Ast.AstProject2{} -> Ast.AstConvert c a+  Ast.AstApply (AstLambda !var !v) ll ->+    astApply (AstLambda var (astConvertDownKFromS c v)) ll+  Ast.AstVar{} -> Ast.AstConvert c a+  Ast.AstLet var u v -> astLet var u (astConvertDownKFromS c v)+  Ast.AstPrimalPart v -> astPrimalPart $ astConvertDownKFromS c v+  Ast.AstDualPart v -> astDualPart $ astConvertDownKFromS c v+  Ast.AstPlainPart v -> astPlainPart $ astConvertDownKFromS c v+  Ast.AstFromPrimal v -> fromPrimal $ astConvertDownKFromS c v+  Ast.AstFromDual v -> fromDual $ astConvertDownKFromS c v+  Ast.AstFromPlain v -> fromPlain $ astConvertDownKFromS c v+  AstPlusS u v -> astConvertDownKFromS c u `astPlusK` astConvertDownKFromS c v+  AstTimesS u v -> astConvertDownKFromS c u `astTimesK` astConvertDownKFromS c v+  Ast.AstN1S opCode u -> astN1K opCode (astConvertDownKFromS c u)+  Ast.AstR1S opCode u -> astR1K opCode (astConvertDownKFromS c u)+  Ast.AstR2S opCode u v ->+    astR2K opCode (astConvertDownKFromS c u) (astConvertDownKFromS c v)+  Ast.AstI2S QuotOp u v ->+    astI2K QuotOp (astConvertDownKFromS c u) (astConvertDownKFromS c v)+  Ast.AstI2S RemOp u v ->+    astI2K RemOp (astConvertDownKFromS c u) (astConvertDownKFromS c v)+  AstConcreteS v -> AstConcreteK (Nested.sunScalar v)+  Ast.AstFloorS v -> astFloorK (kfromS v)+  Ast.AstFromIntegralS v -> astFromIntegralK (kfromS v)+  Ast.AstCastS v -> astCastK (kfromS v)+  Ast.AstArgMinS v | FTKS (_ :$$ ZSS) FTKScalar <- ftkAst v -> astArgMinK v+  Ast.AstArgMaxS v | FTKS (_ :$$ ZSS) FTKScalar <- ftkAst v -> astArgMaxK v+  Ast.AstIndexS @shm ZSS v ix | Refl <- lemAppNil @shm -> astIndexK v ix+  Ast.AstCondS b v1 v2 -> astCondK b (astConvertDownKFromS c v1)+                                     (astConvertDownKFromS c v2)+  Ast.AstFromVectorK ZSS l -> l V.! 0+  Ast.AstFromVectorS ZSS l -> astConvertDownKFromS c $ l V.! 0+  Ast.AstSumS _ v | Dict0 <- numFromTKAllNum (Proxy @r) -> astSumK v+  Ast.AstScatterS{} -> Ast.AstConvert c a  -- not a NF+  Ast.AstReplicateK ZSS v -> v+  Ast.AstReplicateS ZSS v -> astConvertDownKFromS c v+  Ast.AstGatherS{} -> Ast.AstConvert c a  -- not a NF+  Ast.AstTransposeS{} -> Ast.AstConvert c a  -- not a NF+  Ast.AstReshapeS{} -> Ast.AstConvert c a  -- not a NF+  Ast.AstDot1InS{} -> Ast.AstConvert c a  -- not a NF+  Ast.AstBoolNotS b -> astBoolNotK $ astConvertDownKFromS c b+  Ast.AstBoolAndS b1 b2 ->+    astBoolAndK (astConvertDownKFromS c b1) (astConvertDownKFromS c b2)+  Ast.AstLeqS _shb _sh arg1 arg2 -> astLeq arg1 arg2+  Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2++astConvertDownKFromR+  :: KnownSpan s+  => TKConversion (TKR 0 r) (TKScalar r)+  -> AstTensor AstMethodLet s (TKR 0 r)+  -> AstTensor AstMethodLet s (TKScalar r)+astConvertDownKFromR c a = case a of+  Ast.AstProject1{} -> Ast.AstConvert c a+  Ast.AstProject2{} -> Ast.AstConvert c a+  Ast.AstApply (AstLambda !var !v) ll ->+    astApply (AstLambda var (astConvertDownKFromR c v)) ll+  Ast.AstVar{} -> Ast.AstConvert c a+  Ast.AstLet var u v -> astLet var u (astConvertDownKFromR c v)+  Ast.AstPrimalPart v -> astPrimalPart $ astConvertDownKFromR c v+  Ast.AstDualPart v -> astDualPart $ astConvertDownKFromR c v+  Ast.AstPlainPart v -> astPlainPart $ astConvertDownKFromR c v+  Ast.AstFromPrimal v -> fromPrimal $ astConvertDownKFromR c v+  Ast.AstFromDual v -> fromDual $ astConvertDownKFromR c v+  Ast.AstFromPlain v -> fromPlain $ astConvertDownKFromR c v+  Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2++astConvertDownKFromX+  :: KnownSpan s+  => TKConversion (TKX '[] r) (TKScalar r)+  -> AstTensor AstMethodLet s (TKX '[] r)+  -> AstTensor AstMethodLet s (TKScalar r)+astConvertDownKFromX c a = case a of+  Ast.AstProject1{} -> Ast.AstConvert c a+  Ast.AstProject2{} -> Ast.AstConvert c a+  Ast.AstApply (AstLambda !var !v) ll ->+    astApply (AstLambda var (astConvertDownKFromX c v)) ll+  Ast.AstVar{} -> Ast.AstConvert c a+  Ast.AstLet var u v -> astLet var u (astConvertDownKFromX c v)+  Ast.AstPrimalPart v -> astPrimalPart $ astConvertDownKFromX c v+  Ast.AstDualPart v -> astDualPart $ astConvertDownKFromX c v+  Ast.AstPlainPart v -> astPlainPart $ astConvertDownKFromX c v+  Ast.AstFromPrimal v -> fromPrimal $ astConvertDownKFromX c v+  Ast.AstFromDual v -> fromDual $ astConvertDownKFromX c v+  Ast.AstFromPlain v -> fromPlain $ astConvertDownKFromX c v+  Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2++-- We are pushing conversions to shaped tensors down, into concrete values+-- and towards variables, so that the conversions often cancel out.+astConvertDownSFromR :: forall sh x s. KnownSpan s+                     => TKConversion (TKR2 (Rank sh) x) (TKS2 sh x)+                     -> ShS sh -> FullShapeTK x+                     -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)+                     -> AstTensor AstMethodLet s (TKS2 sh x)+astConvertDownSFromR c sh x a0 = case a0 of+  Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2+  Ast.AstProject1{} -> Ast.AstConvert c a0+  Ast.AstProject2{} -> Ast.AstConvert c a0+  Ast.AstApply{} -> Ast.AstConvert c a0+  Ast.AstVar{} -> Ast.AstConvert c a0+  Ast.AstBuild1 snat@SNat (STKR @n _ xstk) (var, a) -> case sh of+    snat2@SNat :$$ rest | Just Refl <- sameNat snat snat2+                        , Refl <- lemRankReplicate (Proxy @n) ->+      let c2 = convCmp (ConvXS' (FTKS rest x)) ConvRX+          !a2 = astConvert c2 a+      in Ast.AstBuild1 snat (STKS rest xstk) (var, a2)+    _ -> error "astConvertDownSFromR: impossible shape"+  Ast.AstLet var u v -> astLet var u (astConvertDownSFromR c sh x v)+  Ast.AstPrimalPart a -> astPrimalPart $ astConvertDownSFromR c sh x a+  Ast.AstDualPart a -> astDualPart $ astConvertDownSFromR c sh x a+  Ast.AstPlainPart a -> astPlainPart $ astConvertDownSFromR c sh x a+  Ast.AstFromPrimal a -> fromPrimal $ astConvertDownSFromR c sh x a+  Ast.AstFromDual a -> fromDual $ astConvertDownSFromR c sh x a+  Ast.AstFromPlain a -> fromPlain $ astConvertDownSFromR c sh x a++astConvertDownSFromX :: forall sh shx x s. (KnownSpan s, Rank shx ~ Rank sh)+                     => TKConversion (TKX2 shx x) (TKS2 sh x)+                     -> ShS sh -> FullShapeTK x+                     -> AstTensor AstMethodLet s (TKX2 shx x)+                     -> AstTensor AstMethodLet s (TKS2 sh x)+astConvertDownSFromX c sh x a0 = case a0 of+  Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2+  Ast.AstProject1{} -> Ast.AstConvert c a0+  Ast.AstProject2{} -> Ast.AstConvert c a0+  Ast.AstApply{} -> Ast.AstConvert c a0+  Ast.AstVar{} -> Ast.AstConvert c a0+  Ast.AstBuild1 snat@SNat (STKX @shx2 _ xstk) (var, a) -> case sh of+    (:$$) @_ @rest snat2@SNat rest | Just Refl <- sameNat snat snat2 ->+      -- This is needed only for GHC 9.10.+      gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank shx2) $+      let c2 = ConvXS' (FTKS rest x)+          !a2 = astConvert c2 a+      in Ast.AstBuild1 snat (STKS rest xstk) (var, a2)+    _ -> error "astConvertDownSFromX: impossible shape"+  Ast.AstLet var u v -> astLet var u (astConvertDownSFromX c sh x v)+  Ast.AstPrimalPart a -> astPrimalPart $ astConvertDownSFromX c sh x a+  Ast.AstDualPart a -> astDualPart $ astConvertDownSFromX c sh x a+  Ast.AstPlainPart a -> astPlainPart $ astConvertDownSFromX c sh x a+  Ast.AstFromPrimal a -> fromPrimal $ astConvertDownSFromX c sh x a+  Ast.AstFromDual a -> fromDual $ astConvertDownSFromX c sh x a+  Ast.AstFromPlain a -> fromPlain $ astConvertDownSFromX c sh x a++-- We are pulling conversions from shaped tensors up.+-- z is shaped or scalar or a product (presumably with some shaped components,+-- but it's not checked; the other components are supposed to be+-- converted identically, which is not checked in c, either).+astConvertUp+  :: KnownSpan s+  => TKConversion y z -> FullShapeTK z -> AstTensor AstMethodLet s y+  -> AstTensor AstMethodLet s z+astConvertUp c zftk t = case (ftkAst t, zftk, t) of+  (yftk, _, _) | Just Refl <- matchingFTK yftk zftk -> t+  -- Rare cases where we don't pull up but push down so that conversions+  -- don't end up interspersed with AstFromPrimal and similar.+  (_, _, Ast.AstFromPrimal v) ->+    fromPrimal $ astConvertUp c zftk v+  (_, _, Ast.AstFromDual v) ->+    fromDual $ astConvertUp c zftk v+  (_, _, Ast.AstFromPlain v )->+    fromPlain $ astConvertUp c zftk v+  (_, _, Ast.AstConvert c2 a2) ->+    astConvert (c `convCmp` c2) a2+  _ -> Ast.AstConvert c t++astConvDown :: forall y z s. KnownSpan s+            => FullShapeTK z -> AstTensor AstMethodLet s y+            -> AstTensor AstMethodLet s z+{-# INLINE astConvDown #-}+astConvDown zftk t = case convDownMaybe (ftkAst t) (ftkToSTK zftk) of+  Just c -> astConvertDown c zftk t+  Nothing -> error $ "astConvDown: unexpected types "+                     ++ "(" ++ show (ftkAst t) ++ ", " ++ show zftk ++ ")"++astConvDownSFromR :: forall sh s x. KnownSpan s+                  => ShS sh -> FullShapeTK x+                  -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)+                  -> AstTensor AstMethodLet s (TKS2 sh x)+astConvDownSFromR sh x | Refl <- lemRankReplicate (Proxy @(Rank sh)) =+  astConvertDownSFromR (ConvCmp (ConvXS' (FTKS sh x)) ConvRX) sh x++astConvDownSFromX :: forall sh sh' s x. (KnownSpan s, Rank sh ~ Rank sh')+                  => ShS sh -> FullShapeTK x+                  -> AstTensor AstMethodLet s (TKX2 sh' x)+                  -> AstTensor AstMethodLet s (TKS2 sh x)+astConvDownSFromX sh x = astConvertDownSFromX (ConvXS' (FTKS sh x)) sh x++astConvUp :: forall y z s. KnownSpan s+          => FullShapeTK z -> AstTensor AstMethodLet s y+          -> AstTensor AstMethodLet s z+{-# INLINE astConvUp #-}+astConvUp zftk t = case convUpMaybe (ftkAst t) zftk of+  Just c -> astConvertUp c zftk t+  Nothing -> error $ "astConvUp: unexpected types "+                     ++ "(" ++ show (ftkAst t) ++ ", " ++ show zftk ++ ")"++astConvUpRFromS :: forall sh x s. KnownSpan s+                => ShS sh -> FullShapeTK x+                -> AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)+astConvUpRFromS sh x t | Refl <- lemRankMapJust sh =+  astConvertUp+    (ConvCmp (ConvXR (ftkToSTK x)) ConvSX) (FTKR (shrFromShS sh) x) t++astConvUpXFromS :: forall sh sh' x s. (Rank sh ~ Rank sh', KnownSpan s)+                => IShX sh' -> FullShapeTK x+                -> AstTensor AstMethodLet s (TKS2 sh x)+                -> AstTensor AstMethodLet s (TKX2 sh' x)+astConvUpXFromS sh' x =+  gcastWith (unsafeCoerceRefl :: Rank (MapJust sh) :~: Rank sh) $+  astConvertUp (ConvCmp (ConvXX' (FTKX sh' x)) ConvSX) (FTKX sh' x)++astSumKContract :: forall shm r s. (NumScalar r, KnownSpan s)+                => AstTensor AstMethodLet s (TKS shm r)+                -> AstTensor AstMethodLet s (TKScalar r)+astSumKContract t0 = case t0 of+  Ast.AstLet var u v -> astLet var u (astSumKContract v)+  AstTimesS t1 t2 -> astDot0 t1 t2+  Ast.AstDot1InS _ _ t1 t2 -> astDot0 t1 t2+  Ast.AstMatmul2S m@SNat SNat p@SNat m1 m2 ->+    astDot0 (astTransposeS (Permutation.makePerm @'[1, 0])+                           (astReplicateS (p :$$ ZSS) m1))+            (astTransposeS (Permutation.makePerm @'[0, 2, 1])+                           (astReplicateS (m :$$ ZSS) m2))+  _ -> astSumK t0++astDot0 :: (NumScalar r, KnownSpan s)+         => AstTensor AstMethodLet s (TKS sh r)+         -> AstTensor AstMethodLet s (TKS sh r)+         -> AstTensor AstMethodLet s (TKScalar r)+astDot0 t1 t2 = case (t1, t2) of+  _ | Just u1 <- unRepl t1 ->+      kfromS u1 `astTimesK` astSumKContract t2+  _ | Just u2 <- unRepl t2 ->+      kfromS u2 `astTimesK` astSumKContract t1+  _ | FTKS (snat :$$ _) _ <- ftkAst t1+    , Just u1 <- unRepl1 t1+    , Just u2 <- unRepl1 t2 ->+      astTimesK (astDot0 u1 u2)+                (fromPlain $ AstConcreteK $ fromIntegral $ fromSNat' snat)+  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->+    fromPrimal $ astDot0 u1 u2+  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->+    fromDual $ astDot0 u1 u2+  (Ast.AstFromPlain u1, Ast.AstFromPlain u2) ->+    fromPlain $ astDot0 u1 u2+  (Ast.AstN1S NegateOp u1, Ast.AstN1S NegateOp u2) -> astDot0 u1 u2+  (u, v) | Just u0 <- unAstS u+         , Just v0 <- unAstS v ->+    withKnownShS (Nested.sshape u0) $+    fromPlain $ astConcreteK $ tsdot0 (Concrete u0) (Concrete v0)+  (Ast.AstFromVectorK _ l1, Ast.AstFromVectorK _ l2) | V.length l1 < 100 ->+    V.foldr (\(a1, a2) acc -> acc + a1 `astTimesK` a2) 0 $ V.zip l1 l2+  -- TODO: instead take the longest common prefix (sufix?) of shm1 and shm2:+  (Ast.AstSumS shm1 u1, Ast.AstSumS shm2 u2)+    | Just Refl <- testEquality shm1 shm2 ->+      astDot0 u1 u2+  {- KnownNat would be needed (or SNat):+  (Ast.AstAppendS @m1 u1 v1, Ast.AstAppendS @m2 u2 v2)+    | Just Refl <- sameNat (SNat @m1) (SNat @m2) ->+      astDot0 u1 u2 + astDot0 v1 v2 -}+  (Ast.AstReverseS u1, Ast.AstReverseS u2) -> astDot0 u1 u2+  (Ast.AstTransposeS @_ @sh1 perm1 u1, Ast.AstTransposeS @_ @sh2 perm2 u2)+    | Just Refl <- testEquality perm1 perm2 ->+      gcastWith (unsafeCoerceRefl :: sh1 :~: sh2) $+      astDot0 u1 u2+  _ -> case ftkAst t1 of+    FTKS ZSS FTKScalar -> kfromS t1 `astTimesK` kfromS t2+    FTKS (SNat' @1 :$$ ZSS) FTKScalar ->+      astTimesK (astIndexK t1 (0 :.$ ZIS)) (astIndexK t2 (0 :.$ ZIS))+    _ -> Ast.AstDot0 t1 t2++astDot1InS :: forall sh n r s. (NumScalar r, KnownSpan s)+           => ShS sh -> SNat n+           -> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)+           -> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)+           -> AstTensor AstMethodLet s (TKS sh r)+astDot1InS sh n@SNat t1 t2 = case (t1, t2) of+  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->+    fromPrimal $ astDot1InS sh n u1 u2+  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->+    fromDual $ astDot1InS sh n u1 u2+  (Ast.AstFromPlain u1, Ast.AstFromPlain u2) ->+    fromPlain $ astDot1InS sh n u1 u2+  (Ast.AstN1S NegateOp u1, Ast.AstN1S NegateOp u2) -> astDot1InS sh n u1 u2+  {- TODO: instead take the longest common prefix (sufix?) of shm1 and shm2.+  -- TODO: possibly this increases work, so wait until we have a generalization+  -- of tsdot1In.+  (Ast.AstSumS @shm1 @shn1 shm1 u1, Ast.AstSumS shm2 u2)+    | Just Refl <- testEquality shm1 shm2 ->+      gcastWith (unsafeCoerceRefl :: shm1 ++ sh ++ '[n] :~: shm1 ++ shn1) $+      astSumS shm1 $ astDot1InS (shm1 `shsAppend` sh) n u1 u2 -}+  (u, v) | Just u0 <- unAstS u+         , Just v0 <- unAstS v ->+    withKnownShS sh $+    fromPlain+    $ astConcreteS $ tsdot1In @_ @sh (SNat @n) (Concrete u0) (Concrete v0)+  _ -> case sh of+    ZSS -> sfromK $ astDot0 t1 t2+    snat :$$ shRest -> case (t1, t2) of+      _ | Just u1 <- unRepl1 t1+        , Just u2 <- unRepl1 t2 ->+          astReplicateS (snat :$$ ZSS) $ astDot1InS shRest n u1 u2+      (Ast.AstReverseS u1, Ast.AstReverseS u2) ->+        astReverseS $ astDot1InS sh n u1 u2+      _ -> Ast.AstDot1InS sh n t1 t2++astMatmul2S :: (NumScalar r, KnownSpan s)+            => SNat m -> SNat n -> SNat p+            -> AstTensor AstMethodLet s (TKS '[m, n] r)+            -> AstTensor AstMethodLet s (TKS '[n, p] r)+            -> AstTensor AstMethodLet s (TKS '[m, p] r)+astMatmul2S m@SNat n@SNat p@SNat t1 t2 = case (t1, t2) of+  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->+    fromPrimal $ astMatmul2S m n p u1 u2+  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->+    fromDual $ astMatmul2S m n p u1 u2+  (Ast.AstFromPlain u1, Ast.AstFromPlain u2) ->+    fromPlain $ astMatmul2S m n p u1 u2+  (u, v) | Just u0 <- unAstS u+         , Just v0 <- unAstS v ->+    fromPlain $ astConcreteS $ tsmatmul2 (Concrete u0) (Concrete v0)+  _ -> Ast.AstMatmul2S m n p t1 t2++astBoolNotK :: AstBool AstMethodLet -> AstBool AstMethodLet+astBoolNotK = \case+  Ast.AstLet var n b -> astLet var n (astBoolNotK b)+  AstConcreteK b -> AstConcreteK $ not b+  Ast.AstCondK b v w -> astCondK b (astBoolNotK v) (astBoolNotK w)+  AstLeqInt n k -> astLeqK (1 + k) n+  Ast.AstBoolNotK b -> b+  b -> Ast.AstBoolNotK b++astBoolNotS :: AstTensor AstMethodLet PlainSpan (TKS sh Bool)+            -> AstTensor AstMethodLet PlainSpan (TKS sh Bool)+astBoolNotS = Ast.AstBoolNotS++astBoolAndK :: AstBool AstMethodLet -> AstBool AstMethodLet+            -> AstBool AstMethodLet+astBoolAndK = \cases+  u v | eqY u v -> u+  (Ast.AstLet var n b) c -> astLetRefresh var n b $ \b' -> astBoolAndK b' c+  b (Ast.AstLet var n c) -> astLetRefresh var n c $ \c' -> astBoolAndK b c'+  (AstConcreteK True) b -> b+  (AstConcreteK False) _b -> AstConcreteK False+  b (AstConcreteK True) -> b+  _b (AstConcreteK False) -> AstConcreteK False+  (Ast.AstBoolAndK b c) d -> b `astBoolAndK` (c `astBoolAndK` d)+  -- Normal forms:+  b@(Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) c ->+    Ast.AstBoolAndK b c+  b@(Ast.AstLeqK AstConcreteK{} (Ast.AstN1K NegateOp Ast.AstVar{})) c ->+    Ast.AstBoolAndK b c+  b@(Ast.AstBoolNotK+       (Ast.AstBoolAndK (Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) _)) c ->+    Ast.AstBoolAndK b c+  b@(Ast.AstBoolNotK+       (Ast.AstBoolAndK+          (Ast.AstLeqK AstConcreteK{}+                       (Ast.AstN1K NegateOp Ast.AstVar{})) _)) c ->+    Ast.AstBoolAndK b c+  -- Rewriting to the normal forms by commutativity:+  b c@(Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) ->+     Ast.AstBoolAndK c b+  b c@(Ast.AstLeqK AstConcreteK{} (Ast.AstN1K NegateOp Ast.AstVar{})) ->+     Ast.AstBoolAndK c b+  b c@(Ast.AstBoolNotK+         (Ast.AstBoolAndK (Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) _)) ->+    Ast.AstBoolAndK c b+  b c@(Ast.AstBoolNotK+         (Ast.AstBoolAndK+            (Ast.AstLeqK AstConcreteK{}+                         (Ast.AstN1K NegateOp Ast.AstVar{})) _)) ->+    Ast.AstBoolAndK c b+  -- Rewriting to the normal forms by commutativity and associativity:+  b (Ast.AstBoolAndK+       c@(Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) d) ->+    astBoolAndK c (b `astBoolAndK` d)+  b (Ast.AstBoolAndK+       c@(Ast.AstLeqK AstConcreteK{}+                      (Ast.AstN1K NegateOp Ast.AstVar{})) d) ->+    astBoolAndK c (b `astBoolAndK` d)+  b (Ast.AstBoolAndK+       c@(Ast.AstBoolNotK+            (Ast.AstBoolAndK (Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) _)) d) ->+    astBoolAndK c (b `astBoolAndK` d)+  b (Ast.AstBoolAndK+       c@(Ast.AstBoolNotK+            (Ast.AstBoolAndK+               (Ast.AstLeqK AstConcreteK{}+                            (Ast.AstN1K NegateOp Ast.AstVar{})) _)) d) ->+    astBoolAndK c (b `astBoolAndK` d)+  b c -> Ast.AstBoolAndK b c++astBoolAndS :: AstTensor AstMethodLet PlainSpan (TKS sh Bool)+            -> AstTensor AstMethodLet PlainSpan (TKS sh Bool)+            -> AstTensor AstMethodLet PlainSpan (TKS sh Bool)+astBoolAndS = Ast.AstBoolAndS++astLeqK :: forall r. NumScalar r+        => AstTensor AstMethodLet PlainSpan (TKScalar r)+        -> AstTensor AstMethodLet PlainSpan (TKScalar r)+        -> AstBool AstMethodLet+astLeqK = \cases+  u v | eqY u v -> true+  u v | Just (u1, u2) <- bounds u+      , Just (v1, v2) <- bounds v+      , u2 <= v1 || u1 > v2 -> AstConcreteK (u2 <= v1)+  -- This is wrong, because LHS is a particular valuation and RHS is all.+  -- Ast.AstLet _ _ u <=. Ast.AstLet _ _ v -> u <=. v+  (Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astLeqK u' v+  u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astLeqK u v'+  u (AstPlusK (AstConcreteK v) w) ->+    astLeqK (astPlusK u (AstConcreteK $ negate v)) w+  (AstPlusK (AstConcreteK u) w) v ->+    astLeqK (AstConcreteK u) (astPlusK (negate w) v)+  u (AstConcreteK v) ->+    astLeqK (AstConcreteK (negate v)) (negate u)+  (AstConcreteK u) (AstTimesK (AstConcreteK v) w)+    | v > 0 && u >= 0+    , Just Refl <- testEquality (typeRep @r) (typeRep @Int) ->+      astLeqK (AstConcreteK ((u + v - 1) `quotH` v)) w+  (AstConcreteK u) (AstTimesK (AstConcreteK v) w)+    | v > 0 && u < 0+    , Just Refl <- testEquality (typeRep @r) (typeRep @Int) ->+      astLeqK (AstConcreteK (u `quotH` v)) w+  (AstConcreteK u) (AstTimesK (AstConcreteK v) w)+    | v < 0+    , Just Refl <- testEquality (typeRep @r) (typeRep @Int) ->+      astLeqK (AstConcreteK u)+              (AstTimesK (AstConcreteK $ negate v)  -- not astTimesK!+                         (negate w))+  u v -> Ast.AstLeqK u v++astLeq :: forall sh r. NumScalar r+       => AstTensor AstMethodLet PlainSpan (TKS sh r)+       -> AstTensor AstMethodLet PlainSpan (TKS sh r)+       -> AstBool AstMethodLet+astLeq = \cases+  u v | eqY u v -> true+  (AstConcreteS u) (AstConcreteS v) ->+    AstConcreteK $ Shaped.stoPrimitive u <= Shaped.stoPrimitive v+  u (AstPlusS (AstConcreteS v) w) ->+    astLeq (astPlusS u (AstConcreteS $ negate v)) w+  (AstPlusS (AstConcreteS u) w) v ->+    astLeq (AstConcreteS u) (astPlusS (negate w) v)+  u (AstConcreteS v) ->+    astLeq (AstConcreteS (negate v)) (negate u)+  (AstConvUpSFromK w) (AstConvUpSFromK v) -> astLeqK w v+  w (AstConvUpSFromK v) -> astLeqK (kfromS w) v+  (AstConvUpSFromK v) w -> astLeqK v (kfromS w)+  u v -> Ast.AstLeq u v++astLeqS :: forall shb sh r. NumScalar r+        => ShS shb -> ShS sh+        -> AstTensor AstMethodLet PlainSpan (TKS (shb ++ sh) r)+        -> AstTensor AstMethodLet PlainSpan (TKS (shb ++ sh) r)+        -> AstTensor AstMethodLet PlainSpan (TKS shb Bool)+astLeqS = Ast.AstLeqS+++-- * Scalar and shaped AstMethodLet AST instances++-- This instance is needed for unwinding in astConcrete.+instance KnownSpan s => ConvertTensor (AstTensor AstMethodLet s) where+  tconvert c _astk = astConvert c++  -- These are somewhat faster than their default implementations.+  kfromR = astConvertDownKFromR (ConvCmp ConvX0 ConvRX)+  kfromS = astConvertDownKFromS (ConvCmp ConvX0 ConvSX)+  kfromX = astConvertDownKFromX ConvX0+  sfromK = astConvertUp (ConvCmp ConvXS (Conv0X STKScalar)) (FTKS ZSS FTKScalar)++  rfromK = astConvertUp (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))+                        (FTKR ZSR FTKScalar)+  rfromS t = case ftkAst t of+    FTKS sh x -> astConvUpRFromS sh x t+  rfromX t = case ftkAst t of+    FTKX sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        astConvUpRFromS sh x $ astConvDownSFromX sh x t+  xfromR t = case ftkAst t of+    FTKR shr x ->+      withShsFromShR shr $ \(sh :: ShS sh) ->+        astConvUpXFromS (shCastSX knownShX sh) x+        $ astConvDownSFromR sh x t+  sfromR t = case ftkAst t of+    FTKR _ x -> astConvDownSFromR knownShS x t+  sfromX t = case ftkAst t of+    FTKX _ x -> astConvDownSFromX knownShS x t+  xfromK = astConvertUp (Conv0X STKScalar) (FTKX ZSX FTKScalar)+  xfromS t = case ftkAst t of+    FTKS sh x -> astConvUpXFromS (shCastSX knownShX sh) x t++  rzip @_ @_ @n a+   | Refl <- lemRankReplicate (Proxy @n) = case ftkAst a of+    FTKProduct (FTKR _sh y) (FTKR _ z) ->+      let c = convCmp+                (ConvXR (ftkToSTK (FTKProduct y z)))+                (convCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvRX ConvRX))+      in astConvert c a+  runzip @_ @_ @n a+   | Refl <- lemRankReplicate (Proxy @n) = case ftkAst a of+    FTKR _sh (FTKProduct y z) ->+      let c = convCmp+                (ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))+                (convCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvRX)+      in astConvert c a+  szip a = case ftkAst a of+    FTKProduct (FTKS _sh y) (FTKS _ z) ->+      let c = convCmp+                ConvXS+                (convCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvSX ConvSX))+      in astConvert c a+  sunzip a = case ftkAst a of+    FTKS _sh (FTKProduct y z) ->+      let c = convCmp+                (ConvT2 ConvXS ConvXS)+                (convCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvSX)+      in astConvert c a+  xzip a = case ftkAst a of+    FTKProduct (FTKX _sh y) (FTKX _ z) ->+      let c = ConvZip (ftkToSTK y) (ftkToSTK z)+      in astConvert c a+  xunzip a = case ftkAst a of+    FTKX _sh (FTKProduct y z) ->+      let c = ConvUnzip (ftkToSTK y) (ftkToSTK z)+      in astConvert c a++  xnestR @sh1 @m @x sh1 a+    | Refl <- lemRankReplicate (Proxy @m) =+      let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)+                            (TKX2 sh1 (TKR2 m x))+          c = convCmp+                (ConvXX (ConvXR (knownSTK @x)))+                (ConvNest @_ @_ @(Replicate m Nothing)+                          (STKX sh1 (knownSTK @x)))+      in astConvert c a+  xnestS @_ @_ @x sh1 a =+    let c = convCmp (ConvXX ConvXS)+                    (ConvNest (STKX sh1 (knownSTK @x)))+    in astConvert c a+  xnest @_ @_ @x sh1 a =+    let c = ConvNest (STKX sh1 (knownSTK @x))+    in astConvert c a+  xunNestR a =+    let c = convCmp ConvUnnest+                    (ConvXX ConvRX)+    in astConvert c a+  xunNestS a =+    let c = convCmp ConvUnnest+                    (ConvXX ConvSX)+    in astConvert c a+  xunNest a =+    let c = ConvUnnest+    in astConvert c a++  tpairConv = astPair+  tunpairConv t = (astProject1 t, astProject2 t)++-- These instances improve readability of rules.+instance (NumScalar r, KnownSpan s)+         => Num (AstTensor AstMethodLet s (TKScalar r)) where+  (+) = astPlusK+  (*) = astTimesK+  negate = astN1K NegateOp+  abs = astN1K AbsOp+  signum = astN1K SignumOp+  {-# INLINE fromInteger #-}+  fromInteger i = fromPlain $ AstConcreteK (fromInteger i)++instance (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)+         => IntegralH (AstTensor AstMethodLet s (TKScalar r)) where+  quotH = astI2K QuotOp+  remH = astI2K RemOp++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Fractional (AstTensor AstMethodLet s (TKScalar r)) where+  (/) = astR2K DivideOp+  recip = astR1K RecipOp+  {-# INLINE fromRational #-}+  fromRational r = fromPlain $ AstConcreteK (fromRational r)++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Floating (AstTensor AstMethodLet s (TKScalar r)) where+  pi = error "pi is not defined for tensors"+  exp = astR1K ExpOp+  log = astR1K LogOp+  sqrt = astR1K SqrtOp+  (**) = astR2K PowerOp+  logBase = astR2K LogBaseOp+  sin = astR1K SinOp+  cos = astR1K CosOp+  tan = astR1K TanOp+  asin = astR1K AsinOp+  acos = astR1K AcosOp+  atan = astR1K AtanOp+  sinh = astR1K SinhOp+  cosh = astR1K CoshOp+  tanh = astR1K TanhOp+  asinh = astR1K AsinhOp+  acosh = astR1K AcoshOp+  atanh = astR1K AtanhOp++instance (NumScalar r, Differentiable r, KnownSpan s)+         => RealFloatH (AstTensor AstMethodLet s (TKScalar r)) where+  atan2H = astR2K Atan2Op++instance (NumScalar r, KnownSpan s)+         => Num (AstTensor AstMethodLet s (TKS sh r)) where+  (+) = astPlusS+  (*) = astTimesS+  negate = astN1S NegateOp+  abs = astN1S AbsOp+  signum = astN1S SignumOp+  fromInteger i = error $ "fromInteger is not defined for shaped tensors: "+                          ++ show i++instance (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)+         => IntegralH (AstTensor AstMethodLet s (TKS sh r)) where+  quotH = astI2S QuotOp+  remH = astI2S RemOp++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Fractional (AstTensor AstMethodLet s (TKS sh r)) where+  (/) = astR2S DivideOp+  recip = astR1S RecipOp+  fromRational r = error $ "fromRational is not defined for shaped tensors: "+                           ++ show r++instance (NumScalar r, Differentiable r, KnownSpan s)+         => Floating (AstTensor AstMethodLet s (TKS sh r)) where+  pi = error "pi is not defined for tensors"+  exp = astR1S ExpOp+  log = astR1S LogOp+  sqrt = astR1S SqrtOp+  (**) = astR2S PowerOp+  logBase = astR2S LogBaseOp+  sin = astR1S SinOp+  cos = astR1S CosOp+  tan = astR1S TanOp+  asin = astR1S AsinOp+  acos = astR1S AcosOp+  atan = astR1S AtanOp+  sinh = astR1S SinhOp+  cosh = astR1S CoshOp+  tanh = astR1S TanhOp+  asinh = astR1S AsinhOp+  acosh = astR1S AcoshOp+  atanh = astR1S AtanhOp++instance (NumScalar r, Differentiable r, KnownSpan s)+         => RealFloatH (AstTensor AstMethodLet s (TKS sh r)) where+  atan2H = astR2S Atan2Op++instance Boolean (AstBool AstMethodLet) where+  true = AstConcreteK True+  false = AstConcreteK False+  notB = astBoolNotK+  (&&*) = astBoolAndK+  b ||* c = notB (notB b &&* notB c)++-- Since u and v are duplicated here, they need to be shared.+-- We share their difference, which would most likely appear in the+-- inequalities once they are rewritten, to ensure it's shared and whatever+-- vectorization substitutes into it is shared as well.+-- Otherwise, if u and v are variables, the sharing would vanish+-- before vectoriation complicates the expression a bit, making it+-- worth sharing.+instance (KnownSpan s, NumScalar r)+         => EqH (AstTensor AstMethodLet s) (TKScalar r) where+  u ==. v | eqY u v = true+  uUnshared ==. vUnshared = astLetFun (uUnshared - vUnshared) $ \uv ->+    0 <=. uv &&* uv <=. 0++instance (KnownSpan s, NumScalar r)+         => EqH (AstTensor AstMethodLet s) (TKS sh r) where+  u ==. v | eqY u v = true+  uUnshared ==. vUnshared = astLetFun (uUnshared - vUnshared) $ \uv ->+    let zero = fromPlain $ AstConcreteS $ defTargetRep $ ftkAst uUnshared+    in zero <=. uv &&* uv <=. zero++instance (KnownSpan s, NumScalar r)+         => OrdH (AstTensor AstMethodLet s) (TKScalar r) where+  u <=. v = fromPlain $ astLeqK (astPlainPart u) (astPlainPart v)++instance (KnownSpan s, NumScalar r)+         => OrdH (AstTensor AstMethodLet s) (TKS sh r) where+  u <=. v = fromPlain $ astLeq (astPlainPart u) (astPlainPart v)+++-- * Helper combinators++-- All but the last case are shortcuts for common forms.+astConcrete :: FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PlainSpan y+{-# INLINE astConcrete #-}+astConcrete ftk v = case ftk of+  FTKScalar -> astConcreteK v+  FTKR ZSR FTKScalar ->+    rfromK $ AstConcreteK $ Nested.runScalar $ unConcrete v+  FTKR sh' FTKScalar ->+    withShsFromShR sh' $ \sh ->+      withKnownShS sh $+      astConvUpRFromS sh FTKScalar $ astConcreteS $ sfromR v+  FTKS ZSS FTKScalar ->+    sfromK $ AstConcreteK $ Nested.sunScalar $ unConcrete v+  FTKS _ FTKScalar -> astConcreteS v+  FTKX ZSX FTKScalar ->+    xfromK $ AstConcreteK $ Nested.munScalar $ unConcrete v+  FTKX sh' FTKScalar ->+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      withKnownShS sh $+      astConvUpXFromS sh' FTKScalar+      $ astConcreteS $ sfromX @_ @sh v+  FTKProduct ftk1 ftk2 ->+    astPair (astConcrete ftk1 (tproject1 v))+            (astConcrete ftk2 (tproject2 v))+  _ -> concreteTarget astConcreteK astConcreteS+                      (\sh -> astConvUpRFromS sh FTKScalar)+                      (\sh' -> astConvUpXFromS sh' FTKScalar)+                      (ftkToSTK ftk) v++-- INLINE here would bloat the binary a lot, probably negating any+-- gains from directly calling the function. Also, this is not a bottleneck.+-- It would also trigger specialization of many other functions from this+-- module, which would fail, because their unfolding is not exposed.+-- To some extend, this is even happening with no pragma (wrappers and+-- similar functions derived from this one fail to specialize in other modules).+astLetFun :: forall y z s s2. (KnownSpan s, KnownSpan s2)+          => AstTensor AstMethodLet s y+          -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)+          -> AstTensor AstMethodLet s2 z+{-# NOINLINE astLetFun #-}+astLetFun a f | astIsSmall True a = f a+astLetFun a f = case a of+  AstConvUp c ftkz v ->+    unsafePerformIO $ do+      var <- funToAstNoBoundsIO (ftkAst v)+      pure $! astLet var v (f $ astConvertUp c ftkz $ astVar var)+  Ast.AstFromPrimal v -> astLetFun v (f . fromPrimal)+  Ast.AstFromDual v -> astLetFun v (f . fromDual)+  Ast.AstFromPlain v -> astLetFun v (f . fromPlain)+  _ -> unsafePerformIO $ case ftkAst a of+    ftk@FTKScalar -> do+        var <- funToAstAutoBoundsIO ftk a+        pure $! astLet var a (f $ astVar var)+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> do+        let v = astConvDownSFromR sh x a+        var <- funToAstNoBoundsIO (FTKS sh x)+        pure $! astLet var v (f $ astConvUpRFromS sh x $ astVar var)+          -- safe, because subsitution ruled out above+    FTKX sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> do+        let v = astConvDownSFromX sh x a+        var <- funToAstNoBoundsIO (FTKS sh x)+        pure $! astLet var v (f $ astConvUpXFromS sh' x $ astVar var)+    FTKS ZSS x@FTKScalar -> do+        let v = kfromS a+        var <- funToAstAutoBoundsIO x v+        pure $! astLet var v (f $ sfromK $ astVar var)+    -- calling recursively for product may be not worth it+    ftk -> do+        var <- funToAstNoBoundsIO ftk+        pure $! astLet var a (f $ astVar var)++astLetRefresh :: forall y x z s s2 s3. (KnownSpan s2, KnownSpan s3)+              => AstVarName '(s, y) -> AstTensor AstMethodLet s y+              -> AstTensor AstMethodLet s3 x+              -> (AstTensor AstMethodLet s3 x -> AstTensor AstMethodLet s2 z)+              -> AstTensor AstMethodLet s2 z+astLetRefresh var u v f =+  withKnownSpan (varNameToSpan var) $+  astLetFun u $ \ !ast1 -> f (substituteAst ast1 var v)++unRepl1 :: forall n sh x s. KnownSpan s+        => AstTensor AstMethodLet s (TKS2 (n ': sh) x)+        -> Maybe (AstTensor AstMethodLet s (TKS2 sh x))+unRepl1 (Ast.AstLet var u t) =+  Ast.AstLet var u <$> unRepl1 t  -- we may be before inlining+unRepl1 (Ast.AstPrimalPart t) = astPrimalPart <$> unRepl1 t+unRepl1 (Ast.AstDualPart t) = astDualPart <$> unRepl1 t+unRepl1 (Ast.AstPlainPart t) = astPlainPart <$> unRepl1 t+unRepl1 (Ast.AstFromPrimal t) = fromPrimal <$> unRepl1 t+unRepl1 (Ast.AstFromDual t) = fromDual <$> unRepl1 t+unRepl1 (Ast.AstFromPlain t) = fromPlain <$> unRepl1 t+unRepl1 (AstConcreteS a) = AstConcreteS <$> sunReplicate1 a+unRepl1 (Ast.AstReplicateK (_ :$$ shmRest) u) = Just $ astReplicateK shmRest u+unRepl1 (Ast.AstReplicateS (_ :$$ shmRest) u) = Just $ astReplicateS shmRest u+unRepl1 _ = Nothing++-- The result must not be equal to the argument.+unRepl :: forall sh x s. KnownSpan s+       => AstTensor AstMethodLet s (TKS2 sh x)+       -> Maybe (AstTensor AstMethodLet s (TKS2 '[] x))+-- This is too costly and not needed in all the places where unRepl is used,+-- hence the restriction to different result and argument:+-- unRepl t | FTKS ZSS _ <- ftkAst t = Just t+unRepl (Ast.AstLet var u t) = Ast.AstLet var u <$> unRepl t+unRepl (Ast.AstPrimalPart t) = astPrimalPart <$> unRepl t+unRepl (Ast.AstDualPart t) = astDualPart <$> unRepl t+unRepl (Ast.AstPlainPart t) = astPlainPart <$> unRepl t+unRepl (Ast.AstFromPrimal t) = fromPrimal <$> unRepl t+unRepl (Ast.AstFromDual t) = fromDual <$> unRepl t+unRepl (Ast.AstFromPlain t) = fromPlain <$> unRepl t+unRepl (AstConcreteS a) | _ :$$ _ <- Nested.sshape a =+  AstConcreteS . Nested.sscalar <$> sunReplicate a+unRepl (Ast.AstReplicateK _ u) = Just $ sfromK u+unRepl (Ast.AstReplicateS _ u) = unRepl u+unRepl _ = Nothing++-- The result must not be equal to the argument.+unReplN :: forall shm shn x s. KnownSpan s+        => ShS shm -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+        -> Maybe (AstTensor AstMethodLet s (TKS2 shn x))+unReplN ZSS a = Just a+unReplN (_ :$$ sh) (Ast.AstReplicateK (_ :$$ shmRest) u) =+  unReplN sh (astReplicateK shmRest u)+unReplN (_ :$$ sh) (Ast.AstReplicateS (_ :$$ shmRest) u) =+  unReplN sh (astReplicateS shmRest u)+unReplN shm (Ast.AstLet var u t) = Ast.AstLet var u <$> unReplN shm t+unReplN shm (Ast.AstPrimalPart t) = astPrimalPart <$> unReplN shm t+unReplN shm (Ast.AstDualPart t) = astDualPart <$> unReplN shm t+unReplN shm (Ast.AstPlainPart t) = astPlainPart <$> unReplN shm t+unReplN shm (Ast.AstFromPrimal t) = fromPrimal <$> unReplN shm t+unReplN shm (Ast.AstFromDual t) = fromDual <$> unReplN shm t+unReplN shm (Ast.AstFromPlain t) = fromPlain <$> unReplN shm t+unReplN shm (AstConcreteS a) = AstConcreteS <$> sunReplicateN shm a+unReplN _ _ = Nothing+++-- * Substitution wrappers++-- | We assume no variable is shared between a binding and its nested binding+-- and nobody substitutes into variables that are bound.+-- This keeps the substitution code simple, because we never need to compare+-- variables to any variable in the bindings.+substituteAst :: forall s s2 y z. KnownSpan s+              => AstTensor AstMethodLet s2 z -> AstVarName '(s2, z)+              -> AstTensor AstMethodLet s y+              -> AstTensor AstMethodLet s y+substituteAst i var v1 =+  fromMaybe v1 $ substitute1Ast i var v1++substituteAstIxS+  :: AstTensor AstMethodLet s y -> AstVarName '(s, y) -> AstIxS AstMethodLet sh+  -> AstIxS AstMethodLet sh+substituteAstIxS i var ix =+  fromMaybe ix $ substitute1AstIxS i var ix+++-- * Substitution workers++-- Invariant: if the variable has bounds, the expression can only have+-- values within the bounds (regardless of what the `bounds` call would say).+substitute1Ast :: forall s s2 y z. KnownSpan s+               => AstTensor AstMethodLet s2 z -> AstVarName '(s2, z)+               -> AstTensor AstMethodLet s y+               -> Maybe (AstTensor AstMethodLet s y)+substitute1Ast i var = subst where+ substIxS :: AstIxS AstMethodLet sh3 -> Maybe (AstIxS AstMethodLet sh3)+ substIxS ix =+   let mix = fmap subst ix+   in if any isJust mix+      then Just $ ixsZipWith fromMaybe ix mix+      else Nothing+ subst :: forall s3 y2. KnownSpan s3+       => AstTensor AstMethodLet s3 y2 -> Maybe (AstTensor AstMethodLet s3 y2)+ subst = \case+  Ast.AstPair u v ->+    case (subst u, subst v) of+      (Nothing, Nothing) -> Nothing+      (mu, mv) -> Just $ astPair (fromMaybe u mu) (fromMaybe v mv)+  Ast.AstProject1 a -> astProject1 <$> subst a+  Ast.AstProject2 a -> astProject2 <$> subst a+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+      case ( substitute1AstHFun i var f, substitute1AstHFun i var df+           , substitute1AstHFun i var rf, subst acc0+           , subst es ) of+        (Nothing, Nothing, Nothing, Nothing, Nothing) -> Nothing+        (mf, mdf, mrf, macc0, mes) ->+          Just $ astMapAccumLDer k bftk eftk+                                 (fromMaybe f mf)+                                 (fromMaybe df mdf)+                                 (fromMaybe rf mrf)+                                 (fromMaybe acc0 macc0)+                                 (fromMaybe es mes)+  Ast.AstApply t ll ->+    case ( substitute1AstHFun i var t+         , subst ll ) of+      (Nothing, Nothing) -> Nothing+      (mt, mll) -> Just $ astApply (fromMaybe t mt) (fromMaybe ll mll)+  Ast.AstVar var2 ->+    -- We can't assert anything here about bounds, because only runtime concrete+    -- concrete values need to be in bounds and bounds approximations+    -- don't have to agree.+    if varNameToAstVarId var == varNameToAstVarId var2+    then case testEquality var var2 of+      Just Refl -> case (i, var2) of+        (Ast.AstVar var3, AstVarName _ (FtkAndBoundsBounds lb2 ub2)) ->+          let bs = case var3 of+                AstVarName _ (FtkAndBoundsBounds lb3 ub3) ->+                  (max lb2 lb3, min ub2 ub3)+                    -- We know all bounds approximations have to be correct+                    -- so we can intersect them. We ignore approximations+                    -- from @var, because they should be equal to those of var2,+                    -- because we never shadow variables nor update their types.+                _ -> (lb2, ub2)+          in Just $ astVar $ reboundsVarName bs var3+        _ -> Just i+      _ -> error $ "substitute1Ast: kind of the variable "+                   ++ show var2 ++ ": " ++ show (varNameToFTK var)+                   ++ ", payload kind: " ++ show (varNameToFTK var2)+                   ++ ", payload: " ++ show i+    else Nothing+  Ast.AstBuild1 k stk (var2, v) ->+    assert (varNameToAstVarId var2 /= varNameToAstVarId var) $+    case subst v of+      Just !v2 -> Just $ Ast.AstBuild1 k stk (var2, v2)+      Nothing -> Nothing++  Ast.AstLet var2 u v ->+    assert (varNameToAstVarId var2 /= varNameToAstVarId var) $+    case (withKnownSpan (varNameToSpan var2) $ subst u, subst v) of+      (Nothing, Nothing) -> Nothing+      (mu, mv) -> Just $ astLet var2 (fromMaybe u mu) (fromMaybe v mv)++  Ast.AstPrimalPart a -> astPrimalPart <$> subst a+  Ast.AstDualPart a -> astDualPart <$> subst a+  Ast.AstPlainPart a -> astPlainPart <$> subst a+  Ast.AstFromPrimal a -> fromPrimal <$> subst a+  Ast.AstFromDual a -> fromDual <$> subst a+  Ast.AstFromPlain a -> fromPlain <$> subst a++  AstPlusK u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astPlusK (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  AstTimesK u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astTimesK (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstN1K opCode u -> astN1K opCode <$> subst u+  Ast.AstR1K opCode u -> astR1K opCode <$> subst u+  Ast.AstR2K opCode u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astR2K opCode (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstI2K QuotOp u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astI2K QuotOp (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstI2K RemOp u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astI2K RemOp (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  AstConcreteK{} -> Nothing+  Ast.AstFloorK a -> astFloorK <$> subst a+  Ast.AstFromIntegralK v -> astFromIntegralK <$> subst v+  Ast.AstCastK v -> astCastK <$> subst v+  Ast.AstArgMinK v -> astArgMinK <$> subst v+  Ast.AstArgMaxK v -> astArgMaxK <$> subst v+  Ast.AstIndexK v ix ->+    case (subst v, substIxS ix) of+      (Nothing, Nothing) -> Nothing+      (mv, mix) -> Just $ astIndexK (fromMaybe v mv) (fromMaybe ix mix)++  AstPlusS u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astPlusS (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  AstTimesS u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astTimesS (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstN1S opCode u -> astN1S opCode <$> subst u+  Ast.AstR1S opCode u -> astR1S opCode <$> subst u+  Ast.AstR2S opCode u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astR2S opCode (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstI2S QuotOp u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astI2S QuotOp (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstI2S RemOp u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astI2S RemOp (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  AstConcreteS{} -> Nothing+  Ast.AstFloorS a -> astFloorS <$> subst a+  Ast.AstFromIntegralS a -> astFromIntegralS <$> subst a+  Ast.AstCastS v -> astCastS <$> subst v+  Ast.AstArgMinS a -> Ast.AstArgMinS <$> subst a+  Ast.AstArgMaxS a -> Ast.AstArgMaxS <$> subst a+  Ast.AstIndexS shn t ix ->+    case (subst t, substIxS ix) of+      (Nothing, Nothing) -> Nothing+      (mt, mix) -> Just $ astIndexS shn (fromMaybe t mt) (fromMaybe ix mix)++  Ast.AstCondK b v w ->+    case (subst b, subst v, subst w) of+      (Nothing, Nothing, Nothing) -> Nothing+      (mb, mv, mw) ->+        Just+        $ astCondKInitial (fromMaybe b mb) (fromMaybe v mv) (fromMaybe w mw)+  Ast.AstCondS b v w ->+    case (subst b, subst v, subst w) of+      (Nothing, Nothing, Nothing) -> Nothing+      (mb, mv, mw) ->+        Just+        $ astCondSInitial (fromMaybe b mb) (fromMaybe v mv) (fromMaybe w mw)+  Ast.AstFromVectorK shm args ->+    let margs = V.map subst args+    in if V.any isJust margs+       then Just $ astFromVectorK shm $ V.zipWith fromMaybe args margs+       else Nothing+  Ast.AstFromVectorS shm args ->+    let margs = V.map subst args+    in if V.any isJust margs+       then Just $ astFromVectorS shm $ V.zipWith fromMaybe args margs+       else Nothing+  Ast.AstSumK v -> astSumK <$> subst v+  Ast.AstSumS shm v -> astSumS shm <$> subst v+  Ast.AstScatterS shm shn shp t (vars, ix) ->+    assert (all (\v -> varNameToAstVarId var /= varNameToAstVarId v) vars) $+    case (subst t, substIxS ix) of+      (Nothing, Nothing) -> Nothing+      (mt, mix) -> Just $ astScatterS shm shn shp+                                      (fromMaybe t mt)+                                      (vars, fromMaybe ix mix)+  Ast.AstReplicateK shm v -> astReplicateK shm <$> subst v+  Ast.AstReplicateS shm v -> astReplicateS shm <$> subst v+  Ast.AstGatherS shm shn shp t (vars, ix) ->+    assert (all (\v -> varNameToAstVarId var /= varNameToAstVarId v) vars) $+    case (subst t, substIxS ix) of+      (Nothing, Nothing) -> Nothing+      (mt, mix) -> Just $ astGatherS shm shn shp+                                     (fromMaybe t mt)+                                     (vars, fromMaybe ix mix)+  Ast.AstIotaS{} -> Nothing+  Ast.AstAppendS x y ->+    case (subst x, subst y) of+      (Nothing, Nothing) -> Nothing+      (mx, my) -> Just $ astAppendS (fromMaybe x mx) (fromMaybe y my)+  Ast.AstSliceS i2 n k v -> astSliceS i2 n k <$> subst v+  Ast.AstReverseS v -> astReverseS <$> subst v+  Ast.AstTransposeS perm v -> astTransposeS perm <$> subst v+  Ast.AstReshapeS sh v -> astReshapeS sh <$> subst v++  Ast.AstConvert c v -> astConvert c <$> subst v++  Ast.AstDot0 u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astDot0 (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstDot1InS sh n u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astDot1InS sh n (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstMatmul2S m n p u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astMatmul2S m n p (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstBoolNotK arg -> astBoolNotK <$> subst arg+  Ast.AstBoolNotS arg -> astBoolNotS <$> subst arg+  Ast.AstBoolAndK arg1 arg2 ->+    let mb1 = subst arg1+        mb2 = subst arg2+    in if isJust mb1 || isJust mb2+       then Just $ astBoolAndK (fromMaybe arg1 mb1) (fromMaybe arg2 mb2)+       else Nothing+  Ast.AstBoolAndS arg1 arg2 ->+    let mb1 = subst arg1+        mb2 = subst arg2+    in if isJust mb1 || isJust mb2+       then Just $ astBoolAndS (fromMaybe arg1 mb1) (fromMaybe arg2 mb2)+       else Nothing+  Ast.AstLeqK arg1 arg2 ->+    let mr1 = subst arg1+        mr2 = subst arg2+    in if isJust mr1 || isJust mr2+       then Just $ astLeqK (fromMaybe arg1 mr1) (fromMaybe arg2 mr2)+       else Nothing+  Ast.AstLeq arg1 arg2 ->+    let mr1 = subst arg1+        mr2 = subst arg2+    in if isJust mr1 || isJust mr2+       then Just $ astLeq (fromMaybe arg1 mr1) (fromMaybe arg2 mr2)+       else Nothing+  Ast.AstLeqS shb sh arg1 arg2 ->+    let mr1 = subst arg1+        mr2 = subst arg2+    in if isJust mr1 || isJust mr2+       then Just $ astLeqS shb sh (fromMaybe arg1 mr1) (fromMaybe arg2 mr2)+       else Nothing++substitute1AstIxS+  :: AstTensor AstMethodLet s2 y -> AstVarName '(s2, y)+  -> AstIxS AstMethodLet sh+  -> Maybe (AstIxS AstMethodLet sh)+substitute1AstIxS i var ix =+  let mix = fmap (substitute1Ast i var) ix+  in if any isJust mix+     then Just $ ixsZipWith fromMaybe ix mix+     else Nothing++substitute1AstHFun+  :: forall s s3 x y z.+     AstTensor AstMethodLet s3 z -> AstVarName '(s3, z) -> AstHFun s x y+  -> Maybe (AstHFun s x y)+substitute1AstHFun _i var (AstLambda var2 _) =+  assert (varNameToAstVarId var2 /= varNameToAstVarId var)+  Nothing  -- no outside free variables
src/HordeAd/Core/AstTools.hs view
@@ -1,44 +1,43 @@-{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE LambdaCase, ViewPatterns #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} -- | An assortment of operations working on AST of the code to be differentiated -- or the code resulting from differentiation. module HordeAd.Core.AstTools   ( -- * Full tensor kind derivation-    ftkAst, isTensorInt+    ftkAst, stkAstX     -- * Variable occurrence detection-  , varInAst, varInAstBool, varInIxS, varNameInAst, varNameInIxS-    -- * Determining if a term is too small to require sharing-  , astIsSmall, ixIsSmall+  , varInAst, varInIxS, varNameInAst, varNameInIxS+    -- * Tools related to sharing+  , astIsSmall, ixIsSmall, astLetDown, astVar, astShare     -- * Odds and ends-  , bounds-  , liftRFromS1, liftRFromS2, liftXFromS1, liftXFromS2-  , cAstConvert, cAstSFromR, cAstSFromX, cAstXFromS-  , pattern AstFromS', checkAstFromS, checkFtkAstFromS, checkFtkAstSFrom-  , cAstFromS, cAstSFrom, convFromS, convSFrom+  , bounds, intBounds+  , pattern AstConvUpSFromK, pattern AstConvUp, AstConvUpMaybe(..)+  , convDownMaybe, convUpMaybe   , setTotalSharing   ) where -import Prelude hiding (foldl')+import Prelude  import Control.Exception.Assert.Sugar-import Data.Int (Int64)+import Data.Bifunctor (second) import Data.IORef-import Data.Maybe (isJust)+import Data.Maybe (fromMaybe) import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (testEquality, (:~:) (Refl))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl)) import Data.Vector.Generic qualified as V import System.IO.Unsafe (unsafePerformIO)-import Type.Reflection (typeRep)+import Type.Reflection (Typeable, typeRep) +import Data.Array.Nested (type (++)) import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX) import Data.Array.Nested.Lemmas import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (snatPlus)+import Data.Array.Nested.Types (fromSNat', snatPlus, unsafeCoerceRefl)  import HordeAd.Core.Ast+import HordeAd.Core.Conversion import HordeAd.Core.TensorKind import HordeAd.Core.Types @@ -56,18 +55,10 @@     FTKProduct ftk1 _ -> ftk1   AstProject2 v -> case ftkAst v of     FTKProduct _ ftk2 -> ftk2-  AstFromVector snat _ l -> case V.uncons l of-    Nothing -> error "ftkAst: empty vector"-    Just (v, _) -> buildFTK snat (ftkAst v)-  AstSum snat stk v -> razeFTK snat stk (ftkAst v)-  AstReplicate snat _ v -> buildFTK snat (ftkAst v)-  AstMapAccumRDer k bftk _eftk _f _df _rf acc0 _es ->-    FTKProduct (ftkAst acc0) (buildFTK k bftk)   AstMapAccumLDer k bftk _eftk _f _df _rf acc0 _es ->     FTKProduct (ftkAst acc0) (buildFTK k bftk)-  AstApply (AstLambda !_ !l) _ -> ftkAst l+  AstApply (AstLambda !_ !u) _ -> ftkAst u   AstVar var -> varNameToFTK var-  AstCond _b v _w -> ftkAst v   AstBuild1 snat _ (_var, v) -> buildFTK snat (ftkAst v)    AstLet _ _ v -> ftkAst v@@ -76,8 +67,10 @@    AstPrimalPart a -> ftkAst a   AstDualPart a -> ftkAst a+  AstPlainPart a -> ftkAst a   AstFromPrimal a -> ftkAst a   AstFromDual a -> ftkAst a+  AstFromPlain a -> ftkAst a    AstPlusK{} -> FTKScalar   AstTimesK{} -> FTKScalar@@ -89,6 +82,9 @@   AstFloorK{} -> FTKScalar   AstFromIntegralK{} -> FTKScalar   AstCastK{} -> FTKScalar+  AstArgMinK{} -> FTKScalar+  AstArgMaxK{} -> FTKScalar+  AstIndexK{} -> FTKScalar    AstPlusS v _ -> ftkAst v   AstTimesS v _ -> ftkAst v@@ -103,17 +99,31 @@     FTKS sh FTKScalar -> FTKS sh FTKScalar   AstCastS v -> case ftkAst v of     FTKS sh FTKScalar -> FTKS sh FTKScalar--  AstIndexS shn v _ix -> case ftkAst v of-    FTKS _ x -> FTKS shn x-  AstScatterS shn v (_ , ix) -> case ftkAst v of-    FTKS _ x -> FTKS (shsFromIxS ix `shsAppend` shn) x-  AstGatherS shn v (vars, _) -> case ftkAst v of-    FTKS _ x -> FTKS (shsFromListS vars `shsAppend` shn) x-  AstMinIndexS v -> case ftkAst v of+  AstArgMinS v -> case ftkAst v of     FTKS sh FTKScalar -> FTKS (shsInit sh) FTKScalar-  AstMaxIndexS v -> case ftkAst v of+  AstArgMaxS v -> case ftkAst v of     FTKS sh FTKScalar -> FTKS (shsInit sh) FTKScalar+  AstIndexS shn v _ix -> case ftkAst v of+    FTKS _ x -> FTKS shn x++  AstCondK _b v _w -> ftkAst v+  AstCondS _b v _w -> ftkAst v+  AstFromVectorK shm _ -> FTKS shm FTKScalar+  AstFromVectorS shm l -> case V.uncons l of+    Just (v, _) | FTKS shn x <- ftkAst v -> FTKS (shm `shsAppend` shn) x+    Nothing -> error "ftkAst: empty vector in AstFromVectorS"+  AstSumK{} -> FTKScalar+  AstSumS @shm @shn shm v -> case ftkAst v of+    FTKS shmshn x | SNat <- shsRank shm ->+      gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $+      FTKS (shsDrop @(Rank shm) shmshn) x+  AstScatterS _ shn shp v _ -> case ftkAst v of+    FTKS _ x -> FTKS (shp `shsAppend` shn) x+  AstReplicateK shm _v -> FTKS shm FTKScalar+  AstReplicateS shm v -> case ftkAst v of+    FTKS shn x -> FTKS (shm `shsAppend` shn) x+  AstGatherS shm shn _shp v _ -> case ftkAst v of+    FTKS _ x -> FTKS (shm `shsAppend` shn) x   AstIotaS n@SNat -> FTKS (n :$$ ZSS) FTKScalar   AstAppendS a b -> case (ftkAst a, ftkAst b) of     (FTKS (m :$$ sh) x, FTKS (n :$$ _) _) -> FTKS (snatPlus m n :$$ sh) x@@ -127,23 +137,24 @@    AstConvert c u -> convertFTK c $ ftkAst u -  AstSum0S v ->  case ftkAst v of-    FTKS _ x -> FTKS ZSS x-  AstDot0S _u _v -> FTKS ZSS FTKScalar+  AstDot0{} -> FTKScalar   AstDot1InS sh _ _u _v -> FTKS sh FTKScalar   AstMatmul2S m@SNat _ p@SNat _u _v -> FTKS (m :$$ p :$$ ZSS) FTKScalar -isTensorInt :: forall s y ms. AstSpan s-            => Proxy s -> FullShapeTK y-            -> Maybe (AstTensor ms s y :~: AstInt ms)-isTensorInt _ ftk = case ftk of-  FTKScalar @r -> case ( testEquality (typeRep @r) (typeRep @Int64)-                       , sameAstSpan @s @PrimalSpan ) of-                    (Just Refl, Just Refl) -> Just Refl-                    _ -> Nothing-  _ -> Nothing+  AstBoolNotK{} -> FTKScalar+  AstBoolNotS a -> ftkAst a+  AstBoolAndK{} -> FTKScalar+  AstBoolAndS a _ -> ftkAst a+  AstLeqK{} -> FTKScalar+  AstLeq{} -> FTKScalar+  AstLeqS shb _ _ _ -> FTKS shb FTKScalar +stkAstX :: forall s x ms sh. AstTensor ms s (TKS2 sh x) -> SingletonTK x+{-# INLINE stkAstX #-}+stkAstX t = case ftkAst t of+  FTKS _ x -> ftkToSTK x + -- * Variable occurrence detection  -- | We assume no variable is shared between a binding and its nested binding@@ -155,28 +166,26 @@   AstPair t1 t2 -> varInAst var t1 || varInAst var t2   AstProject1 t -> varInAst var t   AstProject2 t -> varInAst var t-  AstFromVector _ _ vl -> any (varInAst var) vl-  AstSum _ _ v -> varInAst var v-  AstReplicate _ _ v -> varInAst var v-  AstMapAccumRDer _k _bftk _eftk _f _df _rf acc0 es ->-    varInAst var acc0 || varInAst var es   AstMapAccumLDer _k _bftk _eftk _f _df _rf acc0 es ->     varInAst var acc0 || varInAst var es   AstApply t ll -> varInAstHFun var t || varInAst var ll   AstVar var2 -> var == varNameToAstVarId var2-  AstCond b v w -> varInAstBool var b || varInAst var v || varInAst var w   AstBuild1 _ _ (var2, v) ->     assert (varNameToAstVarId var2 /= var) $     varInAst var v -  AstLet _ u v -> varInAst var u || varInAst var v+  AstLet var2 u v ->+    assert (varNameToAstVarId var2 /= var) $+    varInAst var u || varInAst var v   AstShare _ v -> varInAst var v   AstToShare v -> varInAst var v    AstPrimalPart a -> varInAst var a   AstDualPart a -> varInAst var a+  AstPlainPart a -> varInAst var a   AstFromPrimal v -> varInAst var v   AstFromDual v -> varInAst var v+  AstFromPlain v -> varInAst var v    AstPlusK t u -> varInAst var t || varInAst var u   AstTimesK t u -> varInAst var t || varInAst var u@@ -188,6 +197,9 @@   AstFloorK a -> varInAst var a   AstFromIntegralK t -> varInAst var t   AstCastK t -> varInAst var t+  AstArgMinK t -> varInAst var t+  AstArgMaxK t -> varInAst var t+  AstIndexK v ix -> varInAst var v || varInIxS var ix    AstPlusS t u -> varInAst var t || varInAst var u   AstTimesS t u -> varInAst var t || varInAst var u@@ -199,12 +211,24 @@   AstFloorS a -> varInAst var a   AstFromIntegralS a -> varInAst var a   AstCastS t -> varInAst var t-+  AstArgMinS a -> varInAst var a+  AstArgMaxS a -> varInAst var a   AstIndexS _ v ix -> varInAst var v || varInIxS var ix-  AstScatterS _ v (_vars, ix) -> varInIxS var ix || varInAst var v-  AstGatherS _ v (_vars, ix) -> varInIxS var ix || varInAst var v-  AstMinIndexS a -> varInAst var a-  AstMaxIndexS a -> varInAst var a++  AstCondK b v w -> varInAst var b || varInAst var v || varInAst var w+  AstCondS b v w -> varInAst var b || varInAst var v || varInAst var w+  AstFromVectorK _ l -> any (varInAst var) l+  AstFromVectorS _ l -> any (varInAst var) l+  AstSumK v -> varInAst var v+  AstSumS _ v -> varInAst var v+  AstScatterS _ _ _ t (vars, ix) ->+    assert (all (\v -> var /= varNameToAstVarId v) vars) $+    varInIxS var ix || varInAst var t+  AstReplicateK _ v -> varInAst var v+  AstReplicateS _ v -> varInAst var v+  AstGatherS _ _ _ t (vars, ix) ->+    assert (all (\v -> var /= varNameToAstVarId v) vars) $+    varInIxS var ix || varInAst var t   AstIotaS{} -> False   AstAppendS v u -> varInAst var v || varInAst var u   AstSliceS _ _ _ v -> varInAst var v@@ -214,34 +238,34 @@    AstConvert _ v -> varInAst var v -  AstSum0S v -> varInAst var v-  AstDot0S u v -> varInAst var u || varInAst var v+  AstDot0 u v -> varInAst var u || varInAst var v   AstDot1InS _ _ u v -> varInAst var u || varInAst var v   AstMatmul2S _ _ _ u v -> varInAst var u || varInAst var v +  AstBoolNotK b -> varInAst var b+  AstBoolNotS b -> varInAst var b+  AstBoolAndK arg1 arg2 -> varInAst var arg1 || varInAst var arg2+  AstBoolAndS arg1 arg2 -> varInAst var arg1 || varInAst var arg2+  AstLeqK arg1 arg2 -> varInAst var arg1 || varInAst var arg2+  AstLeq arg1 arg2 -> varInAst var arg1 || varInAst var arg2+  AstLeqS _ _ arg1 arg2 -> varInAst var arg1 || varInAst var arg2+ varInIxS :: AstVarId -> AstIxS ms sh -> Bool varInIxS var = any (varInAst var) -varInAstHFun :: AstVarId -> AstHFun s s2 x y -> Bool-varInAstHFun _var AstLambda{} =+varInAstHFun :: AstVarId -> AstHFun s x y -> Bool+varInAstHFun var (AstLambda var2 _) =+  assert (varNameToAstVarId var2 /= var)   False  -- we take advantage of the term being closed -varInAstBool :: AstVarId -> AstBool ms -> Bool-varInAstBool var = \case-  AstBoolConst{} -> False-  AstBoolNot b -> varInAstBool var b-  AstBoolAnd arg1 arg2 -> varInAstBool var arg1 || varInAstBool var arg2-  AstLeqK arg1 arg2 -> varInAst var arg1 || varInAst var arg2-  AstLeqS arg1 arg2 -> varInAst var arg1 || varInAst var arg2--varNameInAst :: AstVarName f y -> AstTensor ms s2 y2 -> Bool+varNameInAst :: AstVarName '(s, y) -> AstTensor ms s2 y2 -> Bool varNameInAst var = varInAst (varNameToAstVarId var) -varNameInIxS :: AstVarName f y -> AstIxS ms sh -> Bool+varNameInIxS :: AstVarName '(s, y) -> AstIxS ms sh -> Bool varNameInIxS var = varInIxS (varNameToAstVarId var)  --- * Determining if a term requires sharing+-- * Tools related to sharing  -- Turns off all but the most trivial cases of astIsSmall. -- For tests only. Affects all simplification and inlining taking place@@ -251,7 +275,7 @@ unsafeTotalSharingRef = unsafePerformIO $ newIORef False  setTotalSharing :: Bool -> IO ()-setTotalSharing b = atomicWriteIORef unsafeTotalSharingRef b+setTotalSharing = atomicWriteIORef unsafeTotalSharingRef  -- | A term requires sharing if it's too large as a term and so duplicating -- it could affect the performance of simplification@@ -261,7 +285,7 @@ astIsSmall _ AstVar{} = True astIsSmall _ AstShare{} = True astIsSmall _ AstConcreteK{} = True-astIsSmall _ (AstConcreteS a) | sNatValue (Nested.srank a) == 0 = True+astIsSmall _ (AstConcreteS a) | fromSNat' (Nested.srank a) == 0 = True astIsSmall lax t = unsafePerformIO $ do   unsafeTotalSharing <- readIORef unsafeTotalSharingRef   return $! if | unsafeTotalSharing -> False@@ -282,22 +306,27 @@   AstPair t1 t2 -> astIsSmallN (astIsSmallN (n - 1) t1) t2   AstProject1 t -> astIsSmallN (n - 1) t   AstProject2 t -> astIsSmallN (n - 1) t-  AstFromVector (SNat' @1) _ v -> astIsSmallN (n - 1) $ v V.! 0-  AstSum (SNat' @1) _ v -> astIsSmallN (n - 1) v-  AstReplicate _ _ v ->-    astIsSmallN (n - 1) v  -- a really good redex and often in series-      -- executed as a metadata change, which is however not free   AstVar{} -> n-  AstCond b u v -> astIsSmallN (astIsSmallN (astBoolIsSmallN (n - 1) b) u) v-  AstConcreteK _ -> n-  AstConcreteS _ -> n  -- small term with zero interpretation cost;-                       -- the physical arrays is shared on GHC heap   AstShare{} -> n   AstPrimalPart v -> astIsSmallN (n - 1) v   AstDualPart v -> astIsSmallN (n - 1) v+  AstPlainPart v -> astIsSmallN (n - 1) v   AstFromPrimal v -> astIsSmallN (n - 1) v   AstFromDual v -> astIsSmallN (n - 1) v-+  AstFromPlain v -> astIsSmallN (n - 1) v+  AstPlusK AstConcreteK{} AstVar{} -> n - 1  -- likely index offset+  AstTimesK AstConcreteK{} AstVar{} -> n - 1  -- likely index manipulation+  AstConcreteK{} -> n  -- small terms with zero interpretation cost;+  AstConcreteS{} -> n  -- the physical arrays are shared on GHC heap+  -- This often appears from user writing (-1), often reduces away+  -- and it has only one argument.+  AstN1K NegateOp v -> astIsSmallN (n - 1) v+  AstCondK b u v -> astIsSmallN (astIsSmallN (astIsSmallN (n - 1) b) u) v+  AstCondS b u v -> astIsSmallN (astIsSmallN (astIsSmallN (n - 1) b) u) v+  -- This is a really good redex, often nested, executed as a metadata change,+  -- but not completely free, hence non-zero cost.+  AstReplicateK _ v -> astIsSmallN (n - 1) v+  AstReplicateS _ v -> astIsSmallN (n - 1) v   AstIotaS{} -> n   AstSliceS _ _ _ v ->     if n <= 20 then 0 else astIsSmallN (n - 1) v  -- executed as metadata change@@ -305,219 +334,375 @@     astIsSmallN (n - 1) v  -- executed as a cheap metadata change   AstTransposeS _perm v ->     if n <= 20 then 0 else astIsSmallN (n - 1) v  -- executed as metadata change-   AstConvert _ v -> astIsSmallN (n - 1) v--  _ -> 0--astBoolIsSmallN :: Int -> AstBool ms -> Int-astBoolIsSmallN n _ | n <= 0 = 0-astBoolIsSmallN n t0 = case t0 of-  AstBoolConst{} -> n-  AstBoolNot v -> astBoolIsSmallN (n - 1) v-  AstBoolAnd u v -> astBoolIsSmallN (astBoolIsSmallN (n - 1) u) v+  AstBoolNotK v -> astIsSmallN (n - 1) v+  AstBoolAndK u v -> astIsSmallN (astIsSmallN (n - 1) u) v   AstLeqK u v -> astIsSmallN (astIsSmallN (n - 1) u) v-  AstLeqS u v -> astIsSmallN (astIsSmallN (n - 1) u) v+  AstLeq u v -> astIsSmallN (astIsSmallN (n - 1) u) v+  _ -> 0  ixIsSmall :: AstIxS ms sh -> Bool ixIsSmall = all (astIsSmall True) --- * Odds and ends+-- | Try to limit the scope of the let cheaply.+--+-- Note that this doesn't inline lets into indexes, so it can be safely+-- performed on user code without the risk of removing the workaround lets for+-- big non-constant values in indexes that prevent the loss of sharing occurring+-- when differentiating indexing. gathers and scatters.+astLetDown :: forall y z s s2. KnownSpan s2+           => AstVarName '(s, y) -> AstTensor AstMethodLet s y+           -> AstTensor AstMethodLet s2 z+           -> AstTensor AstMethodLet s2 z+astLetDown var u v@(AstVar var2) =+  if varNameToAstVarId var2 == varNameToAstVarId var+  then case testEquality var var2 of+    Just Refl -> u+    _ -> error "astLetDown: wrong variable types at AstVar"+  else v+astLetDown var u v = case v of+  -- Normaly the type bounds pair nesting, so the check is cheap.+  AstPair t1 t2 ->+    if | not (varNameInAst var t1) -> AstPair t1 (astLetDown var u t2)+       | not (varNameInAst var t2) -> AstPair (astLetDown var u t1) t2+       | otherwise -> AstLet var u v+  AstProject1 v2 -> AstProject1 (astLetDown var u v2)+  AstProject2 v2 -> AstProject2 (astLetDown var u v2)+  -- Plausibly, accumulators are small and mapAccums are rare,+  -- so the check is cheap.+  AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    if varNameInAst var acc0+    then AstLet var u v+    else AstMapAccumLDer k bftk eftk f df rf acc0 (astLetDown var u es)+  AstApply f t -> AstApply f (astLetDown var u t)+  -- handled above: AstVar+  AstBuild1 k stk (var2, v2) ->+    let !v3 = astLetDown var u v2+    in AstBuild1 k stk (var2, v3) --- An approximation: lower and upper bound.--- TODO: extend, e.g., to general quot and rem.-bounds :: GoodScalar r => AstTensor ms s (TKScalar r) -> (r, r)-bounds (AstConcreteK u) = (u, u)-bounds (AstVar var) = case varNameToBounds var of-  Nothing -> (-1000000000, 1000000000)-  Just (u1, u2) -> (fromIntegral u1, fromIntegral u2)-bounds (AstFromPrimal u) = bounds u-bounds (AstPrimalPart u) = bounds u-bounds (AstCond _b u v) = let (u1, u2) = bounds u-                              (v1, v2) = bounds v-                          in (min u1 v1, max u2 v2)-bounds (AstLet _ _ u) = bounds u  -- TODO: substitute?-bounds (AstPlusK u v) = let (u1, u2) = bounds u-                            (v1, v2) = bounds v-                        in (u1 + v1, u2 + v2)-bounds (AstN1K NegateOp u) = let (u1, u2) = bounds u in (- u2, - u1)-bounds (AstTimesK u v) =-  let (u1, u2) = bounds u-      (v1, v2) = bounds v-      l = [u1 * v1, u1 * v2, u2 * v1, u2 * v2]-  in (minimum l, maximum l)-bounds (AstI2K QuotOp u (AstConcreteK v)) | v > 0 =  -- a common case-  let (u1, u2) = bounds u-  in (u1 `quotH` v, u2 `quotH` v)-bounds (AstI2K RemOp u (AstConcreteK v)) | v > 0 =-  let (u1, u2) = bounds u-  in if | u1 >= 0 -> (0, min u2 (v - 1))  -- very crude-        | u2 <= 0 -> (max u1 (- v + 1), 0)-        | otherwise -> (- v + 1, v - 1)-bounds _ = (-1000000000, 1000000000)+  AstLet{} -> AstLet var u v -liftRFromS1 :: forall n x ms s.-               (forall sh.-                   AstTensor ms s (TKS2 sh x)-                -> AstTensor ms s (TKS2 sh x))-            -> AstTensor ms s (TKR2 n x)-            -> AstTensor ms s (TKR2 n x)-liftRFromS1 f a = case ftkAst a of-  ftk@(FTKR sh' _) ->-    withShsFromShR sh' $ \(sh :: ShS sh) ->-      cAstFromS @(TKS2 sh x) ftk-      $ f (cAstSFromR @sh sh a)+  AstPrimalPart v2 -> AstPrimalPart (astLetDown var u v2)+  AstDualPart v2 -> AstDualPart (astLetDown var u v2)+  AstPlainPart v2 -> AstPlainPart (astLetDown var u v2)+  AstFromPrimal v2 -> fromPrimal (astLetDown var u v2)+  AstFromDual v2 -> fromDual (astLetDown var u v2)+  AstFromPlain v2 -> fromPlain (astLetDown var u v2) -liftRFromS2 :: forall n x ms s.-               (forall sh.-                   AstTensor ms s (TKS2 sh x) -> AstTensor ms s (TKS2 sh x)-                -> AstTensor ms s (TKS2 sh x))-            -> AstTensor ms s (TKR2 n x) -> AstTensor ms s (TKR2 n x)-            -> AstTensor ms s (TKR2 n x)-liftRFromS2 f a b  = case ftkAst a of-  ftk@(FTKR sh' _) ->-    withShsFromShR sh' $ \(sh :: ShS sh) ->-      cAstFromS @(TKS2 sh x) ftk-      $ f (cAstSFromR @sh sh a) (cAstSFromR @sh sh b)+  AstPlusK{} -> AstLet var u v+  AstTimesK{} -> AstLet var u v+  AstN1K op u2 -> AstN1K op (astLetDown var u u2)+  AstR1K op u2 -> AstR1K op (astLetDown var u u2)+  AstR2K{} -> AstLet var u v+  AstI2K{} -> AstLet var u v+  AstConcreteK{} -> v+  AstFloorK a -> AstFloorK (astLetDown var u a)+  AstFromIntegralK v2 -> AstFromIntegralK (astLetDown var u v2)+  AstCastK v2 -> AstCastK (astLetDown var u v2)+  AstArgMinK v2 -> AstArgMinK (astLetDown var u v2)+  AstArgMaxK v2 -> AstArgMaxK (astLetDown var u v2)+  AstIndexK v2 ix ->+    if varNameInIxS var ix+    then AstLet var u v+    else AstIndexK (astLetDown var u v2) ix -liftXFromS1 :: forall sh' x ms s.-               (forall sh.-                   AstTensor ms s (TKS2 sh x)-                -> AstTensor ms s (TKS2 sh x))-            -> AstTensor ms s (TKX2 sh' x)-            -> AstTensor ms s (TKX2 sh' x)-liftXFromS1 f a = case ftkAst a of-  ftk@(FTKX sh' _) ->-    withShsFromShX sh' $ \(sh :: ShS sh) ->-      cAstFromS @(TKS2 sh x) ftk-      $ f (cAstSFromX @sh @sh' sh a)+  AstPlusS{} -> AstLet var u v+  AstTimesS{} -> AstLet var u v+  AstN1S op u2 -> AstN1S op (astLetDown var u u2)+  AstR1S op u2 -> AstR1S op (astLetDown var u u2)+  AstR2S{} -> AstLet var u v+  AstI2S{} -> AstLet var u v+  AstConcreteS{} -> v+  AstFloorS a -> AstFloorS (astLetDown var u a)+  AstFromIntegralS v2 -> AstFromIntegralS (astLetDown var u v2)+  AstCastS v2 -> AstCastS (astLetDown var u v2)+  AstArgMinS a -> AstArgMinS (astLetDown var u a)+  AstArgMaxS a -> AstArgMaxS (astLetDown var u a)+  -- In these three, index terms are usually small, so the check is cheap.+  -- Also, this undoes precisely the pushing of the lets up that rules+  -- for these three perform when simplifying. Note that we never push+  -- lets down indexes, which is important for user workarounds+  -- and is a legal inlining only when the iteration is trivial and so+  -- the number of dynamic occurrences of the inlined variable is trivial.+  AstIndexS shn v2 ix ->+    if varNameInIxS var ix+    then AstLet var u v+    else AstIndexS shn (astLetDown var u v2) ix -liftXFromS2 :: forall sh' x ms s.-               (forall sh.-                   AstTensor ms s (TKS2 sh x) -> AstTensor ms s (TKS2 sh x)-                -> AstTensor ms s (TKS2 sh x))-            -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKX2 sh' x)-            -> AstTensor ms s (TKX2 sh' x)-liftXFromS2 f a b = case ftkAst a of-  ftk@(FTKX sh' _) ->-    withShsFromShX sh' $ \(sh :: ShS sh) ->-      cAstFromS @(TKS2 sh x) ftk-      $ f (cAstSFromX @sh @sh' sh a) (cAstSFromX @sh @sh' sh b)+  AstCondK{} -> AstLet var u v+  AstCondS{} -> AstLet var u v+  AstFromVectorK{} -> AstLet var u v+  AstFromVectorS{} -> AstLet var u v+  AstSumK v2 -> AstSumK (astLetDown var u v2)+  AstSumS shm v2 -> AstSumS shm (astLetDown var u v2)+  AstScatterS shm shn shp v2 (vars, ix) ->+    if varNameInIxS var ix+    then AstLet var u v+    else AstScatterS shm shn shp (astLetDown var u v2) (vars, ix)+  AstReplicateK shm v2 -> AstReplicateK shm (astLetDown var u v2)+  AstReplicateS shm v2 -> AstReplicateS shm (astLetDown var u v2)+  AstGatherS shm shn shp v2 (vars, ix) ->+    if varNameInIxS var ix+    then AstLet var u v+    else AstGatherS shm shn shp (astLetDown var u v2) (vars, ix)+  AstIotaS{} -> v+  AstAppendS{} -> AstLet var u v+  AstSliceS i n k v2 -> AstSliceS i n k (astLetDown var u v2)+  AstReverseS v2 -> AstReverseS (astLetDown var u v2)+  AstTransposeS perm v2 -> AstTransposeS perm (astLetDown var u v2)+  AstReshapeS sh v2 -> AstReshapeS sh (astLetDown var u v2) -cAstConvert :: TKConversion x z -> AstTensor ms s x -> AstTensor ms s z-cAstConvert c t-  | Just Refl <- matchingFTK (ftkAst t) (convertFTK c (ftkAst t)) = t-cAstConvert c1 (AstConvert c2 t2) = cAstConvert (ConvCmp c1 c2) t2-cAstConvert c1 (AstFromPrimal (AstConvert c2 t2)) =-  AstFromPrimal (cAstConvert (ConvCmp c1 c2) t2)-cAstConvert c t = AstConvert c t+  AstConvert c v2 -> AstConvert c (astLetDown var u v2) -cAstSFromR :: forall sh x ms s.-              ShS sh -> AstTensor ms s (TKR2 (Rank sh) x)-           -> AstTensor ms s (TKS2 sh x)-cAstSFromR sh v = case ftkAst v of-  FTKR _ x | Refl <- lemRankReplicate (Proxy @(Rank sh)) ->-    let c2 = ConvCmp (ConvXS' (FTKS sh x)) ConvRX-    in cAstConvert c2 v+  AstDot0{} -> AstLet var u v+  AstDot1InS{} -> AstLet var u v+  AstMatmul2S{} -> AstLet var u v -cAstSFromX :: forall sh sh' x ms s. Rank sh ~ Rank sh'-           => ShS sh -> AstTensor ms s (TKX2 sh' x)-           -> AstTensor ms s (TKS2 sh x)-cAstSFromX sh v = case ftkAst v of-  FTKX _ x -> let c2 = ConvXS' (FTKS sh x)-              in cAstConvert c2 v+  AstBoolNotK arg -> AstBoolNotK (astLetDown var u arg)+  AstBoolNotS arg -> AstBoolNotS (astLetDown var u arg)+  AstBoolAndK{} -> AstLet var u v+  AstBoolAndS{} -> AstLet var u v+  AstLeqK{} -> AstLet var u v+  AstLeq{} -> AstLet var u v+  AstLeqS{} -> AstLet var u v -cAstXFromS :: forall sh sh' x ms s. Rank sh ~ Rank sh'-           => StaticShX sh' -> AstTensor ms s (TKS2 sh x)-           -> AstTensor ms s (TKX2 sh' x)-cAstXFromS ssx v-  | FTKS sh x <- ftkAst v-  , let shx = shCastSX ssx sh-  , Refl <- lemRankMapJust sh =-    let c2 = ConvCmp (ConvXX' (FTKX shx x)) ConvSX-    in cAstConvert c2 v+astVar :: AstVarName '(s, y) -> AstTensor ms s y+astVar (AstVarName _ (FtkAndBoundsBounds lb ub)) | lb == ub =+  AstConcreteK lb+astVar var = AstVar var -pattern AstFromS' :: forall {z1} {ms1} {s1}.-                     forall y z ms s. (z ~ z1, ms ~ ms1, s ~ s1)-                  => FullShapeTK z -> AstTensor ms s y-                  -> AstTensor ms1 s1 z1-pattern AstFromS' zftk a <--  AstConvert c (checkPatternAstFromS c -> Just (zftk, a))+astShare :: AstVarName '(s, y) -> AstTensor AstMethodShare s y+         -> AstTensor AstMethodShare s y+astShare (AstVarName _ (FtkAndBoundsBounds lb ub)) _ | lb == ub =+  AstConcreteK lb+astShare var t = AstShare var t -checkPatternAstFromS :: TKConversion y z -> AstTensor ms s y-                     -> Maybe (FullShapeTK z, AstTensor ms s y)-checkPatternAstFromS c t =-  let zftk = convertFTK c (ftkAst t)-  in if checkFtkAstFromS (ftkAst t) zftk then Just (zftk, t) else Nothing -checkAstFromS :: TKConversion a b -> AstTensor ms s a -> Bool-checkAstFromS c t = isJust $ checkPatternAstFromS c t+-- * Odds and ends --- TODO: this is too lax, since it accepts nests/unnests. Add tests and fix.-checkFtkAstFromS :: FullShapeTK y -> FullShapeTK z -> Bool-checkFtkAstFromS yftk zftk | Just Refl <- matchingFTK yftk zftk = True-checkFtkAstFromS FTKS{} FTKS{} = False-checkFtkAstFromS FTKS{} _ = True-checkFtkAstFromS (FTKProduct yftk1 yftk2) (FTKProduct zftk1 zftk2) =-  checkFtkAstFromS yftk1 zftk1 && checkFtkAstFromS yftk2 zftk2-checkFtkAstFromS _ _ = False+bounds :: forall r s ms. (Typeable r, KnownSpan s)+       => AstTensor ms s (TKScalar r) -> Maybe (r, r)+bounds t | SPlainSpan <- knownSpan @s+         , Just Refl <- testEquality (typeRep @r) (typeRep @Int) = intBounds t+bounds _ = Nothing -checkFtkAstSFrom :: FullShapeTK y -> FullShapeTK z -> Bool-checkFtkAstSFrom yftk zftk | Just Refl <- matchingFTK yftk zftk = True-checkFtkAstSFrom FTKS{} FTKS{} = False-checkFtkAstSFrom _ FTKS{} = True-checkFtkAstSFrom (FTKProduct yftk1 yftk2) (FTKProduct zftk1 zftk2) =-  checkFtkAstSFrom yftk1 zftk1 && checkFtkAstSFrom yftk2 zftk2-checkFtkAstSFrom _ _ = False+-- An approximation: lower and upper bound.+-- TODO: extend, e.g., to general quot and rem.+intBounds :: AstTensor ms PlainSpan (TKScalar Int) -> Maybe (Int, Int)+intBounds (AstConcreteK u) = Just (u, u)+intBounds (AstApply (AstLambda _ u) _) = intBounds u+intBounds (AstVar var) = varNameToBounds var+intBounds (AstLet _ _ u) = intBounds u  -- TODO: substitute?+intBounds (AstShare var _) = varNameToBounds var+intBounds (AstToShare u) = intBounds u+intBounds (AstPlusK u v) = do+  (u1, u2) <- intBounds u+  (v1, v2) <- intBounds v+  pure (u1 + v1, u2 + v2)+intBounds (AstN1K NegateOp u) = do+  (u1, u2) <- intBounds u+  pure (- u2, - u1)+intBounds (AstTimesK u v) = case (intBounds u, intBounds v) of+  (Nothing, Nothing) -> Nothing+  (mu, mv) ->  -- multiplication by zero makes even one operand enough;+               -- TODO: but this is, in theory, unsound, because the below+               -- are not real infinities+    let (u1, u2) = fromMaybe (-1000000000, 1000000000) mu+        (v1, v2) = fromMaybe (-1000000000, 1000000000) mv+        l = [u1 * v1, u1 * v2, u2 * v1, u2 * v2]+        (lb, ub) = (minimum l, maximum l)+    in if lb == -1000000000 && ub == 1000000000+       then Nothing+       else Just (lb, ub)+intBounds (AstI2K QuotOp u (AstConcreteK v)) | v > 0 = do  -- a common case+  (u1, u2) <- intBounds u+  pure (u1 `quotH` v, u2 `quotH` v)+intBounds (AstI2K RemOp u (AstConcreteK v)) | v > 0 = do+  (u1, u2) <- intBounds u+  pure $ if | u1 >= 0 -> (0, min u2 (v - 1))  -- very crude+            | u2 <= 0 -> (max u1 (- v + 1), 0)+            | otherwise -> (- v + 1, v - 1)+intBounds (AstCondK _b u v) = do+  (u1, u2) <- intBounds u+  (v1, v2) <- intBounds v+  pure (min u1 v1, max u2 v2)+intBounds _ = Nothing -cAstFromS :: forall y z ms s.-             FullShapeTK z -> AstTensor ms s y-          -> AstTensor ms s z-cAstFromS zftk t = cAstConvert (convFromS (ftkAst t) zftk) t+pattern AstConvUpSFromK :: forall r sh s ms. () => sh ~ '[]+                        => AstTensor ms s (TKScalar r)+                        -> AstTensor ms s (TKS sh r)+pattern AstConvUpSFromK t <- (matchAstConvUpSFromK -> Just (Refl, t)) -cAstSFrom :: forall y z ms s.-             FullShapeTK z -> AstTensor ms s y-          -> AstTensor ms s z-cAstSFrom zftk t = cAstConvert (convSFrom (ftkAst t) (ftkToSTK zftk)) t+matchAstConvUpSFromK :: AstTensor ms s (TKS sh r)+                     -> Maybe ( sh :~: '[]+                              , AstTensor ms s (TKScalar r) )+matchAstConvUpSFromK = \case+  AstConvert c t+    | FTKScalar @ry <- ftkAst t+    , FTKS ZSS (FTKScalar @r) <- convertFTK c (ftkAst t)+    , Just Refl <- testEquality (typeRep @ry) (typeRep @r) ->+      Just (Refl, t)+  AstConvert c t+    | FTKR ZSR (FTKScalar @ry) <- ftkAst t+    , FTKS ZSS (FTKScalar @r) <- convertFTK c (ftkAst t)+    , Just Refl <- testEquality (typeRep @ry) (typeRep @r) ->+      Just (Refl, AstConvert (ConvCmp ConvX0 ConvRX) t)+  AstConvert c t+    | FTKX ZSX (FTKScalar @ry) <- ftkAst t+    , FTKS ZSS (FTKScalar @r) <- convertFTK c (ftkAst t)+    , Just Refl <- testEquality (typeRep @ry) (typeRep @r) ->+      Just (Refl, AstConvert ConvX0 t)+  AstConcreteS a | ZSS <- Nested.sshape a ->+    Just (Refl, AstConcreteK $ Nested.sunScalar a)+  AstFromPrimal t -> second AstFromPrimal <$> matchAstConvUpSFromK t+  AstFromDual t -> second AstFromDual <$> matchAstConvUpSFromK t+  AstFromPlain t -> second AstFromPlain <$> matchAstConvUpSFromK t+  _ -> Nothing -convFromS :: FullShapeTK y0 -> FullShapeTK z0 -> TKConversion y0 z0-convFromS yftk0 zftk0 = case (yftk0, zftk0) of-  _ | Just Refl <- matchingFTK yftk0 zftk0 -> ConvId-  (FTKS ZSS (FTKScalar @ry), FTKScalar @rz)+-- TODO: simplify this monstrosity, if possible+pattern AstConvUp :: forall {z1} {ms1} {s1}.+                     forall y z ms s. (z ~ z1, ms ~ ms1, s ~ s1)+                  => TKConversion y z -> FullShapeTK z -> AstTensor ms s y+                  -> AstTensor ms1 s1 z1+pattern AstConvUp c zftk a <- (matchAstConvUp -> AstConvUpJust c zftk a)++type role AstConvUpMaybe nominal nominal nominal+data AstConvUpMaybe z ms s =+    forall y.+    AstConvUpJust (TKConversion y z) (FullShapeTK z) (AstTensor ms s y)+  | AstConvUpNothing++matchAstConvUp :: AstTensor ms s z -> AstConvUpMaybe z ms s+matchAstConvUp = \case+  AstConvert c t+    | FTKR ZSR (FTKScalar @ry) <- ftkAst t+    , let zftk = convertFTK c (ftkAst t)+    , FTKS ZSS (FTKScalar @r) <- zftk+    , Just Refl <- testEquality (typeRep @ry) (typeRep @r) ->+      AstConvUpJust (ConvCmp ConvXS (Conv0X STKScalar))+                    zftk+                    (AstConvert (ConvCmp ConvX0 ConvRX) t)+  AstConvert c t+    | FTKX ZSX (FTKScalar @ry) <- ftkAst t+    , let zftk = convertFTK c (ftkAst t)+    , FTKS ZSS (FTKScalar @r) <- zftk+    , Just Refl <- testEquality (typeRep @ry) (typeRep @r) ->+      AstConvUpJust (ConvCmp ConvXS (Conv0X STKScalar))+                    zftk+                    (AstConvert ConvX0 t)+  AstConvert c t+    | FTKS ZSS (FTKScalar @r) <- ftkAst t+    , let zftk = convertFTK c (ftkAst t)+    , FTKR ZSR (FTKScalar @ry) <- zftk+    , Just Refl <- testEquality (typeRep @ry) (typeRep @r) ->+      AstConvUpJust (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))+                    zftk+                    (AstConvert (ConvCmp ConvX0 ConvSX) t)+  AstConvert c t+    | FTKS ZSS (FTKScalar @r) <- ftkAst t+    , let zftk = convertFTK c (ftkAst t)+    , FTKX ZSX (FTKScalar @ry) <- zftk+    , Just Refl <- testEquality (typeRep @ry) (typeRep @r) ->+      AstConvUpJust (Conv0X STKScalar)+                    zftk+                    (AstConvert (ConvCmp ConvX0 ConvSX) t)+  AstConvert c t ->+    let yftk = ftkAst t+        zftk = convertFTK c yftk+    in case convUpMaybe (ftkAst t) zftk of+      Just c2 -> AstConvUpJust c2 zftk t+      Nothing -> AstConvUpNothing+  AstConcreteS a | ZSS <- Nested.sshape a ->+    AstConvUpJust (ConvCmp ConvXS (Conv0X STKScalar))+                  (FTKS ZSS FTKScalar)+                  (AstConcreteK $ Nested.sunScalar a)+  AstFromPrimal t -> case matchAstConvUp t of+    AstConvUpJust c zftk u -> AstConvUpJust c zftk (AstFromPrimal u)+    AstConvUpNothing -> AstConvUpNothing+  AstFromDual t -> case matchAstConvUp t of+    AstConvUpJust c zftk u -> AstConvUpJust c zftk (AstFromDual u)+    AstConvUpNothing -> AstConvUpNothing+  AstFromPlain t -> case matchAstConvUp t of+    AstConvUpJust c zftk u -> AstConvUpJust c zftk (AstFromPlain u)+    AstConvUpNothing -> AstConvUpNothing+  _ -> AstConvUpNothing++convDownMaybe :: FullShapeTK y0 -> SingletonTK z0 -> Maybe (TKConversion y0 z0)+convDownMaybe = \cases+  yftk0 zstk0 | Just Refl <- sameSTK (ftkToSTK yftk0) zstk0 -> Just ConvId+  (FTKS ZSS (FTKScalar @ry)) (STKScalar @rz)     | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->-      ConvCmp ConvX0 ConvSX-  (FTKS sh x, FTKR rsh rx)-    | Just Refl <- matchingFTK x rx+      Just $ convCmp ConvX0 ConvSX+  (FTKR ZSR (FTKScalar @ry)) (STKScalar @rz)+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      Just $ convCmp ConvX0 ConvRX+  (FTKX ZSX (FTKScalar @ry)) (STKScalar @rz)+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      Just ConvX0+  (FTKR rsh rx) (STKS @sh sh@(_ :$$ _) x)+    | Just Refl <- sameSTK x (ftkToSTK rx)     , Just Refl <- testEquality (shsRank sh) (shrRank rsh)-    , Refl <- lemRankMapJust sh ->-      ConvCmp (ConvXR (ftkToSTK x)) ConvSX-  (FTKS sh x, FTKX xsh xx)-    | Just Refl <- matchingFTK x xx+    , Refl <- lemRankReplicate (Proxy @(Rank sh)) ->+      Just $ convCmp (ConvXS' (FTKS sh rx)) ConvRX+  (FTKX xsh xx) (STKS sh@(_ :$$ _) x)+    | Just Refl <- sameSTK x (ftkToSTK xx)     , Just Refl <- testEquality (shsRank sh) (shxRank xsh)     , Refl <- lemRankMapJust sh ->-      ConvCmp (ConvXX' zftk0) ConvSX-  (FTKProduct yftk1 yftk2, FTKProduct zftk1 zftk2) ->-    ConvT2 (convFromS yftk1 zftk1) (convFromS yftk2 zftk2)-  _ -> error $ "convFromS': unexpected types "  -- TODO: try nevertheless-               ++ "(" ++ show yftk0 ++ ", " ++ show zftk0 ++ ")"+      Just $ ConvXS' (FTKS sh xx)+  (FTKProduct yftk1 yftk2) (STKProduct zstk1 zstk2) -> do+    c1 <- convDownMaybe yftk1 zstk1+    c2 <- convDownMaybe yftk2 zstk2+    Just $ ConvT2 c1 c2+  (FTKProduct (FTKS sh' yftk1) (FTKS sh'' yftk2)) (STKS sh (STKProduct ystk1+                                                                       ystk2))+    | Just Refl <- testEquality sh sh'+    , Just Refl <- testEquality sh sh''+    , Just Refl <- sameSTK ystk1 (ftkToSTK yftk1)+    , Just Refl <- sameSTK ystk2 (ftkToSTK yftk2) ->+      Just+      $ convCmp+          ConvXS+          (convCmp+             (ConvZip ystk1 ystk2)+             (ConvT2 ConvSX ConvSX))+  _ _ -> Nothing -convSFrom :: FullShapeTK y0 -> SingletonTK z0 -> TKConversion y0 z0-convSFrom yftk0 zstk0 = case (zstk0, yftk0) of-  _ | Just Refl <- sameSTK (ftkToSTK yftk0) zstk0 -> ConvId-  (STKS ZSS (STKScalar @ry), FTKScalar @rz)+convUpMaybe :: FullShapeTK y0 -> FullShapeTK z0 -> Maybe (TKConversion y0 z0)+convUpMaybe = \cases+  yftk0 zftk0 | Just Refl <- matchingFTK yftk0 zftk0 -> Just ConvId+  (FTKScalar @rz) (FTKS ZSS (FTKScalar @ry))     | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->-      ConvCmp ConvXS (Conv0X STKScalar)-  (STKS @sh sh x, FTKR rsh rx)-    | Just Refl <- sameSTK x (ftkToSTK rx)+      Just $ convCmp ConvXS (Conv0X STKScalar)+  (FTKScalar @rz) (FTKR ZSR (FTKScalar @ry))+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      Just $ convCmp (ConvXR STKScalar) (Conv0X STKScalar)+  (FTKScalar @rz) (FTKX ZSX (FTKScalar @ry))+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      Just $ Conv0X STKScalar+  (FTKS sh@(_ :$$ _) x) (FTKR rsh rx)+    | Just Refl <- matchingFTK x rx     , Just Refl <- testEquality (shsRank sh) (shrRank rsh)-    , Refl <- lemRankReplicate (Proxy @(Rank sh)) ->-      ConvCmp (ConvXS' (FTKS sh rx)) ConvRX-  (STKS sh x, FTKX xsh xx)-    | Just Refl <- sameSTK x (ftkToSTK xx)+    , Refl <- lemRankMapJust sh ->+      Just $ convCmp (ConvXR (ftkToSTK x)) ConvSX+  (FTKS sh@(_ :$$ _) x) zftk0@(FTKX xsh xx)+    | Just Refl <- matchingFTK x xx     , Just Refl <- testEquality (shsRank sh) (shxRank xsh)     , Refl <- lemRankMapJust sh ->-      ConvXS' (FTKS sh xx)-  (STKProduct zstk1 zstk2, FTKProduct yftk1 yftk2) ->-    ConvT2 (convSFrom yftk1 zstk1) (convSFrom yftk2 zstk2)-  _ -> error $ "convSFrom': unexpected types "  -- TODO: try nevertheless-               ++ "(" ++ show yftk0 ++ ", " ++ show zstk0 ++ ")"+      Just $ convCmp (ConvXX' zftk0) ConvSX+  (FTKProduct yftk1 yftk2) (FTKProduct zftk1 zftk2) -> do+    c1 <- convUpMaybe yftk1 zftk1+    c2 <- convUpMaybe yftk2 zftk2+    Just $ ConvT2 c1 c2+  (FTKS sh (FTKProduct yftk1 yftk2)) (FTKProduct (FTKS sh' yftk1')+                                                 (FTKS sh'' yftk2'))+    | Just Refl <- testEquality sh sh'+    , Just Refl <- testEquality sh sh''+    , Just Refl <- matchingFTK yftk1 yftk1'+    , Just Refl <- matchingFTK yftk2 yftk2' ->+      Just+      $ convCmp+          (ConvT2 ConvXS ConvXS)+          (convCmp+             (ConvUnzip (ftkToSTK yftk1) (ftkToSTK yftk2))+             ConvSX)+  _ _ -> Nothing
src/HordeAd/Core/AstTraverse.hs view
@@ -1,3 +1,7 @@+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif -- | This module implements a few complete bottom-up simplifying passes -- over any AST expression. module HordeAd.Core.AstTraverse@@ -7,34 +11,41 @@   , simplifyAst     -- * The contraction (e.g., from gather expressions) bottom-up pass   , contractAst+    -- * The let down (reducing the scope of lets cheaply) bottom-up pass+  , letDownAst   ) where  import Prelude -import Data.Int (Int64)+import Data.Foldable qualified as Foldable+import Data.Int (Int16, Int32, Int64, Int8) import Data.Maybe (fromMaybe)+import Data.Proxy (Proxy (Proxy)) import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl)) import Data.Vector.Generic qualified as V import Foreign.C (CInt)+import Foreign.Storable (sizeOf) import GHC.TypeLits (KnownNat) import Type.Reflection (typeRep)  import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Permutation (Perm (..)) import Data.Array.Nested.Permutation qualified as Permutation import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (unsafeCoerceRefl)+import Data.Array.Nested.Types (Tail, fromSNat', unsafeCoerceRefl)  import HordeAd.Core.Ast-  ( AstBool (AstBoolConst)-  , AstTensor (AstConcreteK, AstConcreteS, AstPlusK, AstPlusS, AstTimesK, AstTimesS)+  ( AstTensor (AstConcreteK, AstConcreteS, AstPlusK, AstPlusS, AstTimesK, AstTimesS)   )-import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))+import HordeAd.Core.Ast hiding (AstTensor (..)) import HordeAd.Core.Ast qualified as Ast import HordeAd.Core.AstSimplify import HordeAd.Core.AstTools import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Conversion+import HordeAd.Core.ConvertTensor import HordeAd.Core.Ops import HordeAd.Core.TensorKind import HordeAd.Core.Types@@ -50,22 +61,12 @@ -- fusion rules that would be immediately counteracted by expansion rules -- if applied earlier. expandAst-  :: forall s y. AstSpan s+  :: forall s y. KnownSpan s   => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y expandAst t = case t of   Ast.AstPair t1 t2 -> astPair (expandAst t1) (expandAst t2)   Ast.AstProject1 v -> astProject1 (expandAst v)   Ast.AstProject2 v -> astProject2 (expandAst v)-  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map expandAst l)-  Ast.AstSum snat stk v -> astSum snat stk (expandAst v)-  Ast.AstReplicate snat stk v -> astReplicate snat stk (expandAst v)-  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->-    astMapAccumRDer k bftk eftk-                    (expandAstHFun f)-                    (expandAstHFun df)-                    (expandAstHFun rf)-                    (expandAst acc0)-                    (expandAst es)   Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->     astMapAccumLDer k bftk eftk                     (expandAstHFun f)@@ -74,130 +75,94 @@                     (expandAst acc0)                     (expandAst es)   Ast.AstApply v ll -> astApply (expandAstHFun v) (expandAst ll)-  Ast.AstVar var -> astVar var-  Ast.AstCond b a2 a3 ->-    astCond (expandAstBool b) (expandAst a2) (expandAst a3)+  Ast.AstVar{} -> t   Ast.AstBuild1 k stk (var, v) ->     let !v2 = expandAst v     in Ast.AstBuild1 k stk (var, v2) -  Ast.AstLet var u v -> astLet var (expandAst u) (expandAst v)+  Ast.AstLet var u v ->+    astLetDown var (withKnownSpan (varNameToSpan var) $ expandAst u)+                   (expandAst v)    Ast.AstPrimalPart v -> astPrimalPart (expandAst v)   Ast.AstDualPart v -> astDualPart (expandAst v)-  Ast.AstFromPrimal v -> Ast.AstFromPrimal (expandAst v)-  Ast.AstFromDual v -> Ast.AstFromDual (expandAst v)+  Ast.AstPlainPart v -> astPlainPart (expandAst v)+  Ast.AstFromPrimal v -> fromPrimal (expandAst v)+  Ast.AstFromDual v -> fromDual (expandAst v)+  Ast.AstFromPlain v -> fromPlain (expandAst v)    AstPlusK u v -> expandAst u + expandAst v   AstTimesK u v -> expandAst u * expandAst v-  Ast.AstN1K NegateOp u -> negate (expandAst u)-  Ast.AstN1K AbsOp u -> abs (expandAst u)-  Ast.AstN1K SignumOp u -> signum (expandAst u)-  Ast.AstR1K opCode u -> Ast.AstR1K opCode (expandAst u)-  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (expandAst u) (expandAst v)-  Ast.AstI2K QuotOp u v -> quotH (expandAst u) (expandAst v)-  Ast.AstI2K RemOp u v -> remH (expandAst u) (expandAst v)-  AstConcreteK k -> AstConcreteK k+  Ast.AstN1K opCode u -> astN1K opCode (expandAst u)+  Ast.AstR1K opCode u -> astR1K opCode (expandAst u)+  Ast.AstR2K opCode u v -> astR2K opCode (expandAst u) (expandAst v)+  Ast.AstI2K opCode u v -> astI2K opCode (expandAst u) (expandAst v)+  AstConcreteK{} -> t   Ast.AstFloorK a -> astFloorK (expandAst a)   Ast.AstFromIntegralK v -> astFromIntegralK $ expandAst v   Ast.AstCastK v -> astCastK $ expandAst v+  Ast.AstArgMinK v -> astArgMinK $ expandAst v+  Ast.AstArgMaxK v -> astArgMaxK $ expandAst v+  Ast.AstIndexK v ix ->+    astIndexKnobsK (defaultKnobs {knobPhase = PhaseExpansion})+                   (expandAst v) (expandAstIxS ix)    AstPlusS u v -> expandAst u + expandAst v   AstTimesS u v -> expandAst u * expandAst v-  Ast.AstN1S NegateOp u -> negate (expandAst u)-  Ast.AstN1S AbsOp u -> abs (expandAst u)-  Ast.AstN1S SignumOp u -> signum (expandAst u)-  Ast.AstR1S opCode u -> Ast.AstR1S opCode (expandAst u)-  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (expandAst u) (expandAst v)-  Ast.AstI2S QuotOp u v -> quotH (expandAst u) (expandAst v)-  Ast.AstI2S RemOp u v -> remH (expandAst u) (expandAst v)-  AstConcreteS a -> AstConcreteS a+  Ast.AstN1S opCode u -> astN1S opCode (expandAst u)+  Ast.AstR1S opCode u -> astR1S opCode (expandAst u)+  Ast.AstR2S opCode u v -> astR2S opCode (expandAst u) (expandAst v)+  Ast.AstI2S opCode u v -> astI2S opCode (expandAst u) (expandAst v)+  AstConcreteS{} -> t   Ast.AstFloorS a -> astFloorS (expandAst a)   Ast.AstFromIntegralS v -> astFromIntegralS $ expandAst v   Ast.AstCastS v -> astCastS $ expandAst v-+  Ast.AstArgMinS a -> Ast.AstArgMinS (expandAst a)+  Ast.AstArgMaxS a -> Ast.AstArgMaxS (expandAst a)   Ast.AstIndexS shn v ix ->     astIndexKnobsS (defaultKnobs {knobPhase = PhaseExpansion})                    shn (expandAst v) (expandAstIxS ix)-  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->-    astScatterS @shm @shn @shp shn (expandAst v) (vars, expandAstIxS ix)-  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->-    astGatherKnobsS @shm @shn @shp-                    (defaultKnobs {knobPhase = PhaseExpansion})-                    shn (expandAst v) (vars, expandAstIxS ix)-  Ast.AstMinIndexS a -> Ast.AstMinIndexS (expandAst a)-  Ast.AstMaxIndexS a -> Ast.AstMaxIndexS (expandAst a)++  Ast.AstCondK b a2 a3 -> astCondK (expandAst b) (expandAst a2) (expandAst a3)+  Ast.AstCondS b a2 a3 -> astCondS (expandAst b) (expandAst a2) (expandAst a3)+  Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map expandAst l)+  Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map expandAst l)+  Ast.AstSumK a -> astSumK (expandAst a)+  Ast.AstSumS shm a -> astSumS shm (expandAst a)+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterKnobsS (defaultKnobs {knobPhase = PhaseExpansion})+                     shm shn shp (expandAst v) (vars, expandAstIxS ix)+  Ast.AstReplicateK shm v -> astReplicateK shm (expandAst v)+  Ast.AstReplicateS shm v -> astReplicateS shm (expandAst v)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherKnobsS (defaultKnobs {knobPhase = PhaseExpansion})+                    shm shn shp (expandAst v) (vars, expandAstIxS ix)   Ast.AstIotaS{} -> t   Ast.AstAppendS x y -> astAppendS (expandAst x) (expandAst y)   Ast.AstSliceS i n k v -> astSliceS i n k (expandAst v)   Ast.AstReverseS v -> astReverseS (expandAst v)-  Ast.AstTransposeS perm v -> {--   -- disabled until we can reliably fuse back to transpose-   case expandAst v of-    Ast.AstVar{} -> t  -- normal form-    Ast.AstPrimalPart Ast.AstVar{} -> t  -- normal form-    Ast.AstDualPart Ast.AstVar{} -> t  -- normal form-    Ast.AstFromPrimal Ast.AstVar{} -> t  -- normal form-    Ast.AstFromDual Ast.AstVar{} -> t  -- normal form-    Ast.AstProject1{} -> t  -- normal form-    Ast.AstProject2{} -> t  -- normal form-    Ast.AstFromIntegralS{} -> t  -- normal form-    Ast.AstCastS{} -> t  -- normal form-    Ast.AstReplicate{} -> t  -- normal form-      -- TODO: this nf is silly, but right now transposes of replicates-      -- are small arrays and equivalent gathers are large terms and arrays,-      -- so this has to stay. Maybe we should contract gathers back-      -- to transposes of replicates (not only to replicates). Or maybe-      -- we should extend orthotope to any gather schemes, not only-      -- the simple ones.-      -- TODO: review also other nfs here and for AstReshapeS below-    Ast.AstScatterS _ _ (_, ix)-     | gcompare (Permutation.permRank perm) (ixsRank ix) == GGT -> t  -- nf-    v2 ->  -- not nf, let's express all as a gather-      astTransposeAsGatherS (defaultKnobs {knobExpand = True})-                            perm v2  -- TODO: (normalizePermutation perm)-        -- this is expensive but the only way to guarantee full simplification-    -} astTransposeS perm (expandAst v)-  Ast.AstReshapeS sh v -> {-  -- too hard to fuse back to reshape-   case expandAst v of-    Ast.AstVar{} -> t  -- normal form-    Ast.AstPrimalPart Ast.AstVar{} -> t  -- normal form-    Ast.AstDualPart Ast.AstVar{} -> t  -- normal form-    Ast.AstFromPrimal Ast.AstVar{} -> t  -- normal form-    Ast.AstFromDual Ast.AstVar{} -> t  -- normal form-    Ast.AstProject1{} -> t  -- normal form-    Ast.AstProject2{} -> t  -- normal form-    Ast.AstFromIntegralS{} -> t  -- normal form-    Ast.AstCastS{} -> t  -- normal form-    AstPlusS{} -> t  -- normal form-    AstTimesS{} -> t  -- normal form-    Ast.AstR2S{} -> t  -- normal form-    Ast.AstScatterS{} -> t  -- normal form-    v2 ->  -- not nf, let's express all as a gather-      astReshapeAsGatherS (defaultKnobs {knobExpand = True})-                          sh v2-        -- this is expensive but the only way to guarantee full simplification-    -} astReshapeS sh (expandAst v)+  Ast.AstTransposeS perm v -> astTransposeS perm (expandAst v)+  Ast.AstReshapeS sh v -> astReshapeS sh (expandAst v)    Ast.AstConvert c v -> astConvert c $ expandAst v    -- These should not appear in this context unless via wacky tests.-  Ast.AstSum0S{} -> t-  Ast.AstDot0S{} -> t+  Ast.AstDot0{} -> t   Ast.AstDot1InS{} -> t   Ast.AstMatmul2S{} -> t -expandAstHFun :: AstSpan s2-              => AstHFun s s2 x y -> AstHFun s s2 x y-expandAstHFun (AstLambda var l) = AstLambda var (expandAst l)+  Ast.AstBoolNotK arg -> notB $ expandAst arg+  Ast.AstBoolNotS arg -> astBoolNotS $ expandAst arg+  Ast.AstBoolAndK arg1 arg2 -> expandAst arg1 &&* expandAst arg2+  Ast.AstBoolAndS arg1 arg2 -> astBoolAndS (expandAst arg1) (expandAst arg2)+  Ast.AstLeqK arg1 arg2 -> fromPlain $ expandAst arg1 <=. expandAst arg2+  Ast.AstLeq arg1 arg2 -> fromPlain $ expandAst arg1 <=. expandAst arg2+  Ast.AstLeqS shb sh arg1 arg2 ->+    fromPlain $ astLeqS shb sh (expandAst arg1) (expandAst arg2) -expandAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet-expandAstBool t = case t of-  AstBoolConst{} -> t-  Ast.AstBoolNot arg -> notB $ expandAstBool arg-  Ast.AstBoolAnd arg1 arg2 -> expandAstBool arg1 &&* expandAstBool arg2-  Ast.AstLeqK arg1 arg2 -> expandAst arg1 <=. expandAst arg2-  Ast.AstLeqS arg1 arg2 -> expandAst arg1 <=. expandAst arg2+expandAstHFun :: KnownSpan s+              => AstHFun s x y -> AstHFun s x y+expandAstHFun (AstLambda var l) = AstLambda var (expandAst l)   -- * The simplifying bottom-up pass@@ -210,22 +175,12 @@ -- is visited and each combinator applied. The most exhaustive and costly -- variants of each combinator are used, e.g., astIndexR. simplifyAst-  :: forall s y. AstSpan s+  :: forall s y. KnownSpan s   => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y simplifyAst t = case t of   Ast.AstPair t1 t2 -> astPair (simplifyAst t1) (simplifyAst t2)   Ast.AstProject1 v -> astProject1 (simplifyAst v)   Ast.AstProject2 v -> astProject2 (simplifyAst v)-  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map simplifyAst l)-  Ast.AstSum snat stk v -> astSum snat stk (simplifyAst v)-  Ast.AstReplicate snat stk v ->  astReplicate snat stk (simplifyAst v)-  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->-    astMapAccumRDer k bftk eftk-                    (simplifyAstHFun f)-                    (simplifyAstHFun df)-                    (simplifyAstHFun rf)-                    (simplifyAst acc0)-                    (simplifyAst es)   Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->     astMapAccumLDer k bftk eftk                     (simplifyAstHFun f)@@ -233,60 +188,71 @@                     (simplifyAstHFun rf)                     (simplifyAst acc0)                     (simplifyAst es)-  Ast.AstApply v ll -> astApply (simplifyAstHFun v) (simplifyAst ll)-  Ast.AstVar var -> astVar var-  Ast.AstCond b a2 a3 ->-    astCond (simplifyAstBool b) (simplifyAst a2) (simplifyAst a3)+  Ast.AstApply f a -> astApply (simplifyAstHFun f) (simplifyAst a)+  Ast.AstVar{} -> t   Ast.AstBuild1 k stk (var, v) ->     let !v2 = simplifyAst v     in Ast.AstBuild1 k stk (var, v2) -  Ast.AstLet var u v -> astLet var (simplifyAst u) (simplifyAst v)+  Ast.AstLet var u v ->+    astLet var (withKnownSpan (varNameToSpan var) $ simplifyAst u)+               (simplifyAst v)    Ast.AstPrimalPart v -> astPrimalPart (simplifyAst v)   Ast.AstDualPart v -> astDualPart (simplifyAst v)-  Ast.AstFromPrimal v -> Ast.AstFromPrimal (simplifyAst v)-  Ast.AstFromDual v -> Ast.AstFromDual (simplifyAst v)+  Ast.AstPlainPart v -> astPlainPart (simplifyAst v)+  Ast.AstFromPrimal v -> fromPrimal (simplifyAst v)+  Ast.AstFromDual v -> fromDual (simplifyAst v)+  Ast.AstFromPlain v -> fromPlain (simplifyAst v)    AstPlusK u v -> simplifyAst u + simplifyAst v   AstTimesK u v -> simplifyAst u * simplifyAst v-  Ast.AstN1K NegateOp u -> negate (simplifyAst u)-  Ast.AstN1K AbsOp u -> abs (simplifyAst u)-  Ast.AstN1K SignumOp u -> signum (simplifyAst u)-  Ast.AstR1K opCode u -> Ast.AstR1K opCode (simplifyAst u)-  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (simplifyAst u) (simplifyAst v)-  Ast.AstI2K QuotOp u v -> quotH (simplifyAst u) (simplifyAst v)-  Ast.AstI2K RemOp u v -> remH (simplifyAst u) (simplifyAst v)-  AstConcreteK k -> AstConcreteK k+  Ast.AstN1K opCode u -> astN1K opCode (simplifyAst u)+  Ast.AstR1K opCode u -> astR1K opCode (simplifyAst u)+  Ast.AstR2K opCode u v -> astR2K opCode (simplifyAst u) (simplifyAst v)+  Ast.AstI2K opCode u v -> astI2K opCode (simplifyAst u) (simplifyAst v)+  AstConcreteK{} -> t   Ast.AstFloorK a -> astFloorK (simplifyAst a)   Ast.AstFromIntegralK v -> astFromIntegralK $ simplifyAst v   Ast.AstCastK v -> astCastK $ simplifyAst v+  Ast.AstArgMinK v -> astArgMinK $ simplifyAst v+  Ast.AstArgMaxK v -> astArgMaxK $ simplifyAst v+  Ast.AstIndexK v ix ->+    astIndexKnobsK (defaultKnobs {knobPhase = PhaseSimplification})+                   (simplifyAst v) (simplifyAstIxS ix)    AstPlusS u v -> simplifyAst u + simplifyAst v   AstTimesS u v -> simplifyAst u * simplifyAst v-  Ast.AstN1S NegateOp u -> negate (simplifyAst u)-  Ast.AstN1S AbsOp u -> abs (simplifyAst u)-  Ast.AstN1S SignumOp u -> signum (simplifyAst u)-  Ast.AstR1S opCode u -> Ast.AstR1S opCode (simplifyAst u)-  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (simplifyAst u) (simplifyAst v)-  Ast.AstI2S QuotOp u v -> quotH (simplifyAst u) (simplifyAst v)-  Ast.AstI2S RemOp u v -> remH (simplifyAst u) (simplifyAst v)-  AstConcreteS a -> AstConcreteS a+  Ast.AstN1S opCode u -> astN1S opCode (simplifyAst u)+  Ast.AstR1S opCode u -> astR1S opCode (simplifyAst u)+  Ast.AstR2S opCode u v -> astR2S opCode (simplifyAst u) (simplifyAst v)+  Ast.AstI2S opCode u v -> astI2S opCode (simplifyAst u) (simplifyAst v)+  AstConcreteS{} -> t   Ast.AstFloorS a -> astFloorS (simplifyAst a)   Ast.AstFromIntegralS v -> astFromIntegralS $ simplifyAst v   Ast.AstCastS v -> astCastS $ simplifyAst v-+  Ast.AstArgMinS a -> Ast.AstArgMinS (simplifyAst a)+  Ast.AstArgMaxS a -> Ast.AstArgMaxS (simplifyAst a)   Ast.AstIndexS shn v ix ->     astIndexKnobsS (defaultKnobs {knobPhase = PhaseSimplification})                    shn (simplifyAst v) (simplifyAstIxS ix)-  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->-    astScatterS @shm @shn @shp shn (simplifyAst v) (vars, simplifyAstIxS ix)-  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->-    astGatherKnobsS @shm @shn @shp-                    (defaultKnobs {knobPhase = PhaseSimplification})-                    shn (simplifyAst v) (vars, simplifyAstIxS ix)-  Ast.AstMinIndexS a -> Ast.AstMinIndexS (simplifyAst a)-  Ast.AstMaxIndexS a -> Ast.AstMaxIndexS (simplifyAst a)++  Ast.AstCondK b a2 a3 ->+    astCondK (simplifyAst b) (simplifyAst a2) (simplifyAst a3)+  Ast.AstCondS b a2 a3 ->+    astCondS (simplifyAst b) (simplifyAst a2) (simplifyAst a3)+  Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map simplifyAst l)+  Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map simplifyAst l)+  Ast.AstSumK a -> astSumK (simplifyAst a)+  Ast.AstSumS shm a -> astSumS shm (simplifyAst a)+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterKnobsS (defaultKnobs {knobPhase = PhaseSimplification})+                     shm shn shp (simplifyAst v) (vars, simplifyAstIxS ix)+  Ast.AstReplicateK shm v -> astReplicateK shm (simplifyAst v)+  Ast.AstReplicateS shm v -> astReplicateS shm (simplifyAst v)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherKnobsS (defaultKnobs {knobPhase = PhaseSimplification})+                    shm shn shp (simplifyAst v) (vars, simplifyAstIxS ix)   Ast.AstIotaS{} -> t   Ast.AstAppendS x y -> astAppendS (simplifyAst x) (simplifyAst y)   Ast.AstSliceS i n k v -> astSliceS i n k (simplifyAst v)@@ -297,22 +263,22 @@   Ast.AstConvert c v -> astConvert c $ simplifyAst v    -- These should not appear in this context unless via wacky tests.-  Ast.AstSum0S{} -> t-  Ast.AstDot0S{} -> t+  Ast.AstDot0{} -> t   Ast.AstDot1InS{} -> t   Ast.AstMatmul2S{} -> t -simplifyAstHFun :: AstSpan s2-                => AstHFun s s2 x y -> AstHFun s s2 x y-simplifyAstHFun (AstLambda var l) = AstLambda var (simplifyAst l)+  Ast.AstBoolNotK arg -> notB $ simplifyAst arg+  Ast.AstBoolNotS arg -> astBoolNotS $ simplifyAst arg+  Ast.AstBoolAndK arg1 arg2 -> simplifyAst arg1 &&* simplifyAst arg2+  Ast.AstBoolAndS arg1 arg2 -> astBoolAndS (simplifyAst arg1) (simplifyAst arg2)+  Ast.AstLeqK arg1 arg2 -> fromPlain $ simplifyAst arg1 <=. simplifyAst arg2+  Ast.AstLeq arg1 arg2 -> fromPlain $ simplifyAst arg1 <=. simplifyAst arg2+  Ast.AstLeqS shb sh arg1 arg2 ->+    fromPlain $ astLeqS shb sh (simplifyAst arg1) (simplifyAst arg2) -simplifyAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet-simplifyAstBool t = case t of-  AstBoolConst{} -> t-  Ast.AstBoolNot arg -> notB $ simplifyAstBool arg-  Ast.AstBoolAnd arg1 arg2 -> simplifyAstBool arg1 &&* simplifyAstBool arg2-  Ast.AstLeqK arg1 arg2 -> simplifyAst arg1 <=. simplifyAst arg2-  Ast.AstLeqS arg1 arg2 -> simplifyAst arg1 <=. simplifyAst arg2+simplifyAstHFun :: KnownSpan s+                => AstHFun s x y -> AstHFun s x y+simplifyAstHFun (AstLambda var l) = AstLambda var (simplifyAst l)   -- * The contraction (e.g., from gather expressions) bottom-up pass@@ -330,33 +296,142 @@ -- is not written in a compositional style nor close to it, -- but it's instead defined in an ad-hoc way based on benchmarks. contractAst-  :: forall s y. AstSpan s+  :: forall s y. KnownSpan s   => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y contractAst t0 = case t0 of   Ast.AstPair t1 t2 -> astPair (contractAst t1) (contractAst t2)   Ast.AstProject1 v -> astProject1 (contractAst v)   Ast.AstProject2 v -> astProject2 (contractAst v)-  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map contractAst l)-  Ast.AstSum _ (STKS ZSS _) t2 -> astSum0S (contractAst t2)-  Ast.AstSum _ STKScalar t2 -> astKFromS' $ astSum0S (contractAst t2)-  Ast.AstSum-    snat@(SNat @m2)-    stk@(STKS (SNat @n2 :$$ SNat @p2 :$$ ZSS) STKScalar)-    v@(AstTimesS (Ast.AstTransposeS @permt permt-                    (Ast.AstReplicate (SNat @kt) (STKS @sht _ _) t2))-                 (Ast.AstTransposeS @permu permu-                    (Ast.AstReplicate (SNat @ku) (STKS @shu _ _) u2))) ->+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    astMapAccumLDer k bftk eftk+                    (contractAstHFun f)+                    (contractAstHFun df)+                    (contractAstHFun rf)+                    (contractAst acc0)+                    (contractAst es)+  Ast.AstApply v ll -> astApply (contractAstHFun v) (contractAst ll)+  Ast.AstVar{} -> t0+  -- These are only needed for tests that don't vectorize Ast. In fact,+  -- they seem totally unused ATM.+  Ast.AstBuild1 snat stk@(STKS ZSS STKScalar)  -- generalize+                (var, v@(AstConvUpSFromK (Ast.AstSumK+                           (AstTimesS+                              t2+                              (Ast.AstIndexS @shm @shn (n :$$ ZSS)+                                 u ix@(AstIntVar var2 :.$ ZIS))))))+    | var == var2+    , not (varNameInAst var t2), not (varNameInAst var u)+    , FTKS shmshn _ <- ftkAst u ->+      withKnownShS (shsTakeIx @shm @shn Proxy ix shmshn) $+      case knownShS @shm of+        snat2 :$$ _ | Just Refl <- testEquality snat snat2 ->+          astDot1InS (snat :$$ ZSS) n+                     (contractAst u)+                     (contractAst $ Ast.AstReplicateS (snat :$$ ZSS) t2)+        _ ->+          let !v2 = contractAst v+          in Ast.AstBuild1 snat stk (var, v2)+  Ast.AstBuild1 snat stk@STKScalar+                (var, v@(Ast.AstSumK+                           (AstTimesS+                              t2+                              (Ast.AstIndexS @shm @shn (n :$$ ZSS)+                                 u ix@(AstIntVar var2 :.$ ZIS)))))+    | var == var2+    , not (varNameInAst var t2), not (varNameInAst var u)+    , FTKS shmshn _ <- ftkAst u ->+      withKnownShS (shsTakeIx @shm @shn Proxy ix shmshn) $+      case knownShS @shm of+        snat2 :$$ _ | Just Refl <- testEquality snat snat2 ->+          astDot1InS (snat :$$ ZSS) n+                     (contractAst u)+                     (contractAst $ Ast.AstReplicateS (snat :$$ ZSS) t2)+        _ ->+          let !v2 = contractAst v+          in Ast.AstBuild1 snat stk (var, v2)+  Ast.AstBuild1 snat stk (var, v) ->+    let !v2 = contractAst v+    in Ast.AstBuild1 snat stk (var, v2)++  Ast.AstLet var u v ->+    astLet var (withKnownSpan (varNameToSpan var) $ contractAst u)+               (contractAst v)++  Ast.AstPrimalPart v -> astPrimalPart (contractAst v)+  Ast.AstDualPart v -> astDualPart (contractAst v)+  Ast.AstPlainPart v -> astPlainPart (contractAst v)+  Ast.AstFromPrimal v -> fromPrimal (contractAst v)+  Ast.AstFromDual v -> fromDual (contractAst v)+  Ast.AstFromPlain v -> fromPlain (contractAst v)++  AstPlusK u v -> contractAst u + contractAst v+  AstTimesK u v -> contractAst u * contractAst v+  Ast.AstN1K opCode u -> astN1K opCode (contractAst u)+  Ast.AstR1K opCode u -> astR1K opCode (contractAst u)+  Ast.AstR2K opCode u v -> astR2K opCode (contractAst u) (contractAst v)+  Ast.AstI2K opCode u v -> astI2K opCode (contractAst u) (contractAst v)+  AstConcreteK{} -> t0+  Ast.AstFloorK a -> astFloorK (contractAst a)+  Ast.AstFromIntegralK v -> astFromIntegralK $ contractAst v+  Ast.AstCastK v -> astCastK $ contractAst v+  Ast.AstArgMinK v -> astArgMinK $ contractAst v+  Ast.AstArgMaxK v -> astArgMaxK $ contractAst v+  Ast.AstIndexK v ix ->+    astIndexKnobsK (defaultKnobs {knobPhase = PhaseContraction})+                   (contractAst v) (contractAstIxS ix)++  AstPlusS u v -> contractAst u + contractAst v+  AstTimesS u v -> contractAst u * contractAst v+  Ast.AstN1S opCode u -> astN1S opCode (contractAst u)+  Ast.AstR1S opCode u -> astR1S opCode (contractAst u)+  Ast.AstR2S opCode u v -> astR2S opCode (contractAst u) (contractAst v)+  Ast.AstI2S opCode u v -> astI2S opCode (contractAst u) (contractAst v)+  AstConcreteS{} -> t0+  Ast.AstFloorS @r1 @r2 t -> case contractAst t of+    AstConcreteS a | sizeOf (undefined :: r1) >= sizeOf (undefined :: r2) ->+      fromPlain $ astConcreteS $ tsfloor $ Concrete a+    t2 -> astFloorS t2+  Ast.AstFromIntegralS @r1 @r2 t -> case contractAst t of+    AstConcreteS a | sizeOf (undefined :: r1) >= sizeOf (undefined :: r2) ->+      fromPlain $ astConcreteS $ tsfromIntegral $ Concrete a+    t2 -> astFromIntegralS t2+  Ast.AstCastS @r1 @r2 t -> case contractAst t of+    AstConcreteS a | sizeOf (undefined :: r1) >= sizeOf (undefined :: r2) ->+      astConcreteS (tscast $ Concrete a)+    t2 -> astCastS t2+  Ast.AstArgMinS a -> Ast.AstArgMinS (contractAst a)+  Ast.AstArgMaxS a -> Ast.AstArgMaxS (contractAst a)+  Ast.AstIndexS shn v ix ->+    astIndexKnobsS (defaultKnobs {knobPhase = PhaseContraction})+                   shn (contractAst v) (contractAstIxS ix)++  Ast.AstCondK b a2 a3 ->+    astCondK (contractAst b) (contractAst a2) (contractAst a3)+  Ast.AstCondS b a2 a3 ->+    astCondS (contractAst b) (contractAst a2) (contractAst a3)+  Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map contractAst l)+  Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map contractAst l)+  Ast.AstSumK v -> astSumKContract $ contractAst v+  Ast.AstSumS (n :$$ shmRest@(_ :$$ _)) v@AstTimesS{} ->+    contractAst (Ast.AstSumS shmRest (Ast.AstSumS (n :$$ ZSS) v))+  Ast.AstSumS+    (_ :$$ ZSS)+    v@(AstTimesS (Ast.AstTransposeS @permt @nsht permt t1)+                 (Ast.AstTransposeS @permu @nshu permu u1))+    | FTKS (SNat @m2 :$$ SNat @n2 :$$ SNat @p2 :$$ ZSS) FTKScalar <- ftkAst v+    , FTKS ((:$$) @_ @sht _ _) FTKScalar <- ftkAst t1+    , Just t2 <- unRepl1 t1+    , FTKS ((:$$) @_ @shu _ _) FTKScalar <- ftkAst u1+    , Just u2 <- unRepl1 u1 ->     let perm10 = Permutation.makePerm @'[1, 0]-    in fromMaybe (astSum snat stk (contractAst v))+    in fromMaybe (astSumS (SNat @m2 :$$ ZSS) (contractAst v))        $ case (permt, permu) of       ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil        ,SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->         gcastWith (unsafeCoerceRefl-                   :: Permutation.PermutePrefix permt (kt ': sht)-                      :~: [m2, n2, p2]) $+                   :: Permutation.PermutePrefix permt nsht :~: [m2, n2, p2]) $         gcastWith (unsafeCoerceRefl-                   :: Permutation.PermutePrefix permu (ku ': shu)-                      :~: [m2, n2, p2]) $+                   :: Permutation.PermutePrefix permu nshu :~: [m2, n2, p2]) $         -- Sadly, the casts below, though implied by the permutations         -- (as redundantly spelled out by the casts above) are required         -- to make it type-check and they easily mask bugs, too.@@ -402,184 +477,178 @@         attemptMatmul2 (astTransposeS perm10 u2)                        (astTransposeS perm10 t2)       _ -> Nothing-  Ast.AstSum n@(SNat @n) (STKS @sh sh _) (AstTimesS t2 u) ->+  Ast.AstSumS+    (_ :$$ ZSS)+    v@(AstTimesS (Ast.AstFromPlain+                    (Ast.AstTransposeS @permt @nsht permt t1'))+                 (Ast.AstTransposeS @permu @nshu permu u1))+    | FTKS (SNat @m2 :$$ SNat @n2 :$$ SNat @p2 :$$ ZSS) FTKScalar <- ftkAst v+    , FTKS ((:$$) @_ @sht _ _) FTKScalar <- ftkAst t1'+    , Just t2' <- unRepl1 t1'+    , FTKS ((:$$) @_ @shu _ _) FTKScalar <- ftkAst u1+    , Just u2 <- unRepl1 u1 ->+    let perm10 = Permutation.makePerm @'[1, 0]+        t2 = fromPlain @s t2'+    in fromMaybe (astSumS (SNat @m2 :$$ ZSS) (contractAst v))+       $ case (permt, permu) of+      ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix permt nsht :~: [m2, n2, p2]) $+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix permu nshu :~: [m2, n2, p2]) $+        -- Sadly, the casts below, though implied by the permutations+        -- (as redundantly spelled out by the casts above) are required+        -- to make it type-check and they easily mask bugs, too.+        -- In the result, this is as type-unsafe as ranked code would be.+        gcastWith (unsafeCoerceRefl :: sht :~: [n2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, p2]) $+        attemptMatmul2 t2 u2+      ( SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, p2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [n2, m2]) $+        attemptMatmul2 u2 t2+      ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [n2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [p2, m2]) $+        attemptMatmul2 t2 (astTransposeS perm10 u2)+      ( SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [p2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [n2, m2]) $+        attemptMatmul2 u2 (astTransposeS perm10 t2)+      ( SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, n2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, p2]) $+        attemptMatmul2 (astTransposeS perm10 t2) u2+      ( SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, p2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, n2]) $+        attemptMatmul2 (astTransposeS perm10 u2) t2+      ( SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, n2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [p2, m2]) $+        attemptMatmul2 (astTransposeS perm10 t2)+                       (astTransposeS perm10 u2)+      ( SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [p2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, n2]) $+        attemptMatmul2 (astTransposeS perm10 u2)+                       (astTransposeS perm10 t2)+      _ -> Nothing+  Ast.AstSumS+    (_ :$$ ZSS)+    v@(AstTimesS (Ast.AstTransposeS @permt @nsht permt t1)+                 (Ast.AstFromPlain+                    (Ast.AstTransposeS @permu @nshu permu u1')))+    | FTKS (SNat @m2 :$$ SNat @n2 :$$ SNat @p2 :$$ ZSS) FTKScalar <- ftkAst v+    , FTKS ((:$$) @_ @sht _ _) FTKScalar <- ftkAst t1+    , Just t2 <- unRepl1 t1+    , FTKS ((:$$) @_ @shu _ _) FTKScalar <- ftkAst u1'+    , Just u2' <- unRepl1 u1' ->+    let perm10 = Permutation.makePerm @'[1, 0]+        u2 = fromPlain @s u2'+    in fromMaybe (astSumS (SNat @m2 :$$ ZSS) (contractAst v))+       $ case (permt, permu) of+      ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix permt nsht :~: [m2, n2, p2]) $+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix permu nshu :~: [m2, n2, p2]) $+        -- Sadly, the casts below, though implied by the permutations+        -- (as redundantly spelled out by the casts above) are required+        -- to make it type-check and they easily mask bugs, too.+        -- In the result, this is as type-unsafe as ranked code would be.+        gcastWith (unsafeCoerceRefl :: sht :~: [n2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, p2]) $+        attemptMatmul2 t2 u2+      ( SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, p2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [n2, m2]) $+        attemptMatmul2 u2 t2+      ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [n2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [p2, m2]) $+        attemptMatmul2 t2 (astTransposeS perm10 u2)+      ( SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [p2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [n2, m2]) $+        attemptMatmul2 u2 (astTransposeS perm10 t2)+      ( SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, n2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, p2]) $+        attemptMatmul2 (astTransposeS perm10 t2) u2+      ( SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, p2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, n2]) $+        attemptMatmul2 (astTransposeS perm10 u2) t2+      ( SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, n2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [p2, m2]) $+        attemptMatmul2 (astTransposeS perm10 t2)+                       (astTransposeS perm10 u2)+      ( SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [p2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, n2]) $+        attemptMatmul2 (astTransposeS perm10 u2)+                       (astTransposeS perm10 t2)+      _ -> Nothing+  Ast.AstSumS (n@(SNat @n) :$$ ZSS) (AstTimesS t2 u)+    | FTKS @sh1 (_ :$$ sh) _ <- ftkAst t2 ->     let cpermR = backpermCycle $ 1 + shsLength sh-    in Permutation.permFromList cpermR $ \(cperm :: Permutation.Perm cperm) ->-         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (n : sh)) $+    in Permutation.permFromListCont cpermR $ \(cperm+                                               :: Permutation.Perm cperm) ->+         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank sh1) $          gcastWith (unsafeCoerceRefl-                    :: Permutation.PermutePrefix cperm (n : sh)-                       :~: sh ++ '[n]) $+                    :: Permutation.PermutePrefix cperm sh1+                       :~: Tail sh1 ++ '[n]) $          fromMaybe (error "contractAst: impossible non-permutation")          $ Permutation.permCheckPermutation cperm          $ astDot1InS sh n (contractAst $ Ast.AstTransposeS cperm t2)                            (contractAst $ Ast.AstTransposeS cperm u)-  Ast.AstSum snat stk (AstTimesS (Ast.AstLet vart vt t2) u) ->-    {- TODO: do we really need this check?-    | not (varNameInAst vart u) ->-        -- The varNameInAst check is needed, because although variable-        -- capture is impossible, because we don't create nested lets-        -- with the same variable, we could create such nested lets-        -- if we omitted this check.-    -}-    astLet vart-           (contractAst vt)-           (contractAst $ Ast.AstSum snat stk  -- the crucial exposed redex-                                     (AstTimesS t2 u))-  Ast.AstSum snat stk (AstTimesS t2 (Ast.AstLet varu vu u)) ->-    astLet varu-           (contractAst vu)-           (contractAst $ Ast.AstSum snat stk (AstTimesS t2 u))-  Ast.AstSum snat stk (Ast.AstLet var v t2) ->-    astLet var (contractAst v) (contractAst (Ast.AstSum snat stk t2))-  Ast.AstSum snat stk (Ast.AstSum snat2 stk2 (Ast.AstLet var v t2)) ->-    astLet var (contractAst v)-               (contractAst (Ast.AstSum snat stk (Ast.AstSum snat2 stk2 t2)))-  Ast.AstSum snat stk v -> astSum snat stk (contractAst v)-  Ast.AstReplicate snat stk v -> astReplicate snat stk (contractAst v)-  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->-    astMapAccumRDer k bftk eftk-                    (contractAstHFun f)-                    (contractAstHFun df)-                    (contractAstHFun rf)-                    (contractAst acc0)-                    (contractAst es)-  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->-    astMapAccumLDer k bftk eftk-                    (contractAstHFun f)-                    (contractAstHFun df)-                    (contractAstHFun rf)-                    (contractAst acc0)-                    (contractAst es)-  Ast.AstApply v ll -> astApply (contractAstHFun v) (contractAst ll)-  Ast.AstVar var -> astVar var-  Ast.AstCond b a2 a3 ->-    astCond (contractAstBool b) (contractAst a2) (contractAst a3)-  -- These are only needed for tests that don't vectorize Ast.-  Ast.AstBuild1 snat (STKS ZSS _)  -- generalize-                (var, Ast.AstSum-                        n _-                        (AstTimesS-                           t2-                           (Ast.AstIndexS _shn-                              u (((:.$) @m (AstIntVar var2) ZIS)))))-    | Just Refl <- testEquality snat (SNat @m)-    , var == var2-    , not (varNameInAst var t2), not (varNameInAst var u) ->-        astDot1InS (snat :$$ ZSS) n-                   (contractAst u)-                   (contractAst-                    $ Ast.AstReplicate snat (ftkToSTK (ftkAst t2)) t2)-  Ast.AstBuild1 snat (STKS ZSS _)-                (var, Ast.AstSum _ _-                        (Ast.AstReshapeS-                           _sh (AstTimesS-                                  t2-                                  (Ast.AstIndexS _shn-                                     u (((:.$) @m (AstIntVar var2) ZIS))))))-    | ftk2@(FTKS (n :$$ ZSS) _) <- ftkAst t2-    , Just Refl <- testEquality snat (SNat @m)-    , var == var2-    , not (varNameInAst var t2), not (varNameInAst var u) ->-        astDot1InS (snat :$$ ZSS) n-                   (contractAst u)-                   (contractAst $ Ast.AstReplicate snat (ftkToSTK ftk2) t2)-  Ast.AstBuild1 k stk (var, v) ->-    let !v2 = contractAst v-    in Ast.AstBuild1 k stk (var, v2)--  Ast.AstLet var u v -> astLet var (contractAst u) (contractAst v)--  Ast.AstPrimalPart v -> astPrimalPart (contractAst v)-  Ast.AstDualPart v -> astDualPart (contractAst v)-  Ast.AstFromPrimal v -> Ast.AstFromPrimal (contractAst v)-  Ast.AstFromDual v -> Ast.AstFromDual (contractAst v)--  AstPlusK u v -> contractAst u + contractAst v-  AstTimesK u v -> contractAst u * contractAst v-  Ast.AstN1K NegateOp u -> negate (contractAst u)-  Ast.AstN1K AbsOp u -> abs (contractAst u)-  Ast.AstN1K SignumOp u -> signum (contractAst u)-  Ast.AstR1K RecipOp u -> recip (contractAst u)-  Ast.AstR1K ExpOp u -> exp (contractAst u)-  Ast.AstR1K LogOp u -> log (contractAst u)-  Ast.AstR1K SqrtOp u -> sqrt (contractAst u)-  Ast.AstR1K SinOp u -> sin (contractAst u)-  Ast.AstR1K CosOp u -> cos (contractAst u)-  Ast.AstR1K TanOp u -> tan (contractAst u)-  Ast.AstR1K AsinOp u -> asin (contractAst u)-  Ast.AstR1K AcosOp u -> acos (contractAst u)-  Ast.AstR1K AtanOp u -> atan (contractAst u)-  Ast.AstR1K SinhOp u -> sinh (contractAst u)-  Ast.AstR1K CoshOp u -> cosh (contractAst u)-  Ast.AstR1K TanhOp u -> tanh (contractAst u)-  Ast.AstR1K AsinhOp u -> asinh (contractAst u)-  Ast.AstR1K AcoshOp u -> acosh (contractAst u)-  Ast.AstR1K AtanhOp u -> atanh (contractAst u)-  Ast.AstR2K DivideOp u v -> contractAst u / contractAst v-  Ast.AstR2K PowerOp u v -> contractAst u ** contractAst v-  Ast.AstR2K LogBaseOp u v -> logBase (contractAst u) (contractAst v)-  Ast.AstR2K Atan2Op u v -> atan2H (contractAst u) (contractAst v)-  Ast.AstI2K QuotOp u v -> quotH (contractAst u) (contractAst v)-  Ast.AstI2K RemOp u v -> remH (contractAst u) (contractAst v)-  AstConcreteK k -> AstConcreteK k-  Ast.AstFloorK a -> astFloorK (contractAst a)-  Ast.AstFromIntegralK v -> astFromIntegralK $ contractAst v-  Ast.AstCastK v -> astCastK $ contractAst v--  AstPlusS u v -> contractAst u + contractAst v-  AstTimesS u v -> contractAst u * contractAst v-  Ast.AstN1S NegateOp u -> negate (contractAst u)-  Ast.AstN1S AbsOp u -> abs (contractAst u)-  Ast.AstN1S SignumOp u -> signum (contractAst u)-  Ast.AstR1S RecipOp u -> recip (contractAst u)-  Ast.AstR1S ExpOp u -> exp (contractAst u)-  Ast.AstR1S LogOp u -> log (contractAst u)-  Ast.AstR1S SqrtOp u -> sqrt (contractAst u)-  Ast.AstR1S SinOp u -> sin (contractAst u)-  Ast.AstR1S CosOp u -> cos (contractAst u)-  Ast.AstR1S TanOp u -> tan (contractAst u)-  Ast.AstR1S AsinOp u -> asin (contractAst u)-  Ast.AstR1S AcosOp u -> acos (contractAst u)-  Ast.AstR1S AtanOp u -> atan (contractAst u)-  Ast.AstR1S SinhOp u -> sinh (contractAst u)-  Ast.AstR1S CoshOp u -> cosh (contractAst u)-  Ast.AstR1S TanhOp u -> tanh (contractAst u)-  Ast.AstR1S AsinhOp u -> asinh (contractAst u)-  Ast.AstR1S AcoshOp u -> acosh (contractAst u)-  Ast.AstR1S AtanhOp u -> atanh (contractAst u)-  Ast.AstR2S DivideOp u v -> contractAst u / contractAst v-  Ast.AstR2S PowerOp u v -> contractAst u ** contractAst v-  Ast.AstR2S LogBaseOp u v -> logBase (contractAst u) (contractAst v)-  Ast.AstR2S Atan2Op u v -> atan2H (contractAst u) (contractAst v)-  Ast.AstI2S QuotOp u v -> quotH (contractAst u) (contractAst v)-  Ast.AstI2S RemOp u v -> remH (contractAst u) (contractAst v)-  AstConcreteS a -> AstConcreteS a-  Ast.AstFloorS t -> case contractAst t of-    AstConcreteS a -> astConcreteS (tsfloor $ Concrete a)-    t2 -> astFloorS t2-  Ast.AstFromIntegralS t -> case contractAst t of-    AstConcreteS a -> astConcreteS (tsfromIntegral $ Concrete a)-    t2 -> astFromIntegralS t2-  Ast.AstCastS t -> case contractAst t of-    AstConcreteS a -> astConcreteS (tscast $ Concrete a)-    t2 -> astCastS t2--  Ast.AstIndexS shn v ix ->-    astIndexKnobsS (defaultKnobs {knobPhase = PhaseContraction})-                   shn (contractAst v) (contractAstIxS ix)-  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->-    astScatterS @shm @shn @shp shn (contractAst v) (vars, contractAstIxS ix)+  Ast.AstSumS shm (AstTimesS (Ast.AstLet vart vt t2) u) ->+    astLetRefresh+      vart+      (withKnownSpan (varNameToSpan vart) $ contractAst vt)+      t2+      $ \t2' ->+        contractAst $ Ast.AstSumS shm (AstTimesS t2' u)+                                         -- the crucial exposed redex+  Ast.AstSumS shm (AstTimesS t2 (Ast.AstLet varu vu u)) ->+    astLetRefresh+      varu+      (withKnownSpan (varNameToSpan varu) $ contractAst vu)+      u+      $ contractAst . Ast.AstSumS shm . AstTimesS t2+  Ast.AstSumS shm v -> astSumS shm $ contractAst v+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    astScatterKnobsS (defaultKnobs {knobPhase = PhaseContraction})+                     shm shn shp (contractAst v) (vars, contractAstIxS ix)+  Ast.AstReplicateK shm v -> astReplicateK shm (contractAst v)+  Ast.AstReplicateS shm v -> astReplicateS shm (contractAst v)   -- This rule is reverted in vectorization, so contraction phase may be fine.-  Ast.AstGatherS shn v (vars, Ast.AstCond b i1 i2 :.$ prest)-    | not $ any ((`varInAstBool` b) . varNameToAstVarId) (listsToList vars) ->+  Ast.AstGatherS shm shn shp v (vars, Ast.AstCondK b i1 i2 :.$ prest)+    | not $ Foldable.any ((`varInAst` b) . varNameToAstVarId) vars ->       contractAst-      $ Ast.AstCond b-                    (Ast.AstGatherS shn v (vars, i1 :.$ prest))-                    (Ast.AstGatherS shn v (vars, i2 :.$ prest))-  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->-    astGatherKnobsS @shm @shn @shp-                    (defaultKnobs {knobPhase = PhaseContraction})-                    shn (contractAst v) (vars, contractAstIxS ix)+      $ Ast.AstCondS b (Ast.AstGatherS shm shn shp v (vars, i1 :.$ prest))+                       (Ast.AstGatherS shm shn shp v (vars, i2 :.$ prest))+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    astGatherKnobsS (defaultKnobs {knobPhase = PhaseContraction})+                    shm shn shp (contractAst v) (vars, contractAstIxS ix) {- TODO, but sbuild is tricky, so only if benchmarks show it's worth it:   AstGatherS @shm AstIotaS (vars, i :.$ ZIS) | Refl <- lemAppNil @shm ->     gcastWith (unsafeCoerceRefl :: Drop (Rank shm) shm :~: '[]) $@@ -587,10 +656,10 @@     sbuild @_ @_ @(Rank shm)            (interpretLambdaIndexS               interpretAst env-              (vars, fromPrimal @s $ AstFromIntegralS $ AstSFromK i)) -}-  Ast.AstMinIndexS a -> Ast.AstMinIndexS (contractAst a)-  Ast.AstMaxIndexS a -> Ast.AstMaxIndexS (contractAst a)-  Ast.AstIotaS (SNat @n) -> astConcreteS $ tsiota @_ @n+              (vars, fromPrimal @s $ AstFromIntegralS $ AstConvUpSFromK i)) -}+  Ast.AstIotaS snat@(SNat @n) | fromSNat' snat < 100 ->+    astConcreteS $ tsiota @_ @n  -- likely not to be O(data size)+  Ast.AstIotaS{} -> t0  -- tough trade-offs here   Ast.AstAppendS x y -> astAppendS (contractAst x) (contractAst y)   Ast.AstSliceS i n k t -> astSliceS i n k (contractAst t)   Ast.AstReverseS t -> astReverseS (contractAst t)@@ -599,48 +668,178 @@     AstConcreteS v -> astConcreteS (tsreshape sh2 $ Concrete v)     t2 -> astReshapeS sh2 t2 -  Ast.AstConvert c v -> astConvert c $ contractAst v+  Ast.AstConvert c v -> astConvertConcrete c $ contractAst v -  -- These should not appear in this context unless via wacky tests.-  Ast.AstSum0S{} -> t0-  Ast.AstDot0S{} -> t0-  Ast.AstDot1InS{} -> t0-  Ast.AstMatmul2S{} -> t0+  -- These can appear in this context, e.g., if contractAst is run twice.+  Ast.AstDot0 u v -> astDot0 u v+  Ast.AstDot1InS sh n u v -> astDot1InS sh n u v+  Ast.AstMatmul2S m n p u v -> astMatmul2S m n p u v +  Ast.AstBoolNotK arg -> notB $ contractAst arg+  Ast.AstBoolNotS arg -> astBoolNotS $ contractAst arg+  Ast.AstBoolAndK arg1 arg2 -> contractAst arg1 &&* contractAst arg2+  Ast.AstBoolAndS arg1 arg2 -> astBoolAndS (contractAst arg1) (contractAst arg2)+  Ast.AstLeqK arg1 arg2 -> fromPlain $ contractAst arg1 <=. contractAst arg2+  Ast.AstLeq arg1 arg2 -> fromPlain $ contractAst arg1 <=. contractAst arg2+  Ast.AstLeqS shb sh arg1 arg2 ->+    fromPlain $ astLeqS shb sh (contractAst arg1) (contractAst arg2)++contractAstHFun :: KnownSpan s+                => AstHFun s x y -> AstHFun s x y+contractAstHFun (AstLambda var l) = AstLambda var (contractAst l)++astConvertConcrete :: forall y z s. KnownSpan s+                   => TKConversion y z+                   -> AstTensor AstMethodLet s y+                   -> AstTensor AstMethodLet s z+astConvertConcrete c a0 = case a0 of+  AstConcreteK a -> astConcreteKeepShaped (convertFTK c FTKScalar)+                    $ tconvert c STKScalar $ Concrete a+  AstConcreteS a -> let ftk = FTKS (Nested.sshape a) FTKScalar+                    in astConcreteKeepShaped (convertFTK c ftk)+                       $ tconvert c (ftkToSTK ftk) $ Concrete a+  Ast.AstPrimalPart a -> astPrimalPart $ astConvertConcrete c a+  Ast.AstDualPart a -> astDualPart $ astConvertConcrete c a+  Ast.AstPlainPart a -> astPlainPart $ astConvertConcrete c a+  Ast.AstFromPrimal a -> fromPrimal $ astConvertConcrete c a+  Ast.AstFromDual a -> fromDual $ astConvertConcrete c a+  Ast.AstFromPlain a -> fromPlain $ astConvertConcrete c a+  Ast.AstConvert c2 a2 -> astConvertConcrete (c `convCmp` c2) a2+  _ -> astConvert c a0++astConcreteKeepShaped :: FullShapeTK y -> Concrete y+                      -> AstTensor AstMethodLet PlainSpan y+astConcreteKeepShaped ftk v = case ftk of+  FTKS _ FTKScalar -> astConcreteS v+  _ -> astConcrete ftk v+ attemptMatmul2   :: forall m n p r s.-     (KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s)+     (KnownNat m, KnownNat n, KnownNat p, GoodScalar r, KnownSpan s)   => AstTensor AstMethodLet s (TKS '[m, n] r)   -> AstTensor AstMethodLet s (TKS '[n, p] r)   -> Maybe (AstTensor AstMethodLet s (TKS '[m, p] r)) attemptMatmul2 t3 u3 = Just $   let t4 = contractAst t3       u4 = contractAst u3-  in case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl ->-      astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl ->-        astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4-      _ -> case testEquality (typeRep @r) (typeRep @Int64) of-        Just Refl ->-          astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4-        _ -> case testEquality (typeRep @r) (typeRep @CInt) of-          Just Refl ->-            astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4-          _ -> case testEquality (typeRep @r) (typeRep @Z1) of-            Just Refl ->-              astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4-            _ -> error "attemptMatmul2: unexpected scalar"+  in case typeRep @r of+    Is @Int -> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    Is @Double -> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    Is @Float -> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    Is @Int8 -> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    Is @Int16 -> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    Is @Int32-> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    Is @Int64 -> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    Is @CInt -> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    Is @Z1 -> astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    _ -> error "attemptMatmul2: unexpected scalar" -contractAstHFun :: AstSpan s2-                => AstHFun s s2 x y -> AstHFun s s2 x y-contractAstHFun (AstLambda var l) = AstLambda var (contractAst l) -contractAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet-contractAstBool t = case t of-  AstBoolConst{} -> t-  Ast.AstBoolNot arg -> notB $ contractAstBool arg-  Ast.AstBoolAnd arg1 arg2 -> contractAstBool arg1 &&* contractAstBool arg2-  Ast.AstLeqK arg1 arg2 -> contractAst arg1 <=. contractAst arg2-  Ast.AstLeqS arg1 arg2 -> contractAst arg1 <=. contractAst arg2+-- * The let down (reducing the scope of lets cheaply) bottom-up pass++letDownAstIxS :: AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh+letDownAstIxS = fmap letDownAst++letDownAst+  :: forall s y. KnownSpan s+  => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y+letDownAst t = case t of+  Ast.AstPair t1 t2 -> Ast.AstPair (letDownAst t1) (letDownAst t2)+  Ast.AstProject1 v -> Ast.AstProject1 (letDownAst v)+  Ast.AstProject2 v -> Ast.AstProject2 (letDownAst v)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    Ast.AstMapAccumLDer k bftk eftk+                        (letDownAstHFun f)+                        (letDownAstHFun df)+                        (letDownAstHFun rf)+                        (letDownAst acc0)+                        (letDownAst es)+  Ast.AstApply v ll -> Ast.AstApply (letDownAstHFun v) (letDownAst ll)+  Ast.AstVar{} -> t+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = letDownAst v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v ->+    astLetDown var (withKnownSpan (varNameToSpan var) $ letDownAst u)+                   (letDownAst v)++  Ast.AstPrimalPart v -> Ast.AstPrimalPart (letDownAst v)+  Ast.AstDualPart v -> Ast.AstDualPart (letDownAst v)+  Ast.AstPlainPart v -> Ast.AstPlainPart (letDownAst v)+  Ast.AstFromPrimal v -> fromPrimal (letDownAst v)+  Ast.AstFromDual v -> fromDual (letDownAst v)+  Ast.AstFromPlain v -> fromPlain (letDownAst v)++  AstPlusK u v -> AstPlusK (letDownAst u) (letDownAst v)+  AstTimesK u v -> AstTimesK (letDownAst u) (letDownAst v)+  Ast.AstN1K op u -> Ast.AstN1K op (letDownAst u)+  Ast.AstR1K op u -> Ast.AstR1K op (letDownAst u)+  Ast.AstR2K op u v -> Ast.AstR2K op (letDownAst u) (letDownAst v)+  Ast.AstI2K op u v -> Ast.AstI2K op (letDownAst u) (letDownAst v)+  AstConcreteK{} -> t+  Ast.AstFloorK a -> Ast.AstFloorK (letDownAst a)+  Ast.AstFromIntegralK v -> Ast.AstFromIntegralK (letDownAst v)+  Ast.AstCastK v -> Ast.AstCastK (letDownAst v)+  Ast.AstArgMinK v -> Ast.AstArgMinK (letDownAst v)+  Ast.AstArgMaxK v -> Ast.AstArgMaxK (letDownAst v)+  Ast.AstIndexK v ix -> Ast.AstIndexK (letDownAst v) (letDownAstIxS ix)++  AstPlusS u v -> AstPlusS (letDownAst u) (letDownAst v)+  AstTimesS u v -> AstTimesS (letDownAst u) (letDownAst v)+  Ast.AstN1S op u -> Ast.AstN1S op (letDownAst u)+  Ast.AstR1S op u -> Ast.AstR1S op (letDownAst u)+  Ast.AstR2S op u v -> Ast.AstR2S op (letDownAst u) (letDownAst v)+  Ast.AstI2S op u v -> Ast.AstI2S op (letDownAst u) (letDownAst v)+  AstConcreteS{} -> t+  Ast.AstFloorS a -> Ast.AstFloorS (letDownAst a)+  Ast.AstFromIntegralS v -> Ast.AstFromIntegralS (letDownAst v)+  Ast.AstCastS v -> Ast.AstCastS (letDownAst v)+  Ast.AstArgMinS a -> Ast.AstArgMinS (letDownAst a)+  Ast.AstArgMaxS a -> Ast.AstArgMaxS (letDownAst a)+  Ast.AstIndexS shn v ix -> Ast.AstIndexS shn (letDownAst v) (letDownAstIxS ix)++  Ast.AstCondK b a2 a3 ->+    Ast.AstCondK (letDownAst b) (letDownAst a2) (letDownAst a3)+  Ast.AstCondS b a2 a3 ->+    Ast.AstCondS (letDownAst b) (letDownAst a2) (letDownAst a3)+  Ast.AstFromVectorK shm l -> Ast.AstFromVectorK shm (V.map letDownAst l)+  Ast.AstFromVectorS shm l -> Ast.AstFromVectorS shm (V.map letDownAst l)+  Ast.AstSumK v -> Ast.AstSumK (letDownAst v)+  Ast.AstSumS shm v -> Ast.AstSumS shm (letDownAst v)+  Ast.AstScatterS shm shn shp v (vars, ix) ->+    let !ix2 = letDownAstIxS ix+    in Ast.AstScatterS shm shn shp (letDownAst v) (vars, ix2)+  Ast.AstReplicateK shm v -> Ast.AstReplicateK shm (letDownAst v)+  Ast.AstReplicateS shm v -> Ast.AstReplicateS shm (letDownAst v)+  Ast.AstGatherS shm shn shp v (vars, ix) ->+    let !ix2 = letDownAstIxS ix+    in Ast.AstGatherS shm shn shp (letDownAst v) (vars, ix2)+  Ast.AstIotaS{} -> t+  Ast.AstAppendS x y -> Ast.AstAppendS (letDownAst x) (letDownAst y)+  Ast.AstSliceS i n k v -> Ast.AstSliceS i n k (letDownAst v)+  Ast.AstReverseS v -> Ast.AstReverseS (letDownAst v)+  Ast.AstTransposeS perm v -> Ast.AstTransposeS perm (letDownAst v)+  Ast.AstReshapeS sh v -> Ast.AstReshapeS sh (letDownAst v)++  Ast.AstConvert c v -> Ast.AstConvert c (letDownAst v)++  Ast.AstDot0 u v -> Ast.AstDot0 (letDownAst u) (letDownAst v)+  Ast.AstDot1InS sh n u v -> Ast.AstDot1InS sh n (letDownAst u) (letDownAst v)+  Ast.AstMatmul2S m n p u v ->+    Ast.AstMatmul2S m n p (letDownAst u) (letDownAst v)++  Ast.AstBoolNotK arg -> Ast.AstBoolNotK (letDownAst arg)+  Ast.AstBoolNotS arg -> Ast.AstBoolNotS (letDownAst arg)+  Ast.AstBoolAndK arg1 arg2 ->+    Ast.AstBoolAndK (letDownAst arg1) (letDownAst arg2)+  Ast.AstBoolAndS arg1 arg2 ->+    Ast.AstBoolAndS (letDownAst arg1) (letDownAst arg2)+  Ast.AstLeqK arg1 arg2 -> Ast.AstLeqK (letDownAst arg1) (letDownAst arg2)+  Ast.AstLeq arg1 arg2 -> Ast.AstLeq (letDownAst arg1) (letDownAst arg2)+  Ast.AstLeqS shb sh arg1 arg2 ->+    Ast.AstLeqS shb sh (letDownAst arg1) (letDownAst arg2)++letDownAstHFun :: KnownSpan s+               => AstHFun s x y -> AstHFun s x y+letDownAstHFun (AstLambda var l) = AstLambda var (letDownAst l)
src/HordeAd/Core/AstVectorize.hs view
@@ -1,3 +1,7 @@+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} -- | BOT (bulk-operation transformation) of the AST, which is a kind@@ -12,7 +16,6 @@  import Control.Exception.Assert.Sugar import Control.Monad (when)-import Data.Functor.Const import Data.IORef import Data.Maybe (fromMaybe) import Data.Proxy (Proxy (Proxy))@@ -24,19 +27,22 @@  import Data.Array.Nested (type (++)) import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)+import Data.Array.Nested.Lemmas import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Permutation qualified as Permutation import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape import Data.Array.Nested.Types (Tail, unsafeCoerceRefl) -import HordeAd.AstEngine import HordeAd.Core.Ast (AstTensor)-import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))+import HordeAd.Core.Ast hiding (AstTensor (..)) import HordeAd.Core.Ast qualified as Ast import HordeAd.Core.AstFreshId import HordeAd.Core.AstSimplify import HordeAd.Core.AstTools+import HordeAd.Core.Conversion+import HordeAd.Core.ConvertTensor+import HordeAd.Core.PPEngine import HordeAd.Core.TensorKind import HordeAd.Core.Types @@ -50,11 +56,11 @@ -- If no @AstBuild1@ terms occur in @v@, the resulting term won't -- have any, either. build1Vectorize-  :: forall y k s. AstSpan s+  :: forall y k s. KnownSpan s   => SNat k -> SingletonTK y -> (IntVarName, AstTensor AstMethodLet s y)-  -> AstTensor AstMethodLet s (BuildTensorKind k y)-{-# NOINLINE build1Vectorize #-}-build1Vectorize snat@SNat stk (!var, !v0) = unsafePerformIO $ do+  -> IO (AstTensor AstMethodLet s (BuildTensorKind k y))+{-# INLINE build1Vectorize #-}+build1Vectorize snat@SNat stk (!var, !v0) = do   enabled <- readIORef traceRuleEnabledRef   let width = 1000 * traceWidth       startTerm = Ast.AstBuild1 snat stk (var, v0)@@ -73,9 +79,9 @@       ++ ellipsisString width (printAstSimple endTerm)       ++ "\n"   let !_A = assert (ftkAst startTerm == ftkAst endTerm-                   `blame` "build1Vectorize: term shape changed"-                   `swith`( ftkAst startTerm-                          , ftkAst endTerm) ) ()+                    `blame` "build1Vectorize: term shape changed"+                    `swith` ( ftkAst startTerm+                            , ftkAst endTerm )) ()   return endTerm  @@ -85,13 +91,14 @@ -- term @AstBuild1 k (var, v)@, where it's unknown whether -- @var@ occurs in @v@. build1VOccurrenceUnknown-  :: forall y k s. AstSpan s+  :: forall y k s. KnownSpan s   => SNat k -> (IntVarName, AstTensor AstMethodLet s y)   -> AstTensor AstMethodLet s (BuildTensorKind k y) build1VOccurrenceUnknown snat@SNat (!var, !v0)   | stk0 <- ftkToSTK (ftkAst v0) =     let traceRule =-          mkTraceRule "build1VOcc" (Ast.AstBuild1 snat stk0 (var, v0)) v0 1+          mkTraceRule "build1VOcc" (Ast.AstBuild1 snat stk0 (var, v0))+                      (buildFTK snat (ftkAst v0)) v0 1     in if varNameInAst var v0        then build1V snat (var, v0)        else traceRule $@@ -104,12 +111,12 @@ -- and break our invariants that we need for simplified handling of bindings -- when rewriting terms. build1VOccurrenceUnknownRefresh-  :: forall y k s. AstSpan s+  :: forall y k s. KnownSpan s   => SNat k -> (IntVarName, AstTensor AstMethodLet s y)   -> AstTensor AstMethodLet s (BuildTensorKind k y) {-# NOINLINE build1VOccurrenceUnknownRefresh #-}-build1VOccurrenceUnknownRefresh snat@SNat (!var, !v0) =-  funToAstIntVar (varNameToBounds var) $ \ (!varFresh, !astVarFresh) ->+build1VOccurrenceUnknownRefresh snat@SNat (var, v0) =+  funToAstIntMaybe (varNameToBounds var) $ \ (!varFresh, astVarFresh) ->     let !v2 = substituteAst astVarFresh var v0                 -- cheap subst, because only a renaming     in build1VOccurrenceUnknown snat (varFresh, v2)@@ -118,12 +125,12 @@ -- @AstBuild1 k (var, v)@, where it's known that- @var@ occurs in @v@, -- see above for what it means precisely. build1V-  :: forall y k s. AstSpan s+  :: forall y k s. KnownSpan s   => SNat k -> (IntVarName, AstTensor AstMethodLet s y)   -> AstTensor AstMethodLet s (BuildTensorKind k y) build1V snat@SNat (!var, !v0) | ftk0 <- ftkAst v0 =   let bv = Ast.AstBuild1 snat (ftkToSTK ftk0) (var, v0)-      traceRule = mkTraceRule "build1V" bv v0 1+      traceRule = mkTraceRule "build1V" bv (buildFTK snat (ftkAst v0)) v0 1   in case v0 of     Ast.AstPair t1 t2 -> traceRule $       astPair (build1VOccurrenceUnknown snat (var, t1))@@ -132,34 +139,6 @@       astProject1 (build1V snat (var, t))     Ast.AstProject2 t -> traceRule $       astProject2 (build1V snat (var, t))-    Ast.AstFromVector snat1@(SNat @k1) stk l -> traceRule $-      astTrBuild (SNat @k1) (SNat @k) stk-      $ astFromVector snat1 (buildSTK snat stk) (V.map (\v ->-          build1VOccurrenceUnknown snat (var, v)) l)-    Ast.AstSum (SNat @k1) stk v -> traceRule $-      astSum (SNat @k1) (buildSTK snat stk)-      $ astTrBuild (SNat @k) (SNat @k1) stk $ build1V snat (var, v)-    Ast.AstReplicate snat2@(SNat @k2) stk v -> traceRule $-      astTrBuild (SNat @k2) SNat stk-      $ astReplicate snat2 (buildSTK snat stk) $ build1V snat (var, v)-    Ast.AstMapAccumRDer k5@(SNat @k5) bftk eftk f df rf acc0 es-     | Refl <- lemBuildOfAD snat (ftkToSTK (ftkAst acc0))-     , Refl <- lemBuildOfAD snat (ftkToSTK bftk)-     , Refl <- lemBuildOfAD snat (ftkToSTK eftk) -> traceRule $-      astLetFun-        (astMapAccumRDer-           k5-           (buildFTK snat bftk)-           (buildFTK snat eftk)-           (build1VHFun snat (var, f))-           (build1VHFun snat (var, df))-           (build1VHFun snat (var, rf))-           (build1VOccurrenceUnknown snat (var, acc0))-           (astTrBuild (SNat @k) (SNat @k5) (ftkToSTK eftk)-            $ build1VOccurrenceUnknown snat (var, es)))-        (\x1bs1 -> astPair (astProject1 x1bs1)-                           (astTrBuild (SNat @k5) (SNat @k)-                                       (ftkToSTK bftk) (astProject2 x1bs1)))     Ast.AstMapAccumLDer k5@(SNat @k5) bftk eftk f df rf acc0 es      | Refl <- lemBuildOfAD snat (ftkToSTK (ftkAst acc0))      , Refl <- lemBuildOfAD snat (ftkToSTK bftk)@@ -183,14 +162,10 @@                (build1VOccurrenceUnknown snat (var, ll))     Ast.AstVar var2 -> traceRule $       if varNameToAstVarId var2 == varNameToAstVarId var-      then case isTensorInt (Proxy @s) (varNameToFTK var2) of-        Just Refl -> fromPrimal @s $ Ast.AstIotaS (SNat @k)+      then case var2 of+        AstVarName _ FtkAndBoundsBounds{} -> fromPlain @s $ Ast.AstIotaS snat         _ -> error "build1V: build variable is not an index variable"       else error "build1V: AstVar can't contain other free variables"-    Ast.AstCond b u v -> traceRule $-      let uv = astFromVector (SNat @2) (ftkToSTK (ftkAst u)) (V.fromList [u, v])-          t = astIndexBuild (SNat @2) ftk0 uv (astCond b 0 1)-      in build1VOccurrenceUnknown snat (var, t)     Ast.AstBuild1 snat2 _ (var2, v2) -> traceRule $       assert (var2 /= var) $       build1VOccurrenceUnknown@@ -199,7 +174,8 @@      Ast.AstLet var1 u v -> traceRule $       let (var3, v2) = substProjRep snat var var1 v-      in astLet var3 (build1VOccurrenceUnknown snat (var, u))+      in astLet var3 (withKnownSpan (varNameToSpan var1) $+                      build1VOccurrenceUnknown snat (var, u))                      (build1VOccurrenceUnknownRefresh snat (var, v2))            -- ensures no duplicated bindings, see below @@ -207,10 +183,14 @@       astPrimalPart $ build1V snat (var, v)     Ast.AstDualPart v -> traceRule $       astDualPart $ build1V snat (var, v)+    Ast.AstPlainPart v -> traceRule $+      astPlainPart $ build1V snat (var, v)     Ast.AstFromPrimal v -> traceRule $-      Ast.AstFromPrimal $ build1V snat (var, v)+      fromPrimal $ build1V snat (var, v)     Ast.AstFromDual v -> traceRule $-      Ast.AstFromDual $ build1V snat (var, v)+      fromDual $ build1V snat (var, v)+    Ast.AstFromPlain v -> traceRule $+      fromPlain $ build1V snat (var, v)      Ast.AstPlusK u v -> traceRule $       build1VOccurrenceUnknown snat (var, u)@@ -239,6 +219,12 @@       astFromIntegralS $ build1V snat (var, v)     Ast.AstCastK v -> traceRule $       astCastS $ build1V snat (var, v)+    Ast.AstArgMinK v -> traceRule $+      Ast.AstArgMinS $ build1V snat (var, v)+    Ast.AstArgMaxK v -> traceRule $+      Ast.AstArgMaxS $ build1V snat (var, v)+    Ast.AstIndexK @shm v ix | Refl <- lemAppNil @shm -> traceRule $+      build1VIndexS snat ZSS (var, v, ix)  -- @var@ is in @v@ or @ix@      Ast.AstPlusS u v -> traceRule $       build1VOccurrenceUnknown snat (var, u)@@ -267,23 +253,111 @@       astFromIntegralS $ build1V snat (var, v)     Ast.AstCastS v -> traceRule $       astCastS $ build1V snat (var, v)-+    Ast.AstArgMinS v -> traceRule $+      Ast.AstArgMinS $ build1V snat (var, v)+    Ast.AstArgMaxS v -> traceRule $+      Ast.AstArgMaxS $ build1V snat (var, v)     Ast.AstIndexS shn v ix -> traceRule $       build1VIndexS snat shn (var, v, ix)  -- @var@ is in @v@ or @ix@-    Ast.AstScatterS @shm @shn @shp shn v (vars, ix) -> traceRule $++    Ast.AstCondK b u v -> traceRule $+      let uv = astFromVectorK (SNat @2 :$$ ZSS) (V.fromListN 2 [u, v])+          t = astIndexKnobsK (defaultKnobs {knobPhase = PhaseVectorization})+                             uv (astCondK b 0 1 :.$ ZIS)+      in build1VOccurrenceUnknown snat (var, t)+    Ast.AstCondS b u v | FTKS sh _ <- ftk0 -> traceRule $+      let uv = astFromVectorS (SNat @2 :$$ ZSS) (V.fromListN 2 [u, v])+          t = astIndexS sh uv (astCondK b 0 1 :.$ ZIS)+      in build1VOccurrenceUnknown snat (var, t)+    Ast.AstFromVectorK @shm shm l -> traceRule $+      let perm1 = permCycle $ shsLength shm + 1+      in Permutation.permFromListCont perm1 $ \(perm+                                                :: Permutation.Perm perm) ->+           gcastWith (unsafeCoerceRefl+                      :: (Rank perm <=? Rank (shm ++ '[k])) :~: True) $+           gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm (shm ++ '[k])+                         :~: k : shm) $+           fromMaybe (error "build1V: impossible non-permutation")+           $ Permutation.permCheckPermutation perm+           $ astTransposeS perm+           $ astFromVectorS shm (V.map (\v ->+               build1VOccurrenceUnknown snat (var, v)) l)+    Ast.AstFromVectorS @shm @shn shm l -> traceRule $+      let perm1 = permCycle $ shsLength shm + 1+      in Permutation.permFromListCont perm1 $ \(perm+                                                :: Permutation.Perm perm) ->+           gcastWith (unsafeCoerceRefl+                      :: (Rank perm <=? Rank (shm ++ (k : shn))) :~: True) $+           gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm (shm ++ (k : shn))+                         :~: k : shm ++ shn) $+           fromMaybe (error "build1V: impossible non-permutation")+           $ Permutation.permCheckPermutation perm+           $ astTransposeS perm+           $ astFromVectorS shm (V.map (\v ->+               build1VOccurrenceUnknown snat (var, v)) l)+    Ast.AstSumK @shm v -> traceRule $ case ftkAst v of+     FTKS shm _ ->+      let perm1 = backpermCycle $ shsLength shm + 1+      in Permutation.permFromListCont perm1 $ \(perm+                                                :: Permutation.Perm perm) ->+           gcastWith (unsafeCoerceRefl+                      :: Rank perm :~: Rank (k : shm)) $+           gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm (k : shm)+                         :~: shm ++ '[k]) $+           fromMaybe (error "build1V: impossible non-permutation")+           $ Permutation.permCheckPermutation perm+           $ astSumS shm $ astTransposeS perm $ build1V snat (var, v)+    Ast.AstSumS @shm @shn shm v -> traceRule $+      let perm1 = backpermCycle $ shsLength shm + 1+      in Permutation.permFromListCont perm1 $ \(perm+                                                :: Permutation.Perm perm) ->+           gcastWith (unsafeCoerceRefl+                      :: (Rank perm <=? Rank (k : shm ++ shn)) :~: True) $+           gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm (k : shm ++ shn)+                         :~: shm ++ (k : shn)) $+           fromMaybe (error "build1V: impossible non-permutation")+           $ Permutation.permCheckPermutation perm+           $ astSumS shm $ astTransposeS perm $ build1V snat (var, v)+    Ast.AstScatterS shm shn shp v (vars, ix) -> traceRule $       let (varFresh, astVarFresh, ix2) = intBindingRefreshS (var, ix)-      in astScatterS @(k ': shm) @shn @(k ': shp) shn+      in astScatterS (SNat @k :$$ shm) shn (SNat @k :$$ shp)                      (build1VOccurrenceUnknown snat (var, v))-                     (Const varFresh ::$ vars, astVarFresh :.$ ix2)-    Ast.AstGatherS @shm @shn @shp shn v (vars, ix) -> traceRule $+                     (varFresh ::$ vars, astVarFresh :.$ ix2)+    Ast.AstReplicateK @shm shm v -> traceRule $+      let perm1 = permCycle $ shsLength shm + 1+      in Permutation.permFromListCont perm1 $ \(perm+                                                :: Permutation.Perm perm) ->+           gcastWith (unsafeCoerceRefl+                      :: (Rank perm <=? Rank (shm ++ '[k])) :~: True) $+           gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm (shm ++ '[k])+                         :~: k : shm) $+           fromMaybe (error "build1V: impossible non-permutation")+           $ Permutation.permCheckPermutation perm+           $ astTransposeS perm+           $ astReplicateS shm $ build1V snat (var, v)+    Ast.AstReplicateS @shm @shn shm v -> traceRule $+      let perm1 = permCycle $ shsLength shm + 1+      in Permutation.permFromListCont perm1 $ \(perm+                                                :: Permutation.Perm perm) ->+           gcastWith (unsafeCoerceRefl+                      :: (Rank perm <=? Rank (shm ++ (k : shn))) :~: True) $+           gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm (shm ++ (k : shn))+                         :~: k : shm ++ shn) $+           fromMaybe (error "build1V: impossible non-permutation")+           $ Permutation.permCheckPermutation perm+           $ astTransposeS perm+           $ astReplicateS shm $ build1V snat (var, v)+    Ast.AstGatherS shm shn shp v (vars, ix) -> traceRule $       let (varFresh, astVarFresh, ix2) = intBindingRefreshS (var, ix)-      in astGatherS @(k ': shm) @shn @(k ': shp) shn+      in astGatherS (SNat @k :$$ shm) shn (SNat @k :$$ shp)                     (build1VOccurrenceUnknown snat (var, v))-                    (Const varFresh ::$ vars, astVarFresh :.$ ix2)-    Ast.AstMinIndexS v -> traceRule $-      Ast.AstMinIndexS $ build1V snat (var, v)-    Ast.AstMaxIndexS v -> traceRule $-      Ast.AstMaxIndexS $ build1V snat (var, v)+                    (varFresh ::$ vars, astVarFresh :.$ ix2)     Ast.AstIotaS{} ->       error "build1V: AstIotaS can't have free variables"     Ast.AstAppendS v w -> traceRule $@@ -314,18 +388,40 @@       astConvert (buildTKConversion snat (ftkAst v) c)       $ build1V snat (var, v) -    Ast.AstSum0S{} -> error "build1V: term not accessible from user API"-    Ast.AstDot0S{} -> error "build1V: term not accessible from user API"+    Ast.AstDot0{} -> error "build1V: term not accessible from user API"     Ast.AstDot1InS{} -> error "build1V: term not accessible from user API"     Ast.AstMatmul2S{} -> error "build1V: term not accessible from user API" +    Ast.AstBoolNotK b -> traceRule $+      Ast.AstBoolNotS $ build1V snat (var, b)+    Ast.AstBoolNotS b -> traceRule $+      Ast.AstBoolNotS $ build1V snat (var, b)+    Ast.AstBoolAndK b c -> traceRule $+      Ast.AstBoolAndS (build1VOccurrenceUnknown snat (var, b))+                      (build1VOccurrenceUnknown snat (var, c))+    Ast.AstBoolAndS b c -> traceRule $+      Ast.AstBoolAndS (build1VOccurrenceUnknown snat (var, b))+                      (build1VOccurrenceUnknown snat (var, c))+    Ast.AstLeqK u v -> traceRule $+      Ast.AstLeqS (snat :$$ ZSS) ZSS+                  (build1VOccurrenceUnknown snat (var, u))+                  (build1VOccurrenceUnknown snat (var, v))+    Ast.AstLeq u v | FTKS sh _ <- ftkAst u -> traceRule $+      Ast.AstLeqS (snat :$$ ZSS) sh+                  (build1VOccurrenceUnknown snat (var, u))+                  (build1VOccurrenceUnknown snat (var, v))+    Ast.AstLeqS shb sh u v -> traceRule $+      Ast.AstLeqS (snat :$$ shb) sh+                  (build1VOccurrenceUnknown snat (var, u))+                  (build1VOccurrenceUnknown snat (var, v))+ -- This refreshes an index variable in a list of index expressions. intBindingRefreshS   :: (IntVarName, AstIxS AstMethodLet sh)   -> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh) {-# NOINLINE intBindingRefreshS #-}-intBindingRefreshS (!var, !ix) =-  funToAstIntVar (varNameToBounds var) $ \ (!varFresh, !astVarFresh) ->+intBindingRefreshS (var, ix) =+  funToAstIntMaybe (varNameToBounds var) $ \ (!varFresh, astVarFresh) ->     let !ix2 = substituteAstIxS astVarFresh var ix                  -- cheap subst, because only a renaming     in (varFresh, astVarFresh, ix2)@@ -352,7 +448,7 @@ -- eventually proven unnecessary. The rule changes the index to a gather -- and pushes the build down the gather, getting the vectorization unstuck. build1VIndexS-  :: forall k shm shn x s. AstSpan s+  :: forall k shm shn x s. KnownSpan s   => SNat k -> ShS shn   -> ( IntVarName  -- bounds (0, k - 1)      , AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)@@ -360,48 +456,43 @@   -> AstTensor AstMethodLet s (TKS2 (k ': shn) x) build1VIndexS k _ (!var, !v0, ZIS) =   build1VOccurrenceUnknown k (var, v0)-build1VIndexS k@SNat shn (var, v0, ix) | STKS _ x <- ftkToSTK (ftkAst v0) =-  let vTrace = Ast.AstBuild1 k (STKS shn x) (var, Ast.AstIndexS shn v0 ix)-      traceRule = mkTraceRule "build1VIndexS" vTrace v0 1+build1VIndexS k@SNat shn (var, v0, ix) | FTKS shmshn x' <- ftkAst v0 =+  let x = ftkToSTK x'+      vTrace = Ast.AstBuild1 k (STKS shn x) (var, Ast.AstIndexS shn v0 ix)+      traceRule = mkTraceRule "build1VIndexS" vTrace+                              (buildFTK k (FTKS shn x')) v0 1   in if varNameInAst var v0      then case astIndexKnobsS (defaultKnobs {knobPhase = PhaseVectorization})                               shn v0 ix of  -- push deeper        Ast.AstIndexS _ v1 ZIS -> traceRule $          build1VOccurrenceUnknown k (var, v1)-       v@(Ast.AstIndexS shn1 v1 ix1) -> traceRule $+       v@(Ast.AstIndexS @shm1 @shn1 shn1 v1 ix1) -> traceRule $          let (varFresh, astVarFresh, ix2) = intBindingRefreshS (var, ix1)              ruleD :: AstTensor AstMethodLet s (TKS2 (k ': shn) x)-             ruleD = astGatherS-                       shn1 (build1VOccurrenceUnknown k (var, v1))-                       (Const varFresh ::$ ZS, astVarFresh :.$ ix2)-             len = ixsLength ix1+             ruleD | FTKS shmshn1 _ <- ftkAst v1 =+               astGatherS+                 (SNat @k :$$ ZSS)+                 shn1+                 (SNat @k :$$ shsTakeIx @shm1 @shn1 Proxy ix2 shmshn1)+                 (build1VOccurrenceUnknown k (var, v1))+                 (varFresh ::$ ZS, astVarFresh :.$ ix2)          in if varNameInAst var v1-            then case v1 of  -- try to avoid ruleD if not a normal form-              Ast.AstFromVector{} | len == 1 -> ruleD-              Ast.AstScatterS{} -> ruleD-              Ast.AstGatherS{} -> ruleD-              -- These can only be simplified to the AstFromVector NF above.-              Ast.AstReplicate{} -> ruleD-              Ast.AstAppendS{} -> ruleD-              -- These, in general, simplify to gathers, so as bad as ruleD.-              Ast.AstTransposeS{} -> ruleD-              Ast.AstReshapeS{} -> ruleD-              -- Rarely these don't simplify enough; left as an escape hatch:-              -- (TODO: simplify better)-              Ast.AstConvert{} -> ruleD-              _ -> build1VOccurrenceUnknown k (var, v)  -- not a normal form-            else build1VOccurrenceUnknown k (var, v)  -- shortcut+            then ruleD+            else build1VOccurrenceUnknown k (var, v)  -- var may be in ix1        v -> traceRule $          build1VOccurrenceUnknown k (var, v)            -- peel off yet another constructor      else traceRule $-            astGatherS shn v0 (Const var ::$ ZS, ix)+            astGatherS (SNat @k :$$ ZSS)+                       shn+                       (shsTakeIx @shm @shn Proxy ix shmshn)+                       v0 (var ::$ ZS, ix)  build1VHFun-  :: forall k x z s s2. (AstSpan s, AstSpan s2)-  => SNat k -> (IntVarName, AstHFun s s2 x z)-  -> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)-build1VHFun snat@SNat (!var, !v0) = case v0 of+  :: forall k x z s. KnownSpan s+  => SNat k -> (IntVarName, AstHFun s x z)+  -> AstHFun s (BuildTensorKind k x) (BuildTensorKind k z)+build1VHFun snat@SNat (var, v0) = case v0 of   Ast.AstLambda var1 t ->     -- This handles the case of l having free variables beyond var1,     -- which is not possible for lambdas used in folds, etc.@@ -414,26 +505,50 @@  -- * Auxiliary operations -astTr :: forall n s r. AstSpan s+astReplicate+  :: forall k s y. KnownSpan s+  => SNat k -> SingletonTK y+  -> AstTensor AstMethodLet s y+  -> AstTensor AstMethodLet s (BuildTensorKind k y)+astReplicate k stk t = case stk of+  STKScalar -> astReplicateK (k :$$ ZSS) t+  STKR{} -> case ftkAst t of+    FTKR shn0 x ->+      withShsFromShR shn0 $ \(shn :: ShS shn) ->+        astConvUpRFromS (k :$$ shn) x . astReplicateS (k :$$ ZSS)+        . astConvDownSFromR shn x $ t+  STKS{} -> astReplicateS (k :$$ ZSS) t+  STKX{} -> case ftkAst t of+    FTKX shn0 x ->+      withShsFromShX shn0 $ \(shn :: ShS shn) ->+        astConvUpXFromS @(k : shn) (SKnown k :$% shn0) x+        . astReplicateS (k :$$ ZSS) . astConvDownSFromX shn x $ t+  STKProduct stk1 stk2 ->+    astLetFun t $ \ !tShared ->+      let (u1, u2) = (astProject1 tShared, astProject2 tShared)+      in astPair (astReplicate k stk1 u1) (astReplicate k stk2 u2)++astTr :: forall n s r. KnownSpan s       => AstTensor AstMethodLet s (TKR2 (2 + n) r)       -> AstTensor AstMethodLet s (TKR2 (2 + n) r) astTr a = case Permutation.makePerm @'[1, 0] of   (perm :: Permutation.Perm perm) -> case ftkAst a of-    FTKR sh'@(k :$: m :$: shr) x | SNat <- shrRank sh' ->+    FTKR sh' x | SNat <- shrRank sh' ->       withShsFromShR sh' $ \(sh :: ShS sh) ->         gcastWith (unsafeCoerceRefl :: (Rank perm <=? Rank sh) :~: True) $-        astFromS' (FTKR (m :$: k :$: shr) x)-        . astTransposeS perm . astSFromR' sh $ a-    _ -> error "astTr: impossible shape"+        gcastWith (unsafeCoerceRefl+                   :: Rank (Permutation.PermutePrefix perm sh) :~: Rank sh) $+        astConvUpRFromS (shsPermutePrefix perm sh) x+        . astTransposeS perm . astConvDownSFromR sh x $ a -astTrS :: forall n m sh s r. AstSpan s+astTrS :: forall n m sh s r. KnownSpan s        => AstTensor AstMethodLet s (TKS2 (n ': m ': sh) r)        -> AstTensor AstMethodLet s (TKS2 (m ': n ': sh) r) astTrS a | FTKS (_ :$$ _ :$$ sh) _ <- ftkAst a-         , SNat <- shsRank sh =  -- why on Earth is this needed?+         , SNat <- shsRank sh =  -- needed to use auto-instances to infer 2 <= 2   astTransposeS (Permutation.makePerm @'[1, 0]) a -astTrX :: forall n m shx s r. AstSpan s+astTrX :: forall n m shx s r. KnownSpan s        => AstTensor AstMethodLet s (TKX2 (Just n ': Just m ': shx) r)        -> AstTensor AstMethodLet s (TKX2 (Just m ': Just n ': shx) r) astTrX a = case Permutation.makePerm @'[1, 0] of@@ -441,11 +556,13 @@     FTKX sh'@(mn :$% mm :$% shx) x ->       withShsFromShX sh' $ \(sh :: ShS sh) ->         gcastWith (unsafeCoerceRefl :: (Rank perm <=? Rank sh) :~: True) $-        astFromS' (FTKX (mm :$% mn :$% shx) x)-        . astTransposeS perm . astSFromX' sh $ a+        gcastWith (unsafeCoerceRefl+                   :: Rank (Permutation.PermutePrefix perm sh) :~: Rank sh) $+        astConvUpXFromS (mm :$% mn :$% shx) x+        . astTransposeS perm . astConvDownSFromX sh x $ a  astTrBuild-  :: forall k1 k2 s y. AstSpan s+  :: forall k1 k2 s y. KnownSpan s   => SNat k1 -> SNat k2 -> SingletonTK y   -> AstTensor AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))   -> AstTensor AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))@@ -460,45 +577,59 @@       in astPair (astTrBuild (SNat @k1) (SNat @k2) stk1 u1)                  (astTrBuild (SNat @k1) (SNat @k2) stk2 u2) -astIndexBuild :: forall y k s. AstSpan s+astIndexBuild :: forall y k s. KnownSpan s               => SNat k -> FullShapeTK y               -> AstTensor AstMethodLet s (BuildTensorKind k y)               -> AstInt AstMethodLet               -> AstTensor AstMethodLet s y astIndexBuild snat@SNat ftk u i = case ftk of-  FTKScalar -> astFromS' ftk $ astIndexS ZSS u (i :.$ ZIS)+  FTKScalar -> kfromS $ astIndexS ZSS u (i :.$ ZIS)   FTKR{} -> case ftkAst u of-    FTKR shmshn _ ->+    FTKR shmshn x ->       withShsFromShR shmshn $ \(sh :: ShS sh) ->         gcastWith (unsafeCoerceRefl :: k ': Tail sh :~: sh) $-        astFromS' ftk $ astIndexS (shsTail sh) (astSFromR' sh u) (i :.$ ZIS)+        astConvUpRFromS (shsTail sh) x+        $ astIndexS (shsTail sh) (astConvDownSFromR sh x u) (i :.$ ZIS)   FTKS sh _ -> astIndexS sh u (i :.$ ZIS)-  FTKX{} -> case ftkAst u of-   FTKX shBuild' _->+  FTKX @sh' sh' _ -> case ftkAst u of+   FTKX shBuild' x ->     withShsFromShX shBuild' $ \shBuild -> case shBuild of-      _ :$$ rest ->-        astFromS' ftk $ astIndexS rest (astSFromX' shBuild u) (i :.$ ZIS)+      _ :$$ (rest :: ShS rest) ->+        -- Needed for GHC 9.10 only:+        gcastWith (unsafeCoerceRefl :: Rank sh' :~: Rank rest) $+        astConvUpXFromS sh' x+        $ astIndexS rest (astConvDownSFromX shBuild x u) (i :.$ ZIS)   FTKProduct ftk1 ftk2 ->     astLetFun u $ \ !u3 ->       astPair (astIndexBuild snat ftk1 (astProject1 u3) i)               (astIndexBuild snat ftk2 (astProject2 u3) i)+  {- The following prevents generating duplicated gathers by vectorization+     but it generates big conversions that are even more costly,+     due to transforming HFuns, e.g., the arguments of folds:+  FTKProduct{} ->+    -- This conversion trick, in particular, prevents duplication of a variable+    -- in AstCond, for which we can then substitute indexing, no big deal,+    -- but each of the indexing can subsequently get turned into a gather,+    -- which then makes a performance difference.+    unNestTarget (ftkToSTK ftk)+    $ astIndexS ZSS (nestTargetK snat ftk u) (i :.$ ZIS) -}  substProjRep-  :: forall k s s2 y2 y. (AstSpan s, AstSpan s2)+  :: forall k s s2 y2 y. KnownSpan s   => SNat k -> IntVarName-  -> AstVarName s2 y2 -> AstTensor AstMethodLet s y-  -> (AstVarName s2 (BuildTensorKind k y2), AstTensor AstMethodLet s y)+  -> AstVarName '(s2, y2) -> AstTensor AstMethodLet s y+  -> (AstVarName '(s2, BuildTensorKind k y2), AstTensor AstMethodLet s y) substProjRep snat@SNat var var1 v =-  let ftk3 = buildFTK snat $ varNameToFTK var1-      var3 :: AstVarName s2 (BuildTensorKind k y2)-      var3 = mkAstVarName ftk3 (varNameToBounds var1) (varNameToAstVarId var1)+  let ftk1 = varNameToFTK var1+      ftk3 = buildFTK snat ftk1+      var3 :: AstVarName '(s2, BuildTensorKind k y2)+      var3 = reshapeVarName ftk3 var1       astVar3 = astVar var3-      v2 = substituteAst-             (astIndexBuild snat (varNameToFTK var1)-                            astVar3 (astVar var))-             var1 v-        -- The subsitutions of projections don't break sharing,-        -- because they don't duplicate variables and the added var+      indexing = withKnownSpan (varNameToSpan var1) $+                 astIndexBuild snat ftk1 astVar3 (astVar var)+      v2 = substituteAst indexing var1 v+        -- The subsitutions of projections and indexing don't break sharing+        -- too much, because they don't duplicate variables and the added var         -- is eventually being eliminated instead of substituted for.   in (var3, v2) @@ -528,21 +659,27 @@                                then cropped                                else take (width - 3) cropped ++ "..." -mkTraceRule :: forall y z s. AstSpan s+-- We can't force @to@, because we want the debug info displayed before+-- @to@ is evaluated (or diverges).+-- TODO: switch away from IORefs to ensure correct blocking and then+-- really display the first part of the message before @to@ diverges.+mkTraceRule :: forall y z s. KnownSpan s             => String             -> AstTensor AstMethodLet s y+            -> FullShapeTK y             -> AstTensor AstMethodLet s z             -> Int             -> AstTensor AstMethodLet s y             -> AstTensor AstMethodLet s y {-# NOINLINE mkTraceRule #-}-mkTraceRule !prefix !from !caseAnalysed !nwords ~to = unsafePerformIO $ do+mkTraceRule prefix from !fromFTK caseAnalysed nwords to = unsafePerformIO $ do   enabled <- readIORef traceRuleEnabledRef   let width = traceWidth       constructorName =         unwords $ take nwords $ words $ take 21         $ case caseAnalysed of           Ast.AstVar{} -> "variable"+          Ast.AstIndexK{} -> "sindex0"           Ast.AstIndexS{} -> "sindex"           _ -> printAstSimple caseAnalysed       ruleName = prefix ++ "." ++ constructorName@@ -558,10 +695,10 @@                             ++ " sends " ++ padString width stringFrom                             ++ " to " ++ padString width stringTo     modifyIORef' traceNestingLevel pred-  let !_A = assert (ftkAst from == ftkAst to+  let !_A = assert (fromFTK == ftkAst to                     `blame` "mkTraceRule: term shape changed"-                    `swith`( ftkAst from, ftkAst to-                           , from, to )) ()+                    `swith` ( fromFTK, ftkAst to+                            , from, to )) ()   return $! to  hPutStrLnFlush :: Handle -> String -> IO ()
src/HordeAd/Core/CarriersADVal.hs view
@@ -9,7 +9,7 @@ -- in case of reverse derivatives). module HordeAd.Core.CarriersADVal   ( -- * The dual number type-    ADVal, pattern D, dD, dDnotShared+    ADVal, pattern D, dD, dDnotShared, fromPrimalFTK     -- * Auxiliary definitions   , unDeltaPair, unDeltaPairUnshared, dScale, dAdd   , dSFromR, dSFromX, dXFromS@@ -19,7 +19,6 @@  import Prelude -import Data.Int (Int64) import Data.Proxy (Proxy (Proxy)) import Data.Type.Equality ((:~:) (Refl)) @@ -27,7 +26,7 @@ import Data.Array.Nested.Shaped.Shape  import Data.Array.Nested.Lemmas-import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Conversion import HordeAd.Core.Delta import HordeAd.Core.DeltaFreshId import HordeAd.Core.Ops@@ -72,7 +71,10 @@ dDnotShared :: f z -> Delta f z -> ADVal f z dDnotShared = ADVal +fromPrimalFTK :: FullShapeTK z -> f z -> ADVal f z+fromPrimalFTK ftk a = dDnotShared a (DeltaZero ftk) + -- * Auxiliary definitions  -- TODO: maybe create a separate module of Delta smart constructors@@ -136,7 +138,7 @@   , Just Refl <- matchingFTK (ftkDelta d) (FTKS sh x) = d dSFromR sh d | FTKR _ x <- ftkDelta d              , Refl <- lemRankReplicate (Proxy @(Rank sh)) =-  let c2 = ConvCmp (ConvXS' (FTKS sh x)) ConvRX+  let c2 = convCmp (ConvXS' (FTKS sh x)) ConvRX   in DeltaConvert c2 d  dSFromX :: forall sh sh' x target. Rank sh ~ Rank sh'@@ -160,12 +162,12 @@   | FTKS sh x <- ftkDelta d   , let shx = shCastSX ssx sh   , Refl <- lemRankMapJust sh =-    let c2 = ConvCmp (ConvXX' (FTKX shx x)) ConvSX+    let c2 = convCmp (ConvXX' (FTKX shx x)) ConvSX     in DeltaConvert c2 d  -- This hack is needed to recover shape from tensors, -- in particular in case of numeric literals and also for forward derivative.-intOfShape :: forall z f. ADReadyNoLet f+intOfShape :: forall z f. (TKAllNum z, ADReadyNoLet f)            => Delta f z -> Int -> f z intOfShape tsh c = treplTarget (fromIntegral c) (ftkDelta tsh) @@ -205,8 +207,6 @@  -- * Assorted instances -type instance BoolOf (ADVal f) = BoolOf f- instance EqH f (TKScalar r) => EqH (ADVal f) (TKScalar r) where   D u _ ==. D v _ = u ==. v @@ -231,12 +231,10 @@ instance OrdH f (TKX sh r) => OrdH (ADVal f) (TKX sh r) where   D u _ <=. D v _ = u <=. v -type instance HFunOf (ADVal f) x y = HFun x y-+type instance HFunOf (ADVal f) = HFun type instance PrimalOf (ADVal f) = f- type instance DualOf (ADVal f) = Delta f-+type instance PlainOf (ADVal f) = PlainOf f type instance ShareOf (ADVal f) = ADVal f   -- Maybe this should be ADVal (ShareOf f), but we'd need tests   -- that use this, probably tests with ADVal (AST) nested in ADVal@@ -244,21 +242,19 @@  -- * Numeric instances --- These two instances are required for the numeric tensor instances.+-- These two instances are required for the numeric tensor instances+-- and they are also useful for unrolling recursion in non-symbolic pipelines. -- They can't be made valid for AST, because they require interpretation before--- they can be compared with an instant Bool result, so let's fail early--- also here.-instance Eq (ADVal f z) where-  (==) = error "Eq is not defined for ADVal; please use EqH instead"-  (/=) = error "Eq is not defined for ADVal; please use EqH instead"+-- they can be compared with an instant Bool result.+instance Eq (f z) => Eq (ADVal f z) where+  D a _ == D b _ = a == b -instance Ord (ADVal f z) where-  (<=) = error "Ord is not defined for ADVal; please use OrdH instead"+instance Ord (f z) => Ord (ADVal f z) where+  D a _ <= D b _ = a <= b --- This is copied from below to permit fromInteger for TKScalar.--- This OVERLAPPABLE seems to work 100% reliably for indexes--- and not at all for a variant of rfromListLinear that takes scalars.-instance (GoodScalar r, ShareTensor f, ADReadyNoLet f)+-- This is copied below to permit fromInteger for TKScalar and to forbid+-- TKProduct (also nested) in order to simplify the reverse pass.+instance (NumScalar r, ShareTensor f, ADReadyNoLet f)          => Num (ADVal f (TKScalar r)) where   D u u' + D v v' = dD (u + v) (dAdd u' v')   D u u' - D v v' =@@ -271,17 +267,18 @@   negate (D v v') = dD (negate v) (dScale (intOfShape v' (-1)) v')   abs (D ve v') = let !v = tshare ve                   in dD (abs v) (dScale (signum v) v')-  signum (D v v') = dDnotShared (signum v) (DeltaZero $ ftkDelta v')-  fromInteger i = dDnotShared (fromInteger i) (DeltaZero FTKScalar)+  signum (D v v') = fromPrimalFTK (ftkDelta v') (signum v)+  fromInteger i = fromPrimalFTK FTKScalar (fromInteger i)   -- The constraints in the pragmas below are needed only to avoid-  -- module import cycles.+  -- module import cycles. But the pragmas probably don't work due to+  -- the constraints.+  {-   {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKScalar Double)) #-}   {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKScalar Float)) #-}-  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKScalar Int64)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKScalar Int)) #-} -} -instance {-# OVERLAPPABLE #-}-         (Num (f z), ShareTensor f, ADReadyNoLet f)-         => Num (ADVal f z) where+instance (TKAllNum (TKR n x), Num (f (TKR n x)), ShareTensor f, ADReadyNoLet f)+         => Num (ADVal f (TKR n x)) where   D u u' + D v v' = dD (u + v) (dAdd u' v')   D u u' - D v v' =     dD (u - v) (dAdd u' (dScale (intOfShape v' (-1)) v'))@@ -293,39 +290,74 @@   negate (D v v') = dD (negate v) (dScale (intOfShape v' (-1)) v')   abs (D ve v') = let !v = tshare ve                   in dD (abs v) (dScale (signum v) v')-  signum (D v v') = dDnotShared (signum v) (DeltaZero $ ftkDelta v')+  signum (D v v') = fromPrimalFTK (ftkDelta v') (signum v)   fromInteger = error "fromInteger is not defined for tensors in general"++instance ( TKAllNum (TKS sh x), Num (f (TKS sh x)), ShareTensor f+         , ADReadyNoLet f )+         => Num (ADVal f (TKS sh x)) where+  D u u' + D v v' = dD (u + v) (dAdd u' v')+  D u u' - D v v' =+    dD (u - v) (dAdd u' (dScale (intOfShape v' (-1)) v'))+  D ue u' * D ve v' =+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u * v) (dAdd (dScale v u') (dScale u v'))+  negate (D v v') = dD (negate v) (dScale (intOfShape v' (-1)) v')+  abs (D ve v') = let !v = tshare ve+                  in dD (abs v) (dScale (signum v) v')+  signum (D v v') = fromPrimalFTK (ftkDelta v') (signum v)+  fromInteger = error "fromInteger is not defined for tensors in general"++instance ( TKAllNum (TKX sh x), Num (f (TKX sh x)), ShareTensor f+         , ADReadyNoLet f )+         => Num (ADVal f (TKX sh x)) where+  D u u' + D v v' = dD (u + v) (dAdd u' v')+  D u u' - D v v' =+    dD (u - v) (dAdd u' (dScale (intOfShape v' (-1)) v'))+  D ue u' * D ve v' =+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u * v) (dAdd (dScale v u') (dScale u v'))+  negate (D v v') = dD (negate v) (dScale (intOfShape v' (-1)) v')+  abs (D ve v') = let !v = tshare ve+                  in dD (abs v) (dScale (signum v) v')+  signum (D v v') = fromPrimalFTK (ftkDelta v') (signum v)+  fromInteger = error "fromInteger is not defined for tensors in general"+   -- The constraints in the pragmas below are needed only to avoid   -- module import cycles.+  {-   {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKR n Double)) #-}   {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKR n Float)) #-}-  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKR n Int64)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKR n Int)) #-}   {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKS sh Double)) #-}   {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKS sh Float)) #-}-  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKS sh Int64)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKS sh Int)) #-}   {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKX sh Double)) #-}   {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKX sh Float)) #-}-  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKX sh Int64)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKX sh Int)) #-} -} -instance (Real (f z), ShareTensor f, ADReadyNoLet f)+instance (TKAllNum z, Num (ADVal f z), Real (f z), ADReadyNoLet f)          => Real (ADVal f z) where   toRational (D v _) = toRational v     -- this is most probably not what the user expects, but the type     -- of the result (Rational) doesn't permit any better solution -instance (IntegralH (f z), ShareTensor f, ADReadyNoLet f)+instance (TKAllNum z, Num (ADVal f z), IntegralH (f z), ADReadyNoLet f)          => IntegralH (ADVal f z) where-  quotH (D u _) (D v v') = dDnotShared (quotH u v) (DeltaZero $ ftkDelta v')-  remH (D u _) (D v v') = dDnotShared (remH u v) (DeltaZero $ ftkDelta v')+  quotH (D u _) (D v v') = fromPrimalFTK (ftkDelta v') (quotH u v)+  remH (D u _) (D v v') = fromPrimalFTK (ftkDelta v') (remH u v)   -- The constraints in the pragmas below are needed only to avoid   -- module import cycles.-  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKR n Int64)) #-}-  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKS sh Int64)) #-}-  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKX sh Int64)) #-}+  {-+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKR n Int)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKS sh Int)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKX sh Int)) #-} -}  -- This is copied from below to permit fromRational for TKScalar.-instance ( GoodScalar r, Fractional (f (TKScalar r)), ShareTensor f-         , ADReadyNoLet f )+instance ( TKAllNum (TKScalar r), NumScalar r, Fractional (f (TKScalar r))+         , ShareTensor f, ADReadyNoLet f )          => Fractional (ADVal f (TKScalar r)) where   D ue u' / D ve v' =     let !u = tshare ue in@@ -335,11 +367,13 @@     let !v = tshare ve         minusRecipSq = - recip (v * v)     in dD (recip v) (dScale minusRecipSq v')-  fromRational r = dDnotShared (fromRational r) (DeltaZero FTKScalar)+  fromRational r = fromPrimalFTK FTKScalar (fromRational r) -instance {-# OVERLAPPABLE #-}-         (Fractional (f z), ShareTensor f, ADReadyNoLet f)-         => Fractional (ADVal f z) where+-- This is copied three times, because OVERLAPPABLE either doesn't work+-- across packages or is unreliable.+instance ( TKAllNum (TKR n x), Num (ADVal f (TKR n x)), Fractional (f (TKR n x))+         , ShareTensor f, ADReadyNoLet f )+         => Fractional (ADVal f (TKR n x)) where   D ue u' / D ve v' =     let !u = tshare ue in     let !v = tshare ve@@ -350,7 +384,34 @@     in dD (recip v) (dScale minusRecipSq v')   fromRational = error "fromRational is not defined for tensors in general" -instance (Floating (f z), ShareTensor f, ADReadyNoLet f)+instance ( TKAllNum (TKS sh x), Num (ADVal f (TKS sh x))+         , Fractional (f (TKS sh x)), ShareTensor f, ADReadyNoLet f )+         => Fractional (ADVal f (TKS sh x)) where+  D ue u' / D ve v' =+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u / v) (dAdd (dScale (recip v) u') (dScale ((- u) / (v * v)) v'))+  recip (D ve v') =+    let !v = tshare ve+        minusRecipSq = - recip (v * v)+    in dD (recip v) (dScale minusRecipSq v')+  fromRational = error "fromRational is not defined for tensors in general"++instance ( TKAllNum (TKX sh x), Num (ADVal f (TKX sh x))+         , Fractional (f (TKX sh x)), ShareTensor f, ADReadyNoLet f )+         => Fractional (ADVal f (TKX sh x)) where+  D ue u' / D ve v' =+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u / v) (dAdd (dScale (recip v) u') (dScale ((- u) / (v * v)) v'))+  recip (D ve v') =+    let !v = tshare ve+        minusRecipSq = - recip (v * v)+    in dD (recip v) (dScale minusRecipSq v')+  fromRational = error "fromRational is not defined for tensors in general"++instance ( TKAllNum z, Fractional (ADVal f z), Floating (f z)+         , ShareTensor f, ADReadyNoLet f )          => Floating (ADVal f z) where   pi = error "pi is not defined for tensors"   exp (D ue u') = let !expU = tshare (exp ue)@@ -397,13 +458,14 @@                     in dD (atanh u)                           (dScale (recip (intOfShape u' 1 - u * u)) u') -instance (RealFrac (f z), ShareTensor f, ADReadyNoLet f)+instance (TKAllNum z, Fractional (ADVal f z), RealFrac (f z), ADReadyNoLet f)          => RealFrac (ADVal f z) where   properFraction = error "properFraction is not defined for tensors"     -- The integral type doesn't have a Storable constraint,     -- so we can't implement this (nor RealFracB from Boolean package). -instance (Fractional (f z), RealFloatH (f z), ShareTensor f, ADReadyNoLet f)+instance ( TKAllNum z, Fractional (ADVal f z), RealFloatH (f z)+         , ShareTensor f, ADReadyNoLet f )          => RealFloatH (ADVal f z) where   atan2H (D ue u') (D ve v') =     let !u = tshare ue in@@ -411,7 +473,8 @@     let !t = tshare (recip (u * u + v * v))     in dD (atan2H u v) (dAdd (dScale ((- u) * t) v') (dScale (v * t) u')) -instance (RealFloat (f z), ShareTensor f, ADReadyNoLet f)+instance ( TKAllNum z, Fractional (ADVal f z), RealFloat (f z)+         , ShareTensor f, ADReadyNoLet f )          => RealFloat (ADVal f z) where   atan2 (D ue u') (D ve v') =     let !u = tshare ue in
src/HordeAd/Core/CarriersAst.hs view
@@ -1,1149 +1,342 @@-{-# LANGUAGE UndecidableInstances #-}-{-# OPTIONS_GHC -Wno-orphans #-}--- | Definitions, mostly class instances, needed to make AST a valid--- carrier for a tensor class algebra (instance) defined--- in "HordeAd.Core.OpsAst".--- This algebra permits user programs to be instantiated as AST terms,--- as well as to be interpreted into AST terms and it also permits derivatives--- to be expressed as AST terms.-module HordeAd.Core.CarriersAst-  ( AstRaw(..), AstNoVectorize(..), AstNoSimplify(..)-  , astLetFunNoSimplify, sunReplicateScal, sunReplicate1, sunReplicateN-  ) where--import Prelude hiding (foldl')--import Data.Int (Int64)-import Data.Type.Equality (testEquality, (:~:) (Refl))-import Foreign.C (CInt)-import Type.Reflection (typeRep)--import Data.Array.Nested (type (++))-import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)-import Data.Array.Nested.Mixed qualified as Mixed-import Data.Array.Nested.Mixed.Shape-import Data.Array.Nested.Shaped.Shape--import HordeAd.Core.Ast-import HordeAd.Core.AstFreshId-import HordeAd.Core.AstTools-import HordeAd.Core.CarriersConcrete-import HordeAd.Core.ConvertTensor-import HordeAd.Core.OpsConcrete ()-import HordeAd.Core.TensorKind-import HordeAd.Core.Types---- * Type family instances for AstTensor--type instance PrimalOf (AstTensor ms s) = AstTensor ms PrimalSpan-type instance DualOf (AstTensor ms s) = AstTensor ms DualSpan-type instance ShareOf (AstTensor ms s) = AstRaw s---- This can't be just HFun, because they need to be vectorized--- and vectorization applies such functions to the variable from build1--- and the variable has to be eliminated via vectorization to preserve--- the closed form of the function. Just applying a Haskell closure--- to the build1 variable and then duplicating the result of the function--- would not eliminate the variable and also would likely results--- in more costly computations. Also, that would prevent simplification--- of the instances, especially after applied to arguments that are terms.-type instance HFunOf (AstTensor AstMethodLet s) x z = AstHFun s s x z--type instance BoolOf (AstTensor ms s) = AstBool ms----- * Unlawful numeric instances for AST scalars; they are lawful modulo evaluation---- These are, unfortunately, required by some numeric instances.-instance Eq (AstTensor ms s y) where-  (==) = error "Eq is not defined for AST; please use EqH instead"-  (/=) = error "Eq is not defined for AST; please use EqH instead"--instance Ord (AstTensor ms s y) where-  (<=) = error "Ord is not defined for AST; please use OrdH instead"---- TODO: perhaps aim for a polynomial normal form? but that requires global--- inspection of the whole expression--- TODO: let's aim at SOP (Sum-of-Products) form, just as--- ghc-typelits-natnormalise does. Also, let's associate to the right--- and let's push negation down.------ Not considered are rules that would require comparing non-constant terms--- or that would duplicate a non-constant term, as well as most rules--- informed by inequalities, expressed via max or min, such as--- max n (signum (abs x)) | n <= 0 --> signum (abs x).--- We could use sharing via @tlet@ if terms are duplicated, but it's--- unclear if the term bloat is worth it and also we'd need to restrict--- this extended simplification to AstMethodLet.------ | Integer terms need to be simplified, because large ones are sometimes--- created due to vectorization, e.g., via astTransposeAsGather--- or astReshapeAsGather and can be a deciding factor in whether--- the other tensor terms can be simplified in turn.------ The normal form has AstConcreteK, if any, as the first argument--- of the constructor. No flattening is performed beyond that.-instance (GoodScalar r, AstSpan s)-         => Num (AstTensor ms s (TKScalar r)) where-  AstFromPrimal u + AstFromPrimal v = AstFromPrimal $ u + v-  AstFromDual u + AstFromDual v = AstFromDual $ u + v-  -- TODO: define a pattern synonym that captures the below. Also elsewhere.-  AstConvert c u + AstConvert _ v-    | FTKS ZSS x <- ftkAst u-    , FTKS ZSS y <- ftkAst v-    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u))-    , Just Refl <- matchingFTK x y =-      AstConvert c $ u + v-  AstConcreteK 0 + u = u-  u + AstConcreteK 0 = u-  AstConcreteK n + AstConcreteK k = AstConcreteK (n + k)-  AstConcreteK n + AstPlusK (AstConcreteK k) u = AstConcreteK (n + k) + u-  AstPlusK (AstConcreteK n) u + AstConcreteK k = AstConcreteK (n + k) + u-  AstPlusK (AstConcreteK n) u + AstPlusK (AstConcreteK k) v =-    AstConcreteK (n + k) + AstPlusK u v  -- u and v can cancel, but unlikely--  -- Unfortunately, these only fire if the required subterms are at the top-  -- of the reduced term, which happens rarely except in small terms.-  -- We could keep variables at the top, but they'd compete with AstConcreteK.-  AstN1K NegateOp (AstVar var) + AstVar var'-    | var == var' = 0-  AstN1K NegateOp (AstConvert _ (AstVar var))-    + AstConvert _ (AstVar var')-    | varNameToAstVarId var == varNameToAstVarId var' = 0-  AstN1K NegateOp (AstVar var) + AstPlusK (AstVar var') u-    | var == var' = u-  AstN1K NegateOp (AstConvert _ (AstVar var))-    + AstPlusK (AstConvert _ (AstVar var')) u-    | varNameToAstVarId var == varNameToAstVarId var' = u-  AstVar var' + AstN1K NegateOp (AstVar var)-    | var == var' = 0-  AstConvert _ (AstVar var')-    + AstN1K NegateOp (AstConvert _ (AstVar var))-    | varNameToAstVarId var == varNameToAstVarId var' = 0-  AstVar var' + AstPlusK (AstN1K NegateOp (AstVar var)) u-    | var == var' = u-  AstConvert _ (AstVar var')-    + AstPlusK (AstN1K NegateOp (AstConvert _ (AstVar var))) u-    | varNameToAstVarId var == varNameToAstVarId var' = u--  AstI2K RemOp (AstN1K NegateOp (AstVar var)) (AstConcreteK n)-   + AstI2K RemOp (AstVar var') (AstConcreteK n')-     | var == var' && n == n' = 0-  AstI2K RemOp (AstN1K NegateOp (AstConvert _ (AstVar var)))-               (AstConcreteK n)-   + AstI2K RemOp (AstConvert _ (AstVar var')) (AstConcreteK n')-     | varNameToAstVarId var == varNameToAstVarId var' && n == n' = 0-  AstI2K RemOp (AstN1K NegateOp (AstVar var)) (AstConcreteK n)-   + AstPlusK (AstI2K RemOp (AstVar var') (AstConcreteK n')) u-     | var == var' && n == n' = u-  AstI2K RemOp (AstN1K NegateOp (AstConvert _ (AstVar var)))-               (AstConcreteK n)-   + AstPlusK (AstI2K RemOp (AstConvert _ (AstVar var'))-                            (AstConcreteK n')) u-     | varNameToAstVarId var == varNameToAstVarId var' && n == n' = u-  AstI2K RemOp (AstVar var') (AstConcreteK n')-   + AstI2K RemOp (AstN1K NegateOp (AstVar var)) (AstConcreteK n)-     | var == var' && n == n' = 0-  AstI2K RemOp (AstConvert _ (AstVar var')) (AstConcreteK n')-   + AstI2K RemOp (AstN1K NegateOp (AstConvert _ (AstVar var)))-                                   (AstConcreteK n)-     | varNameToAstVarId var == varNameToAstVarId var' && n == n' = 0-  AstI2K RemOp (AstVar var') (AstConcreteK n')-   + AstPlusK (AstI2K RemOp (AstN1K NegateOp (AstVar var)) (AstConcreteK n)) u-     | var == var' && n == n' = u-  AstI2K RemOp (AstConvert _ (AstVar var')) (AstConcreteK n')-   + AstPlusK (AstI2K RemOp (AstN1K NegateOp (AstConvert _ (AstVar var)))-                            (AstConcreteK n)) u-     | varNameToAstVarId var == varNameToAstVarId var' && n == n' = u--  AstPlusK u@AstConcreteK{} v + w = AstPlusK u (AstPlusK v w)  -- as above-  u + v@AstConcreteK{} = AstPlusK v u-  u + AstPlusK v@AstConcreteK{} w = AstPlusK v (AstPlusK u w)  -- as above-  t1 + t2 | eqK t1 t2 = 2 * t1-  t1 + AstTimesK (AstConcreteK n) t2 | eqK t1 t2 = AstConcreteK (n + 1) * t1-  AstTimesK (AstConcreteK n) t2 + t1 | eqK t1 t2 = AstConcreteK (n + 1) * t1-  AstTimesK (AstConcreteK n1) t1 + AstTimesK (AstConcreteK n2) t2-    | eqK t1 t2 = AstConcreteK (n1 + n2) * t1-  u + v = AstPlusK u v--  AstFromPrimal u * AstFromPrimal v = AstFromPrimal $ u * v-  AstFromDual{} * AstFromDual{} = 0-  AstConcreteK 0 * _ = 0-  _ * AstConcreteK 0 = 0-  AstConcreteK 1 * u = u-  u * AstConcreteK 1 = u-  AstConvert c u * AstConvert _ v-    | FTKS ZSS x <- ftkAst u-    , FTKS ZSS y <- ftkAst v-    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u))-    , Just Refl <- matchingFTK x y =-      AstConvert c $ u * v-  AstConcreteK n * AstConcreteK k = AstConcreteK (n * k)-  AstConcreteK n * AstTimesK (AstConcreteK k) u = AstConcreteK (n * k) * u-  AstTimesK (AstConcreteK n) u * AstConcreteK k = AstConcreteK (n * k) * u-  AstTimesK (AstConcreteK n) u * AstTimesK (AstConcreteK k) v =-    AstConcreteK (n * k) * AstTimesK u v  -- u and v can cancel, but unlikely--  u@AstConcreteK{} * AstPlusK v w = AstPlusK (u * v) (u * w)-  AstTimesK u@AstConcreteK{} x * AstPlusK v w =-    AstTimesK x (AstPlusK (u * v) (u * w))-  AstPlusK v w * u@AstConcreteK{} = AstPlusK (v * u) (w * u)-  AstPlusK v w * AstTimesK u@AstConcreteK{} x =-    AstTimesK (AstPlusK (v * u) (w * u)) x--  AstN1K NegateOp u * AstN1K NegateOp v = AstTimesK u v--  {- TODO: these rules increase the number of occurrences of a variable-     and trade multiplication and quotient for an equally problematic remnant,-     so they are disabled until we find a way to profit from them.-  -- With static shapes, the second argument to QuotOp and RemOp-  -- is often a constant, which makes such rules worth including,-  -- since they are likely to fire. To help them fire, we avoid changing-  -- that constant, if possible, e.g., in rules for NegateOp.-  AstConcreteK n * AstI2K QuotOp (AstVar var) (AstConcreteK n')-    | n == n' =-      AstPlusK-        (AstVar var)-        (negate (AstI2K RemOp (AstVar var) (AstConcreteK n)))-  AstTimesK (AstConcreteK n) x * AstI2K QuotOp (AstVar var)-                                               (AstConcreteK n')-    | n == n' =-      AstTimesK-        x-        (AstPlusK-           (AstVar var)-           (negate (AstI2K RemOp (AstVar var) (AstConcreteK n))))-  AstI2K QuotOp (AstVar var) (AstConcreteK n') * AstConcreteK n-    | n == n' =-      AstPlusK-        (AstVar var)-        (negate (AstI2K RemOp (AstVar var) (AstConcreteK n)))-  AstI2K QuotOp (AstVar var)-                (AstConcreteK n') * AstTimesK (AstConcreteK n) x-    | n == n' =-      AstTimesK-        (AstPlusK-           (AstVar var)-           (negate (AstI2K RemOp (AstVar var) (AstConcreteK n))))-        x-  -}--  AstTimesK u@AstConcreteK{} v * w = AstTimesK u (AstTimesK v w)  -- as above-  u * v@AstConcreteK{} = AstTimesK v u-  u * AstTimesK v@AstConcreteK{} w = AstTimesK v (AstTimesK u w)  -- as above-  u * v = AstTimesK u v--  negate (AstCond b n k) = AstCond b (negate n) (negate k)-  negate (AstLet var n k) = AstLet var n (negate k)-  negate (AstFromPrimal n) = AstFromPrimal (negate n)-  negate (AstFromDual n) = AstFromDual (negate n)-  negate (AstPlusK u v) = AstPlusK (negate u) (negate v)-  negate (AstTimesK u v) = negate u * v-  negate (AstN1K NegateOp u) = u-  negate (AstN1K SignumOp u) = AstN1K SignumOp (negate u)-  negate (AstI2K QuotOp u v) = AstI2K QuotOp (negate u) v-    -- v is likely positive and let's keep it so-  negate (AstI2K RemOp u v) = AstI2K RemOp (negate u) v-    -- v is likely positive and let's keep it so-  negate (AstConcreteK n) = AstConcreteK (negate n)-  negate (AstConvert c u)-    | FTKS ZSS x <- ftkAst u-    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u)) =-      AstConvert c (negate u)-  negate u = AstN1K NegateOp u-  abs (AstFromPrimal n) = AstFromPrimal (abs n)-  abs (AstFromDual n) = AstFromDual (abs n)-  abs (AstConcreteK n) = AstConcreteK (abs n)-  abs (AstN1K AbsOp u) = AstN1K AbsOp u-  abs (AstN1K NegateOp u) = abs u-  abs (AstConvert c u)-    | FTKS ZSS x <- ftkAst u-    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u)) =-      AstConvert c (abs u)-  abs u = AstN1K AbsOp u-  signum (AstFromPrimal n) = AstFromPrimal (signum n)-  signum (AstFromDual n) = AstFromDual (signum n)-  signum (AstConcreteK n) = AstConcreteK (signum n)-  signum (AstN1K SignumOp u) = AstN1K SignumOp u-  signum (AstConvert c u)-    | FTKS ZSS x <- ftkAst u-    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u)) =-      AstConvert c (signum u)-  signum u = AstN1K SignumOp u-  fromInteger i = fromPrimal $ AstConcreteK (fromInteger i)-  {-# SPECIALIZE instance Num (AstTensor ms FullSpan (TKScalar Int64)) #-}-  {-# SPECIALIZE instance Num (AstTensor ms PrimalSpan (TKScalar Int64)) #-}-  {-# SPECIALIZE instance Num (AstTensor ms FullSpan (TKScalar CInt)) #-}-  {-# SPECIALIZE instance Num (AstTensor ms PrimalSpan (TKScalar CInt)) #-}-  {-# SPECIALIZE instance Num (AstTensor ms FullSpan (TKScalar Double)) #-}-  {-# SPECIALIZE instance Num (AstTensor ms PrimalSpan (TKScalar Double)) #-}-  {-# SPECIALIZE instance Num (AstTensor ms FullSpan (TKScalar Float)) #-}-  {-# SPECIALIZE instance Num (AstTensor ms PrimalSpan (TKScalar Float)) #-}---- An approximation. False doesn't imply terms have different semantics,--- but True implies they have equal semantics.-eqK :: AstTensor ms s (TKScalar r) -> AstTensor ms s (TKScalar r) -> Bool-eqK (AstVar var1) (AstVar var2) = var1 == var2-eqK (AstLet @_ @_ @s1 var1 u1 v1) (AstLet @_ @_ @s2 var2 u2 v2)-  | FTKScalar @r1 <- ftkAst u1, FTKScalar @r2 <- ftkAst u2-  , Just Refl <- testEquality (typeRep @r1) (typeRep @r2)-  , Just Refl <- sameAstSpan @s1 @s2 =-    var1 == var2 && eqK u1 u2 && eqK v1 v2-eqK (AstPrimalPart u1) (AstPrimalPart u2) = eqK u1 u2-eqK (AstDualPart u1) (AstDualPart u2) = eqK u1 u2-eqK (AstFromPrimal u1) (AstFromPrimal u2) = eqK u1 u2-eqK (AstFromDual u1) (AstFromDual u2) = eqK u1 u2-eqK (AstPlusK u1 v1) (AstPlusK u2 v2) =-  eqK u1 u2 && eqK v1 v2 || eqK u1 v2 && eqK v1 u2-eqK (AstTimesK u1 v1) (AstTimesK u2 v2) =-  eqK u1 u2 && eqK v1 v2 || eqK u1 v2 && eqK v1 u2-eqK (AstN1K opCode1 u1) (AstN1K opCode2 u2) = opCode1 == opCode2 && eqK u1 u2-eqK (AstR1K opCode1 u1) (AstR1K opCode2 u2) = opCode1 == opCode2 && eqK u1 u2-eqK (AstR2K opCode1 u1 v1) (AstR2K opCode2 u2 v2) =-  opCode1 == opCode2 && eqK u1 u2 && eqK v1 v2-eqK (AstI2K opCode1 u1 v1) (AstI2K opCode2 u2 v2) =-  opCode1 == opCode2 && eqK u1 u2 && eqK v1 v2-eqK (AstConcreteK u1) (AstConcreteK u2) = u1 == u2-eqK (AstFloorK @r1 u1) (AstFloorK @r2 u2)-  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) = eqK u1 u2-eqK (AstFromIntegralK @r1 u1) (AstFromIntegralK @r2 u2)-  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) = eqK u1 u2-eqK (AstCastK @r1 u1) (AstCastK @r2 u2)-  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) = eqK u1 u2-eqK _ _ = False---- Div and mod operations are very costly (simplifying them requires--- constructing conditionals, etc), so they are not included in IntegralH.-instance (GoodScalar r, IntegralH r, Nested.IntElt r, AstSpan s)-         => IntegralH (AstTensor ms s (TKScalar r)) where-  quotH (AstFromPrimal n) (AstFromPrimal k) = AstFromPrimal (quotH n k)-  quotH (AstConvert c n) (AstConvert _ k)-    | FTKS ZSS x <- ftkAst n-    , FTKS ZSS y <- ftkAst k-    , Just Refl <- matchingFTK x (convertFTK c (ftkAst n))-    , Just Refl <- matchingFTK x y =-      AstConvert c (quotH n k)-  quotH (AstConcreteK n) (AstConcreteK k) = AstConcreteK (quotH n k)-  quotH (AstConcreteK 0) _ = 0-  quotH u (AstConcreteK 1) = u-  quotH (AstI2K RemOp _ (AstConcreteK k)) (AstConcreteK k')-    | k' >= k && k >= 0 = 0-  quotH (AstI2K QuotOp u v) w = quotH u (v * w)-  quotH (AstTimesK (AstConcreteK n) v) (AstConcreteK n')-    | n == n' = v-  quotH u v =-    let t = AstI2K QuotOp u v-        (u1, u2) = bounds t-    in if u1 == u2 then fromPrimal $ AstConcreteK u1 else t--  remH (AstFromPrimal n) (AstFromPrimal k) = AstFromPrimal (remH n k)-  remH (AstConvert c n) (AstConvert _ k)-    | FTKS ZSS x <- ftkAst n-    , FTKS ZSS y <- ftkAst k-    , Just Refl <- matchingFTK x (convertFTK c (ftkAst n))-    , Just Refl <- matchingFTK x y =-      AstConvert c (remH n k)-  remH (AstConcreteK n) (AstConcreteK k) = AstConcreteK (remH n k)-  remH (AstConcreteK 0) _ = 0-  remH _ (AstConcreteK 1) = 0-  remH (AstI2K RemOp u (AstConcreteK k)) (AstConcreteK k')-    | k' >= k && k >= 0 = AstI2K RemOp u (AstConcreteK k)-  remH (AstI2K RemOp u (AstConcreteK k)) (AstConcreteK k')-    | remH k k' == 0 && k > 0 = remH u (AstConcreteK k')-  remH (AstTimesK (AstConcreteK n) _) (AstConcreteK n')-    | remH n n' == 0 = 0-  remH u v =-    let t = AstI2K RemOp u v-        (u1, u2) = bounds t-    in if u1 == u2 then fromPrimal $ AstConcreteK u1 else t-  {-# SPECIALIZE instance IntegralH (AstTensor ms FullSpan (TKScalar Int64)) #-}-  {-# SPECIALIZE instance IntegralH (AstTensor ms PrimalSpan (TKScalar Int64)) #-}-  {-# SPECIALIZE instance IntegralH (AstTensor ms FullSpan (TKScalar CInt)) #-}-  {-# SPECIALIZE instance IntegralH (AstTensor ms PrimalSpan (TKScalar CInt)) #-}--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => Fractional (AstTensor ms s (TKScalar r)) where-  AstFromPrimal u / AstFromPrimal v = AstFromPrimal $ u / v-  AstConcreteK u / AstConcreteK v = AstConcreteK $ u / v-  u / v = AstR2K DivideOp u v-  recip (AstFromPrimal u) = AstFromPrimal (recip u)-  recip (AstConcreteK u) = AstConcreteK (recip u)-  recip u = AstR1K RecipOp u-  fromRational r = fromPrimal $ AstConcreteK (fromRational r)--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => Floating (AstTensor ms s (TKScalar r)) where-  pi = error "pi is not defined for tensors"-  exp (AstFromPrimal u) = AstFromPrimal $ exp u-  exp (AstConcreteK u) = AstConcreteK $ exp u-  exp u = AstR1K ExpOp u-  log (AstFromPrimal u) = AstFromPrimal $ log u-  log (AstConcreteK u) = AstConcreteK $ log u-  log u = AstR1K LogOp u-  sqrt (AstFromPrimal u) = AstFromPrimal $ sqrt u-  sqrt (AstConcreteK u) = AstConcreteK $ sqrt u-  sqrt u = AstR1K SqrtOp u-  (AstFromPrimal u) ** (AstFromPrimal v) = AstFromPrimal $ u ** v-  (AstConcreteK u) ** (AstConcreteK v) = AstConcreteK $ u ** v-  u ** v = AstR2K PowerOp u v-  logBase (AstFromPrimal u) (AstFromPrimal v) = AstFromPrimal $ logBase u v-  logBase (AstConcreteK u) (AstConcreteK v) = AstConcreteK $ logBase u v-  logBase u v = AstR2K LogBaseOp u v-  sin (AstFromPrimal u) = AstFromPrimal $ sin u-  sin (AstConcreteK u) = AstConcreteK $ sin u-  sin u = AstR1K SinOp u-  cos (AstFromPrimal u) = AstFromPrimal $ cos u-  cos (AstConcreteK u) = AstConcreteK $ cos u-  cos u = AstR1K CosOp u-  tan (AstFromPrimal u) = AstFromPrimal $ tan u-  tan (AstConcreteK u) = AstConcreteK $ tan u-  tan u = AstR1K TanOp u-  asin (AstFromPrimal u) = AstFromPrimal $ asin u-  asin (AstConcreteK u) = AstConcreteK $ asin u-  asin u = AstR1K AsinOp u-  acos (AstFromPrimal u) = AstFromPrimal $ acos u-  acos (AstConcreteK u) = AstConcreteK $ acos u-  acos u = AstR1K AcosOp u-  atan (AstFromPrimal u) = AstFromPrimal $ atan u-  atan (AstConcreteK u) = AstConcreteK $ atan u-  atan u = AstR1K AtanOp u-  sinh (AstFromPrimal u) = AstFromPrimal $ sinh u-  sinh (AstConcreteK u) = AstConcreteK $ sinh u-  sinh u = AstR1K SinhOp u-  cosh (AstFromPrimal u) = AstFromPrimal $ cosh u-  cosh (AstConcreteK u) = AstConcreteK $ cosh u-  cosh u = AstR1K CoshOp u-  tanh (AstFromPrimal u) = AstFromPrimal $ tanh u-  tanh (AstConcreteK u) = AstConcreteK $ tanh u-  tanh u = AstR1K TanhOp u-  asinh (AstFromPrimal u) = AstFromPrimal $ asinh u-  asinh (AstConcreteK u) = AstConcreteK $ asinh u-  asinh u = AstR1K AsinhOp u-  acosh (AstFromPrimal u) = AstFromPrimal $ acosh u-  acosh (AstConcreteK u) = AstConcreteK $ acosh u-  acosh u = AstR1K AcoshOp u-  atanh (AstFromPrimal u) = AstFromPrimal $ atanh u-  atanh (AstConcreteK u) = AstConcreteK $ atanh u-  atanh u = AstR1K AtanhOp u--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => RealFloatH (AstTensor ms s (TKScalar r)) where-  atan2H (AstFromPrimal u) (AstFromPrimal v) = AstFromPrimal $ atan2H u v-  atan2H (AstConcreteK u) (AstConcreteK v) = AstConcreteK $ atan2H u v-  atan2H u v = AstR2K Atan2Op u v----- * Unlawful numeric instances for ranked AST; lawful modulo evaluation--instance (GoodScalar r, AstSpan s)-         => Num (AstTensor ms s (TKR n r)) where-  (+) = liftRFromS2 (+)-  (-) = liftRFromS2 (-)-  (*) = liftRFromS2 (*)-  negate = liftRFromS1 negate-  abs = liftRFromS1 abs-  signum = liftRFromS1 signum-  fromInteger i = error $ "fromInteger is not defined for ranked tensors: "-                          ++ show i--instance (GoodScalar r, IntegralH r, Nested.IntElt r, AstSpan s)-         => IntegralH (AstTensor ms s (TKR n r)) where-  quotH = liftRFromS2 quotH-  remH = liftRFromS2 remH--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => Fractional (AstTensor ms s (TKR n r)) where-  (/) = liftRFromS2 (/)-  recip = liftRFromS1 recip-  fromRational r = error $ "fromRational is not defined for ranked tensors: "-                           ++ show r--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => Floating (AstTensor ms s (TKR n r)) where-  pi = error "pi is not defined for tensors"-  exp = liftRFromS1 exp-  log = liftRFromS1 log-  sqrt = liftRFromS1 sqrt-  (**) = liftRFromS2 (**)-  logBase = liftRFromS2 logBase-  sin = liftRFromS1 sin-  cos = liftRFromS1 cos-  tan = liftRFromS1 tan-  asin = liftRFromS1 asin-  acos = liftRFromS1 acos-  atan = liftRFromS1 atan-  sinh = liftRFromS1 sinh-  cosh = liftRFromS1 cosh-  tanh = liftRFromS1 tanh-  asinh = liftRFromS1 asinh-  acosh = liftRFromS1 acosh-  atanh = liftRFromS1 atanh--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => RealFloatH (AstTensor ms s (TKR n r)) where-  atan2H = liftRFromS2 atan2H----- * Unlawful numeric instances for shaped AST; lawful modulo evaluation--instance (GoodScalar r, AstSpan s)-         => Num (AstTensor ms s (TKS sh r)) where-  AstReplicate snat stk@STKS{} u + AstReplicate _ STKS{} v =-    AstReplicate snat stk $ u + v-  AstFromPrimal u + AstFromPrimal v = AstFromPrimal $ u + v-  AstFromDual u + AstFromDual v = AstFromDual $ u + v-  AstConvert c u + AstConvert _ v-    | FTKS ZSS x <- convertFTK c (ftkAst u)-    , Just Refl <- matchingFTK x (ftkAst u)-    , Just Refl <- matchingFTK x (ftkAst v) =-      AstConvert c $ u + v-  AstConcreteS z + u | Just 0 <- sunReplicateScal z = u-  u + AstConcreteS z | Just 0 <- sunReplicateScal z = u-  AstConcreteS n + AstConcreteS k = AstConcreteS (n + k)-  AstConcreteS n + AstPlusS (AstConcreteS k) u =-    AstPlusS (AstConcreteS (n + k)) u-  AstPlusS (AstConcreteS n) u + AstConcreteS k =-    AstPlusS (AstConcreteS (n + k)) u-  AstPlusS (AstConcreteS n) u + AstPlusS (AstConcreteS k) v =-    AstPlusS (AstConcreteS (n + k)) (AstPlusS u v)----  AstN1S NegateOp (AstVar var) + AstVar var'---    | var == var' = 0-  AstN1S NegateOp (AstVar var) + AstPlusS (AstVar var') u-    | var == var' = u---  AstVar var' + AstN1S NegateOp (AstVar var)---    | var == var' = 0-  AstVar var' + AstPlusS (AstN1S NegateOp (AstVar var)) u-    | var == var' = u--  AstPlusS u@AstConcreteS{} v + w = AstPlusS u (AstPlusS v w)-  u + v@AstConcreteS{} = AstPlusS v u-  u + AstPlusS v@AstConcreteS{} w = AstPlusS v (AstPlusS u w)-  u + v = AstPlusS u v--  AstReplicate snat stk@STKS{} u * AstReplicate _ STKS{} v =-    AstReplicate snat stk $ u * v-  AstFromPrimal u * AstFromPrimal v = AstFromPrimal $ u * v---  AstFromDual{} * AstFromDual{} = 0-  AstConcreteS z * _ | Just 0 <- sunReplicateScal z = AstConcreteS z-  _ * AstConcreteS z | Just 0 <- sunReplicateScal z = AstConcreteS z-  AstConcreteS s * u | Just 1 <- sunReplicateScal s = u-  u * AstConcreteS s | Just 1 <- sunReplicateScal s = u-  AstConvert c u * AstConvert _ v-    | FTKS ZSS x <- convertFTK c (ftkAst u)-    , Just Refl <- matchingFTK x (ftkAst u)-    , Just Refl <- matchingFTK x (ftkAst v) =-      AstConvert c $ u * v-  AstConcreteS n * AstConcreteS k = AstConcreteS (n * k)-  AstConcreteS n * AstTimesS (AstConcreteS k) u =-    AstTimesS (AstConcreteS (n * k)) u-  AstTimesS (AstConcreteS n) u * AstConcreteS k =-    AstTimesS (AstConcreteS (n * k)) u-  AstTimesS (AstConcreteS n) u * AstTimesS (AstConcreteS k) v =-    AstTimesS (AstConcreteS (n * k)) (AstTimesS u v)--  u@AstConcreteS{} * AstPlusS v w = AstPlusS (u * v) (u * w)-  AstTimesS u@AstConcreteS{} x * AstPlusS v w =-    AstTimesS x (AstPlusS (u * v) (u * w))-  AstPlusS v w * u@AstConcreteS{} = AstPlusS (v * u) (w * u)-  AstPlusS v w * AstTimesS u@AstConcreteS{} x =-    AstTimesS (AstPlusS (v * u) (w * u)) x--  AstN1S NegateOp u * AstN1S NegateOp v = AstTimesS u v--  AstTimesS u@AstConcreteS{} v * w = AstTimesS u (AstTimesS v w)-  u * v@AstConcreteS{} = AstTimesS v u-  u * AstTimesS v@AstConcreteS{} w = AstTimesS v (AstTimesS u w)-  u * v = AstTimesS u v--  negate (AstReplicate snat stk@STKS{} u) = AstReplicate snat stk (negate u)-  negate (AstCond b n k) = AstCond b (negate n) (negate k)--- TODO: negate (AstBuild1 k stk (var, v)) = AstBuild1 k stk (var, negate v)-  negate (AstLet var n k) = AstLet var n (negate k)-  negate (AstFromPrimal n) = AstFromPrimal (negate n)-  negate (AstFromDual n) = AstFromDual (negate n)-  negate (AstPlusS u v) = AstPlusS (negate u) (negate v)-  negate (AstTimesS u v) = AstTimesS (negate u) v-  negate (AstN1S NegateOp u) = u-  negate (AstN1S SignumOp u) = AstN1S SignumOp (negate u)-  negate (AstI2S QuotOp u v) = AstI2S QuotOp (negate u) v-    -- v is likely positive and let's keep it so-  negate (AstI2S RemOp u v) = AstI2S RemOp (negate u) v-    -- v is likely positive and let's keep it so-  negate (AstConcreteS n) = AstConcreteS (negate n)-  negate (AstGatherS @shm @shn @shp shn v (vars, ix)) =-    AstGatherS @shm @shn @shp shn (negate v) (vars, ix)-  negate (AstConvert c n)-    | FTKS ZSS x <- convertFTK c (ftkAst n)-    , Just Refl <- matchingFTK x (ftkAst n) =-      AstConvert c (negate n)-  negate u = AstN1S NegateOp u-  abs (AstReplicate snat stk@STKS{} u) = AstReplicate snat stk (abs u)-  abs (AstFromPrimal n) = AstFromPrimal (abs n)-  abs (AstFromDual n) = AstFromDual (abs n)-  abs (AstN1S AbsOp u) = AstN1S AbsOp u-  abs (AstConcreteS u) = AstConcreteS (abs u)-  abs (AstConvert c n)-    | FTKS ZSS x <- convertFTK c (ftkAst n)-    , Just Refl <- matchingFTK x (ftkAst n) =-      AstConvert c (abs n)-  abs (AstN1S NegateOp u) = abs u-  abs u = AstN1S AbsOp u-  signum (AstReplicate snat stk@STKS{} u) = AstReplicate snat stk (signum u)-  signum (AstFromPrimal n) = AstFromPrimal (signum n)-  signum (AstFromDual n) = AstFromDual (signum n)-  signum (AstN1S SignumOp u) = AstN1S SignumOp u-  signum (AstConcreteS u) = AstConcreteS (signum u)-  signum (AstConvert c n)-    | FTKS ZSS x <- convertFTK c (ftkAst n)-    , Just Refl <- matchingFTK x (ftkAst n) =-      AstConvert c (signum n)-  signum u = AstN1S SignumOp u-  fromInteger i = error $ "fromInteger is not defined for shaped tensors: "-                          ++ show i--instance (GoodScalar r, IntegralH r, Nested.IntElt r, AstSpan s)-         => IntegralH (AstTensor ms s (TKS sh r)) where-  quotH (AstReplicate snat stk@STKS{} u) (AstReplicate _ STKS{} v) =-    AstReplicate snat stk $ quotH u v-  quotH (AstFromPrimal n) (AstFromPrimal k) = AstFromPrimal (quotH n k)-  quotH (AstConvert c n) (AstConvert _ k)-    | FTKS ZSS x <- convertFTK c (ftkAst n)-    , Just Refl <- matchingFTK x (ftkAst n)-    , Just Refl <- matchingFTK x (ftkAst k) =-      AstConvert c (quotH n k)-  quotH (AstConcreteS n) (AstConcreteS k) = AstConcreteS (quotH n k)-  quotH (AstConcreteS z) _ | Just 0 <- sunReplicateScal z = AstConcreteS z-  quotH u (AstConcreteS s) | Just 1 <- sunReplicateScal s = u-  quotH (AstI2S QuotOp u v) w = quotH u (v * w)-  quotH u v = AstI2S QuotOp u v--  remH (AstReplicate snat stk@STKS{} u) (AstReplicate _ STKS{} v) =-    AstReplicate snat stk $ remH u v-  remH (AstFromPrimal n) (AstFromPrimal k) = AstFromPrimal (remH n k)-  remH (AstConvert c n) (AstConvert _ k)-    | FTKS ZSS x <- convertFTK c (ftkAst n)-    , Just Refl <- matchingFTK x (ftkAst n)-    , Just Refl <- matchingFTK x (ftkAst k) =-      AstConvert c (remH n k)-  remH (AstConcreteS n) (AstConcreteS k) = AstConcreteS (remH n k)-  remH (AstConcreteS z) _ | Just 0 <- sunReplicateScal z = AstConcreteS z---  remH _ (AstConcreteS s) | Just 1 <- sunReplicateScal s = AstConcreteS 0-  remH u v = AstI2S RemOp u v--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => Fractional (AstTensor ms s (TKS sh r)) where-  AstFromPrimal u / AstFromPrimal v = AstFromPrimal $ u / v-  AstConcreteS u / AstConcreteS v = AstConcreteS $ u / v-  u / v = AstR2S DivideOp u v-  recip (AstFromPrimal u) = AstFromPrimal (recip u)-  recip (AstConcreteS u) = AstConcreteS (recip u)-  recip u = AstR1S RecipOp u-  fromRational r = error $ "fromRational is not defined for shaped tensors: "-                           ++ show r--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => Floating (AstTensor ms s (TKS sh r)) where-  pi = error "pi is not defined for tensors"-  exp (AstFromPrimal u) = AstFromPrimal $ exp u-  exp (AstConcreteS u) = AstConcreteS $ exp u-  exp u = AstR1S ExpOp u-  log (AstFromPrimal u) = AstFromPrimal $ log u-  log (AstConcreteS u) = AstConcreteS $ log u-  log u = AstR1S LogOp u-  sqrt (AstFromPrimal u) = AstFromPrimal $ sqrt u-  sqrt (AstConcreteS u) = AstConcreteS $ sqrt u-  sqrt u = AstR1S SqrtOp u-  (AstFromPrimal u) ** (AstFromPrimal v) = AstFromPrimal $ u ** v-  (AstConcreteS u) ** (AstConcreteS v) = AstConcreteS $ u ** v-  u ** v = AstR2S PowerOp u v-  logBase (AstFromPrimal u) (AstFromPrimal v) = AstFromPrimal $ logBase u v-  logBase (AstConcreteS u) (AstConcreteS v) = AstConcreteS $ logBase u v-  logBase u v = AstR2S LogBaseOp u v-  sin (AstFromPrimal u) = AstFromPrimal $ sin u-  sin (AstConcreteS u) = AstConcreteS $ sin u-  sin u = AstR1S SinOp u-  cos (AstFromPrimal u) = AstFromPrimal $ cos u-  cos (AstConcreteS u) = AstConcreteS $ cos u-  cos u = AstR1S CosOp u-  tan (AstFromPrimal u) = AstFromPrimal $ tan u-  tan (AstConcreteS u) = AstConcreteS $ tan u-  tan u = AstR1S TanOp u-  asin (AstFromPrimal u) = AstFromPrimal $ asin u-  asin (AstConcreteS u) = AstConcreteS $ asin u-  asin u = AstR1S AsinOp u-  acos (AstFromPrimal u) = AstFromPrimal $ acos u-  acos (AstConcreteS u) = AstConcreteS $ acos u-  acos u = AstR1S AcosOp u-  atan (AstFromPrimal u) = AstFromPrimal $ atan u-  atan (AstConcreteS u) = AstConcreteS $ atan u-  atan u = AstR1S AtanOp u-  sinh (AstFromPrimal u) = AstFromPrimal $ sinh u-  sinh (AstConcreteS u) = AstConcreteS $ sinh u-  sinh u = AstR1S SinhOp u-  cosh (AstFromPrimal u) = AstFromPrimal $ cosh u-  cosh (AstConcreteS u) = AstConcreteS $ cosh u-  cosh u = AstR1S CoshOp u-  tanh (AstFromPrimal u) = AstFromPrimal $ tanh u-  tanh (AstConcreteS u) = AstConcreteS $ tanh u-  tanh u = AstR1S TanhOp u-  asinh (AstFromPrimal u) = AstFromPrimal $ asinh u-  asinh (AstConcreteS u) = AstConcreteS $ asinh u-  asinh u = AstR1S AsinhOp u-  acosh (AstFromPrimal u) = AstFromPrimal $ acosh u-  acosh (AstConcreteS u) = AstConcreteS $ acosh u-  acosh u = AstR1S AcoshOp u-  atanh (AstFromPrimal u) = AstFromPrimal $ atanh u-  atanh (AstConcreteS u) = AstConcreteS $ atanh u-  atanh u = AstR1S AtanhOp u--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => RealFloatH (AstTensor ms s (TKS sh r)) where-  atan2H (AstFromPrimal u) (AstFromPrimal v) = AstFromPrimal $ atan2H u v-  atan2H (AstConcreteS u) (AstConcreteS v) = AstConcreteS $ atan2H u v-  atan2H u v = AstR2S Atan2Op u v----- * Unlawful numeric instances for mixed AST; lawful modulo evaluation--instance (GoodScalar r, AstSpan s)-         => Num (AstTensor ms s (TKX sh r)) where-  (+) = liftXFromS2 (+)-  (-) = liftXFromS2 (-)-  (*) = liftXFromS2 (*)-  negate = liftXFromS1 negate-  abs = liftXFromS1 abs-  signum = liftXFromS1 signum-  fromInteger i = error $ "fromInteger is not defined for mixed tensors: "-                          ++ show i--instance (GoodScalar r, IntegralH r, Nested.IntElt r, AstSpan s)-         => IntegralH (AstTensor ms s (TKX sh r)) where-  quotH = liftXFromS2 quotH-  remH = liftXFromS2 remH--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => Fractional (AstTensor ms s (TKX sh r)) where-  (/) = liftXFromS2 (/)-  recip = liftXFromS1 recip-  fromRational r = error $ "fromRational is not defined for mixed tensors: "-                           ++ show r--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => Floating (AstTensor ms s (TKX sh r)) where-  pi = error "pi is not defined for tensors"-  exp = liftXFromS1 exp-  log = liftXFromS1 log-  sqrt = liftXFromS1 sqrt-  (**) = liftXFromS2 (**)-  logBase = liftXFromS2 logBase-  sin = liftXFromS1 sin-  cos = liftXFromS1 cos-  tan = liftXFromS1 tan-  asin = liftXFromS1 asin-  acos = liftXFromS1 acos-  atan = liftXFromS1 atan-  sinh = liftXFromS1 sinh-  cosh = liftXFromS1 cosh-  tanh = liftXFromS1 tanh-  asinh = liftXFromS1 asinh-  acosh = liftXFromS1 acosh-  atanh = liftXFromS1 atanh--instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)-         => RealFloatH (AstTensor ms s (TKX sh r)) where-  atan2H = liftXFromS2 atan2H----- * Unlawful instances of AST for bool; they are lawful modulo evaluation---- Simple variable comparisons, if any, come first.-instance Boolean (AstBool ms) where-  true = AstBoolConst True-  false = AstBoolConst False-  notB (AstBoolConst b) = AstBoolConst $ not b-  notB (AstBoolNot b) = b-  notB b = AstBoolNot b-  AstBoolConst True &&* b = b-  AstBoolConst False &&* _b = AstBoolConst False-  b &&* AstBoolConst True = b-  _b &&* AstBoolConst False = AstBoolConst False-  AstBoolAnd b c &&* d = b &&* (c &&* d)-  b@(AstLeqK AstConcreteK{} AstVar{}) &&* c = AstBoolAnd b c-  b@(AstLeqK AstConcreteK{} (AstN1K NegateOp-                                    AstVar{})) &&* c = AstBoolAnd b c-  b@(AstBoolNot-       (AstLeqK AstConcreteK{} AstVar{})) &&* c = AstBoolAnd b c-  b@(AstBoolNot-       (AstLeqK AstConcreteK{} (AstN1K NegateOp-                                       AstVar{}))) &&* c = AstBoolAnd b c-  b@(AstBoolNot-       (AstBoolAnd (AstLeqK AstConcreteK{} AstVar{}) _)) &&* c = AstBoolAnd b c-  b@(AstBoolNot-       (AstBoolAnd-          (AstLeqK AstConcreteK{}-                   (AstN1K NegateOp AstVar{})) _)) &&* c = AstBoolAnd b c-  b@(AstBoolNot-       (AstBoolAnd (AstBoolNot (AstLeqK AstConcreteK{}-                                        AstVar{})) _)) &&* c = AstBoolAnd b c-  b@(AstBoolNot-       (AstBoolAnd-          (AstBoolNot-             (AstLeqK AstConcreteK{}-                      (AstN1K NegateOp AstVar{}))) _)) &&* c = AstBoolAnd b c-  b &&* c@(AstLeqK AstConcreteK{} AstVar{}) = AstBoolAnd c b-  b &&* c@(AstLeqK AstConcreteK{} (AstN1K NegateOp-                                          AstVar{})) = AstBoolAnd c b-  b &&* c@(AstBoolNot-             (AstLeqK AstConcreteK{} AstVar{})) = AstBoolAnd c b-  b &&* c@(AstBoolNot-             (AstLeqK AstConcreteK{} (AstN1K NegateOp-                                             AstVar{}))) = AstBoolAnd c b-  b &&* c@(AstBoolNot-             (AstBoolAnd (AstLeqK AstConcreteK{} AstVar{}) _)) = AstBoolAnd c b-  b &&* c@(AstBoolNot-             (AstBoolAnd-                (AstLeqK AstConcreteK{}-                         (AstN1K NegateOp AstVar{})) _)) = AstBoolAnd c b-  b &&* c@(AstBoolNot-             (AstBoolAnd (AstBoolNot (AstLeqK AstConcreteK{}-                                              AstVar{})) _)) = AstBoolAnd c b-  b &&* c@(AstBoolNot-             (AstBoolAnd-                (AstBoolNot-                   (AstLeqK AstConcreteK{}-                            (AstN1K NegateOp AstVar{}))) _)) = AstBoolAnd c b-  b &&* AstBoolAnd-          c@(AstLeqK AstConcreteK{} AstVar{}) d = AstBoolAnd c (b &&* d)-  b &&* AstBoolAnd-          c@(AstLeqK AstConcreteK{}-                     (AstN1K NegateOp AstVar{})) d = AstBoolAnd c (b &&* d)-  b &&* AstBoolAnd-          c@(AstBoolNot (AstLeqK AstConcreteK{}-                                 AstVar{})) d = AstBoolAnd c (b &&* d)-  b &&* AstBoolAnd-          c@(AstBoolNot-               (AstLeqK AstConcreteK{}-                        (AstN1K NegateOp AstVar{}))) d = AstBoolAnd c (b &&* d)-  b &&* AstBoolAnd-          c@(AstBoolNot-               (AstBoolAnd (AstLeqK AstConcreteK{}-                                    AstVar{}) _)) d = AstBoolAnd c (b &&* d)-  b &&* AstBoolAnd-          c@(AstBoolNot-               (AstBoolAnd-                  (AstLeqK AstConcreteK{}-                           (AstN1K NegateOp-                                   AstVar{})) _)) d = AstBoolAnd c (b &&* d)-  b &&* AstBoolAnd-          c@(AstBoolNot-               (AstBoolAnd-                  (AstBoolNot-                     (AstLeqK AstConcreteK{}-                              AstVar{})) _)) d = AstBoolAnd c (b &&* d)-  b &&* AstBoolAnd-          c@(AstBoolNot-               (AstBoolAnd-                  (AstBoolNot-                     (AstLeqK AstConcreteK{}-                              (AstN1K-                                 NegateOp-                                 AstVar{}))) _)) d = AstBoolAnd c (b &&* d)-  b &&* c = AstBoolAnd b c-  b ||* c = notB (notB b &&* notB c)---- TODO: refactor with something like liftRFromS2-instance (AstSpan s, GoodScalar r) => EqH (AstTensor ms s) (TKR n r) where-  v ==. u = case ftkAst v of-    FTKR shv' _ -> case ftkAst u of-      FTKR shu' _ ->-        withShsFromShR shv' $ \shv ->-          withShsFromShR shu' $ \shu ->-            case testEquality shv shu of-              Just Refl ->-                cAstSFromR shu (primalPart v)-                ==. cAstSFromR shv (primalPart u)-              _ -> error $ "(==.): shapes don't match: "-                           ++ show (shu, shv)--instance (AstSpan s, GoodScalar r) => EqH (AstTensor ms s) (TKX sh r) where-  v ==. u = case ftkAst v of-    FTKX shv' _ -> case ftkAst u of-      FTKX shu' _ ->-        withShsFromShX shv' $ \shv ->-          withShsFromShX shu' $ \shu ->-            case testEquality shv shu of-              Just Refl ->-                cAstSFromX shu (primalPart v)-                ==. cAstSFromX shv (primalPart u)-              _ -> error $ "(==.): shapes don't match: "-                           ++ show (shu, shv)--instance (AstSpan s, GoodScalar r) => OrdH (AstTensor ms s) (TKR n r) where-  v <=. u = case ftkAst v of-    FTKR shv' _ -> case ftkAst u of-      FTKR shu' _ ->-        withShsFromShR shv' $ \shv ->-          withShsFromShR shu' $ \shu ->-            case testEquality shv shu of-              Just Refl ->-                cAstSFromR shu (primalPart v)-                <=. cAstSFromR shv (primalPart u)-              _ -> error $ "(<=.): shapes don't match: "-                           ++ show (shu, shv)--instance (AstSpan s, GoodScalar r) => OrdH (AstTensor ms s) (TKX sh r) where-  v <=. u = case ftkAst v of-    FTKX shv' _ -> case ftkAst u of-      FTKX shu' _ ->-        withShsFromShX shv' $ \shv ->-          withShsFromShX shu' $ \shu ->-            case testEquality shv shu of-              Just Refl ->-                cAstSFromX shu (primalPart v)-                <=. cAstSFromX shv (primalPart u)-              _ -> error $ "(<=.): shapes don't match: "-                           ++ show (shu, shv)---- TODO: share u and v, since they are duplicated here-instance (AstSpan s, GoodScalar r)-         => EqH (AstTensor ms s) (TKScalar r) where-  v ==. u = v <=. u &&* u <=. v-  {- TODO: for this to work, booleans have to be first-class:-  vUnshared ==. uUnshared = astLetFunNoSimplify vUnshared $ \v ->-                            astLetFunNoSimplify uUnshared $ \u ->-    v <=. u &&* u <=. v -}--instance (AstSpan s, GoodScalar r)-         => EqH (AstTensor ms s) (TKS sh r) where-  v ==. u = v <=. u &&* u <=. v---- These are common in indexing, so worth optimizing early via AstConcrete.--- We keep AstConcrete on the left, as with AstPlusK and others.-instance (AstSpan s, GoodScalar r)-         => OrdH (AstTensor ms s) (TKScalar r) where-  u <=. v | let (u1, u2) = bounds u-                (v1, v2) = bounds v-          , u2 <= v1 || u1 > v2 = AstBoolConst (u2 <= v1)-  AstFromPrimal u <=. AstFromPrimal v = u <=. v-  AstPrimalPart u <=. AstPrimalPart v = u <=. v-  AstFromDual{} <=. AstFromDual{} = AstBoolConst True-  AstConvert _ u <=. AstConvert _ v-    | FTKS ZSS (FTKScalar @ru) <- ftkAst u-    , FTKS ZSS (FTKScalar @rv) <- ftkAst v-    , Just Refl <- testEquality (typeRep @ru) (typeRep @rv)-    = u <=. v-  AstConcreteK u <=. AstConvert _ v-    | FTKS ZSS (FTKScalar @rv) <- ftkAst v-    , Just Refl <- testEquality (typeRep @rv) (typeRep @r)-    = AstConcreteS (unConcrete $ sfromK $ Concrete u) <=. v-  u <=. AstPlusK (AstConcreteK v) w =-    u - AstConcreteK v <=. w-  AstPlusK (AstConcreteK u) w <=. v =-    AstConcreteK u <=. v - w-  u <=. AstConcreteK v =-    AstConcreteK (negate v) <=. negate u-  AstConcreteK u <=. AstTimesK (AstConcreteK v) w-    | v > 0 && u >= 0-    , Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =-      AstConcreteK ((u + v - 1) `quotH` v) <=. w -- 10 == 5 * 2, 11 > 5 * 2-  AstConcreteK u <=. AstTimesK (AstConcreteK v) w-    | v > 0 && u < 0-    , Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =-      AstConcreteK (u `quotH` v) <=. w  -- -10 == 5 * -2, -9 > 5 * -2-  AstConcreteK u <=. AstTimesK (AstConcreteK v) w-    | v < 0-    , Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =-      AstConcreteK u <=. AstTimesK (AstConcreteK $ negate v) (AstN1K NegateOp w)-  v@AstConcreteK{} <=. u =-    AstLeqK (primalPart v) (primalPart u)-  v <=. u =-    AstConcreteK 0 <=. primalPart u - primalPart v--instance (AstSpan s, GoodScalar r)-         => OrdH (AstTensor ms s) (TKS sh r) where-  AstFromPrimal u <=. AstFromPrimal v = u <=. v-  AstFromDual{} <=. AstFromDual{} = AstBoolConst True-  AstPrimalPart u <=. AstPrimalPart v = u <=. v-  AstConvert c u <=. AstConvert _ v-    | FTKS ZSS x <- convertFTK c (ftkAst u)-    , Just Refl <- matchingFTK x (ftkAst u)-    , Just Refl <- matchingFTK x (ftkAst v) = u <=. v-  AstConcreteS u <=. AstConvert c v-    | FTKS ZSS (FTKScalar @rz) <- convertFTK c (ftkAst v)-    , FTKScalar @ry <- ftkAst v-    , Just Refl <- testEquality (typeRep @ry) (typeRep @rz) =-      AstConcreteK (unConcrete $ kfromS $ Concrete u) <=. v-  AstConcreteS u <=. AstConcreteS v =-    AstBoolConst $ Concrete @(TKS sh r) u <=. Concrete v-  u <=. AstPlusS (AstConcreteS v) w =-    u - AstConcreteS v <=. w-  AstPlusS (AstConcreteS u) w <=. v =-    AstConcreteS u <=. v - w-  u <=. AstConcreteS v =-    AstConcreteS (negate v) <=. negate u-  AstVar u <=. AstVar v | u == v =-    AstBoolConst True-  AstConvert _ (AstVar u) <=. AstConvert _ (AstVar v)-    | varNameToAstVarId u == varNameToAstVarId v =-      AstBoolConst True-  v <=. u = AstLeqS (primalPart v) (primalPart u)----- * AstRaw, AstNoVectorize and AstNoSimplify definitions---- | An AST variant that doesn't vectorize terms and also builds them--- with ordinary, non-simplifying constructors. It's based on sharing--- rather than lets and commonly used as the instance for primals--- inside ADVal and, consequently, used for evaluating delta expressions.-type role AstRaw nominal nominal-newtype AstRaw s y =-  AstRaw {unAstRaw :: AstTensor AstMethodShare s y}- deriving Show---- | An AST variant for testing that doesn't vectorize terms, but still--- builds them using simplifying smart constructors.-type role AstNoVectorize nominal nominal-newtype AstNoVectorize s y =-  AstNoVectorize {unAstNoVectorize :: AstTensor AstMethodLet s y}- deriving Show---- | An AST variant for testing that vectorizes terms, but builds them--- with ordinary, non-simplifying constructors.-type role AstNoSimplify nominal nominal-newtype AstNoSimplify s y =-  AstNoSimplify {unAstNoSimplify :: AstTensor AstMethodLet s y}- deriving Show----- * AstRaw, AstNoVectorize and AstNoSimplify type family instances--type instance PrimalOf (AstRaw s) = AstRaw PrimalSpan-type instance DualOf (AstRaw s) = AstTensor AstMethodShare DualSpan-type instance ShareOf (AstRaw s) = AstRaw s-type instance HFunOf (AstRaw s) x y = AstHFun s s x y-type instance BoolOf (AstRaw s) = AstBool AstMethodShare--type instance PrimalOf (AstNoVectorize s) = AstNoVectorize PrimalSpan-type instance DualOf (AstNoVectorize s) = AstTensor AstMethodLet DualSpan-type instance ShareOf (AstNoVectorize s) = AstRaw s-type instance HFunOf (AstNoVectorize s) x z = AstHFun s s x z-type instance BoolOf (AstNoVectorize s) = AstBool AstMethodLet--type instance PrimalOf (AstNoSimplify s) = AstNoSimplify PrimalSpan-type instance DualOf (AstNoSimplify s) = AstTensor AstMethodLet DualSpan-type instance ShareOf (AstNoSimplify s) = AstRaw s-type instance HFunOf (AstNoSimplify s) x z = AstHFun s s x z-type instance BoolOf (AstNoSimplify s) = AstBool AstMethodLet----- * AstRaw, AstNoVectorize and AstNoSimplify other instances--instance EqH (AstTensor AstMethodShare s) y => EqH (AstRaw s) y where-  AstRaw v ==. AstRaw u = v ==. u-instance OrdH (AstTensor AstMethodShare s) y => OrdH (AstRaw s) y where-  AstRaw v <=. AstRaw u = v <=. u--deriving instance Eq (AstRaw s y)-deriving instance Ord (AstRaw s y)-deriving instance Num (AstTensor AstMethodShare s y) => Num (AstRaw s y)-deriving instance IntegralH (AstTensor AstMethodShare s y)-                  => IntegralH (AstRaw s y)-deriving instance Fractional (AstTensor AstMethodShare s y)-                  => Fractional (AstRaw s y)-deriving instance Floating (AstTensor AstMethodShare s y)-                  => Floating (AstRaw s y)-deriving instance RealFloatH (AstTensor AstMethodShare s y)-                  => RealFloatH (AstRaw s y)--instance EqH (AstTensor AstMethodLet s) y => EqH (AstNoVectorize s) y where-  AstNoVectorize v ==. AstNoVectorize u = v ==. u-instance OrdH (AstTensor AstMethodLet s) y => OrdH (AstNoVectorize s) y where-  AstNoVectorize v <=. AstNoVectorize u = v <=. u-deriving instance Eq (AstNoVectorize s y)-deriving instance Ord (AstNoVectorize s y)-deriving instance Num (AstTensor AstMethodLet s y) => Num (AstNoVectorize s y)-deriving instance (IntegralH (AstTensor AstMethodLet s y))-                  => IntegralH (AstNoVectorize s y)-deriving instance Fractional (AstTensor AstMethodLet s y)-                  => Fractional (AstNoVectorize s y)-deriving instance Floating (AstTensor AstMethodLet s y)-                  => Floating (AstNoVectorize s y)-deriving instance (RealFloatH (AstTensor AstMethodLet s y))-                  => RealFloatH (AstNoVectorize s y)--instance EqH (AstTensor AstMethodLet s) y => EqH (AstNoSimplify s) y where-  AstNoSimplify v ==. AstNoSimplify u = v ==. u-instance OrdH (AstTensor AstMethodLet s) y => OrdH (AstNoSimplify s) y where-  AstNoSimplify v <=. AstNoSimplify u = v <=. u-deriving instance Eq (AstNoSimplify s y)-deriving instance Ord (AstNoSimplify s y)-deriving instance Num (AstTensor AstMethodLet s y) => Num (AstNoSimplify s y)-deriving instance (IntegralH (AstTensor AstMethodLet s y))-                  => IntegralH (AstNoSimplify s y)-deriving instance Fractional (AstTensor AstMethodLet s y)-                  => Fractional (AstNoSimplify s y)-deriving instance Floating (AstTensor AstMethodLet s y)-                  => Floating (AstNoSimplify s y)-deriving instance (RealFloatH (AstTensor AstMethodLet s y))-                  => RealFloatH (AstNoSimplify s y)----- * Misc--astLetFunNoSimplify-  :: forall y z s s2. AstSpan s-  => AstTensor AstMethodLet s y-  -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)-  -> AstTensor AstMethodLet s2 z-astLetFunNoSimplify a f | astIsSmall True a = f a-                            -- too important an optimization to skip-astLetFunNoSimplify a f = case a of-  AstFromS' @y2 ftkz v ->-    let (var, ast) = funToAst2 (ftkAst v) Nothing (f . cAstFromS @y2 ftkz)-    in AstLet var v ast-  AstFromPrimal (AstFromS' @y2 ftkz vRaw) ->-    let v = AstFromPrimal vRaw-        (var, ast) = funToAst2 (ftkAst v) Nothing (f . cAstFromS @y2 ftkz)-    in AstLet var v ast-  _ -> case ftkAst a of-    ftk@(FTKR @_ @x2 sh' x) ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        let (var, ast) =-              funToAst2 (FTKS sh x) Nothing-                        (f . cAstFromS @(TKS2 sh x2) ftk)-        in AstLet var (cAstSFromR @sh sh a) ast-             -- safe, because subsitution ruled out above-    ftk@(FTKX @_ @x sh' x) ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        let (var, ast) =-              funToAst2 (FTKS sh x) Nothing-                        (f . cAstFromS @(TKS2 sh x) ftk)-        in AstLet var (cAstSFromX @sh sh a) ast-    -- processing product recursively may be not worth it-    ftk -> let (var, ast) = funToAst2 ftk Nothing f-           in AstLet var a ast--sunReplicateScal :: Nested.Elt a-                 => Nested.Shaped sh a -> Maybe a-sunReplicateScal (Nested.Shaped arr)-  | all (all (== 0) . take (shxLength (Nested.mshape arr)))-        (Mixed.marrayStrides arr)-  , shxSize (Nested.mshape arr) /= 0 =-    Just $ Nested.mindex arr $ ixxZero' $ Nested.mshape arr-sunReplicateScal _ = Nothing--sunReplicate1 :: Nested.Elt a-              => Nested.Shaped (n ': sh) a -> Maybe (Nested.Shaped sh a)-sunReplicate1 a | (snat :$$ _) <- Nested.sshape a =-  sunReplicateN (snat :$$ ZSS) a--sunReplicateN :: Nested.Elt a-              => ShS shm -> Nested.Shaped (shm ++ shn) a-              -> Maybe (Nested.Shaped shn a)-sunReplicateN shm a@(Nested.Shaped arr)-  | all (all (== 0) . take (shsLength shm)) (Mixed.marrayStrides arr)-  , shsSize shm /= 0 =-    Just $ Nested.sindexPartial a $ ixsZero shm-sunReplicateN _ _ = Nothing+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Definitions, mostly class instances, needed to make AST a valid+-- carrier for a tensor class algebra (instance) defined+-- in "HordeAd.Core.OpsAst".+-- This algebra permits user programs to be instantiated as AST terms,+-- as well as to be interpreted into AST terms and it also permits derivatives+-- to be expressed as AST terms.+module HordeAd.Core.CarriersAst+  ( AstRaw(..), AstNoVectorize(..), AstNoSimplify(..)+  , sunReplicate, sunReplicate1, sunReplicateN, unReplC, unAstK, unAstS+  , eqY, eqUnknownShapes+  ) where++import Prelude++import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Type.Reflection (typeRep)++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed qualified as Mixed+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.Ast+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Conversion+import HordeAd.Core.OpsConcrete ()+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Type family instances for AstTensor++-- This can't be just HFun, because they need to be vectorized+-- and vectorization applies such functions to the variable from build1+-- and the variable has to be eliminated via vectorization to preserve+-- the closed form of the function. Just applying a Haskell closure+-- to the build1 variable and then duplicating the result of the function+-- would not eliminate the variable and also would likely results+-- in more costly computations. Also, that would prevent simplification+-- of the instances, especially after applied to arguments that are terms.+type instance HFunOf (AstTensor AstMethodLet s) = AstHFun s++-- This can't be defined only for FullSpan, because the BaseTensor instance+-- for @AstTensor ms PrimalSpan@ needs it and we need the instance+-- to satisfy ADReady constraints for AST.+type instance PrimalOf (AstTensor ms s) = AstTensor ms (PrimalStepSpan s)+type instance DualOf (AstTensor ms s) = AstTensor ms DualSpan+type instance PlainOf (AstTensor ms s) = AstTensor ms PlainSpan+type instance ShareOf (AstTensor ms s) = AstRaw s++-- These are, unfortunately, required by some numeric instances.+instance Eq (AstTensor ms s y) where+  (==) = error "Eq is not defined for AST; please use EqH instead"+  (/=) = error "Eq is not defined for AST; please use EqH instead"++instance Ord (AstTensor ms s y) where+  (<=) = error "Ord is not defined for AST; please use OrdH instead"+++-- * AstRaw, AstNoVectorize and AstNoSimplify definitions++-- | An AST variant that doesn't vectorize terms and also builds them+-- with ordinary, non-simplifying constructors. It's based on sharing+-- rather than lets and commonly used as the instance for primals+-- inside ADVal and, consequently, used for evaluating delta expressions.+type role AstRaw nominal nominal+newtype AstRaw s y =+  AstRaw {unAstRaw :: AstTensor AstMethodShare s y}+ deriving Show++-- | An AST variant for testing that doesn't vectorize terms, but still+-- builds them using simplifying smart constructors.+type role AstNoVectorize nominal nominal+newtype AstNoVectorize s y =+  AstNoVectorize {unAstNoVectorize :: AstTensor AstMethodLet s y}+ deriving Show++-- | An AST variant for testing that vectorizes terms, but builds them+-- with ordinary, non-simplifying constructors.+type role AstNoSimplify nominal nominal+newtype AstNoSimplify s y =+  AstNoSimplify {unAstNoSimplify :: AstTensor AstMethodLet s y}+ deriving Show+++-- * AstRaw, AstNoVectorize and AstNoSimplify type family instances++type instance PrimalOf (AstRaw s) = AstRaw (PrimalStepSpan s)+type instance DualOf (AstRaw s) = AstTensor AstMethodShare DualSpan+type instance PlainOf (AstRaw s) = AstRaw PlainSpan+type instance ShareOf (AstRaw s) = AstRaw s+type instance HFunOf (AstRaw s) = AstHFun s++type instance PrimalOf (AstNoVectorize s) = AstNoVectorize (PrimalStepSpan s)+type instance DualOf (AstNoVectorize s) = AstTensor AstMethodLet DualSpan+type instance PlainOf (AstNoVectorize s) = AstNoVectorize PlainSpan+type instance ShareOf (AstNoVectorize s) = AstRaw s+type instance HFunOf (AstNoVectorize s) = AstHFun s++type instance PrimalOf (AstNoSimplify s) = AstNoSimplify (PrimalStepSpan s)+type instance DualOf (AstNoSimplify s) = AstTensor AstMethodLet DualSpan+type instance PlainOf (AstNoSimplify s) = AstNoSimplify PlainSpan+type instance ShareOf (AstNoSimplify s) = AstRaw s+type instance HFunOf (AstNoSimplify s) = AstHFun s+++-- * Helper functions++sunReplicate :: Nested.Elt a+             => Nested.Shaped sh a -> Maybe a+{-# INLINE sunReplicate #-}+sunReplicate (Nested.Shaped arr)+  | all (all (== 0) . take (shxLength (Nested.mshape arr)))+        (Mixed.marrayStrides arr)+  , shxSize (Nested.mshape arr) /= 0 =+    Just $ Nested.mindex arr $ ixxZero' $ Nested.mshape arr+sunReplicate arr | ZSS <- Nested.sshape arr = Just $ Nested.sunScalar arr+sunReplicate _ = Nothing++sunReplicate1 :: Nested.Elt a+              => Nested.Shaped (n ': sh) a -> Maybe (Nested.Shaped sh a)+{-# INLINE sunReplicate1 #-}+sunReplicate1 a | (snat :$$ _) <- Nested.sshape a =+  sunReplicateN (snat :$$ ZSS) a++sunReplicateN :: Nested.Elt a+              => ShS shm -> Nested.Shaped (shm ++ shn) a+              -> Maybe (Nested.Shaped shn a)+{-# INLINE sunReplicateN #-}+sunReplicateN shm a@(Nested.Shaped arr)+  | all (all (== 0) . take (shsLength shm)) (Mixed.marrayStrides arr)+  , shsSize shm /= 0 =+    Just $ Nested.sindexPartial a $ ixsZero shm+sunReplicateN _ _ = Nothing++unReplC :: forall sh x s ms.+           AstTensor ms s (TKS2 sh x) -> Maybe (RepConcrete x)+unReplC (AstReplicateK _ a) = unAstK a+unReplC (AstReplicateS _ u) = unReplC u+unReplC (AstLet _ _ t) = unReplC t  -- we may be before inlining+unReplC (AstPrimalPart t) = unReplC t+unReplC (AstDualPart t) = unReplC t+unReplC (AstPlainPart t) = unReplC t+unReplC (AstFromPrimal t) = unReplC t+unReplC (AstFromDual t) = unReplC t+unReplC (AstFromPlain t) = unReplC t+unReplC (AstConcreteS a) = sunReplicate a+unReplC (AstConvert (ConvCmp ConvXS (Conv0X STKScalar)) (AstConcreteK a)) =+  Just a+unReplC _ = Nothing  -- e.g., a variable++-- No cases for, e.g., arithmetic, because it'd get simplified away beforehand.+unAstK :: forall r s ms.+          AstTensor ms s (TKScalar r) -> Maybe r+unAstK (AstLet _ _ t) = unAstK t  -- we may be before inlining+unAstK (AstPrimalPart t) = unAstK t+unAstK (AstDualPart t) = unAstK t+unAstK (AstPlainPart t) = unAstK t+unAstK (AstFromPrimal t) = unAstK t+unAstK (AstFromDual t) = unAstK t+unAstK (AstFromPlain t) = unAstK t+unAstK (AstConcreteK a) = Just a+unAstK (AstConvert (ConvCmp ConvX0 ConvSX) a) = unReplC a+unAstK _ = Nothing++unAstS :: forall sh x s ms.+          AstTensor ms s (TKS2 sh x)+       -> Maybe (Nested.Shaped sh (RepConcrete x))+unAstS (AstReplicateS shm u) | Dict <- eltDictRep (stkAstX u) =+  Nested.sreplicate shm <$> unAstS u+unAstS (AstLet _ _ t) = unAstS t+unAstS (AstPrimalPart t) = unAstS t+unAstS (AstDualPart t) = unAstS t+unAstS (AstPlainPart t) = unAstS t+unAstS (AstFromPrimal t) = unAstS t+unAstS (AstFromDual t) = unAstS t+unAstS (AstFromPlain t) = unAstS t+unAstS (AstConcreteS a) = Just a+unAstS (AstConvert (ConvCmp ConvXS (Conv0X STKScalar)) (AstConcreteK r)) =+  gcastWith (unsafeCoerceRefl :: sh :~: '[]) $+  Just $ Nested.sscalar r+unAstS _ = Nothing++-- | An approximation of equality. `False` doesn't imply terms+-- have different semantics, but `True` implies they have equal semantics.+eqY :: AstTensor ms s y -> AstTensor ms s y -> Bool+eqY t1 t2 = case eqZ t1 t2 of+  Just Refl -> True+  Nothing -> False++-- | An approximation of equality, where `Just Refl` means that+-- the shapes of the terms are equal and the values denoted by them are equal.+eqUnknownShapes :: AstTensor ms s y -> AstTensor ms s z -> Maybe (y :~: z)+eqUnknownShapes = eqZ++-- | An approximation of equality, where `Just Refl` means that+-- the shapes of the terms are equal and if their spans are equal+-- then the values denoted by them are equal.+eqZ :: AstTensor ms s1 y -> AstTensor ms s2 z -> Maybe (y :~: z)+eqZ (AstPair u1 v1) (AstPair u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstProject1 u1) (AstProject1 u2) | Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstProject2 u1) (AstProject2 u2) | Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstVar u1) (AstVar u2)+  | varNameToAstVarId u1 == varNameToAstVarId u2+  , Just Refl <- matchingFTK (varNameToFTK u1) (varNameToFTK u2) = Just Refl+eqZ v1 (AstLet _ _ v2) = eqZ v1 v2+eqZ (AstLet _ _  v1) v2 = eqZ v1 v2+-- We can remove these wrappers from only one side, because if the spans+-- are equal, the values will be, thanks to the careful processing below.+eqZ u1 (AstPrimalPart u2) = eqZ u1 u2+eqZ (AstPrimalPart u1) u2 = eqZ u1 u2+eqZ u1 (AstDualPart u2) = eqZ u1 u2+eqZ (AstDualPart u1) u2 = eqZ u1 u2+eqZ u1 (AstPlainPart u2) = eqZ u1 u2+eqZ (AstPlainPart u1) u2 = eqZ u1 u2+-- We can't remove the following wrappers from only one side,+-- because taking a primal part above may result in a dual-number with zero+-- dual part and if AstFromPrimal would also be applied to only one side,+-- the empty dual part would no longer be reflected in the type (the span).+eqZ (AstFromPrimal u1) (AstFromPrimal u2) = eqZ u1 u2+eqZ (AstFromDual u1) (AstFromDual u2) = eqZ u1 u2+eqZ (AstFromPlain u1) (AstFromPlain u2) = eqZ u1 u2+eqZ (AstPlusK u1 v1) (AstPlusK u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstPlusK u1 v1) (AstPlusK u2 v2)+  | Just Refl <- eqZ u1 v2+  , Just Refl <- eqZ v1 u2 = Just Refl+eqZ (AstTimesK u1 v1) (AstTimesK u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstTimesK u1 v1) (AstTimesK u2 v2)+  | Just Refl <- eqZ u1 v2+  , Just Refl <- eqZ v1 u2 = Just Refl+eqZ (AstN1K opCode1 u1) (AstN1K opCode2 u2) | opCode1 == opCode2 = eqZ u1 u2+eqZ (AstR1K opCode1 u1) (AstR1K opCode2 u2) | opCode1 == opCode2 = eqZ u1 u2+eqZ (AstR2K opCode1 u1 v1) (AstR2K opCode2 u2 v2)+  | opCode1 == opCode2+  , Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstI2K opCode1 u1 v1) (AstI2K opCode2 u2 v2)+  | opCode1 == opCode2+  , Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstConcreteK @r1 u1) (AstConcreteK @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , u1 == u2 = Just Refl+eqZ (AstFloorK @_ @r1 u1) (AstFloorK @_ @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstFromIntegralK @_ @r1 u1) (AstFromIntegralK @_ @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstCastK @_ @r1 u1) (AstCastK @_ @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstPlusS u1 v1) (AstPlusS u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstPlusS u1 v1) (AstPlusS u2 v2)+  | Just Refl <- eqZ u1 v2+  , Just Refl <- eqZ v1 u2 = Just Refl+eqZ (AstTimesS u1 v1) (AstTimesS u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstTimesS u1 v1) (AstTimesS u2 v2)+  | Just Refl <- eqZ u1 v2+  , Just Refl <- eqZ v1 u2 = Just Refl+eqZ (AstN1S opCode1 u1) (AstN1S opCode2 u2) | opCode1 == opCode2 = eqZ u1 u2+eqZ (AstR1S opCode1 u1) (AstR1S opCode2 u2) | opCode1 == opCode2 = eqZ u1 u2+eqZ (AstR2S opCode1 u1 v1) (AstR2S opCode2 u2 v2)+  | opCode1 == opCode2+  , Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstI2S opCode1 u1 v1) (AstI2S opCode2 u2 v2)+  | opCode1 == opCode2+  , Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstFloorS @_ @r1 u1) (AstFloorS @_ @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstFromIntegralS @_ @r1 u1) (AstFromIntegralS @_ @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstCastS @_ @r1 u1) (AstCastS @_ @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstSumK u1) (AstSumK u2)+  | Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstSumS shm1 u1) (AstSumS shm2 u2)+  | Just Refl <- testEquality shm1 shm2+  , Just Refl <- eqZ u1 u2 = Just unsafeCoerceRefl+eqZ (AstReplicateK shm1 u1) (AstReplicateK shm2 u2)+  | Just Refl <- testEquality shm1 shm2+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstReplicateS shm1 u1) (AstReplicateS shm2 u2)+  | Just Refl <- testEquality shm1 shm2+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstIotaS @_ @r1 k1) (AstIotaS @_ @r2 k2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , Just Refl <- testEquality k1 k2 = Just Refl+eqZ (AstReverseS u1) (AstReverseS u2) = eqZ u1 u2+eqZ (AstTransposeS perm1 u1) (AstTransposeS perm2 u2)+  | Just Refl <- testEquality perm1 perm2+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstReshapeS sh1 u1) (AstReshapeS sh2 u2)+  | Just Refl <- testEquality sh1 sh2+  , Just Refl <- eqZ u1 u2 = Just Refl+eqZ (AstConvert c1 u1) (AstConvert c2 u2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- testEquality c1 c2 = Just Refl+eqZ (AstBoolNotK u1) (AstBoolNotK u2) = eqZ u1 u2+eqZ (AstBoolNotS u1) (AstBoolNotS u2) = eqZ u1 u2+eqZ (AstBoolAndK u1 v1) (AstBoolAndK u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstBoolAndK u1 v1) (AstBoolAndK u2 v2)+  | Just Refl <- eqZ u1 v2+  , Just Refl <- eqZ v1 u2 = Just Refl+eqZ (AstBoolAndS u1 v1) (AstBoolAndS u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstBoolAndS u1 v1) (AstBoolAndS u2 v2)+  | Just Refl <- eqZ u1 v2+  , Just Refl <- eqZ v1 u2 = Just Refl+eqZ (AstLeqK u1 v1) (AstLeqK u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ (AstLeq u1 v1) (AstLeq u2 v2)+  | Just Refl <- eqZ u1 u2+  , Just Refl <- eqZ v1 v2 = Just Refl+eqZ _ _ = Nothing
src/HordeAd/Core/CarriersConcrete.hs view
@@ -7,21 +7,29 @@ module HordeAd.Core.CarriersConcrete   ( -- * RepConcrete and its operations     RepConcrete, tftkG, eltDictRep, showDictRep-    -- * Concrete and its operations-  , Concrete(..), rtoVector, stoVector, xtoVector+  , replTargetRep, defTargetRep+    -- * Concrete and its instances+  , Concrete(..), ConcreteFun(..), rtoVector, stoVector, xtoVector+  , DummyDualTarget(..)   ) where -import Prelude hiding (foldl')+import Prelude  import Control.DeepSeq (NFData (..))+import Data.Default+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality ((:~:) (Refl)) import Data.Vector.Storable qualified as VS  import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Lemmas import Data.Array.Nested.Mixed qualified as Mixed import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Ranked qualified as Ranked+import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped qualified as Shaped import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat') import Data.Array.Strided.Orthotope (liftVEltwise1)  import HordeAd.Core.TensorKind@@ -47,27 +55,27 @@   quotH a b = Nested.mquotArray a (mmakeNonZero b)   remH a b = Nested.mremArray a (mmakeNonZero b) -instance GoodScalar r+instance (NumScalar r, Ord (Nested.Ranked n r))          => Real (Nested.Ranked n r) where   toRational = error "toRational is not defined for tensors" -instance GoodScalar r+instance (NumScalar r, Ord (Nested.Shaped sh r))          => Real (Nested.Shaped sh r) where   toRational = error "toRational is not defined for tensors" -instance GoodScalar r+instance (NumScalar r, Ord (Nested.Mixed sh r))          => Real (Nested.Mixed sh r) where   toRational = error "toRational is not defined for tensors" -instance (GoodScalar r, Nested.FloatElt r)+instance (NumScalar r, Ord (Nested.Ranked n r), Nested.FloatElt r)          => RealFrac (Nested.Ranked n r) where   properFraction = error "properFraction is not defined for tensors" -instance (GoodScalar r, RealFrac r, Nested.FloatElt r)+instance (NumScalar r, Ord (Nested.Shaped sh r), Nested.FloatElt r)          => RealFrac (Nested.Shaped sh r) where   properFraction = error "properFraction is not defined for tensors" -instance (GoodScalar r, Nested.FloatElt r)+instance (NumScalar r, Ord (Nested.Mixed sh r), Nested.FloatElt r)          => RealFrac (Nested.Mixed sh r) where   properFraction = error "properFraction is not defined for tensors" @@ -83,7 +91,8 @@          => RealFloatH (Nested.Mixed sh r) where   atan2H = Nested.matan2Array -instance (GoodScalar r, Nested.PrimElt r, RealFloat r, Nested.FloatElt r)+instance ( NumScalar r, Ord (Nested.Ranked n r), Nested.PrimElt r+         , Nested.FloatElt r )          => RealFloat (Nested.Ranked n r) where   atan2 = Nested.ratan2Array   floatRadix = error "operation not defined for tensors"@@ -97,7 +106,8 @@   isNegativeZero = error "operation not defined for tensors"   isIEEE = error "operation not defined for tensors" -instance (GoodScalar r, Nested.PrimElt r, RealFloat r, Nested.FloatElt r)+instance ( NumScalar r, Ord (Nested.Shaped sh r), Nested.PrimElt r+         , Nested.FloatElt r )          => RealFloat (Nested.Shaped sh r) where   atan2 = Nested.satan2Array   floatRadix = error "operation not defined for tensors"@@ -111,7 +121,8 @@   isNegativeZero = error "operation not defined for tensors"   isIEEE = error "operation not defined for tensors" -instance (GoodScalar r, Nested.PrimElt r, RealFloat r, Nested.FloatElt r)+instance ( NumScalar r, Ord (Nested.Mixed sh r), Nested.PrimElt r+         , Nested.FloatElt r )          => RealFloat (Nested.Mixed sh r) where   atan2 = Nested.matan2Array   floatRadix = error "operation not defined for tensors"@@ -130,7 +141,7 @@              => Nested.Mixed sh r -> Nested.Mixed sh r mmakeNonZero =   Mixed.mliftNumElt1-    (`liftVEltwise1` (VS.map (\x -> if x == 0 then 1 else x)))+    (`liftVEltwise1` VS.map (\x -> if x == 0 then 1 else x))   -- * RepConcrete and its operations@@ -143,57 +154,77 @@   RepConcrete (TKX2 sh x) = Nested.Mixed sh (RepConcrete x)   RepConcrete (TKProduct x z) = (RepConcrete x, RepConcrete z) --- | Computing full shape tensor kinds for concrete arrays.+-- TODO: check if we can get errors due to the made-up shapes+-- and if the errors are early and clear. Maybe for nested fully shaped tensors+-- no errors are possible? Statically known dimensions are preserved fine+-- also in mixed arrays.+-- | Compute the full shape tensor kind for a concrete array.+-- If the array is empty and not shaped, shapes can be made up+-- (defaulting to zero dimensions). tftkG :: SingletonTK y -> RepConcrete y -> FullShapeTK y tftkG stk t =   let repackShapeTree :: SingletonTK y                       -> Mixed.ShapeTree (RepConcrete y)                       -> FullShapeTK y+      {-# NOINLINE repackShapeTree #-}       repackShapeTree stk0 tree = case stk0 of         STKScalar -> FTKScalar         STKR _ stk1 -> let (sh, rest) = tree-                       in FTKR sh $ repackShapeTree stk1 rest+                       in FTKR sh $ if shrSize sh == 0  -- rest crashes+                                    then zeroShapes stk1+                                    else repackShapeTree stk1 rest         STKS _ stk1 -> let (sh, rest) = tree-                       in FTKS sh $ repackShapeTree stk1 rest+                       in FTKS sh $ if shsSize sh == 0  -- rest crashes+                                    then zeroShapes stk1+                                    else repackShapeTree stk1 rest         STKX _ stk1 -> let (sh, rest) = tree-                       in FTKX sh $ repackShapeTree stk1 rest+                       in FTKX sh $ if shxSize sh == 0  -- rest crashes+                                    then zeroShapes stk1+                                    else repackShapeTree stk1 rest         STKProduct stk1 stk2 ->                        let (tree1, tree2) = tree                        in FTKProduct (repackShapeTree stk1 tree1)                                      (repackShapeTree stk2 tree2)-  in case stk of+      zeroShapes :: SingletonTK y -> FullShapeTK y+      {-# NOINLINE zeroShapes #-}+      zeroShapes stk0 = case stk0 of+        STKScalar -> FTKScalar+        STKR n stk1 -> FTKR (shrFromList n (replicate (fromSNat' n) 0))+                            (zeroShapes stk1)+        STKS sh stk1 -> FTKS sh $ zeroShapes stk1  -- not made up in this case+        STKX ssx stk1 -> FTKX (shxCompleteZeros ssx) $ zeroShapes stk1+                           -- statically known shapes not made up+        STKProduct stk1 stk2 -> FTKProduct (zeroShapes stk1) (zeroShapes stk2)+  in case stk of  -- this starts with non-recursive shorthands     STKScalar -> FTKScalar-    STKR _ stk1 | Dict <- eltDictRep stk1 ->-      FTKR (Nested.rshape t) $ repackShapeTree stk1-      $ snd $ Mixed.mshapeTree t-    STKS sh stk1 | Dict <- eltDictRep stk1 ->-      FTKS sh $ repackShapeTree stk1-      $ snd $ Mixed.mshapeTree t-    STKX _ stk1 | Dict <- eltDictRep stk1 ->-      FTKX (Nested.mshape t) $ repackShapeTree stk1-      $ snd $ Mixed.mshapeTree t-    STKProduct stk1 stk2 ->-      FTKProduct (tftkG stk1 (fst t))-                 (tftkG stk2 (snd t))+    STKR _ STKScalar -> FTKR (Nested.rshape t) FTKScalar+    STKS sh STKScalar -> FTKS sh FTKScalar+    STKX _ STKScalar -> FTKX (Nested.mshape t) FTKScalar+    _ | Dict <- eltDictRec stk -> repackShapeTree stk (Mixed.mshapeTree t)  eltDictRep :: SingletonTK y -> Dict Nested.KnownElt (RepConcrete y)-eltDictRep = \case-    STKScalar -> Dict-    STKR SNat x | Dict <- eltDictRep x -> Dict-    STKS sh x | Dict <- eltDictRep x -> withKnownShS sh Dict-    STKX sh x | Dict <- eltDictRep x -> withKnownShX sh Dict-    STKProduct stk1 stk2 | Dict <- eltDictRep stk1-                         , Dict <- eltDictRep stk2 -> Dict+eltDictRep STKScalar = Dict  -- the prevalent case+eltDictRep stk = eltDictRec stk +eltDictRec :: SingletonTK y -> Dict Nested.KnownElt (RepConcrete y)+{-# NOINLINE eltDictRec #-}+eltDictRec = \case+  STKScalar -> Dict+  STKR SNat x | Dict <- eltDictRec x -> Dict+  STKS sh x | Dict <- eltDictRec x -> withKnownShS sh Dict+  STKX sh x | Dict <- eltDictRec x -> withKnownShX sh Dict+  STKProduct stk1 stk2 | Dict <- eltDictRec stk1+                       , Dict <- eltDictRec stk2 -> Dict+ showDictRep :: SingletonTK y -> Dict Show (RepConcrete y) showDictRep = \case     STKScalar -> Dict     STKR _ x | Dict <- showDictRep x-             , Dict <- eltDictRep x -> Dict+             , Dict <- eltDictRec x -> Dict     STKS _ x | Dict <- showDictRep x-             , Dict <- eltDictRep x -> Dict+             , Dict <- eltDictRec x -> Dict     STKX _ x | Dict <- showDictRep x-             , Dict <- eltDictRep x -> Dict+             , Dict <- eltDictRec x -> Dict     STKProduct stk1 stk2 | Dict <- showDictRep stk1                          , Dict <- showDictRep stk2 -> Dict @@ -201,15 +232,56 @@ nfdataDictRep = \case     STKScalar -> Dict     STKR _ x | Dict <- nfdataDictRep x-             , Dict <- eltDictRep x -> Dict+             , Dict <- eltDictRec x -> Dict     STKS _ x | Dict <- nfdataDictRep x-             , Dict <- eltDictRep x -> Dict+             , Dict <- eltDictRec x -> Dict     STKX _ x | Dict <- nfdataDictRep x-             , Dict <- eltDictRep x -> Dict+             , Dict <- eltDictRec x -> Dict     STKProduct stk1 stk2 | Dict <- nfdataDictRep stk1                          , Dict <- nfdataDictRep stk2 -> Dict +replTargetRep :: TKAllNum y+              => (forall r. NumScalar r => r) -> FullShapeTK y -> RepConcrete y+replTargetRep r = \case+  FTKScalar @r | Dict0 <- numFromTKAllNum (Proxy @r) -> r+  FTKR sh (FTKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r) ->+    Nested.rreplicatePrim sh r+  FTKS @sh sh (FTKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r)+                             , Refl <- lemAppNil @sh ->+    Nested.sreplicatePrim sh r+  FTKX @sh sh (FTKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r)+                             , Refl <- lemAppNil @sh ->+    Nested.mreplicatePrim sh r+  FTKR sh x | Dict <- eltDictRec (ftkToSTK x) ->+    Nested.rreplicate sh $ Nested.rscalar $ replTargetRep r x+  FTKS @sh sh x | Dict <- eltDictRec (ftkToSTK x)+                , Refl <- lemAppNil @sh ->+    Nested.sreplicate sh $ Nested.sscalar $ replTargetRep r x+  FTKX @sh sh x | Dict <- eltDictRec (ftkToSTK x)+                , Refl <- lemAppNil @sh ->+    Nested.mreplicate sh $ Nested.mscalar $ replTargetRep r x+  FTKProduct ftk1 ftk2 -> (replTargetRep r ftk1, replTargetRep r ftk2) +defTargetRep :: FullShapeTK y -> RepConcrete y+defTargetRep = \case+  FTKScalar -> def+  FTKR sh FTKScalar ->+    Nested.rreplicatePrim sh def+  FTKS @sh sh FTKScalar | Refl <- lemAppNil @sh ->+    Nested.sreplicatePrim sh def+  FTKX @sh sh FTKScalar | Refl <- lemAppNil @sh ->+    Nested.mreplicatePrim sh def+  FTKR sh x | Dict <- eltDictRec (ftkToSTK x) ->+    Nested.rreplicate sh $ Nested.rscalar $ defTargetRep x+  FTKS @sh sh x | Dict <- eltDictRec (ftkToSTK x)+                , Refl <- lemAppNil @sh ->+    Nested.sreplicate sh $ Nested.sscalar $ defTargetRep x+  FTKX @sh sh x | Dict <- eltDictRec (ftkToSTK x)+                , Refl <- lemAppNil @sh ->+    Nested.mreplicate sh $ Nested.mscalar $ defTargetRep x+  FTKProduct ftk1 ftk2 -> (defTargetRep ftk1, defTargetRep ftk2)++ -- * Concrete and its instances  -- | A newtype wrapper over 'RepConcrete'.@@ -226,40 +298,39 @@ instance KnownSTK y => NFData (Concrete y) where   rnf (Concrete t) | Dict <- nfdataDictRep (knownSTK @y) = rnf t -type instance BoolOf Concrete = Bool--type instance HFunOf Concrete x z = RepConcrete x -> RepConcrete z+type role ConcreteFun nominal nominal+newtype ConcreteFun x z = ConcreteFun (RepConcrete x -> RepConcrete z)+type instance HFunOf Concrete = ConcreteFun  type instance PrimalOf Concrete = Concrete- type instance DualOf Concrete = DummyDualTarget-+type instance PlainOf Concrete = Concrete type instance ShareOf Concrete = Concrete  instance GoodScalar r => EqH Concrete (TKScalar r) where-  Concrete u ==. Concrete v = u == v-+  Concrete u ==. Concrete v = Concrete $ u == v instance GoodScalar r => OrdH Concrete (TKScalar r) where-  Concrete u <=. Concrete v = u <= v-+  Concrete u <=. Concrete v = Concrete $ u <= v instance GoodScalar r => EqH Concrete (TKR n r) where-  Concrete u ==. Concrete v = u == v-+  Concrete u ==. Concrete v =+    Concrete $ Ranked.rtoPrimitive u == Ranked.rtoPrimitive v instance GoodScalar r => OrdH Concrete (TKR n r) where-  Concrete u <=. Concrete v = u <= v-+  Concrete u <=. Concrete v =+    Concrete $ Ranked.rtoPrimitive u <= Ranked.rtoPrimitive v instance GoodScalar r => EqH Concrete (TKS sh r) where-  Concrete u ==. Concrete v = u == v-+  Concrete u ==. Concrete v =+    Concrete $ Shaped.stoPrimitive u == Shaped.stoPrimitive v instance GoodScalar r => OrdH Concrete (TKS sh r) where-  Concrete u <=. Concrete v = u <= v-+  Concrete u <=. Concrete v =+    Concrete $ Shaped.stoPrimitive u <= Shaped.stoPrimitive v instance GoodScalar r => EqH Concrete (TKX sh r) where-  Concrete u ==. Concrete v = u == v-+  Concrete u ==. Concrete v =+    Concrete $ Mixed.toPrimitive u == Mixed.toPrimitive v instance GoodScalar r => OrdH Concrete (TKX sh r) where-  Concrete u <=. Concrete v = u <= v+  Concrete u <=. Concrete v =+    Concrete $ Mixed.toPrimitive u <= Mixed.toPrimitive v +deriving instance Boolean (Concrete (TKScalar Bool)) deriving instance Eq (RepConcrete y) => Eq (Concrete y) deriving instance Ord (RepConcrete y) => Ord (Concrete y) deriving instance Num (RepConcrete y) => Num (Concrete y)@@ -279,3 +350,8 @@  xtoVector :: GoodScalar r => Concrete (TKX sh r) -> VS.Vector r xtoVector = Nested.mtoVector . unConcrete++type role DummyDualTarget nominal+type DummyDualTarget :: Target+newtype DummyDualTarget y = DummyDualTarget (FullShapeTK y)+  deriving (Eq, Show)
+ src/HordeAd/Core/Conversion.hs view
@@ -0,0 +1,408 @@+{-# LANGUAGE LambdaCase #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Conversion constructors.+module HordeAd.Core.Conversion+  ( TKConversion(..), convCmp, lemTKAllNumConvert, lemTKAllNumBackConvert+  , convertSTK, convertFTK, buildTKConversion, transposeTKConversion+  ) where++import Prelude hiding ((.))++import Control.Category+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (TestEquality (..), gcastWith, (:~:) (Refl))+import GHC.TypeLits (type (+))++import Data.Array.Nested (MapJust, Replicate, type (++))+import Data.Array.Nested.Convert+  (shrFromShXAnyShape, shsFromSSX, shsFromShX, shxFromShR, shxFromShS)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- | This is copied, with modifications, from ox-arrays.+--+-- This is a recipe for converting arrays, not always followed,+-- and a proof a conversion is possible, with some proof obligations+-- delayed to runtime (in ConvXS' and ConvXX', where not only the ranks+-- of the shapes need to agree, but also the dimensions of the input+-- array and of the output shape, which is not all captured in the type).+-- As in ox-arrays, conversions only change the meta-data, not the underlying+-- vector representation of the array.+type role TKConversion nominal nominal+data TKConversion (a :: TK) (b :: TK) where+  ConvId  :: TKConversion a a+  ConvCmp :: TKConversion b c -> TKConversion a b -> TKConversion a c++  ConvRX  :: TKConversion (TKR2 n a) (TKX2 (Replicate n Nothing) a)+  ConvSX  :: TKConversion (TKS2 sh a) (TKX2 (MapJust sh) a)++  ConvXR  :: SingletonTK a -> TKConversion (TKX2 sh a) (TKR2 (Rank sh) a)+  ConvXS  :: TKConversion (TKX2 (MapJust sh) a) (TKS2 sh a)+  ConvXS' :: Rank sh ~ Rank sh'+          => FullShapeTK (TKS2 sh' a)+          -> TKConversion (TKX2 sh a) (TKS2 sh' a)++  ConvXX' :: Rank sh ~ Rank sh'+          => FullShapeTK (TKX2 sh' a)+          -> TKConversion (TKX2 sh a) (TKX2 sh' a)++  ConvRR  :: TKConversion a b -> TKConversion (TKR2 n a) (TKR2 n b)+  ConvSS  :: TKConversion a b -> TKConversion (TKS2 sh a) (TKS2 sh b)+  ConvXX  :: TKConversion a b -> TKConversion (TKX2 sh a) (TKX2 sh b)+  ConvT2  :: TKConversion a a'+          -> TKConversion b b'+          -> TKConversion (TKProduct a b) (TKProduct a' b')++  Conv0X  :: SingletonTK a -> TKConversion a (TKX2 '[] a)+  ConvX0  :: TKConversion (TKX2 '[] a) a++  ConvNest   :: SingletonTK (TKX2 sh a)+             -> TKConversion (TKX2 (sh ++ sh') a) (TKX2 sh (TKX2 sh' a))+  ConvUnnest :: TKConversion (TKX2 sh (TKX2 sh' a)) (TKX2 (sh ++ sh') a)++  ConvZip   :: SingletonTK a -> SingletonTK b+            -> TKConversion (TKProduct (TKX2 sh a) (TKX2 sh b))+                            (TKX2 sh (TKProduct a b))+  ConvUnzip :: SingletonTK a -> SingletonTK b+            -> TKConversion (TKX2 sh (TKProduct a b))+                            (TKProduct (TKX2 sh a) (TKX2 sh b))++deriving instance Show (TKConversion a b)++instance Category TKConversion where+  id = ConvId+  (.) = convCmp++instance TestEquality (TKConversion a) where+  testEquality = \cases+    ConvId ConvId -> Just Refl+    (ConvCmp c1 d1) (ConvCmp c2 d2)+      | Just Refl <- testEquality d1 d2+      , Just Refl <- testEquality c1 c2 -> Just Refl+    ConvRX ConvRX -> Just Refl+    ConvSX ConvSX -> Just Refl+    (ConvXR stk1) (ConvXR stk2)+      | Just Refl <- sameSTK stk1 stk2 -> Just Refl+    ConvXS ConvXS -> Just unsafeCoerceRefl+                       -- equality with MapJust implies equality without+    (ConvXS' ftk1) (ConvXS' ftk2)+      | Just Refl <- matchingFTK ftk1 ftk2 -> Just Refl+    (ConvXX' ftk1) (ConvXX' ftk2)+      | Just Refl <- matchingFTK ftk1 ftk2 -> Just Refl+    (ConvRR c1) (ConvRR c2)+      | Just Refl <- testEquality c1 c2 -> Just Refl+    (ConvSS c1) (ConvSS c2)+      | Just Refl <- testEquality c1 c2 -> Just Refl+    (ConvXX c1) (ConvXX c2)+      | Just Refl <- testEquality c1 c2 -> Just Refl+    (ConvT2 c1 d1) (ConvT2 c2 d2)+      | Just Refl <- testEquality c1 c2+      , Just Refl <- testEquality d1 d2 -> Just Refl+    (Conv0X stk1) (Conv0X stk2)+      | Just Refl <- sameSTK stk1 stk2 -> Just Refl+    ConvX0 ConvX0 -> Just Refl+    (ConvNest stk1) (ConvNest stk2)+      | Just Refl <- sameSTK stk1 stk2 -> Just unsafeCoerceRefl+                                            -- (sh ++) is morally injective+    ConvUnnest ConvUnnest -> Just Refl+    (ConvZip stk1 stk'1) (ConvZip stk2 stk'2)+      | Just Refl <- sameSTK stk1 stk2+      , Just Refl <- sameSTK stk'1 stk'2 -> Just Refl+    (ConvUnzip stk1 stk'1) (ConvUnzip stk2 stk'2)+      | Just Refl <- sameSTK stk1 stk2+      , Just Refl <- sameSTK stk'1 stk'2 -> Just Refl+    _ _ -> Nothing++-- TODO: expand+convCmp :: TKConversion b c -> TKConversion a b -> TKConversion a c+convCmp a b = case (a, b) of+  (_, ConvId) -> a+  (_, ConvCmp ConvId c) -> convCmp a c+  (ConvId, _) -> b+  (ConvCmp a1 a2, _) -> a1 . (a2 . b)+  (ConvSX, ConvXS) -> ConvId+  (ConvSX, ConvCmp ConvXS c) -> c+  (ConvXR{}, ConvRX @n) | Refl <- lemRankReplicate (Proxy @n) -> ConvId+  (ConvXR{}, ConvCmp (ConvRX @n) c) | Refl <- lemRankReplicate (Proxy @n) -> c+  (ConvXR stk, ConvXX'{}) -> ConvXR stk+  (ConvXR stk, ConvCmp ConvXX'{} c) -> convCmp (ConvXR stk) c+  (ConvXS @sh, ConvSX @sh') ->+    gcastWith (unsafeCoerceRefl :: sh :~: sh')+    ConvId+  (ConvXS @sh, ConvCmp (ConvSX @sh') c) ->+    gcastWith (unsafeCoerceRefl :: sh :~: sh')+    c+  (ConvXS, ConvXX' (FTKX sh x)) | Refl <- lemRankMapJust (shsFromShX sh) ->+    ConvXS' (FTKS (shsFromShX sh) x)+  (ConvXS, ConvCmp (ConvXX' (FTKX sh x)) c)+    | Refl <- lemRankMapJust (shsFromShX sh) ->+      convCmp (ConvXS' (FTKS (shsFromShX sh) x)) c+  (ConvXS' @_ @sh' _, ConvSX @sh) ->+    gcastWith (unsafeCoerceRefl :: sh :~: sh')+    ConvId+  (ConvXS' @_ @sh' _, ConvCmp (ConvSX @sh) c) ->+    gcastWith (unsafeCoerceRefl :: sh :~: sh')+    c+  (ConvXS' ftk, ConvXX'{}) -> ConvXS' ftk+  (ConvXS' ftk, ConvCmp ConvXX'{} c) -> convCmp (ConvXS' ftk) c+  (ConvXS' (FTKS ZSS _), Conv0X stk) -> convCmp ConvXS (Conv0X stk)+  (ConvXS' (FTKS ZSS _), ConvCmp (Conv0X stk) c) ->+    convCmp (convCmp ConvXS (Conv0X stk)) c+  (ConvXX' ftk, ConvXX'{}) -> ConvXX' ftk+  (ConvXX' ftk, ConvCmp ConvXX'{} c) -> convCmp (ConvXX' ftk) c+  (ConvXX' (FTKX ZSX _), Conv0X stk) -> Conv0X stk+  (ConvXX' (FTKX ZSX _), ConvCmp (Conv0X stk) c) -> convCmp (Conv0X stk) c+  (ConvRR a', ConvRR b') -> ConvRR (convCmp a' b')+  (ConvRR a', ConvCmp (ConvRR b') c) -> convCmp (ConvRR (convCmp a' b')) c+  (ConvSS a', ConvSS b') -> ConvSS (convCmp a' b')+  (ConvSS a', ConvCmp (ConvSS b') c) -> convCmp (ConvSS (convCmp a' b')) c+  (ConvXX a', ConvXX b') -> convXX (convCmp a' b')+  (ConvXX a', ConvCmp (ConvXX b') c) -> convCmp (convXX (convCmp a' b')) c+  (Conv0X{}, ConvX0) -> ConvId+  (Conv0X{}, ConvCmp ConvX0 c) -> c+  (ConvX0, ConvXX' @sh (FTKX ZSX _)) ->+    gcastWith (unsafeCoerceRefl :: sh :~: '[])+    ConvX0+  (ConvX0, ConvCmp (ConvXX' @sh (FTKX ZSX _)) c) ->+    gcastWith (unsafeCoerceRefl :: sh :~: '[]) $+    convCmp ConvX0 c+  (ConvX0, Conv0X{}) -> ConvId+  (ConvX0, ConvCmp Conv0X{} c) -> c+  (ConvT2 a1 a2, ConvT2 b1 b2) -> convT2 (convCmp a1 b1) (convCmp a2 b2)+  (ConvT2 a1 a2, ConvCmp (ConvT2 b1 b2) c) ->+    convCmp (convT2 (convCmp a1 b1) (convCmp a2 b2)) c+  {- Not enough singletons to decide the equality here:+  (ConvNest @sh @_ @sh' _, ConvUnnest @sh2 @sh2')+    | Just Refl <- testEquality sh :~: sh2) ->+      gcastWith (unsafeCoerceRefl :: sh' :~: sh2') $+      ConvId+  (ConvNest @sh @_ @sh' _, ConvCmp (ConvUnnest @sh2 @sh2') c)+    | Just Refl <- testEquality sh :~: sh2) ->+      gcastWith (unsafeCoerceRefl :: sh' :~: sh2') $+      c -}+-- not enough type info in the AST:+-- (ConvNest (STKX sh x), ConvXX d) ->+--   convCmp (ConvXX (ConvXX d)) (ConvNest (STKX sh (convertSTKBack d x)))+  (ConvUnnest, ConvNest{}) -> ConvId+  (ConvUnnest, ConvCmp ConvNest{} c) -> c+  (ConvXX d, ConvUnnest) ->+    convCmp ConvUnnest (convXX (convXX d))+  (ConvXX d, ConvCmp ConvUnnest c) ->+    convCmp ConvUnnest (convCmp (convXX (convXX d)) c)+  (ConvZip{}, ConvUnzip{}) -> ConvId+  (ConvZip{}, ConvCmp (ConvUnzip{}) c) -> c+  (ConvXX' (FTKX sh (FTKProduct c1 c2)), ConvZip stk1 stk2) ->+    convCmp (ConvZip stk1 stk2)+            (convT2 (ConvXX' (FTKX sh c1)) (ConvXX' (FTKX sh c2)))+  (ConvXX' (FTKX sh (FTKProduct c1 c2)), ConvCmp (ConvZip stk1 stk2) c) ->+    convCmp (ConvZip stk1 stk2)+            (convCmp (convT2 (ConvXX' (FTKX sh c1)) (ConvXX' (FTKX sh c2))) c)+  (ConvUnzip{}, ConvZip{}) -> ConvId+  (ConvUnzip{}, ConvCmp (ConvZip{}) c) -> c+  (ConvUnzip stk1 stk2, ConvXX' (FTKX sh (FTKProduct c1 c2))) ->+    convCmp (convT2 (ConvXX' (FTKX sh c1)) (ConvXX' (FTKX sh c2)))+            (ConvUnzip stk1 stk2)+  (ConvUnzip stk1 stk2, ConvCmp (ConvXX' (FTKX sh (FTKProduct c1 c2))) c) ->+    convCmp (convT2 (ConvXX' (FTKX sh c1)) (ConvXX' (FTKX sh c2)))+            (convCmp (ConvUnzip stk1 stk2) c)+  _ -> ConvCmp a b++convT2  :: TKConversion a a'+        -> TKConversion b b'+        -> TKConversion (TKProduct a b) (TKProduct a' b')+convT2 ConvId ConvId = ConvId+convT2 a b = ConvT2 a b++convXX  :: TKConversion a b -> TKConversion (TKX2 sh a) (TKX2 sh b)+convXX ConvId = ConvId+convXX a = ConvXX a++lemTKAllNumConvert :: TKAllNum a+                   => TKConversion a b -> Dict0 (TKAllNum b)+lemTKAllNumConvert = \case+  ConvId -> Dict0+  ConvCmp c1 c2 | Dict0 <- lemTKAllNumConvert c2 ->+    lemTKAllNumConvert c1+  ConvRX -> Dict0+  ConvSX -> Dict0+  ConvXR{}  -> Dict0+  ConvXS -> Dict0+  ConvXS'{} -> Dict0+  ConvXX'{} -> Dict0+  ConvRR c | Dict0 <- lemTKAllNumConvert c -> Dict0+  ConvSS c | Dict0 <- lemTKAllNumConvert c -> Dict0+  ConvXX c | Dict0 <- lemTKAllNumConvert c -> Dict0+  ConvT2 c1 c2 | Dict0 <- lemTKAllNumConvert c1+               , Dict0 <- lemTKAllNumConvert c2 -> Dict0+  Conv0X{} -> Dict0+  ConvX0 -> Dict0+  ConvNest{} -> Dict0+  ConvUnnest -> Dict0+  ConvZip{} -> Dict0+  ConvUnzip{} -> Dict0++lemTKAllNumBackConvert :: TKAllNum b+                       => TKConversion a b -> Dict0 (TKAllNum a)+lemTKAllNumBackConvert = \case+  ConvId -> Dict0+  ConvCmp c1 c2 | Dict0 <- lemTKAllNumBackConvert c1 ->+    lemTKAllNumBackConvert c2+  ConvRX -> Dict0+  ConvSX -> Dict0+  ConvXR{}  -> Dict0+  ConvXS -> Dict0+  ConvXS'{} -> Dict0+  ConvXX'{} -> Dict0+  ConvRR c | Dict0 <- lemTKAllNumBackConvert c -> Dict0+  ConvSS c | Dict0 <- lemTKAllNumBackConvert c -> Dict0+  ConvXX c | Dict0 <- lemTKAllNumBackConvert c -> Dict0+  ConvT2 c1 c2 | Dict0 <- lemTKAllNumBackConvert c1+               , Dict0 <- lemTKAllNumBackConvert c2 -> Dict0+  Conv0X{} -> Dict0+  ConvX0 -> Dict0+  ConvNest{} -> Dict0+  ConvUnnest -> Dict0+  ConvZip{} -> Dict0+  ConvUnzip{} -> Dict0++convertSTK :: TKConversion a b -> SingletonTK a -> SingletonTK b+convertSTK = \cases+  ConvId astk -> astk+  (ConvCmp c1 c2) astk -> convertSTK c1 (convertSTK c2 astk)+  ConvRX (STKR n a) -> STKX (ssxReplicate n) a+  ConvSX (STKS sh a) -> STKX (ssxFromShX $ shxFromShS sh) a+  (ConvXR _stk) (STKX ssx a) -> STKR (ssxRank ssx) a+  ConvXS (STKX ssx a) -> STKS (shsFromSSX ssx) a+  (ConvXS' (FTKS sh _x)) (STKX _ssx2 a) -> STKS sh a+  (ConvXX' (FTKX shx _x)) (STKX _ssx2 a) -> STKX (ssxFromShX shx) a+  (ConvRR c) (STKR n a) -> STKR n (convertSTK c a)+  (ConvSS c) (STKS sh a) -> STKS sh (convertSTK c a)+  (ConvXX c) (STKX ssx a) -> STKX ssx (convertSTK c a)+  (ConvT2 c1 c2) (STKProduct stk1 stk2) ->+    STKProduct (convertSTK c1 stk1) (convertSTK c2 stk2)+  (Conv0X _stk) stk -> STKX ZKX stk+  ConvX0 (STKX ZKX stk) -> stk+  (ConvNest (STKX ssx x)) (STKX shsh' _x) ->+    STKX ssx (STKX (ssxDropSSX ssx shsh') x)+  ConvUnnest (STKX sh (STKX sh' x)) -> STKX (sh `ssxAppend` sh') x+  (ConvZip _ _) (STKProduct (STKX sh a1) (STKX _sh a2)) ->+    STKX sh (STKProduct a1 a2)+  (ConvUnzip _ _) (STKX sh (STKProduct a1 a2)) ->+    STKProduct (STKX sh a1) (STKX sh a2)++convertFTK :: TKConversion a b -> FullShapeTK a -> FullShapeTK b+convertFTK = \cases+  ConvId aftk -> aftk+  (ConvCmp c1 c2) aftk -> convertFTK c1 (convertFTK c2 aftk)+  ConvRX (FTKR shr a) -> FTKX (shxFromShR shr) a+  ConvSX (FTKS sh a) -> FTKX (shxFromShS sh) a+  (ConvXR _stk) (FTKX shx a) -> FTKR (shrFromShXAnyShape shx) a+  ConvXS (FTKX shx a) -> FTKS (shsFromShX shx) a+  (ConvXS' ftk) _ -> ftk+  (ConvXX' ftk) _ -> ftk+  (ConvRR c) (FTKR shr a) -> FTKR shr (convertFTK c a)+  (ConvSS c) (FTKS sh a) -> FTKS sh (convertFTK c a)+  (ConvXX c) (FTKX shx a) -> FTKX shx (convertFTK c a)+  (ConvT2 c1 c2) (FTKProduct ftk1 ftk2) ->+    FTKProduct (convertFTK c1 ftk1) (convertFTK c2 ftk2)+  (Conv0X _stk) ftk -> FTKX ZSX ftk+  ConvX0 (FTKX ZSX ftk) -> ftk+  (ConvNest @_ @_ @sh' (STKX ssx _x)) (FTKX shsh' x) ->+    let (t, d) = shxSplitApp (Proxy @sh') ssx shsh'+    in FTKX t (FTKX d x)+  ConvUnnest (FTKX sh (FTKX sh' x)) -> FTKX (sh `shxAppend` sh') x+  (ConvZip _ _) (FTKProduct (FTKX sh a1) (FTKX _sh a2)) ->+    FTKX sh (FTKProduct a1 a2)+  (ConvUnzip _ _) (FTKX sh (FTKProduct a1 a2)) ->+    FTKProduct (FTKX sh a1) (FTKX sh a2)++buildTKConversion :: SNat k -> FullShapeTK a+                  -> TKConversion a b+                  -> TKConversion (BuildTensorKind k a) (BuildTensorKind k b)+buildTKConversion k aftk c0 = case c0 of+  ConvId -> ConvId+  ConvCmp c1 c2 -> convCmp (buildTKConversion k (convertFTK c2 aftk) c1)+                           (buildTKConversion k aftk c2)+  ConvRX | FTKR @n shr xstk <- aftk+         , Refl <- lemRankReplicate (Proxy @n)+         , Refl <- lemRankReplicate (Proxy @(1 + n)) ->+    convCmp (ConvXX' (FTKX (SKnown k :$% shxFromShR shr) xstk)) ConvRX+  ConvSX -> ConvSX+  ConvXR stk -> ConvXR stk+  ConvXS -> ConvXS+  ConvXS' ftk -> ConvXS' (buildFTK k ftk)+  ConvXX' ftk -> ConvXX' (buildFTK k ftk)+  ConvRR c -> ConvRR c+  ConvSS c -> ConvSS c+  ConvXX c -> ConvXX c+  ConvT2 c1 c2 | FTKProduct ftk1 ftk2 <- aftk ->+    ConvT2 (buildTKConversion k ftk1 c1) (buildTKConversion k ftk2 c2)+  Conv0X _astk -> case aftk of+    FTKScalar -> ConvSX+    FTKR @n shr x | Refl <- lemRankReplicate (Proxy @n)+                  , Refl <- lemRankReplicate (Proxy @(1 + n)) ->+      convCmp (ConvXX (ConvXR (ftkToSTK x)))+              (convCmp (ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x)))+                       (convCmp+                          (ConvXX' (FTKX (SKnown k :$% shxFromShR shr) x))+                          ConvRX))+    FTKS _sh x ->+      convCmp (ConvXX ConvXS)+              (convCmp (ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x)))+                       ConvSX)+    FTKX _ssx x -> ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x))+    FTKProduct aftk1 aftk2 ->+      buildTKConversion+        k aftk (convCmp (ConvZip (ftkToSTK aftk1) (ftkToSTK aftk2))+                        (ConvT2 (Conv0X (ftkToSTK aftk1))+                                (Conv0X (ftkToSTK aftk2))))+  ConvX0 -> case aftk of+    FTKX ZSX FTKScalar -> ConvXS+    FTKX ZSX (FTKR @n _n x) | Refl <- lemRankReplicate (Proxy @n) ->+      convCmp (ConvXR (ftkToSTK x))+              (convCmp ConvUnnest (ConvXX ConvRX))+    FTKX ZSX FTKS{} ->+      convCmp ConvXS+              (convCmp ConvUnnest (ConvXX ConvSX))+    FTKX ZSX FTKX{} -> ConvUnnest+    FTKX ZSX (FTKProduct aftk1 aftk2) ->+      buildTKConversion+        k aftk (convCmp (ConvT2 ConvX0 ConvX0)+                        (ConvUnzip (ftkToSTK aftk1) (ftkToSTK aftk2)))+  ConvNest (STKX sh x) -> ConvNest (STKX (SKnown k :!% sh) x)+  ConvUnnest -> ConvUnnest+  ConvZip astk1 astk2 -> ConvZip astk1 astk2+  ConvUnzip astk1 astk2 -> ConvUnzip astk1 astk2++transposeTKConversion :: FullShapeTK a -> TKConversion a b -> TKConversion b a+transposeTKConversion aftk c0 = case c0 of+  ConvId -> ConvId+  ConvCmp c1 c2 -> transposeTKConversion aftk c2+                   `convCmp` transposeTKConversion (convertFTK c2 aftk) c1+  ConvRX | FTKR @n _ x <- aftk+         , Refl <- lemRankReplicate (Proxy @n) ->+    ConvXR (ftkToSTK x)+  ConvSX -> ConvXS+  ConvXR @_ @sh _stk | Refl <- lemRankReplicate (Proxy @(Rank sh)) ->+    convCmp (ConvXX' aftk) ConvRX+  ConvXS -> ConvSX+  ConvXS' (FTKS sh _) | Refl <- lemRankMapJust sh ->+    convCmp (ConvXX' aftk) ConvSX+  ConvXX' _ftk -> ConvXX' aftk+  ConvRR c | FTKR _ x <- aftk -> ConvRR (transposeTKConversion x c)+  ConvSS c | FTKS _ x <- aftk -> ConvSS (transposeTKConversion x c)+  ConvXX c | FTKX _ x <- aftk -> ConvXX (transposeTKConversion x c)+  ConvT2 c1 c2 | FTKProduct x1 x2 <- aftk ->+    ConvT2 (transposeTKConversion x1 c1) (transposeTKConversion x2 c2)+  Conv0X _stk -> ConvX0+  ConvX0 | FTKX ZSX x <- aftk -> Conv0X (ftkToSTK x)+  ConvNest _stk -> ConvUnnest+  ConvUnnest | (FTKX shx (FTKX _ x)) <- aftk ->+    ConvNest (STKX (ssxFromShX shx) (ftkToSTK x))+  ConvZip stk1 stk2 -> ConvUnzip stk1 stk2+  ConvUnzip stk1 stk2 -> ConvZip stk1 stk2
src/HordeAd/Core/ConvertTensor.hs view
@@ -3,6 +3,7 @@ -- the same data but varying amounts of shape information. module HordeAd.Core.ConvertTensor   ( ConvertTensor(..)+  , nestTarget, nestTargetK, unNestTarget, unNestTargetK   ) where  import Prelude@@ -17,6 +18,7 @@ import Data.Array.Nested.Shaped.Shape import Data.Array.Nested.Types (unsafeCoerceRefl) +import HordeAd.Core.Conversion import HordeAd.Core.TensorKind import HordeAd.Core.Types @@ -35,7 +37,7 @@   kfromR = kfromS . sfromR   -- | The conversion from an empty shape shaped tensor to a scalar.   kfromS :: GoodScalar r => target (TKS '[] r) -> target (TKScalar r)-  kfromS = let c = ConvCmp ConvX0 ConvSX+  kfromS = let c = convCmp ConvX0 ConvSX            in tconvert c (STKS ZSS STKScalar)   kfromX :: GoodScalar r => target (TKX '[] r) -> target (TKScalar r)   kfromX = kfromS . sfromX@@ -46,12 +48,12 @@   rfromS :: forall sh x. (KnownShS sh, KnownSTK x)          => target (TKS2 sh x) -> target (TKR2 (Rank sh) x)   rfromS | Refl <- lemRankMapJust (knownShS @sh) =-    let c = ConvCmp (ConvXR knownSTK) ConvSX+    let c = convCmp (ConvXR knownSTK) ConvSX     in tconvert c (STKS knownShS knownSTK)   rfromX :: forall sh x. KnownSTK x          => target (TKX2 sh x) -> target (TKR2 (Rank sh) x)   sfromK :: GoodScalar r => target (TKScalar r) -> target (TKS '[] r)-  sfromK = let c = ConvCmp ConvXS (Conv0X STKScalar)+  sfromK = let c = convCmp ConvXS (Conv0X STKScalar)            in tconvert c STKScalar   -- | The conversion from a ranked tensor to the corresponding shaped tensor   -- of the same rank.@@ -232,3 +234,31 @@   -- | A clone of tunpair, if @ShareTensor@ is available, or an implementation   -- that duplicates the argument, otherwise.   tunpairConv :: target (TKProduct x z) -> (target x, target z)++-- | Convert a tensor into a trivial array with the tensor as the only element.+nestTarget :: forall y target. ConvertTensor target+           => SingletonTK y -> target y -> target (TKS2 '[] y)+nestTarget stk = tconvert (convCmp ConvXS (Conv0X stk)) stk++--- | Convert similarly as in @nestTarget@.+nestTargetK :: forall k y target. ConvertTensor target+            => SNat k -> FullShapeTK y+            -> target (BuildTensorKind k y)+            -> target (BuildTensorKind k (TKS2 '[] y))+nestTargetK k ftk =+  tconvert (buildTKConversion k ftk $ convCmp ConvXS (Conv0X (ftkToSTK ftk)))+           (buildSTK k $ ftkToSTK ftk)++-- | Convert a tensor from a trivial array with the tensor as the only element.+unNestTarget :: forall y target. ConvertTensor target+             => SingletonTK y -> target (TKS2 '[] y) -> target y+unNestTarget stk = tconvert (convCmp ConvX0 ConvSX) (STKS ZSS stk)++-- | Convert similarly as in @unNestTarget@.+unNestTargetK :: forall k y target. ConvertTensor target+              => SNat k -> FullShapeTK y+              -> target (BuildTensorKind k (TKS2 '[] y))+              -> target (BuildTensorKind k y)+unNestTargetK k@SNat ftk =+  tconvert (buildTKConversion k (FTKS ZSS ftk) $ convCmp ConvX0 ConvSX)+           (STKS (k :$$ ZSS) (ftkToSTK ftk))
src/HordeAd/Core/Delta.hs view
@@ -15,7 +15,7 @@ -- -- The \'sparsity\' is less obvious when a delta expression -- contains big concrete tensors, e.g., via the `DeltaScale` constructor.--- However, via 'DeltaReplicate' and other constructors, the tensors+-- However, via 'DeltaReplicateR' and other constructors, the tensors -- can be enlarged much beyond what's embedded in the delta term. -- Also, if the expression refers to unknown inputs ('DeltaInput') -- it may denote, after evaluation, a still larger tensor.@@ -42,7 +42,6 @@ import Control.Exception.Assert.Sugar import Data.Dependent.EnumMap.Strict qualified as DMap import Data.Kind (Type)-import Data.Some import Data.Type.Equality (TestEquality (..), gcastWith, testEquality, (:~:)) import Data.Vector.Generic qualified as V import Data.Vector.Strict qualified as Data.Vector@@ -50,13 +49,13 @@ import Text.Show.Functions ()  import Data.Array.Nested (type (++))-import Data.Array.Nested qualified as Nested import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Permutation qualified as Permutation import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape import Data.Array.Nested.Types (snatPlus, unsafeCoerceRefl) +import HordeAd.Core.Conversion import HordeAd.Core.Ops import HordeAd.Core.TensorKind import HordeAd.Core.Types@@ -75,9 +74,9 @@     showsPrec d n  -- less verbose, more readable  instance DMap.Enum1 (NodeId target) where-  type Enum1Info (NodeId target) = Some FullShapeTK-  fromEnum1 (NodeId ftk n) = (n, Some ftk)-  toEnum1 n (Some ftk) = Some $ NodeId ftk n+  type Enum1Info (NodeId target) = FullShapeTK+  fromEnum1 (NodeId ftk n) = (n, ftk)+  toEnum1 n ftk = NodeId ftk n  instance TestEquality (NodeId target) where   testEquality (NodeId ftk1 _) (NodeId ftk2 _) = matchingFTK ftk1 ftk2@@ -103,9 +102,9 @@       . shows n  instance DMap.Enum1 (InputId target) where-  type Enum1Info (InputId target) = Some FullShapeTK-  fromEnum1 (InputId ftk n) = (n, Some ftk)-  toEnum1 n (Some ftk) = Some $ InputId ftk n+  type Enum1Info (InputId target) = FullShapeTK+  fromEnum1 (InputId ftk n) = (n, ftk)+  toEnum1 n ftk = InputId ftk n  instance TestEquality (InputId target) where   testEquality (InputId ftk1 _) (InputId ftk2 _) = matchingFTK ftk1 ftk2@@ -165,36 +164,6 @@                    Delta target (TKProduct y z) -> Delta target y   DeltaProject2 :: forall y z target.                    Delta target (TKProduct y z) -> Delta target z-  DeltaFromVector :: forall y k target.-                     SNat k -> SingletonTK y-                  -> Data.Vector.Vector (Delta target y)-                  -> Delta target (BuildTensorKind k y)-  DeltaSum :: forall y k target.-              SNat k -> SingletonTK y-           -> Delta target (BuildTensorKind k y)-           -> Delta target y-  DeltaReplicate :: forall y k target.-                    SNat k -> SingletonTK y-                 -> Delta target y-                 -> Delta target (BuildTensorKind k y)-  DeltaMapAccumR-    :: forall target k accy by ey.-       ( Show (target (BuildTensorKind k accy))-       , Show (target (BuildTensorKind k ey)) )-    => SNat k-    -> FullShapeTK by-    -> FullShapeTK ey-    -> target (BuildTensorKind k accy)-    -> target (BuildTensorKind k ey)-    -> HFun (TKProduct (ADTensorKind (TKProduct accy ey))-                       (TKProduct accy ey))-            (ADTensorKind (TKProduct accy by))-    -> HFun (TKProduct (ADTensorKind (TKProduct accy by))-                       (TKProduct accy ey))-            (ADTensorKind (TKProduct accy ey))-    -> Delta target accy-    -> Delta target (BuildTensorKind k ey)-    -> Delta target (TKProduct accy (BuildTensorKind k by))   DeltaMapAccumL     :: forall target k accy by ey.        ( Show (target (BuildTensorKind k accy))@@ -222,28 +191,52 @@            => Delta target y -> Delta target y -> Delta target y    -- Scalar arithmetic-  DeltaCastK :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+  DeltaCastK :: ( NumScalar r1, Differentiable r1+                , NumScalar r2, Differentiable r2 )              => Delta target (TKScalar r1) -> Delta target (TKScalar r2)    -- Ranked tensor operations-  DeltaCastR :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+  DeltaCastR :: ( NumScalar r1, Differentiable r1+                , NumScalar r2, Differentiable r2 )              => Delta target (TKR n r1) -> Delta target (TKR n r2)-  DeltaSum0R :: Delta target (TKR2 n r) -> Delta target (TKR2 0 r)-  DeltaDot0R :: (GoodScalar r, Show (target (TKR n r)))+  DeltaDot0R :: (NumScalar r, Show (target (TKR n r)))              => target (TKR n r) -> Delta target (TKR n r)-             -> Delta target (TKR 0 r)+             -> Delta target (TKScalar r)+  DeltaIndex0R :: forall m r target. GoodScalar r+               => Delta target (TKR m r) -> IxROf target m+               -> Delta target (TKScalar r)   DeltaIndexR :: forall m n r target.                  SNat n               -> Delta target (TKR2 (m + n) r) -> IxROf target m               -> Delta target (TKR2 n r)+  DeltaFromVectorLinearR :: forall m r target. GoodScalar r+                         => IShR m+                         -> Data.Vector.Vector (Delta target (TKScalar r))+                         -> Delta target (TKR m r)+  DeltaFromVectorR :: forall m n x target.+                      IShR m+                   -> Data.Vector.Vector (Delta target (TKR2 n x))+                   -> Delta target (TKR2 (m + n) x)+  DeltaSum0R :: forall m r target. NumScalar r+             => Delta target (TKR m r)+             -> Delta target (TKScalar r)+  DeltaSumR :: forall m n x target.+               SNat m -> Delta target (TKR2 (m + n) x)+            -> Delta target (TKR2 n x)   DeltaScatterR :: forall m n p r target.                    SNat m -> SNat n -> SNat p-                -> IShR (p + n) -> Delta target (TKR2 (m + n) r)+                -> IShR p -> Delta target (TKR2 (m + n) r)                 -> (IxROf target m -> IxROf target p)                 -> Delta target (TKR2 (p + n) r)+  DeltaReplicate0NR :: forall m r target. GoodScalar r+                    => IShR m -> Delta target (TKScalar r)+                    -> Delta target (TKR m r)+  DeltaReplicateR :: forall m n x target.+                     IShR m -> Delta target (TKR2 n x)+                  -> Delta target (TKR2 (m + n) x)   DeltaGatherR :: forall m n p r target.                   SNat m -> SNat n -> SNat p-               -> IShR (m + n) -> Delta target (TKR2 (p + n) r)+               -> IShR m -> Delta target (TKR2 (p + n) r)                -> (IxROf target m -> IxROf target p)                -> Delta target (TKR2 (m + n) r)   DeltaAppendR :: Delta target (TKR2 (1 + n) r)@@ -259,21 +252,44 @@                 -> Delta target (TKR2 m r)    -- Shaped tensor operations-  DeltaCastS :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+  DeltaCastS :: ( NumScalar r1, Differentiable r1+                , NumScalar r2, Differentiable r2 )              => Delta target (TKS sh r1) -> Delta target (TKS sh r2)-  DeltaSum0S :: Delta target (TKS2 sh r) -> Delta target (TKS2 '[] r)-  DeltaDot0S :: (GoodScalar r, Show (target (TKS sh r)))+  DeltaDot0S :: (NumScalar r, Show (target (TKS sh r)))              => target (TKS sh r) -> Delta target (TKS sh r)-             -> Delta target (TKS '[] r)+             -> Delta target (TKScalar r)+  DeltaIndex0S :: forall shm r target. GoodScalar r+               => Delta target (TKS shm r) -> IxSOf target shm+               -> Delta target (TKScalar r)   DeltaIndexS :: forall shm shn r target.                  ShS shn               -> Delta target (TKS2 (shm ++ shn) r) -> IxSOf target shm               -> Delta target (TKS2 shn r)+  DeltaFromVectorLinearS :: forall shm r target. GoodScalar r+                         => ShS shm+                         -> Data.Vector.Vector (Delta target (TKScalar r))+                         -> Delta target (TKS shm r)+  DeltaFromVectorS :: forall shm shn x target.+                      ShS shm+                   -> Data.Vector.Vector (Delta target (TKS2 shn x))+                   -> Delta target (TKS2 (shm ++ shn) x)+  DeltaSum0S :: forall shm r target. NumScalar r+             => Delta target (TKS shm r)+             -> Delta target (TKScalar r)+  DeltaSumS :: forall shm shn x target.+               ShS shm -> Delta target (TKS2 (shm ++ shn) x)+            -> Delta target (TKS2 shn x)   DeltaScatterS :: forall shm shn shp r target.                    ShS shm -> ShS shn -> ShS shp                 -> Delta target (TKS2 (shm ++ shn) r)                 -> (IxSOf target shm -> IxSOf target shp)                 -> Delta target (TKS2 (shp ++ shn) r)+  DeltaReplicate0NS :: forall shm r target. GoodScalar r+                    => ShS shm -> Delta target (TKScalar r)+                    -> Delta target (TKS shm r)+  DeltaReplicateS :: forall shm shn x target.+                     ShS shm -> Delta target (TKS2 shn x)+                  -> Delta target (TKS2 (shm ++ shn) x)   DeltaGatherS :: forall shm shn shp r target.                   ShS shm -> ShS shn -> ShS shp                -> Delta target (TKS2 (shp ++ shn) r)@@ -299,30 +315,54 @@                 -> Delta target (TKS2 sh2 r)    -- Mixed tensor operations-  DeltaCastX :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+  DeltaCastX :: ( NumScalar r1, Differentiable r1+                , NumScalar r2, Differentiable r2 )              => Delta target (TKX sh r1) -> Delta target (TKX sh r2)-  DeltaSum0X :: Delta target (TKX2 sh r) -> Delta target (TKX2 '[] r)-  DeltaDot0X :: (GoodScalar r, Show (target (TKX sh r)))+  DeltaDot0X :: (NumScalar r, Show (target (TKX sh r)))              => target (TKX sh r) -> Delta target (TKX sh r)-             -> Delta target (TKX '[] r)+             -> Delta target (TKScalar r)+  DeltaIndex0X :: forall shm r target. GoodScalar r+               => Delta target (TKX shm r) -> IxXOf target shm+               -> Delta target (TKScalar r)   DeltaIndexX :: forall shm shn r target.                  StaticShX shn               -> Delta target (TKX2 (shm ++ shn) r) -> IxXOf target shm               -> Delta target (TKX2 shn r)+  DeltaFromVectorLinearX :: forall shm r target. GoodScalar r+                         => IShX shm+                         -> Data.Vector.Vector (Delta target (TKScalar r))+                         -> Delta target (TKX shm r)+  DeltaFromVectorX :: forall shm shn x target.+                      IShX shm+                   -> Data.Vector.Vector (Delta target (TKX2 shn x))+                   -> Delta target (TKX2 (shm ++ shn) x)+  DeltaSum0X :: forall shm r target. NumScalar r+             => Delta target (TKX shm r)+             -> Delta target (TKScalar r)+  DeltaSumX :: forall shm shn x target.+               StaticShX shm -> Delta target (TKX2 (shm ++ shn) x)+            -> Delta target (TKX2 shn x)   DeltaScatterX :: StaticShX shm -> StaticShX shn -> StaticShX shp-                -> IShX (shp ++ shn) -> Delta target (TKX2 (shm ++ shn) r)+                -> IShX shp -> Delta target (TKX2 (shm ++ shn) r)                 -> (IxXOf target shm -> IxXOf target shp)                 -> Delta target (TKX2 (shp ++ shn) r)+  DeltaReplicate0NX :: forall shm r target. GoodScalar r+                    => IShX shm -> Delta target (TKScalar r)+                    -> Delta target (TKX shm r)+  DeltaReplicateX :: forall shm shn x target.+                     IShX shm -> Delta target (TKX2 shn x)+                  -> Delta target (TKX2 (shm ++ shn) x)   DeltaGatherX :: StaticShX shm -> StaticShX shn -> StaticShX shp-               -> IShX (shm ++ shn) -> Delta target (TKX2 (shp ++ shn) r)+               -> IShX shm -> Delta target (TKX2 (shp ++ shn) r)                -> (IxXOf target shm -> IxXOf target shp)                -> Delta target (TKX2 (shm ++ shn) r)-  DeltaAppendX :: Delta target (TKX2 (Just m ': sh) r)-               -> Delta target (TKX2 (Just n ': sh) r)-               -> Delta target (TKX2 (Just (m + n) ': sh) r)-  DeltaSliceX :: SNat i -> SNat n -> SNat k-              -> Delta target (TKX2 (Just (i + n + k) ': sh) r)-              -> Delta target (TKX2 (Just n ': sh) r)+  DeltaAppendX :: forall m n sh r target.+                  Delta target (TKX2 (m ': sh) r)+               -> Delta target (TKX2 (n ': sh) r)+               -> Delta target (TKX2 (AddMaybe m n ': sh) r)+  DeltaSliceX :: SMayNat Int i -> SMayNat Int n -> SMayNat Int k+              -> Delta target (TKX2 (AddMaybe (AddMaybe i n) k ': sh) r)+              -> Delta target (TKX2 (n ': sh) r)   DeltaReverseX :: Delta target (TKX2 (mn ': sh) r)                 -> Delta target (TKX2 (mn ': sh) r)   DeltaTransposeX :: forall perm sh r target.@@ -368,13 +408,6 @@     FTKProduct ftk1 _ -> ftk1   DeltaProject2 v -> case ftkDelta v of     FTKProduct _ ftk2 -> ftk2-  DeltaFromVector snat _ l -> case V.uncons l of-    Nothing -> error "ftkDelta: empty vector"-    Just (d, _) -> buildFTK snat (ftkDelta d)-  DeltaSum snat stk d -> razeFTK snat stk (ftkDelta d)-  DeltaReplicate snat _ d -> buildFTK snat (ftkDelta d)-  DeltaMapAccumR k bftk _eftk _q _es _df _rf acc0' _es' ->-    FTKProduct (ftkDelta acc0') (buildFTK k bftk)   DeltaMapAccumL k bftk _eftk _q _es _df _rf acc0' _es' ->     FTKProduct (ftkDelta acc0') (buildFTK k bftk) @@ -387,19 +420,28 @@    DeltaCastR d -> case ftkDelta d of     FTKR sh _ -> FTKR sh FTKScalar-  DeltaSum0R d -> case ftkDelta d of-    FTKR _ x -> FTKR ZSR x-  DeltaDot0R{} -> FTKR ZSR FTKScalar-  DeltaIndexR SNat d ix | SNat <- ixrRank ix -> case ftkDelta d of+  DeltaDot0R{} -> FTKScalar+  DeltaIndex0R _ _ -> FTKScalar+  DeltaIndexR _ d ix | SNat <- ixrRank ix -> case ftkDelta d of     FTKR sh x -> FTKR (shrDrop sh) x-  DeltaScatterR _ _ _ sh d _ -> case ftkDelta d of-    FTKR _ x -> FTKR sh x-  DeltaGatherR _ _ _ sh d _ -> case ftkDelta d of-    FTKR _ x -> FTKR sh x+  DeltaFromVectorLinearR shm _ -> FTKR shm FTKScalar+  DeltaFromVectorR shm l -> case V.uncons l of+    Nothing -> error "ftkDelta: empty vector"+    Just (d, _) | FTKR shn x <- ftkDelta d -> FTKR (shm `shrAppend` shn) x+  DeltaSum0R{} -> FTKScalar+  DeltaSumR @m SNat d -> case ftkDelta d of+    FTKR shmshn x -> FTKR (shrDrop @m shmshn) x+  DeltaScatterR (SNat @m) _ _ shp d _ -> case ftkDelta d of+    FTKR sh x -> FTKR (shp `shrAppend` shrDrop @m sh) x+  DeltaReplicate0NR shm _ -> FTKR shm FTKScalar+  DeltaReplicateR shm d -> case ftkDelta d of+    FTKR shn x -> FTKR (shm `shrAppend` shn) x+  DeltaGatherR _ _ (SNat @p) shm d _ -> case ftkDelta d of+    FTKR sh x -> FTKR (shm `shrAppend` shrDrop @p sh) x+  -- Depite the warning, the pattern match is exhaustive and if a dummy+  -- pattern is added, GHC 9.14.1 complains about that, in turn.   DeltaAppendR a b -> case ftkDelta a of-    FTKR ZSR _ -> error "ftkDelta: impossible pattern needlessly required"     FTKR (ai :$: ash) x -> case ftkDelta b of-      FTKR ZSR _ -> error "ftkDelta: impossible pattern needlessly required"       FTKR (bi :$: _) _ -> FTKR (ai + bi :$: ash) x   DeltaSliceR _ n d -> case ftkDelta d of     FTKR sh x -> FTKR (n :$: shrTail sh) x@@ -411,13 +453,24 @@    DeltaCastS d -> case ftkDelta d of     FTKS sh FTKScalar -> FTKS sh FTKScalar-  DeltaSum0S d -> case ftkDelta d of-    FTKS _ x -> FTKS ZSS x-  DeltaDot0S{} -> FTKS ZSS FTKScalar+  DeltaDot0S{} -> FTKScalar+  DeltaIndex0S _ _ -> FTKScalar   DeltaIndexS shn d _ix -> case ftkDelta d of     FTKS _ x -> FTKS shn x+  DeltaFromVectorLinearS shm _ -> FTKS shm FTKScalar+  DeltaFromVectorS shm l -> case V.uncons l of+    Nothing -> error "ftkDelta: empty vector"+    Just (d, _) | FTKS shn x <- ftkDelta d -> FTKS (shm `shsAppend` shn) x+  DeltaSum0S{} -> FTKScalar+  DeltaSumS @shm @shn shm d | SNat <- shsRank shm -> case ftkDelta d of+    FTKS shmshn x ->+      gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $+      FTKS (shsDrop @(Rank shm) shmshn) x   DeltaScatterS _shm shn shp d _ -> case ftkDelta d of     FTKS _ x -> FTKS (shp `shsAppend` shn) x+  DeltaReplicate0NS shm _ -> FTKS shm FTKScalar+  DeltaReplicateS shm d -> case ftkDelta d of+    FTKS shn x -> FTKS (shm `shsAppend` shn) x   DeltaGatherS shm shn _shp d _ -> case ftkDelta d of     FTKS _ x -> FTKS (shm `shsAppend` shn) x   DeltaAppendS a b -> case (ftkDelta a, ftkDelta b) of@@ -432,24 +485,33 @@    DeltaCastX d -> case ftkDelta d of     FTKX sh FTKScalar -> FTKX sh FTKScalar-  DeltaSum0X d -> case ftkDelta d of-    FTKX _ x -> FTKX ZSX x-  DeltaDot0X{} -> FTKX ZSX FTKScalar-  DeltaIndexX @shm @shn shn d ix -> case ftkDelta d of+  DeltaDot0X{} -> FTKScalar+  DeltaIndex0X _ _ -> FTKScalar+  DeltaIndexX @shm @shn _shn d ix -> case ftkDelta d of     FTKX sh x | SNat @len <- ixxRank ix ->       gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $-      withKnownShX (ssxFromShX sh) $-      withKnownShX shn $-      FTKX (shxDrop @len sh) x  -- TODO: (shxDropSSX sh (ssxFromIxX ix)) x-  DeltaScatterX _ _ _ sh d _ -> case ftkDelta d of-    FTKX _ x -> FTKX sh x-  DeltaGatherX _ _ _ sh d _ -> case ftkDelta d of-    FTKX _ x -> FTKX sh x+      FTKX (shxDrop @len sh) x+  DeltaFromVectorLinearX shm _ -> FTKX shm FTKScalar+  DeltaFromVectorX shm l -> case V.uncons l of+    Nothing -> error "ftkDelta: empty vector"+    Just (d, _) | FTKX shn x <- ftkDelta d -> FTKX (shm `shxAppend` shn) x+  DeltaSum0X{} -> FTKScalar+  DeltaSumX @shm @shn shm d | SNat <- ssxRank shm -> case ftkDelta d of+    FTKX shmshn x ->+      gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $+      FTKX (shxDrop @(Rank shm) shmshn) x+  DeltaScatterX @_ @shn ssm _ _ shp d _ -> case ftkDelta d of+    FTKX sh x -> FTKX (shp `shxAppend` shxDropSSX @_ @shn ssm sh) x+  DeltaReplicate0NX shm _ -> FTKX shm FTKScalar+  DeltaReplicateX shm d -> case ftkDelta d of+    FTKX shn x -> FTKX (shm `shxAppend` shn) x+  DeltaGatherX @_ @shn _ _ ssp shm d _ -> case ftkDelta d of+    FTKX sh x -> FTKX (shm `shxAppend` shxDropSSX @_ @shn ssp sh) x   DeltaAppendX a b -> case (ftkDelta a, ftkDelta b) of-    (FTKX (Nested.SKnown m :$% sh) x, FTKX (Nested.SKnown n :$% _) _) ->-      FTKX (Nested.SKnown (snatPlus m n) :$% sh) x-  DeltaSliceX _ n@SNat _ d -> case ftkDelta d of-    FTKX (_ :$% sh) x -> FTKX (Nested.SKnown n :$% sh) x+    (FTKX (m :$% sh) x, FTKX (n :$% _) _) ->+      FTKX (smnAddMaybe m n :$% sh) x+  DeltaSliceX _ n _ d -> case ftkDelta d of+    FTKX (_ :$% sh) x -> FTKX (n :$% sh) x   DeltaReverseX d -> ftkDelta d   DeltaTransposeX perm d -> case ftkDelta d of     FTKX sh x -> FTKX (shxPermutePrefix perm sh) x
src/HordeAd/Core/DeltaEval.hs view
@@ -1,1098 +1,1120 @@-{-# LANGUAGE AllowAmbiguousTypes, QuantifiedConstraints #-}-{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}-{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}--- | Evaluation of delta expressions, that is, transpose of the linear--- maps of which the delta expressions are sparse representations.--- See comments in "HordeAd.Core.Delta".-module HordeAd.Core.DeltaEval-  ( -- * Delta expression evaluation-    gradientFromDelta, derivativeFromDelta-    -- * Exported to be specialized elsewhere-  , evalRev, evalRevFTK, evalRevSame, evalRevFromnMap, EvalState-  ) where--import Prelude--import Control.Arrow (second)-import Control.Exception.Assert.Sugar-import Data.Dependent.EnumMap.Strict (DEnumMap)-import Data.Dependent.EnumMap.Strict qualified as DMap-import Data.Dependent.Sum (DSum (..))-import Data.Proxy (Proxy (Proxy))-import Data.Traversable (mapAccumL)-import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))-import Data.Vector.Generic qualified as V-import Text.Show (showListWith)-import Text.Show.Functions ()-import Type.Reflection (typeRep)--import Data.Array.Nested (type (++))-import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Lemmas-import Data.Array.Nested.Mixed.Shape-import Data.Array.Nested.Permutation (permInverse)-import Data.Array.Nested.Permutation qualified as Permutation-import Data.Array.Nested.Ranked.Shape-import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (unsafeCoerceRefl)--import HordeAd.Core.ConvertTensor-import HordeAd.Core.Delta-import HordeAd.Core.Ops-import HordeAd.Core.TensorKind-import HordeAd.Core.Types-import HordeAd.Core.Unwind---- * Top-level functions for computing derivatives from delta expressions---- | The top-level function for computing a gradient of an objective function.------ Delta expressions naturally denote forward derivatives, as encoded--- in function 'derivativeFromDelta'. However, we are usually more--- interested in computing gradients, which is what @gradientFromDelta@ does.--- The two functions are bound by the equation from Lemma 5 from the paper--- "Provably correct, asymptotically efficient, higher-order reverse-mode--- automatic differentiation":------ > dt <.> derivativeFromDelta d ds = gradientFromDelta d dt <.> ds------ where @\<.\>@ denotes generalized dot product (multiplying--- all tensors element-wise and summing the results), @d@ is the top level--- delta expression from translation of the objective function @f@ to dual--- numbers, @ds@ belongs to the domain of @f@ and @dt@ to the codomain.--- In other words, @ds@ is a perturbation (small change) of the arguments--- of @f@, for which we compute the derivative, and @dt@ is a sensitivity--- of the result of @f@, for which we compute the gradient.--- Nota bene, this property is checked for many example objective functions--- (and perturbations and sensitivities) in the horde-ad testsuite.-gradientFromDelta-  :: forall x z target. (ADReadyNoLet target, ShareTensor target)-  => FullShapeTK x-  -> FullShapeTK z-  -> target (ADTensorKind z)-  -> Delta target z-  -> target (ADTensorKind x)-gradientFromDelta !xftk !zftk !dt deltaTopLevel =-  let s0 = initEvalState xftk-      s1 = evalRev zftk s0 dt deltaTopLevel-      s2 = evalRevFromnMap s1-      (res, remainder) =-        rebuildInputs @(ADTensorKind x) (DMap.toAscList $ iMap s2) s2-        $ adFTK xftk-  in assert (null remainder) res---- | The top-level function for computing a (forward) derivative--- of an objective function.-derivativeFromDelta-  :: forall x z target. (ADReadyNoLet target, ShareTensor target)-  => Delta target z -> FullShapeTK (ADTensorKind x)-  -> target (ADTensorKind x)-  -> target (ADTensorKind z)-derivativeFromDelta deltaTopLevel ftk ds =-    let iMap = DMap.fromDistinctAscList $ fst $ generateDSums 0 ftk ds-        s0 = DMap.empty-        !(!_s2, !c) = evalFwd iMap s0 deltaTopLevel-    in c----- * Auxiliary datatypes for Delta evaluation--type ADMap target = DEnumMap (NodeId target) (Cotangent target)--type IMap target = DEnumMap (InputId target) (TensorOrZero target)--showsPrec_IMap-  :: (forall y. KnownSTK y => Show (TensorOrZero target y))-  => Int -> IMap target -> ShowS-showsPrec_IMap d demap =-  showParen (d > 10) $-    showString "fromList "-    . showListWith-        (\(k :=> target) ->-          withKnownSTK (ftkToSTK $ inputIdToFTK k) $-          showsPrec 2 k . showString " :=> " . showsPrec 1 target)-        (DMap.toList demap)--show_IMap-  :: (forall y. KnownSTK y => Show (TensorOrZero target y))-  => IMap target -> String-show_IMap iMap = showsPrec_IMap 0 iMap ""--type role Cotangent nominal nominal-newtype Cotangent target y =-  Cotangent {unCotangent :: target (ADTensorKind y)}---- | This is a tensor representation where zero tensors are marked specially--- at construction, when it's cheap to do so (as opposed to, later on.--- requiring a traversal of large terms or checking that each cell--- of a huge concrete array is zero). It also makes the computation--- of a special case of addTensorOrZero cheaper.-type role TensorOrZero nominal nominal-data TensorOrZero target y =-    TOTensor (target y)-  | TOZero (FullShapeTK y)-  deriving Show--evalTensorOrZero :: forall target x. ADReadyNoLet target-                 => TensorOrZero target x -> target x-evalTensorOrZero = \case-  TOTensor t -> t-  TOZero ftk -> tdefTarget ftk---- The ShareTensor constraint is needed, despite what GHC says,--- in order not to require duplicable arguments.-addTensorOrZero :: forall target y. (ADReadyNoLet target, ShareTensor target)-                => SingletonTK y-                -> TensorOrZero target y -> TensorOrZero target y-                -> TensorOrZero target y-addTensorOrZero stk a b = case (a, b) of-  (TOTensor ta, TOTensor tb) -> TOTensor $ taddTarget stk ta tb-    -- target has a ShareTensor instance, so ta and tb don't need-    -- to be duplicable-  (TOZero{}, _) -> b-  (_, TOZero{}) -> a---- Matches generateDSumsDummy.-rebuildInputs :: forall ady target. ADReadyNoLet target-              => [DSum (InputId target) (TensorOrZero target)]-              -> EvalState target  -- original state; only for error messages-              -> FullShapeTK ady-              -> (target ady, [DSum (InputId target) (TensorOrZero target)])-rebuildInputs els s2 ftk = case ftk of-  FTKProduct ftk1 ftk2 ->-      let (t1, rest1) = rebuildInputs els s2 ftk1-          (t2, rest2) = rebuildInputs rest1 s2 ftk2-          !t = tpair t1 t2-      in (t, rest2)-  _ | differentiableFTK ftk -> case els of-    (n :=> tz@(TOTensor t)) : rest ->-      case matchingFTK (inputIdToFTK n) ftk of-        Just Refl ->-          (t, rest)-        _ | Dict <- lemKnownSTK (ftkToSTK $ inputIdToFTK n) ->-          error $ "rebuildInputs: wrong Tensor type: "-                  ++ show (n, tz, show_IMap (iMap s2))-    (n :=> tz@(TOZero ftk2)) : rest ->-      case matchingFTK ftk2 ftk of-        Just Refl ->-          let !zero = tdefTarget ftk-          in (zero, rest)-          -- TODO: actually pass this ZERO through to optimizers-          -- and use there to avoid updating the gradient-          -- and maybe use elsewhere, too, to manage sparsity a bit.-          -- We'd probably need TKProduct to take TensorOrZero.-        _ | Dict <- lemKnownSTK (ftkToSTK $ inputIdToFTK n) ->-          error $ "rebuildInputs: wrong Zero type: "-                  ++ show (n, tz, show_IMap (iMap s2))-    _ -> error $ "rebuildInputs: illegal TensorOrZero: "-                 ++ show_IMap (iMap s2)-  _ -> (tdefTarget ftk, els)---- Matches generateDeltaInputs.-generateDSumsDummy :: Int -> FullShapeTK y-                   -> ([DSum (InputId target) (TensorOrZero target)], Int)-generateDSumsDummy j ftk  = case ftk of-  FTKProduct ftk1 ftk2 ->-    let (ds1, j1) = generateDSumsDummy j ftk1-        (ds2, j2) = generateDSumsDummy j1 ftk2-    in (ds1 ++ ds2, j2)-  _ | differentiableFTK ftk -> ([mkInputId ftk j :=> TOZero ftk], j + 1)-  _ -> ([], j)---- Matches generateDeltaInputs.-generateDSums :: ShareTensor target-              => Int -> FullShapeTK y -> target y-              -> ([DSum (InputId target) (TensorOrZero target)], Int)-generateDSums j ftk t = case ftk of-  FTKProduct ftk1 ftk2 ->-    let (t1, t2) = tunpair t-        (ds1, j1) = generateDSums j ftk1 t1-        (ds2, j2) = generateDSums j1 ftk2 t2-    in (ds1 ++ ds2, j2)-  _ | differentiableFTK ftk -> ([mkInputId ftk j :=> TOTensor t], j + 1)-  _ -> ([], j)----- * Delta evaluation state---- | The state of evaluation. It consists of several maps.--- The maps indexed by input identifiers and node identifiers--- eventually store cotangents for their respective nodes.--- The cotangents are built gradually during the evaluation,--- by summing cotangent contributions.------ Data invariant: keys nMap == keys dMap.-type role EvalState nominal-data EvalState target = EvalState-  { iMap :: IMap target-      -- ^ eventually, cotangents of objective function inputs-      -- (eventually copied to the vector representing the gradient-      -- of the objective function);-      -- the identifiers need to be contiguous and start at 0-  , dMap :: ADMap target-      -- ^ eventually, cotangents of non-input subterms indexed-      -- by their node identifiers-  , nMap :: DEnumMap (NodeId target) (Delta target)-      -- ^ nodes left to be evaluated;-      -- we can't evaluate them at once, because their other shared copies-      -- may still not be processed, so we'd not take advantage of the sharing-      -- and/or not take into account the whole summed context when finally-      -- evaluating-  }---- | Initialization of the evalutation state, which consists of--- creating the finite maps that hold values associated with inputs--- and with (possibly shared) term tree nodes.--- The former are usually initialized with dummy values so that it's cheap--- to check if any update has already been performed to a cell--- (allocating big vectors filled with zeros is too costly,--- especially if never used in an iteration, and adding to such vectors--- and especially using them as cotangent accumulators is wasteful).-initEvalState :: FullShapeTK x -> EvalState target-initEvalState ftk0 =-  let iMap = DMap.fromDistinctAscList $ fst $ generateDSumsDummy 0 $ adFTK ftk0-      dMap = DMap.empty-      nMap = DMap.empty-  in EvalState {..}----- * Reverse pass, transpose/evaluation of the delta expressions--evalRevScalarRuntimeSpecialized-  :: forall r target.-     (GoodScalar r, ADReadyNoLet target, ShareTensor target)-  => EvalState target -> target (ADTensorKind (TKScalar r))-  -> Delta target (TKScalar r)-  -> EvalState target-{-# INLINE evalRevScalarRuntimeSpecialized #-}-evalRevScalarRuntimeSpecialized !s !c =-  case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> evalRevSame @(TKScalar Double) s c-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> evalRevSame @(TKScalar Float) s c-      _ -> const s--evalRevRRuntimeSpecialized-  :: forall n r target.-     (GoodScalar r, ADReadyNoLet target, ShareTensor target)-  => EvalState target -> target (ADTensorKind (TKR n r))-  -> Delta target (TKR n r)-  -> EvalState target-{-# INLINE evalRevRRuntimeSpecialized #-}-evalRevRRuntimeSpecialized !s !c =-  -- We dispatch on all expected underyling scalar types, which is-  -- necessary to run the correct specialization when unpacking-  -- an existential type. All IfDifferentiable instances should-  -- be included in the list of expected underlying scalar types.-  -- If the scalar type is not on the list, performance suffers greatly.-  case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> evalRevSame @(TKR n Double) s c-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> evalRevSame @(TKR n Float) s c-      _ -> const s--evalSRuntimeSpecialized-  :: forall sh r target.-     (GoodScalar r, ADReadyNoLet target, ShareTensor target)-  => EvalState target -> target (ADTensorKind (TKS sh r))-  -> Delta target (TKS sh r)-  -> EvalState target-{-# INLINE evalSRuntimeSpecialized #-}-evalSRuntimeSpecialized !s !c =-  case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> evalRevSame @(TKS sh Double) s c-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> evalRevSame @(TKS sh Float) s c-      _ -> const s--evalXRuntimeSpecialized-  :: forall sh r target.-     (GoodScalar r, ADReadyNoLet target, ShareTensor target)-  => EvalState target -> target (ADTensorKind (TKX sh r))-  -> Delta target (TKX sh r)-  -> EvalState target-{-# INLINE evalXRuntimeSpecialized #-}-evalXRuntimeSpecialized !s !c =-  case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> evalRevSame @(TKX sh Double) s c-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> evalRevSame @(TKX sh Float) s c-      _ -> const s---- | Reverse pass, that is, transpose/evaluation of the delta expressions--- in order to produce the gradient for the objective function runtime--- trace represented by the delta expression.------ The first argument is the tensor kind that constrains the shapes--- of the contangent accumulator and the delta expression arguments.--- The second is the evaluation state being modified.--- The third is the cotangent accumulator that will become an actual--- cotangent contribution when complete (see below for an explanation).--- The fourth is the delta expression node to evaluate.------ Obtaining the gradient amounts to transposing the linear map--- that is straightforwardly represented by the delta expression.--- The @evalRev@ function transposes the linear map and,--- at the same time, evaluates the transposed map on the cotangent accumulator--- value contained in the third argument. If the cotangent and the tensor--- operations are symbolic, the resulting value represents the transposed--- map itself, if its free variables are treated as the map's inputs.-evalRev-  :: forall y target.-     (ADReadyNoLet target, ShareTensor target)-  => FullShapeTK y-  -> EvalState target -> target (ADTensorKind y) -> Delta target y-  -> EvalState target-evalRev ftk !s !c d = case ftk of-  FTKScalar @r -> evalRevScalarRuntimeSpecialized @r s c d-  FTKR @n _ (FTKScalar @r) -> evalRevRRuntimeSpecialized @n @r s c d-  FTKS @sh _ (FTKScalar @r) -> evalSRuntimeSpecialized @sh @r s c d-  FTKX @sh _ (FTKScalar @r) -> evalXRuntimeSpecialized @sh @r s c d-  _ -> evalRevFTK s c d---- | A helper function to `evalRev`. The @FTK@ suffix denotes it doesn't get--- an FTK as an argument but reconstructs it as needed.------ All constructors that can have a type with TKProduct kind--- need to be handled here,--- as opposed to in 'evalRevSame', except for DeltaInput that is always--- constructed only in basic kinds even though its type permits others.-evalRevFTK-  :: forall y target.-     (ADReadyNoLet target, ShareTensor target)-  => EvalState target -> target (ADTensorKind y) -> Delta target y-  -> EvalState target-evalRevFTK !s !c d0 = case d0 of-  DeltaShare n d ->-    -- In this context, by construction, @d@ is the dual component-    -- of a dual number term. Let's say that, at this point, evaluation-    -- considers position (node) p out of possibly multiple positions-    -- at which that dual number resides in the whole term tree-    -- of the dual number representation of the objective function.-    -- (Equivalently, considers edge p, one of many leading to the only-    -- node with identifier @n@ in the DAG representing the term).-    -- If so, the @c@ argument of @eval0@ is the cotangent-    -- contribution for position p, that is, the partial derivative-    -- of the objective function with respect to position p.-    ---    -- If there are indeed multiple such positions-    -- (the term is non-trivially shared) then,-    -- over the course of evaluation, cotangent contributions-    -- of them all are gradually accumulated in the finite-    -- maps and eventually their total sum represents the total-    -- influence of the objective function's subcomputation-    -- (more precisely, subgraph of the data flow graph in question)-    -- corresponding to the shared term @DeltaShare n d@. This total-    -- influence over the objective function's behaviour is called-    -- in short the cotangent of the node identifier @n@.-    -- In other words, the cotangent of @n@ is the sum,-    -- over all positions (edges) q in the global delta-expression DAG-    -- that are a reference to node @n@, of the partial derivative-    -- of the objective function with respect to the subcomputation-    -- corresponding to @q@ (meaning, subcomputations denoted by-    -- Haskell terms whose dual components are @Share n ...@).-    ---    -- For @Input@ terms, the eventual lists of cotangents end up-    -- in the cells of the gradient vectors that are the final-    -- result of the evaluation.-    assert (case d of  -- should match shareDelta-              DeltaZero{} -> False-              DeltaPair{} -> False-              DeltaInput{} -> False-              DeltaShare{} -> False-              _ -> True)-    $ if DMap.member n $ nMap s-      then let addc x =-                 Cotangent $ taddTarget (adSTK $ ftkToSTK $ nodeIdToFTK n)-                                        c (unCotangent x)-             -- target has a ShareTensor instance, so taddTarget arguments-             -- don't need to be duplicable-           in s {dMap = DMap.adjust addc n $ dMap s}-      else let cd = Cotangent c-           in s { nMap = DMap.insert n d $ nMap s-                , dMap = DMap.insert n cd $ dMap s }--  DeltaPair d1 d2 ->-    let (c1, c2) = tunpair c-    in evalRevFTK (evalRevFTK s c1 d1) c2 d2-  DeltaProject1 d -> case ftkDelta d of-    FTKProduct _ ftk2 ->-      let zero = tdefTarget $ adFTK ftk2-      in evalRevFTK s (tpair c zero) d-    -- if y is, e.g., TKR 0 Int64, we eval this delta anyway, even though-    -- we could ignore it at the price of complicating or duplicating-    -- the code slightly more-  DeltaProject2 d -> case ftkDelta d of-    FTKProduct ftk1 _ ->-      let zero = tdefTarget $ adFTK ftk1-      in evalRevFTK s (tpair zero c) d-  DeltaFromVector snat stk ld | Refl <- lemBuildOfAD snat stk ->-    let cxs = tunravelToListShare snat (adSTK stk) c-    in foldl' (\ !s2 (cx, d2) -> evalRevFTK s2 cx d2) s-       $ zip cxs (V.toList ld)-  DeltaSum snat stk d | Refl <- lemBuildOfAD snat stk ->-    evalRevFTK s (treplicate snat (adSTK stk) c) d-  DeltaReplicate snat stk d | Refl <- lemBuildOfAD snat stk ->-    evalRevFTK s (tsum snat (adSTK stk) c) d-  DeltaMapAccumR k bftk eftk q es _df rf acc0' es'-   | Refl <- lemBuildOfAD k (ftkToSTK bftk)-   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->-    let accftk = ftkDelta acc0'-        accftkAD = adFTK accftk-        bftkAD = adFTK bftk-        eftkAD = adFTK eftk-        (c0, crest) = tunpair c-        dacc_des =-          tmapAccumL (Proxy @target)-                     k accftkAD eftkAD (FTKProduct bftkAD-                                                   (FTKProduct accftk eftk))-                     (\dx db_acc_e ->-                        ttlet db_acc_e $ \ !db_acc_e1 ->-                          unHFun rf (tpair (tpair dx (tproject1 db_acc_e1))-                                           (tproject2 db_acc_e1)))-                     c0-                     (tpair crest (tpair q es))-        (dacc, des) = tunpair dacc_des-        s2 = evalRevFTK s dacc acc0'-    in evalRevFTK s2 des es'-  DeltaMapAccumL k bftk eftk q es _df rf acc0' es'-   | Refl <- lemBuildOfAD k (ftkToSTK bftk)-   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->-    let accftk = ftkDelta acc0'-        accftkAD = adFTK accftk-        bftkAD = adFTK bftk-        eftkAD = adFTK eftk-        (c0, crest) = tunpair c-        dacc_des =-          tmapAccumR (Proxy @target)-                     k accftkAD eftkAD (FTKProduct bftkAD-                                                   (FTKProduct accftk eftk))-                     (\dx db_acc_e ->-                        ttlet db_acc_e $ \ !db_acc_e1 ->-                          unHFun rf (tpair (tpair dx (tproject1 db_acc_e1))-                                           (tproject2 db_acc_e1)))-                     c0-                     (tpair crest (tpair q es))-        (dacc, des) = tunpair dacc_des-        s2 = evalRevFTK s dacc acc0'-    in evalRevFTK s2 des es'--  _ -> let y = ftkDelta d0-       in case matchingFTK y (adFTK y) of-         Just Refl -> evalRevSame s c d0-         _ -> s  -- the constructors remaining here have y that is-                 -- a non-TKProduct so if y is equal to ADTensorKind y,-                 -- the latter has the Z1 scalar type and so no influence-                 -- on the derivative.---- | A helper function to `evalRev`. It assumes the scalar underlying--- the tensor kind of its arguments is differentiable.------ All constructors that can only have types with non-TKProduct kinds--- (and the DeltaInput constructor and the vector space constructors)--- can be handled here, where the extra equality constraint makes it easier.-evalRevSame-  :: forall y target.-     (ADReadyNoLet target, ShareTensor target, y ~ ADTensorKind y)-  => EvalState target -> target (ADTensorKind y) -> Delta target y-  -> EvalState target-evalRevSame !s !c = \case-  DeltaInput i ->-    let cs = TOTensor c-    in s {iMap = DMap.adjust (addTensorOrZero (ftkToSTK $ inputIdToFTK i) cs) i-                 $ iMap s}-    -- This and similar don't need to be runtime-specialized,-    -- because the type of c determines the Num instance for (+).-    -- Note that we can't express sharing by inserting DeltaShare constructors-    -- into iMap, because often sharing needs to work across many-    -- iMap keys. That's why global sharing is used.--  -- By placing these here, we force their derivatives to be zeroed-  -- whenever they are called on non-base types, which they should not ever be.-  -- This is ensured by the types of the three constructors, assuming that-  -- no Num instances are defined for the non-base type tensors.-  DeltaZero{} -> s-  DeltaScale (NestedTarget k) d -> evalRevSame s (k * c) d-  DeltaAdd d e ->-    let cShared = tshare c-    in evalRevSame (evalRevSame s cShared d) cShared e--  DeltaCastK @r1 d ->-    evalRevScalarRuntimeSpecialized-    s (toADTensorKindShared (FTKScalar @r1) $ tkcast c) d-  DeltaCastR d -> case ftkDelta d of-    y ->-      evalRevRRuntimeSpecialized-      s (toADTensorKindShared y $ trcast c) d-  DeltaSum0R d -> case ftkDelta d of-    FTKR sh x | SNat <- shrRank sh ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (trreplicate0N sh c) d-  DeltaDot0R v d -> case ftkDelta d of-    FTKR sh x | SNat <- shrRank sh ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (v * trreplicate0N (rshape v) c) d-        -- too slow: evalRevSame s (rmap0N (* (tscalar c)) v) vd-  DeltaIndexR SNat d ix -> case ftkDelta d of-    FTKR sh x | SNat <- ixrRank ix ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (troneHot (shrTake sh) c ix) d  -- TODO: ixrToShR ix-  DeltaScatterR SNat SNat SNat _sh d f -> case ftkDelta d of-    FTKR sh x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (trgather sh c f) d-  DeltaGatherR SNat SNat SNat _sh d f -> case ftkDelta d of-    FTKR sh x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (trscatter sh c f) d-  DeltaAppendR d e -> case (ftkDelta d, ftkDelta e) of-    (FTKR (m :$: _) x, FTKR (n :$: _) _) ->-      withKnownSTK (ftkToSTK x) $-      let cShared = tshare c-          s2 = evalRevSame s (trslice 0 m cShared) d-      in evalRevSame s2 (trslice m n cShared) e-    _ -> error "evalRevSame: impossible pattern needlessly required"-  DeltaSliceR i n d -> case ftkDelta d of-    FTKR (l :$: rest) x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (trappend-                       (tdefTarget (FTKR (i :$: rest) x))-                       (trappend c-                          (tdefTarget (FTKR (l - i - n :$: rest) x)))) d-    FTKR ZSR _ -> error "evalRevSame: impossible pattern needlessly required"-  DeltaReverseR d -> case ftkDelta d of-    FTKR _ x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (trreverse c) d-  DeltaTransposeR perm d -> case ftkDelta d of-    FTKR _ x ->-      withKnownSTK (ftkToSTK x) $-      let permR = permRInverse perm-      in evalRevSame s (trtranspose permR c) d-  DeltaReshapeR _sh2 d -> case ftkDelta d of-    FTKR sh x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (trreshape sh c) d--  DeltaCastS d -> case ftkDelta d of-    y ->-      evalSRuntimeSpecialized-      s (toADTensorKindShared y $ tscast c) d-  DeltaSum0S d -> case ftkDelta d of-    FTKS sh x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (tsreplicate0N sh c) d-  DeltaDot0S v d -> case ftkDelta d of-    FTKS sh FTKScalar ->-      evalRevSame s (v * tsreplicate0N sh c) d-  DeltaIndexS shn d ix -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShS shn $-      withKnownShS (shsFromIxS ix) $-      evalRevSame s (tsoneHot c ix) d-  DeltaScatterS @shm @shn shm shn shp d f -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShS shm $-      withKnownShS shn $-      withKnownShS shp $-      evalRevSame s (tsgather @_ @shm @shn c f) d-  DeltaGatherS @shm @shn shm shn shp d f -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShS shm $-      withKnownShS shn $-      withKnownShS shp $-      evalRevSame s (tsscatter @_ @shm @shn c f) d-  DeltaAppendS d e -> case (ftkDelta d, ftkDelta e) of-    (FTKS (msnat :$$ _) x, FTKS (_ :$$ _) _) ->-      withKnownSTK (ftkToSTK x) $-      let cShared = tshare c-          s2 = evalRevSame s (tsslice (SNat @0) SNat SNat cShared) d-      in evalRevSame s2 (tsslice msnat SNat SNat cShared) e-  DeltaSliceS i@SNat _ k@SNat d -> case ftkDelta d of-    FTKS (_ :$$ sh) x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (tsappend-                       (tdefTarget (FTKS (i :$$ sh) x))-                          (tsappend-                             c (tdefTarget (FTKS (k :$$ sh) x)))) d-  DeltaReverseS d -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (tsreverse c) d-  DeltaTransposeS @perm @sh2 perm d -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      permInverse perm $ \(permRev :: Permutation.Perm permR) _ ->-        gcastWith (unsafeCoerceRefl-                   :: Permutation.PermutePrefix-                        permR (Permutation.PermutePrefix perm sh2) :~: sh2)-        $ gcastWith (unsafeCoerceRefl-                     :: Rank (Permutation.PermutePrefix perm sh2) :~: Rank sh2)-        $ gcastWith (unsafeCoerceRefl-                     :: Rank permR :~: Rank perm)-        $ evalRevSame s (tstranspose permRev c) d-  DeltaReshapeS _sh2 d -> case ftkDelta d of-    FTKS sh x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (tsreshape sh c) d--  DeltaCastX d -> case ftkDelta d of-    y ->-      evalXRuntimeSpecialized-      s (toADTensorKindShared y $ txcast c) d-  DeltaSum0X d -> case ftkDelta d of-    FTKX sh x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShX (ssxFromShX sh) $-      evalRevSame s (txreplicate0N sh c) d-  DeltaDot0X v d -> case ftkDelta d of-    FTKX sh FTKScalar ->-      withKnownShX (ssxFromShX sh) $-      evalRevSame s (v * txreplicate0N (xshape v) c) d-  DeltaIndexX @shm @shn shn d ix -> case ftkDelta d of-    FTKX sh x | SNat @len <- ixxRank ix ->-      withKnownSTK (ftkToSTK x) $-      withKnownShX shn $-      withKnownShX (ssxFromShX sh) $-      withKnownShX (ssxTakeIx @shm @shn (ssxFromShX sh) ix) $-      gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $-      evalRevSame s (txoneHot (shxTake @len sh) c ix) d---TODO      evalRevSame s (xoneHot (shxTakeSSX (Proxy @shn) sh---                                         (ssxFromIxX ix)) c ix) d-  DeltaScatterX @shm @shn shm shn shp _sh d f -> case ftkDelta d of-    FTKX sh x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShX shm $-      withKnownShX shn $-      withKnownShX shp $-      evalRevSame s (txgather @_ @shm @shn sh c f) d-  DeltaGatherX @shm @shn shm shn shp _sh d f -> case ftkDelta d of-    FTKX sh x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShX shm $-      withKnownShX shn $-      withKnownShX shp $-      evalRevSame s (txscatter @_ @shm @shn sh c f) d-  DeltaAppendX d e -> case (ftkDelta d, ftkDelta e) of-    (FTKX (Nested.SKnown m@SNat :$% _) x, FTKX (Nested.SKnown SNat :$% _) _) ->-      withKnownSTK (ftkToSTK x) $-      let cShared = tshare c-          s2 = evalRevSame s (txslice (SNat @0) SNat SNat cShared) d-      in evalRevSame s2 (txslice m SNat SNat cShared) e-  DeltaSliceX i@SNat _ k@SNat d -> case ftkDelta d of-    FTKX (_ :$% sh) x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (txappend-                       (tdefTarget (FTKX (Nested.SKnown i :$% sh) x))-                          (txappend-                             c (tdefTarget-                                  (FTKX (Nested.SKnown k :$% sh) x)))) d-  DeltaReverseX d -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (txreverse c) d-  DeltaTransposeX @perm @sh2 perm d -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      permInverse perm $ \(permR :: Permutation.Perm permR) _ ->-        gcastWith (unsafeCoerceRefl-                   :: Permutation.PermutePrefix-                        permR (Permutation.PermutePrefix perm sh2) :~: sh2) $-        gcastWith (unsafeCoerceRefl-                   :: Rank (Permutation.PermutePrefix perm sh2) :~: Rank sh2) $-        gcastWith (unsafeCoerceRefl-                   :: Rank permR :~: Rank perm) $-        evalRevSame s (txtranspose permR c) d-  DeltaReshapeX _sh2 d -> case ftkDelta d of-    FTKX sh x ->-      withKnownSTK (ftkToSTK x) $-      evalRevSame s (txreshape sh c) d--  DeltaConvert @a c1 d -> case ftkDelta d of-    aftk ->-      -- This follows from the same property for @b@ and from @c1@-      -- not changing the underlying scalar types.-      gcastWith (unsafeCoerceRefl :: ADTensorKind a :~: a) $-      evalRevSame-        s (tconvert (transposeTKConversion aftk c1)-                    (convertSTK c1 $ ftkToSTK $ ftkDelta d) c) d--  d -> evalRevFTK s c d-    -- the remaining constructors are already handled in evalRevFTK--transposeTKConversion :: FullShapeTK a -> TKConversion a b -> TKConversion b a-transposeTKConversion aftk c0 = case c0 of-  ConvId -> ConvId-  ConvCmp c1 c2 -> ConvCmp (transposeTKConversion aftk c2)-                           (transposeTKConversion (convertFTK c2 aftk) c1)-  ConvRX | FTKR @n _ x <- aftk-         , Refl <- lemRankReplicate (Proxy @n) ->-    ConvXR (ftkToSTK x)-  ConvSX -> ConvXS-  ConvXR @_ @sh _stk | Refl <- lemRankReplicate (Proxy @(Rank sh)) ->-    ConvCmp (ConvXX' aftk) ConvRX-  ConvXS -> ConvSX-  ConvXS' (FTKS sh _) | Refl <- lemRankMapJust sh ->-    ConvCmp (ConvXX' aftk) ConvSX-  ConvXX' _ftk -> ConvXX' aftk-  ConvRR c | FTKR _ x <- aftk -> ConvRR (transposeTKConversion x c)-  ConvSS c | FTKS _ x <- aftk -> ConvSS (transposeTKConversion x c)-  ConvXX c | FTKX _ x <- aftk -> ConvXX (transposeTKConversion x c)-  ConvT2 c1 c2 | FTKProduct x1 x2 <- aftk ->-    ConvT2 (transposeTKConversion x1 c1) (transposeTKConversion x2 c2)-  Conv0X _stk -> ConvX0-  ConvX0 | FTKX ZSX x <- aftk -> Conv0X (ftkToSTK x)-  ConvNest _stk -> ConvUnnest-  ConvUnnest | (FTKX shx (FTKX _ x)) <- aftk ->-    ConvNest (STKX (ssxFromShX shx) (ftkToSTK x))-  ConvZip stk1 stk2 -> ConvUnzip stk1 stk2-  ConvUnzip stk1 stk2 -> ConvZip stk1 stk2--evalRevFromnMap :: forall target. (ADReadyNoLet target, ShareTensor target)-                => EvalState target -> EvalState target-evalRevFromnMap s@EvalState{nMap, dMap} =-  case DMap.maxViewWithKey nMap of-    Just (n :=> d, nMap2) ->-      let s2 = s {nMap = nMap2}-          s3 = case DMap.lookup n dMap of-            Just (Cotangent c) -> evalRev (nodeIdToFTK n) s2 c d-            Nothing -> error $ "evalRevFromnMap: missing cotangent " ++ show n-      in evalRevFromnMap s3-    Nothing -> s  -- loop ends--{- TODO: optimize similarly?-        -- The general case is given as the last one below,-        -- but for a few constructors it's faster to inline @evalRev@ instead.-        -- BTW, such an optimization doesn't really belong in the simplified-        -- horde-ad and no consistent benefit should be expected here.-        DeltaIndex0 ZeroR{} _ _ -> s  -- shortcut-        DeltaIndex0 (InputR i) ixs' sh ->-          let ixs = indexToList ixs'-              f v = if isTensorDummy v-                    then treplicate0ND sh 0 `OD.update` [(ixs, c)]-                    else v `OD.update` [(ixs, v `rindex0D` ixs + c)]-          in s {iMap = DMap.adjust f i $ iMap s}-        DeltaIndex0 (ShareR n d) ixs' sh ->-          let ixs = indexToList ixs'-          in case DMap.lookup n $ nMap s of-            Just (DynamicRanked _) ->-              let f v = v `OD.update` [(ixs, v `rindex0D` ixs + c)]-              in s {dMap = DMap.adjust f n $ dMap s}-                -- This would be an asymptotic optimization compared to-                -- the general case below, if not for the non-mutable update,-                -- which implies copying the whole @v@ vector,-                -- so it's only several times faster (same allocation,-                -- but not adding to each cell of @v@).-            Nothing ->-              let v = treplicate0ND sh 0 `OD.update` [(ixs, c)]-              in s { nMap = DMap.insert n (DynamicRanked d) $ nMap s-                   , dMap = DMap.insert n v $ dMap s }-            _ -> error "evalRevFromnMap: corrupted nMap"--}----- * Forward derivative computation from the delta expressions---- | Forward derivative computation via forward-evaluation of delta-expressions--- (which is surprisingly competitive to the direct forward method,--- until the allocation of deltas gets large enough to affect cache hits).--- This is the directional derivative, calculated for the point,--- at which the delta expression was computed (which is the point--- represented by the parameters of the objective function and used--- to compute it's dual number result) and along the direction vector--- given by the parameters in the arguments.------ This mimics the reverse derivative code, but in reverse. Perhaps this can be--- simplified, but the obvious simplest formulation does not honour sharing--- and evaluates shared subexpressions repeatedly, so this state-passing--- formulation is adopted.-evalFwd-  :: forall target y. (ADReadyNoLet target, ShareTensor target)-  => IMap target -> ADMap target -> Delta target y-  -> (ADMap target, target (ADTensorKind y))-evalFwd params s d0 = case d0 of-  DeltaShare n d ->-    case DMap.lookup n s of-      Just e1 -> (s, unCotangent e1)-      Nothing ->-        let (s2, cRaw) = evalFwd params s d-            cShared = tshare cRaw-            cd = Cotangent cShared-              -- cRaw is shared, because it's put into the map and then-              -- potentially looked up many times, so it'd get duplicated-            s3 = DMap.insert n cd s2-        in (s3, cShared)-  DeltaInput inputId ->-    case DMap.lookup inputId params of-      Just dtk -> (s, toADTensorKindShared (inputIdToFTK inputId)-                      $ evalTensorOrZero dtk)-      Nothing -> error "evalFwd: missing input"--  DeltaPair d1 d2 ->-    let (s2, t) = evalFwd params s d1-        (s3, u) = evalFwd params s2 d2-    in (s3, tpair t u)-  DeltaProject1 d ->-    let (s2, v) = evalFwd params s d-    in (s2, tproject1 v)-  DeltaProject2 d ->-    let (s2, v) = evalFwd params s d-    in (s2, tproject2 v)-  DeltaFromVector snat stk lsd | Refl <- lemBuildOfAD snat stk ->-    let (s2, l) = mapAccumL (evalFwd params) s lsd-    in (s2, tfromVector snat (adSTK stk) l)-  DeltaSum snat stk d | Refl <- lemBuildOfAD snat stk ->-    let (s2, t) = evalFwd params s d-    in (s2, tsum snat (adSTK stk) t)-  DeltaReplicate snat stk d | Refl <- lemBuildOfAD snat stk ->-    let (s2, t) = evalFwd params s d-    in (s2, treplicate snat (adSTK stk) t)-  DeltaMapAccumR k bftk eftk q es df _rf acc0' es'-   | Refl <- lemBuildOfAD k (ftkToSTK bftk)-   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->-    let accftk = ftkDelta acc0'-        accftkAD = adFTK accftk-        bftkAD = adFTK bftk-        eftkAD = adFTK eftk-        (s2, cacc0) = evalFwd params s acc0'-        (s3, ces) = evalFwd params s2 es'-    in (s3, tmapAccumR (Proxy @target)-                       k accftkAD bftkAD (FTKProduct eftkAD-                                                     (FTKProduct accftk eftk))-                       (\dacc de_acc_e ->-                        ttlet de_acc_e $ \ !de_acc_e1 ->-                          unHFun df (tpair (tpair dacc (tproject1 de_acc_e1))-                                           (tproject2 de_acc_e1)))-                       cacc0-                       (tpair ces (tpair q es)))-  DeltaMapAccumL k bftk eftk q es df _rf acc0' es'-   | Refl <- lemBuildOfAD k (ftkToSTK bftk)-   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->-    let accftk = ftkDelta acc0'-        accftkAD = adFTK accftk-        bftkAD = adFTK bftk-        eftkAD = adFTK eftk-        (s2, cacc0) = evalFwd params s acc0'-        (s3, ces) = evalFwd params s2 es'-    in (s3, tmapAccumL (Proxy @target)-                       k accftkAD bftkAD (FTKProduct eftkAD-                                                     (FTKProduct accftk eftk))-                       (\dacc de_acc_e ->-                        ttlet de_acc_e $ \ !de_acc_e1 ->-                          unHFun df (tpair (tpair dacc (tproject1 de_acc_e1))-                                           (tproject2 de_acc_e1)))-                       cacc0-                       (tpair ces (tpair q es)))--  _ -> let y = ftkDelta d0-           ay = adFTK y-       in case matchingFTK y ay of-         Just Refl -> evalFwdSame params s d0-         _ -> (s, tdefTarget ay)--evalFwdSame-  :: forall target y.-     (ADReadyNoLet target, ShareTensor target, y ~ ADTensorKind y)-  => IMap target -> ADMap target -> Delta target y-  -> (ADMap target, target (ADTensorKind y))-evalFwdSame params s = \case-  DeltaInput inputId ->-    case DMap.lookup inputId params of-      Just dtk -> (s, evalTensorOrZero dtk)-      Nothing -> error "evalFwdSame: missing input"--  -- See the comment about these three in evalRevSame.-  DeltaZero ftk -> (s, tdefTarget $ adFTK ftk)-  DeltaScale (NestedTarget k) d -> second (* k) $ evalFwdSame params s d-  DeltaAdd d e -> let (s2, t) = evalFwdSame params s d-                      (s3, u) = evalFwdSame params s2 e-                  in (s3, t + u)--  d0@(DeltaCastK @r1 d) ->-    case sameSTK (STKScalar @r1) (adSTK (STKScalar @r1)) of-      Just Refl -> second tkcast $ evalFwdSame params s d-      _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)--  d0@(DeltaCastR d) -> case ftkDelta d of-    y -> case matchingFTK y (adFTK y) of-      Just Refl -> second trcast $ evalFwdSame params s d-      _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)-  DeltaSum0R (DeltaZero (FTKR _ x)) -> (s, tdefTarget (FTKR ZSR x))-  DeltaSum0R d -> case ftkDelta d of-    FTKR sh x | SNat <- shrRank sh ->-      withKnownSTK (ftkToSTK x) $-      second trsum0 $ evalFwdSame params s d-  DeltaDot0R _ DeltaZero{} -> (s, trconcrete $ Nested.rscalar 0)-  DeltaDot0R v d -> case ftkDelta d of-    FTKR sh x | SNat <- shrRank sh ->-      withKnownSTK (ftkToSTK x) $-      second (trdot0 v) $ evalFwdSame params s d-  DeltaIndexR SNat d ix -> case ftkDelta d of-    FTKR _ x | SNat <- ixrRank ix ->-      withKnownSTK (ftkToSTK x) $-      second (`trindex` ix) $ evalFwdSame params s d-  DeltaScatterR SNat SNat SNat sh d f -> case ftkDelta d of-    FTKR _ x ->-      withKnownSTK (ftkToSTK x) $-      let (s2, t) = evalFwdSame params s d-      in (s2, trscatter sh t f)-  DeltaGatherR SNat SNat SNat sh d f -> case ftkDelta d of-    FTKR _ x ->-      withKnownSTK (ftkToSTK x) $-      let (s2, t) = evalFwdSame params s d-      in (s2, trgather sh t f)-  DeltaAppendR d e -> case ftkDelta d of-    FTKR _ x ->-      withKnownSTK (ftkToSTK x) $-      let (s2, t) = evalFwdSame params s d-          (s3, u) = evalFwdSame params s2 e-      in (s3, trappend t u)-  DeltaSliceR i n d -> case ftkDelta d of-    FTKR _ x ->-      withKnownSTK (ftkToSTK x) $-      second (trslice i n) $ evalFwdSame params s d-  DeltaReverseR d -> case ftkDelta d of-    FTKR _ x ->-      withKnownSTK (ftkToSTK x) $-      second trreverse $ evalFwdSame params s d-  DeltaTransposeR perm d -> case ftkDelta d of-    FTKR _ x ->-      withKnownSTK (ftkToSTK x) $-      second (trtranspose perm) $ evalFwdSame params s d-  DeltaReshapeR sh2 d -> case ftkDelta d of-    FTKR _sh x ->-      withKnownSTK (ftkToSTK x) $-      second (trreshape sh2) $ evalFwdSame params s d--  d0@(DeltaCastS d) -> case ftkDelta d of-    y -> case matchingFTK y (adFTK y) of-      Just Refl -> second tscast $ evalFwdSame params s d-      _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)-  DeltaSum0S (DeltaZero (FTKS _ x)) -> (s, tdefTarget (FTKS ZSS x))-  DeltaSum0S d -> case ftkDelta d of-    FTKS sh x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShS sh $-      second tssum0 $ evalFwdSame params s d-  DeltaDot0S _ DeltaZero{} -> (s, tsconcrete $ Nested.sscalar 0)-  DeltaDot0S v d -> case ftkDelta d of-    FTKS sh FTKScalar ->-      withKnownShS sh $-      second (tsdot0 v) $ evalFwdSame params s d-  DeltaIndexS shn d ix -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShS shn $-      withKnownShS (shsFromIxS ix) $-      second (`tsindex` ix) $ evalFwdSame params s d-  DeltaScatterS @shm @shn shm shn shp d f -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShS shm $-      withKnownShS shn $-      withKnownShS shp $-      let (s2, t) = evalFwdSame params s d-      in (s2, tsscatter @_ @shm @shn t f)-  DeltaGatherS @shm @shn shm shn shp d f -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShS shm $-      withKnownShS shn $-      withKnownShS shp $-      let (s2, t) = evalFwdSame params s d-      in (s2, tsgather @_ @shm @shn t f)-  DeltaAppendS d e -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      let (s2, t) = evalFwdSame params s d-          (s3, u) = evalFwdSame params s2 e-      in (s3, tsappend t u)-  DeltaSliceS i n k d -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      second (tsslice i n k) $ evalFwdSame params s d-  DeltaReverseS d -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      second tsreverse $ evalFwdSame params s d-  DeltaTransposeS perm d -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      second (tstranspose perm) $ evalFwdSame params s d-  DeltaReshapeS sh2 d -> case ftkDelta d of-    FTKS _ x ->-      withKnownSTK (ftkToSTK x) $-      second (tsreshape sh2) $ evalFwdSame params s d--  d0@(DeltaCastX d) -> case ftkDelta d of-    y -> case matchingFTK y (adFTK y) of-      Just Refl -> second txcast $ evalFwdSame params s d-      _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)-  DeltaSum0X (DeltaZero (FTKX _ x)) -> (s, tdefTarget (FTKX ZSX x))-  DeltaSum0X d -> case ftkDelta d of-    FTKX sh x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShX (ssxFromShX sh) $-      second txsum0 $ evalFwdSame params s d-  DeltaDot0X _ DeltaZero{} -> (s, txconcrete $ Nested.mscalar 0)-  DeltaDot0X v d -> case ftkDelta d of-    FTKX sh FTKScalar ->-      withKnownShX (ssxFromShX sh) $-      second (txdot0 v) $ evalFwdSame params s d-  DeltaIndexX @shm @shn shn d ix -> case ftkDelta d of-    FTKX sh x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShX shn $--- TODO      withKnownShX (ssxFromIxX ix) $-      withKnownShX (ssxTakeIx @shm @shn (ssxFromShX sh) ix) $-      second (`txindex` ix) $ evalFwdSame params s d-  DeltaScatterX @shm @shn shm shn shp sh d f -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShX shm $-      withKnownShX shn $-      withKnownShX shp $-      let (s2, t) = evalFwdSame params s d-      in (s2, txscatter @_ @shm @shn sh t f)-  DeltaGatherX @shm @shn shm shn shp sh d f -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      withKnownShX shm $-      withKnownShX shn $-      withKnownShX shp $-      let (s2, t) = evalFwdSame params s d-      in (s2, txgather @_ @shm @shn sh t f)-  DeltaAppendX d e -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      let (s2, t) = evalFwdSame params s d-          (s3, u) = evalFwdSame params s2 e-      in (s3, txappend t u)-  DeltaSliceX i n k d -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      second (txslice i n k) $ evalFwdSame params s d-  DeltaReverseX d -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      second txreverse $ evalFwdSame params s d-  DeltaTransposeX perm d -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      second (txtranspose perm) $ evalFwdSame params s d-  DeltaReshapeX sh2 d -> case ftkDelta d of-    FTKX _ x ->-      withKnownSTK (ftkToSTK x) $-      second (txreshape sh2) $ evalFwdSame params s d--  DeltaConvert @a c1 d ->-    -- This follows from the same property for @b@ and from @c1@-    -- not changing the underlying scalar types.-    gcastWith (unsafeCoerceRefl :: ADTensorKind a :~: a) $-    second (tconvert c1 (ftkToSTK (ftkDelta d)))-           (evalFwdSame params s d)--  d -> evalFwd params s d+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Evaluation of delta expressions, that is, transpose of the linear+-- maps of which the delta expressions are sparse representations.+-- See comments in "HordeAd.Core.Delta".+module HordeAd.Core.DeltaEval+  ( -- * Delta expression evaluation+    gradientFromDelta, derivativeFromDelta+  ) where++import Prelude++import Control.Arrow (second)+import Control.Exception.Assert.Sugar+import Data.Array.Nested.Lemmas+import Data.Dependent.EnumMap.Strict (DEnumMap)+import Data.Dependent.EnumMap.Strict qualified as DMap+import Data.Dependent.Sum (DSum (..))+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Text.Show (showListWith)+import Text.Show.Functions ()+import Type.Reflection (typeRep)++import Data.Array.Nested (type (++))+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation (permInverse)+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.Conversion+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Delta+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind++-- * Top-level functions for computing derivatives from delta expressions++-- | The top-level function for computing a gradient of an objective function.+--+-- Delta expressions naturally denote forward derivatives, as encoded+-- in function 'derivativeFromDelta'. However, we are usually more+-- interested in computing gradients, which is what @gradientFromDelta@ does.+-- The two functions are bound by the equation from Lemma 5 from the paper+-- "Provably correct, asymptotically efficient, higher-order reverse-mode+-- automatic differentiation":+--+-- > dt <.> derivativeFromDelta d ds = gradientFromDelta d dt <.> ds+--+-- where @\<.\>@ denotes generalized dot product (multiplying+-- all tensors element-wise and summing the results), @d@ is the top level+-- delta expression from translation of the objective function @f@ to dual+-- numbers, @ds@ belongs to the domain of @f@ and @dt@ to the codomain.+-- In other words, @ds@ is a perturbation (small change) of the arguments+-- of @f@, for which we compute the derivative, and @dt@ is a sensitivity+-- of the result of @f@, for which we compute the gradient.+-- Nota bene, this property is checked for many example objective functions+-- (and perturbations and sensitivities) in the horde-ad testsuite.+gradientFromDelta+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => FullShapeTK x+  -> target (ADTensorKind z)+  -> Delta target z+  -> target (ADTensorKind x)+gradientFromDelta !xftk !dt deltaTopLevel =+  let s0 = initEvalState xftk+      s1 = evalRev s0 dt deltaTopLevel+      s2 = evalRevFromnMap s1+      (res, remainder) =+        rebuildInputs @(ADTensorKind x) (DMap.toAscList $ iMap s2) s2+        $ adFTK xftk+  in assert (null remainder) res++-- | The top-level function for computing a (forward) derivative+-- of an objective function.+derivativeFromDelta+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => Proxy x -> Delta target z -> FullShapeTK (ADTensorKind x)+  -> target (ADTensorKind x)+  -> target (ADTensorKind z)+derivativeFromDelta _ deltaTopLevel ftk ds =+    let iMap = DMap.fromDistinctAscList $ fst $ generateDSums 0 ftk ds+        s0 = DMap.empty+        !(!_s2, !c) = evalFwd iMap s0 deltaTopLevel+    in c+++-- * Auxiliary datatypes for Delta evaluation++type ADMap target = DEnumMap (NodeId target) (Cotangent target)++type IMap target = DEnumMap (InputId target) (TensorOrZero target)++showsPrec_IMap+  :: AllTargetShow target+  => Int -> IMap target -> ShowS+showsPrec_IMap d demap =+  showParen (d > 10) $+    showString "fromList "+    . showListWith+        (\(k :=> target) ->+          withKnownSTK (ftkToSTK $ inputIdToFTK k) $+          showsPrec 2 k . showString " :=> " . showsPrec 1 target)+        (DMap.toList demap)++show_IMap+  :: AllTargetShow target+  => IMap target -> String+show_IMap iMap = showsPrec_IMap 0 iMap ""++type role Cotangent nominal nominal+newtype Cotangent target y =+  Cotangent {unCotangent :: target (ADTensorKind y)}++-- | This is a tensor representation where zero tensors are marked specially+-- at construction, when it's cheap to do so (as opposed to, later on.+-- requiring a traversal of large terms or checking that each cell+-- of a huge concrete array is zero). It also makes the computation+-- of a special case of addTensorOrZero cheaper.+type role TensorOrZero nominal nominal+data TensorOrZero target y =+    TOTensor (target y)+  | TOZero (FullShapeTK y)+  deriving Show++evalTensorOrZero :: forall target x. ADReadyNoLet target+                 => TensorOrZero target x -> target x+evalTensorOrZero = \case+  TOTensor t -> t+  TOZero ftk -> tdefTarget ftk++-- The ShareTensor constraint is needed, despite what GHC says,+-- in order not to require duplicable arguments.+addTensorOrZero+  :: forall target y. (TKAllNum y, ADReadyNoLet target, ShareTensor target)+  => SingletonTK y -> TensorOrZero target y -> TensorOrZero target y+  -> TensorOrZero target y+addTensorOrZero stk a b = case (a, b) of+  (TOTensor ta, TOTensor tb) -> TOTensor $ taddTarget stk ta tb+    -- target has a ShareTensor instance, so ta and tb don't need+    -- to be duplicable+  (TOZero{}, _) -> b+  (_, TOZero{}) -> a++-- Matches generateDSumsDummy.+rebuildInputs :: forall ady target. ADReadyNoLet target+              => [DSum (InputId target) (TensorOrZero target)]+              -> EvalState target  -- original state; only for error messages+              -> FullShapeTK ady+              -> (target ady, [DSum (InputId target) (TensorOrZero target)])+rebuildInputs els s2 ftk = case ftk of+  FTKProduct ftk1 ftk2 ->+      let (t1, rest1) = rebuildInputs els s2 ftk1+          (t2, rest2) = rebuildInputs rest1 s2 ftk2+          !t = tpair t1 t2+      in (t, rest2)+  _ | differentiableFTK ftk -> case els of+    (n :=> tz@(TOTensor t)) : rest ->+      case matchingFTK (inputIdToFTK n) ftk of+        Just Refl ->+          (t, rest)+        _ | Dict <- lemKnownSTK (ftkToSTK $ inputIdToFTK n) ->+          error $ "rebuildInputs: wrong Tensor type: "+                  ++ show (n, tz, show_IMap (iMap s2))+    (n :=> tz@(TOZero ftk2)) : rest ->+      case matchingFTK ftk2 ftk of+        Just Refl ->+          let !zero = tdefTarget ftk+          in (zero, rest)+          -- TODO: actually pass this ZERO through to optimizers+          -- and use there to avoid updating the gradient+          -- and maybe use elsewhere, too, to manage sparsity a bit.+          -- We'd probably need TKProduct to take TensorOrZero.+        _ | Dict <- lemKnownSTK (ftkToSTK $ inputIdToFTK n) ->+          error $ "rebuildInputs: wrong Zero type: "+                  ++ show (n, tz, show_IMap (iMap s2))+    _ -> error $ "rebuildInputs: illegal TensorOrZero: "+                 ++ show_IMap (iMap s2)+  _ -> (tdefTarget ftk, els)++-- Matches generateDeltaInputs.+generateDSumsDummy :: Int -> FullShapeTK y+                   -> ([DSum (InputId target) (TensorOrZero target)], Int)+generateDSumsDummy j ftk  = case ftk of+  FTKProduct ftk1 ftk2 ->+    let (ds1, j1) = generateDSumsDummy j ftk1+        (ds2, j2) = generateDSumsDummy j1 ftk2+    in (ds1 ++ ds2, j2)+  _ | differentiableFTK ftk -> ([mkInputId ftk j :=> TOZero ftk], j + 1)+  _ -> ([], j)++-- Matches generateDeltaInputs.+generateDSums :: ShareTensor target+              => Int -> FullShapeTK y -> target y+              -> ([DSum (InputId target) (TensorOrZero target)], Int)+generateDSums j ftk t = case ftk of+  FTKProduct ftk1 ftk2 ->+    let (t1, t2) = tunpair t+        (ds1, j1) = generateDSums j ftk1 t1+        (ds2, j2) = generateDSums j1 ftk2 t2+    in (ds1 ++ ds2, j2)+  _ | differentiableFTK ftk -> ([mkInputId ftk j :=> TOTensor t], j + 1)+  _ -> ([], j)+++-- * Delta evaluation state++-- | The state of evaluation. It consists of several maps.+-- The maps indexed by input identifiers and node identifiers+-- eventually store cotangents for their respective nodes.+-- The cotangents are built gradually during the evaluation,+-- by summing cotangent contributions.+--+-- Data invariant: keys nMap == keys dMap.+type role EvalState nominal+data EvalState target = EvalState+  { iMap :: IMap target+      -- ^ eventually, cotangents of objective function inputs+      -- (eventually copied to the vector representing the gradient+      -- of the objective function);+      -- the identifiers need to be contiguous and start at 0+  , dMap :: ADMap target+      -- ^ eventually, cotangents of non-input subterms indexed+      -- by their node identifiers+  , nMap :: DEnumMap (NodeId target) (Delta target)+      -- ^ nodes left to be evaluated;+      -- we can't evaluate them at once, because their other shared copies+      -- may still not be processed, so we'd not take advantage of the sharing+      -- and/or not take into account the whole summed context when finally+      -- evaluating+  }++-- | Initialization of the evalutation state, which consists of+-- creating the finite maps that hold values associated with inputs+-- and with (possibly shared) term tree nodes.+-- The former are usually initialized with dummy values so that it's cheap+-- to check if any update has already been performed to a cell+-- (allocating big vectors filled with zeros is too costly,+-- especially if never used in an iteration, and adding to such vectors+-- and especially using them as cotangent accumulators is wasteful).+initEvalState :: FullShapeTK x -> EvalState target+initEvalState ftk0 =+  let iMap = DMap.fromDistinctAscList $ fst $ generateDSumsDummy 0 $ adFTK ftk0+      dMap = DMap.empty+      nMap = DMap.empty+  in EvalState {..}+++-- * Reverse pass, transpose/evaluation of the delta expressions++-- | Reverse pass, that is, transpose/evaluation of the delta expressions+-- in order to produce the gradient for the objective function runtime+-- trace represented by the delta expression.+--+-- The first argument is the tensor kind that constrains the shapes+-- of the contangent accumulator and the delta expression arguments.+-- The second is the evaluation state being modified.+-- The third is the cotangent accumulator that will become an actual+-- cotangent contribution when complete (see below for an explanation).+-- The fourth is the delta expression node to evaluate.+--+-- Obtaining the gradient amounts to transposing the linear map+-- that is straightforwardly represented by the delta expression.+-- The @evalRev@ function transposes the linear map and,+-- at the same time, evaluates the transposed map on the cotangent accumulator+-- value contained in the third argument. If the cotangent and the tensor+-- operations are symbolic, the resulting value represents the transposed+-- map itself, if its free variables are treated as the map's inputs.+evalRev+  :: forall y target. (ADReadyNoLet target, ShareTensor target)+  => EvalState target -> target (ADTensorKind y) -> Delta target y+  -> EvalState target+evalRev !s !c d0 = case d0 of+  DeltaShare n d | Dict0 <- lemTKAllNumAD (ftkToSTK $ nodeIdToFTK n) ->+    -- In this context, by construction, @d@ is the dual component+    -- of a dual number term. Let's say that, at this point, evaluation+    -- considers position (node) p out of possibly multiple positions+    -- at which that dual number resides in the whole term tree+    -- of the dual number representation of the objective function.+    -- (Equivalently, considers edge p, one of many leading to the only+    -- node with identifier @n@ in the DAG representing the term).+    -- If so, the @c@ argument of @eval0@ is the cotangent+    -- contribution for position p, that is, the partial derivative+    -- of the objective function with respect to position p.+    --+    -- If there are indeed multiple such positions+    -- (the term is non-trivially shared) then,+    -- over the course of evaluation, cotangent contributions+    -- of them all are gradually accumulated in the finite+    -- maps and eventually their total sum represents the total+    -- influence of the objective function's subcomputation+    -- (more precisely, subgraph of the data flow graph in question)+    -- corresponding to the shared term @DeltaShare n d@. This total+    -- influence over the objective function's behaviour is called+    -- in short the cotangent of the node identifier @n@.+    -- In other words, the cotangent of @n@ is the sum,+    -- over all positions (edges) q in the global delta-expression DAG+    -- that are a reference to node @n@, of the partial derivative+    -- of the objective function with respect to the subcomputation+    -- corresponding to @q@ (meaning, subcomputations denoted by+    -- Haskell terms whose dual components are @Share n ...@).+    --+    -- For @Input@ terms, the eventual lists of cotangents end up+    -- in the cells of the gradient vectors that are the final+    -- result of the evaluation.+    assert (case d of  -- should match shareDelta+              DeltaZero{} -> False+              DeltaPair{} -> False+              DeltaInput{} -> False+              DeltaShare{} -> False+              _ -> True)+    $ if DMap.member n $ nMap s+      then let addc x =+                 Cotangent $ taddTarget (adSTK $ ftkToSTK $ nodeIdToFTK n)+                                        c (unCotangent x)+             -- target has a ShareTensor instance, so taddTarget arguments+             -- don't need to be duplicable+           in s {dMap = DMap.adjust addc n $ dMap s}+      else let cd = Cotangent c+           in s { nMap = DMap.insert n d $ nMap s+                , dMap = DMap.insert n cd $ dMap s }+  DeltaInput i | Dict0 <- lemTKAllNumAD (ftkToSTK $ inputIdToFTK i) ->+    -- This is true for DeltaInput by construction:+    gcastWith (unsafeCoerceRefl :: ADTensorKind y :~: y) $+    let cs = TOTensor c+    in s {iMap = DMap.adjust (addTensorOrZero (ftkToSTK $ inputIdToFTK i) cs) i+                 $ iMap s}+    -- Note that we can't express sharing by inserting DeltaShare constructors+    -- into iMap, because often sharing needs to work across many+    -- iMap keys. That's why global sharing is used.++  DeltaPair d1 d2 ->+    let (c1, c2) = tunpair c+    in evalRev (evalRev s c1 d1) c2 d2+  DeltaProject1 d -> case ftkDelta d of+    FTKProduct _ ftk2 ->+      let zero = tdefTarget $ adFTK ftk2+      in evalRev s (tpair c zero) d+    -- if y is, e.g., TKR 0 Int, we eval this delta anyway, even though+    -- we could ignore it at the price of complicating or duplicating+    -- the code slightly more+  DeltaProject2 d -> case ftkDelta d of+    FTKProduct ftk1 _ ->+      let zero = tdefTarget $ adFTK ftk1+      in evalRev s (tpair zero c) d+  DeltaMapAccumL k (FTKScalar @z1) eftk+                 as es _df rf acc0' es' -- special case to speed up folds+   | Just Refl <- testEquality (typeRep @z1) (typeRep @Z1)+   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->+    let accftk = ftkDelta acc0'+        accftkAD = adFTK accftk+        eftkAD = adFTK eftk+        (c0, _crest) = tunpair c+        dacc_des =+          tmapAccumR (Proxy @target)+                     k accftkAD eftkAD (FTKProduct accftk eftk)+                     (\dx acc_e ->+                        unHFun rf (tpair (tpair dx (tkconcrete Z1)) acc_e))+                     c0+                     (tpair as es)+        (dacc, des) = tunpair dacc_des+        s2 = evalRev s dacc acc0'+    in evalRev s2 des es'+  DeltaMapAccumL k bftk eftk as es _df rf acc0' es'+   | Refl <- lemBuildOfAD k (ftkToSTK bftk)+   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->+    let accftk = ftkDelta acc0'+        accftkAD = adFTK accftk+        bftkAD = adFTK bftk+        eftkAD = adFTK eftk+        (c0, crest) = tunpair c+        dacc_des =+          tmapAccumR (Proxy @target)+                     k accftkAD eftkAD (FTKProduct bftkAD+                                                   (FTKProduct accftk eftk))+                     (\dx db_acc_e ->+                        ttlet db_acc_e $ \ !db_acc_e1 ->+                          unHFun rf (tpair (tpair dx (tproject1 db_acc_e1))+                                           (tproject2 db_acc_e1)))+                     c0+                     (tpair crest (tpair as es))+        (dacc, des) = tunpair dacc_des+        s2 = evalRev s dacc acc0'+    in evalRev s2 des es'++  DeltaZero{} -> s+  DeltaScale (NestedTarget k) d ->+    let z = ftkDelta d+    in case matchingFTK z (adFTK z) of+         Just Refl -> evalRev s (k * c) d+         _ -> s+  DeltaAdd d e ->+    let cShared = tshare c+    in evalRev (evalRev s cShared d) cShared e++  DeltaCastK d -> evalRev s (tkcast c) d++  DeltaCastR d -> evalRev s (trcast c) d+  DeltaDot0R @r v d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKR sh FTKScalar | SNat <- shrRank sh ->+           evalRev s (v * trreplicate0N sh c) d)+      s+  DeltaIndex0R @_ @r d ix ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKR shm FTKScalar | SNat <- shrRank shm ->+           evalRev s (troneHot shm (rfromK c) ix) d)+      s+  DeltaIndexR SNat d ix -> case ftkDelta d of+    FTKR sh x | SNat <- ixrRank ix+              , Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (troneHot (shrTake sh) c ix) d+        -- note that this is the correct derivative also in OOB base,+        -- because then indexing ignores the argument array+        -- and so the array is zeroed in the derivative+  DeltaFromVectorLinearR @_ @r _shm ld ->+    ifDifferentiable @r+      (let cShared = tshare c+           cxs = trtoListLinear cShared+       in foldl' (\ !s2 (cx, d2) -> evalRev s2 cx d2) s+          $ zip cxs (V.toList ld))+      s+  DeltaFromVectorR shm ld -> case V.uncons ld of+    Nothing -> error "evalRev: empty vector"+    Just (d, _) | FTKR shn x <- ftkDelta d+                , SNat <- shrRank shn ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let cShared = tshare c+          cxs = trunravelToListN shm cShared+      in foldl' (\ !s2 (cx, d2) -> evalRev s2 cx d2) s+         $ zip cxs (V.toList ld)+  DeltaSum0R d -> case adFTK $ ftkDelta d of+    FTKR shm FTKScalar | SNat <- shrRank shm ->+      evalRev s (trreplicate0N shm c) d+  DeltaSumR @m SNat d -> case ftkDelta d of+    FTKR shmshn x | SNat <- shrRank shmshn ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (trreplicateN (shrTake @m shmshn) c) d+  DeltaScatterR SNat SNat SNat _sh d f -> case ftkDelta d of+    FTKR sh x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (trgather (shrTake sh) c f) d+  DeltaReplicate0NR @_ @r shm d | SNat <- shrRank shm ->+    ifDifferentiable @r+      (evalRev s (trsum0 c) d)+      s+  DeltaReplicateR shm d -> case ftkDelta d of+    ftk@(FTKR shn x) | SNat <- shrRank shm+                     , SNat <- shrRank shn+                     , Dict0 <- lemTKAllNumAD (ftkToSTK ftk) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (trsumN c) d+  DeltaGatherR SNat SNat SNat _sh d f -> case ftkDelta d of+    FTKR sh x | Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (trscatter (shrTake sh) c f) d+  DeltaAppendR d e -> case (ftkDelta d, ftkDelta e) of+    (FTKR (m :$: _) x, FTKR (n :$: _) _) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let cShared = tshare c+          s2 = evalRev s (trslice 0 m cShared) d+      in evalRev s2 (trslice m n cShared) e+    _ -> error "evalRev: impossible pattern needlessly required"+  -- Depite the warning, the pattern match is exhaustive and if a dummy+  -- pattern is added, GHC 9.14.1 complains about that, in turn.+  DeltaSliceR i n d -> case ftkDelta d of+    FTKR (l :$: rest) x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (trappend+                   (tdefTarget (FTKR (i :$: rest) (adFTK x)))+                   (trappend c+                      (tdefTarget (FTKR (l - i - n :$: rest) (adFTK x))))) d+  DeltaReverseR d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (trreverse c) d+  DeltaTransposeR perm d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let permR = permRInverse perm+      in evalRev s (trtranspose permR c) d+  DeltaReshapeR _sh2 d -> case ftkDelta d of+    FTKR sh x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (trreshape sh c) d++  DeltaCastS d -> evalRev s (tscast c) d+  DeltaDot0S @r v d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKS sh FTKScalar->+           evalRev s (v * tsreplicate0N sh c) d)+      s+  DeltaIndex0S @shm @r d ix | Refl <- lemAppNil @shm ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKS shm FTKScalar ->+           withKnownShS shm $+           evalRev s (tsoneHot (sfromK c) ix) d)+      s+  DeltaIndexS @shm @shn shn d ix -> case ftkDelta d of+    FTKS sh x | Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shn $+      withKnownShS (shsTakeIx @shm @shn Proxy ix sh) $+      evalRev s (tsoneHot c ix) d+  DeltaFromVectorLinearS @_ @r _shm ld ->+    ifDifferentiable @r+      (let cShared = tshare c+           cxs = tstoListLinear cShared+       in foldl' (\ !s2 (cx, d2) -> evalRev s2 cx d2) s+          $ zip cxs (V.toList ld))+      s+  DeltaFromVectorS shm ld -> case V.uncons ld of+    Nothing -> error "evalRev: empty vector"+    Just (d, _) | FTKS shn x <- ftkDelta d ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shn $+      let cShared = tshare c+          cxs = tsunravelToListN shm cShared+      in foldl' (\ !s2 (cx, d2) -> evalRev s2 cx d2) s+         $ zip cxs (V.toList ld)+  DeltaSum0S d -> case adFTK $ ftkDelta d of+    FTKS shm FTKScalar ->+      evalRev s (tsreplicate0N shm c) d+  DeltaSumS shm d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (tsreplicateN shm c) d+  DeltaScatterS @shm @shn shm shn shp d f -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      withKnownShS shp $+      evalRev s (tsgather @_ @shm @shn c f) d+  DeltaReplicate0NS @_ @r shm d ->+    ifDifferentiable @r+      (withKnownShS shm $+       evalRev s (tssum0 c) d)+      s+  DeltaReplicateS @shm shm d -> case ftkDelta d of+    ftk@(FTKS shn x) | Dict0 <- lemTKAllNumAD (ftkToSTK ftk) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      evalRev s (tssumN @_ @shm c) d+  DeltaGatherS @shm @shn shm shn shp d f -> case ftkDelta d of+    FTKS _ x | Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      withKnownShS shp $+      evalRev s (tsscatter @_ @shm @shn c f) d+  DeltaAppendS d e -> case (ftkDelta d, ftkDelta e) of+    (FTKS (msnat :$$ _) x, FTKS (nsnat :$$ _) _) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let cShared = tshare c+          s2 = evalRev s (tsslice (SNat @0) msnat nsnat cShared) d+      in evalRev s2 (tsslice msnat nsnat SNat cShared) e+  DeltaSliceS i@SNat _ k@SNat d -> case ftkDelta d of+    FTKS (_ :$$ sh) x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (tsappend+                       (tdefTarget (FTKS (i :$$ sh) (adFTK x)))+                          (tsappend+                             c (tdefTarget (FTKS (k :$$ sh) (adFTK x))))) d+  DeltaReverseS d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (tsreverse c) d+  DeltaTransposeS @perm @sh2 perm d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      permInverse perm $ \(permRev :: Permutation.Perm permR) _ ->+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix+                        permR (Permutation.PermutePrefix perm sh2) :~: sh2)+        $ gcastWith (unsafeCoerceRefl+                     :: Rank (Permutation.PermutePrefix perm sh2) :~: Rank sh2)+        $ gcastWith (unsafeCoerceRefl+                     :: Rank permR :~: Rank perm)+        $ evalRev s (tstranspose permRev c) d+  DeltaReshapeS _sh2 d -> case ftkDelta d of+    FTKS sh x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (tsreshape sh c) d++  DeltaCastX d -> evalRev s (txcast c) d+  DeltaDot0X @r v d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKX sh FTKScalar ->+           evalRev s (v * txreplicate0N sh c) d)+      s+  DeltaIndex0X @shm @r d ix | Refl <- lemAppNil @shm ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKX shm FTKScalar ->+           withKnownShX (ssxFromShX shm) $+           evalRev s (txoneHot shm (xfromK c) ix) d)+      s+  DeltaIndexX @shm @shn shn d ix -> case ftkDelta d of+    FTKX sh x | SNat @len <- ixxRank ix+              , Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX shn $+      withKnownShX (ssxTakeIx @shm @shn Proxy ix (ssxFromShX sh)) $+      gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+      evalRev s (txoneHot (shxTake @len sh) c ix) d+  DeltaFromVectorLinearX @_ @r _shm ld ->+    ifDifferentiable @r+      (let cShared = tshare c+           cxs = txtoListLinear cShared+       in foldl' (\ !s2 (cx, d2) -> evalRev s2 cx d2) s+          $ zip cxs (V.toList ld))+      s+  DeltaFromVectorX shm ld -> case V.uncons ld of+    Nothing -> error "evalRev: empty vector"+    Just (d, _) | FTKX shn x <- ftkDelta d ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX (ssxFromShX shn) $+      let cShared = tshare c+          cxs = txunravelToListN shm cShared+      in foldl' (\ !s2 (cx, d2) -> evalRev s2 cx d2) s+         $ zip cxs (V.toList ld)+  DeltaSum0X d -> case adFTK $ ftkDelta d of+    FTKX shm FTKScalar ->+      evalRev s (txreplicate0N shm c) d+  DeltaSumX @_ @shn shm d -> case ftkDelta d of+    FTKX shmshn x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (txreplicateN (shxTakeSSX (Proxy @shn) shm shmshn) c) d+  DeltaScatterX @shm @shn shm shn shp _sh d f -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX shm $+      withKnownShX shn $+      withKnownShX shp $+      evalRev s (txgather @_ @shm @shn (shxTakeSSX @_ @shn Proxy shm sh) c f) d+  DeltaReplicate0NX @_ @r shm d ->+    ifDifferentiable @r+      (withKnownShX (ssxFromShX shm) $+       evalRev s (txsum0 c) d)+      s+  DeltaReplicateX @shm shm d -> case ftkDelta d of+    ftk@(FTKX shn x) | Dict0 <- lemTKAllNumAD (ftkToSTK ftk) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX (ssxFromShX shm) $+      withKnownShX (ssxFromShX shn) $+      evalRev s (txsumN @_ @shm c) d+  DeltaGatherX @shm @shn shm shn shp _sh d f -> case ftkDelta d of+    FTKX sh x | Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX shm $+      withKnownShX shn $+      withKnownShX shp $+      evalRev s (txscatter @_ @shm @shn (shxTakeSSX @_ @shn Proxy shp sh) c f) d+  DeltaAppendX @m @n d e -> case (ftkDelta d, ftkDelta e) of+    (FTKX (m :$% _) x, FTKX (n :$% _) _) ->+      gcastWith (unsafeCoerceRefl+                 :: AddMaybe (AddMaybe (Just 0) m) n :~: AddMaybe m n) $+      withKnownSTK (adSTK $ ftkToSTK x) $+      let cShared = tshare c+          s2 =+               evalRev s (txslice (SKnown (SNat @0)) m n cShared) d+      in gcastWith (unsafeCoerceRefl+                    :: AddMaybe (AddMaybe m n) (Just 0) :~: AddMaybe m n) $+         evalRev s2 (txslice m n (SKnown (SNat @0)) cShared) e+  DeltaSliceX @i @n @k i _ k d -> case ftkDelta d of+    FTKX (_ :$% sh) x ->+      gcastWith (unsafeCoerceRefl+                :: AddMaybe (AddMaybe i n) k :~: AddMaybe i (AddMaybe n k)) $+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (txappend+                   (tdefTarget (FTKX (i :$% sh) (adFTK x)))+                   (txappend+                      c (tdefTarget+                           (FTKX (k :$% sh) (adFTK x))))) d+  DeltaReverseX d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (txreverse c) d+  DeltaTransposeX @perm @sh2 perm d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      permInverse perm $ \(permR :: Permutation.Perm permR) _ ->+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix+                        permR (Permutation.PermutePrefix perm sh2) :~: sh2) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (Permutation.PermutePrefix perm sh2) :~: Rank sh2) $+        gcastWith (unsafeCoerceRefl+                   :: Rank permR :~: Rank perm) $+        evalRev s (txtranspose permR c) d+  DeltaReshapeX _sh2 d -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      evalRev s (txreshape sh c) d++  DeltaConvert c1 d ->+    let z = ftkDelta d+    in case matchingFTK z (adFTK z) of+         Just Refl ->+           -- This follows from the property of conversions reflecting+           -- the underlying scalar types unchanged.+           gcastWith (unsafeCoerceRefl :: ADTensorKind y :~: y) $+           evalRev s (tconvert (transposeTKConversion (ftkDelta d) c1)+                   (convertSTK c1 $ ftkToSTK $ ftkDelta d) c) d+         _ -> s++evalRevFromnMap :: forall target. (ADReadyNoLet target, ShareTensor target)+                => EvalState target -> EvalState target+evalRevFromnMap s@EvalState{nMap, dMap} =+  case DMap.maxViewWithKey nMap of+    Just (n :=> d, nMap2) ->+      let s2 = s {nMap = nMap2}+          s3 = case DMap.lookup n dMap of+            Just (Cotangent c) -> evalRev s2 c d+            Nothing -> error $ "evalRevFromnMap: missing cotangent " ++ show n+      in evalRevFromnMap s3+    Nothing -> s  -- loop ends++{- TODO: optimize similarly?+        -- The general case is given as the last one below,+        -- but for a few constructors it's faster to inline @evalRev@ instead.+        -- BTW, such an optimization doesn't really belong in the simplified+        -- horde-ad and no consistent benefit should be expected here.+        DeltaIndex0 ZeroR{} _ _ -> s  -- shortcut+        DeltaIndex0 (InputR i) ixs' sh ->+          let ixs = indexToList ixs'+              f v = if isTensorDummy v+                    then treplicate0ND sh 0 `OD.update` [(ixs, c)]+                    else v `OD.update` [(ixs, v `rindex0D` ixs + c)]+          in s {iMap = DMap.adjust f i $ iMap s}+        DeltaIndex0 (ShareR n d) ixs' sh ->+          let ixs = indexToList ixs'+          in case DMap.lookup n $ nMap s of+            Just (DynamicRanked _) ->+              let f v = v `OD.update` [(ixs, v `rindex0D` ixs + c)]+              in s {dMap = DMap.adjust f n $ dMap s}+                -- This would be an asymptotic optimization compared to+                -- the general case below, if not for the non-mutable update,+                -- which implies copying the whole @v@ vector,+                -- so it's only several times faster (same allocation,+                -- but not adding to each cell of @v@).+            Nothing ->+              let v = treplicate0ND sh 0 `OD.update` [(ixs, c)]+              in s { nMap = DMap.insert n (DynamicRanked d) $ nMap s+                   , dMap = DMap.insert n v $ dMap s }+            _ -> error "evalRevFromnMap: corrupted nMap"+-}+++-- * Forward derivative computation from the delta expressions++-- | Forward derivative computation via forward-evaluation of delta-expressions+-- (which is surprisingly competitive to the direct forward method,+-- until the allocation of deltas gets large enough to affect cache hits).+-- This is the directional derivative, calculated for the point,+-- at which the delta expression was computed (which is the point+-- represented by the parameters of the objective function and used+-- to compute it's dual number result) and along the direction vector+-- given by the parameters in the arguments.+--+-- This mimics the reverse derivative code, but in reverse. Perhaps this can be+-- simplified, but the obvious simplest formulation does not honour sharing+-- and evaluates shared subexpressions repeatedly, so this state-passing+-- formulation is adopted.+evalFwd+  :: forall target y. (ADReadyNoLet target, ShareTensor target)+  => IMap target -> ADMap target -> Delta target y+  -> (ADMap target, target (ADTensorKind y))+evalFwd params !s d0 = case d0 of+  DeltaShare n d ->+    case DMap.lookup n s of+      Just e1 -> (s, unCotangent e1)+      Nothing ->+        let (s2, cRaw) = evalFwd params s d+            cShared = tshare cRaw+            cd = Cotangent cShared+              -- cRaw is shared, because it's put into the map and then+              -- potentially looked up many times, so it'd get duplicated+            s3 = DMap.insert n cd s2+        in (s3, cShared)+  DeltaInput inputId ->+    case DMap.lookup inputId params of+      Just dtk -> (s, toADTensorKindShared (inputIdToFTK inputId)+                      $ evalTensorOrZero dtk)+      Nothing -> error "evalFwd: missing input"++  DeltaPair d1 d2 ->+    let (s2, t) = evalFwd params s d1+        (s3, u) = evalFwd params s2 d2+    in (s3, tpair t u)+  DeltaProject1 d ->+    let (s2, v) = evalFwd params s d+    in (s2, tproject1 v)+  DeltaProject2 d ->+    let (s2, v) = evalFwd params s d+    in (s2, tproject2 v)+  -- No special case to speed up folds, because the concrete tmapAccumLDer+  -- does the very Z1 optimization that could be performed here.+  DeltaMapAccumL k bftk eftk as es df _rf acc0' es'+   | Refl <- lemBuildOfAD k (ftkToSTK bftk)+   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->+    let accftk = ftkDelta acc0'+        accftkAD = adFTK accftk+        bftkAD = adFTK bftk+        eftkAD = adFTK eftk+        (s2, cacc0) = evalFwd params s acc0'+        (s3, ces) = evalFwd params s2 es'+    in (s3, tmapAccumL (Proxy @target)+                       k accftkAD bftkAD (FTKProduct eftkAD+                                                     (FTKProduct accftk eftk))+                       (\dacc de_acc_e ->+                        ttlet de_acc_e $ \ !de_acc_e1 ->+                          unHFun df (tpair (tpair dacc (tproject1 de_acc_e1))+                                           (tproject2 de_acc_e1)))+                       cacc0+                       (tpair ces (tpair as es)))++  DeltaZero ftk -> (s, tdefTarget $ adFTK ftk)+  DeltaScale (NestedTarget k) d ->+    let z = ftkDelta d+    in case matchingFTK z (adFTK z) of+         Just Refl -> second (* k) $ evalFwd params s d+         _ -> (s, tdefTarget $ adFTK z)+  DeltaAdd d e ->+    let y = ftkDelta d0  -- either d or e may be DeltaShare+    in case matchingFTK y (adFTK y) of+         Just Refl ->+           let (s2, t) = evalFwd params s d+               (s3, u) = evalFwd params s2 e+           in (s3, t + u)+         _ -> (s, tdefTarget $ adFTK y)++  DeltaCastK d -> second tkcast $ evalFwd params s d++  DeltaCastR d -> second trcast $ evalFwd params s d+  DeltaDot0R @r _ DeltaZero{} -> (s, toADTensorKindShared (FTKScalar @r) 0)+  DeltaDot0R @r v d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKR sh FTKScalar | SNat <- shrRank sh ->+           second (trdot0 v) $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaIndex0R @_ @r d ix ->+    ifDifferentiable @r+      (second (`trindex0` ix) $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaIndexR SNat d ix -> case ftkDelta d of+    FTKR _ x | SNat <- ixrRank ix ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (`trindex` ix) $ evalFwd params s d+  DeltaFromVectorLinearR shm ld -> case V.uncons ld of+    Nothing -> error "evalFwd: empty vector"+    Just (d, _) | FTKScalar <- adFTK $ ftkDelta d ->+      let (s2, l) = mapAccumL' (evalFwd params) s $ V.toList ld+      in (s2, trfromVectorLinear shm $ V.fromListN (V.length ld) l)+  DeltaFromVectorR shm ld -> case V.uncons ld of+    Nothing -> error "evalFwd: empty vector"+    Just (d, _) | FTKR shn x <- ftkDelta d+                , SNat <- shrRank shn ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let (s2, l) = mapAccumL' (evalFwd params) s $ V.toList ld+      in (s2, trfromVectorN shm $ V.fromListN (V.length ld) l)+  DeltaSum0R @_ @r DeltaZero{} -> (s, toADTensorKindShared (FTKScalar @r) 0)+  DeltaSum0R @_ @r d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKR shm FTKScalar | SNat <- shrRank shm ->+           second trsum0 $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaSumR SNat d -> case ftkDelta d of+    ftk@(FTKR shmshn x) | SNat <- shrRank shmshn+                        , Dict0 <- lemTKAllNumAD (ftkToSTK ftk) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second trsumN $ evalFwd params s d+  DeltaScatterR SNat SNat SNat sh d f -> case ftkDelta d of+    FTKR _ x | Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let (s2, t) = evalFwd params s d+      in (s2, trscatter sh t f)+  DeltaReplicate0NR shm d | SNat <- shrRank shm+                          , FTKScalar <- adFTK $ ftkDelta d ->+    second (trreplicate0N shm) $ evalFwd params s d+  DeltaReplicateR shm d -> case ftkDelta d of+    FTKR shn x | SNat <- shrRank shn ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (trreplicateN shm) $ evalFwd params s d+  DeltaGatherR SNat SNat SNat sh d f -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let (s2, t) = evalFwd params s d+      in (s2, trgather sh t f)+  DeltaAppendR d e -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let (s2, t) = evalFwd params s d+          (s3, u) = evalFwd params s2 e+      in (s3, trappend t u)+  DeltaSliceR i n d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (trslice i n) $ evalFwd params s d+  DeltaReverseR d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second trreverse $ evalFwd params s d+  DeltaTransposeR perm d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (trtranspose perm) $ evalFwd params s d+  DeltaReshapeR sh2 d -> case ftkDelta d of+    FTKR _sh x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (trreshape sh2) $ evalFwd params s d++  DeltaCastS d -> second tscast $ evalFwd params s d+  DeltaDot0S @r _ DeltaZero{} -> (s, toADTensorKindShared (FTKScalar @r) 0)+  DeltaDot0S @r v d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKS sh FTKScalar ->+           withKnownShS sh $+           second (tsdot0 v) $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaIndex0S @_ @r d ix ->+    ifDifferentiable @r+      (second (`tsindex0` ix) $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaIndexS shn d ix -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shn $+      second (`tsindex` ix) $ evalFwd params s d+  DeltaFromVectorLinearS shm ld -> case V.uncons ld of+    Nothing -> error "evalFwd: empty vector"+    Just (d, _) | FTKScalar <- adFTK $ ftkDelta d ->+      let (s2, l) = mapAccumL' (evalFwd params) s $ V.toList ld+      in (s2, tsfromVectorLinear shm $ V.fromListN (V.length ld) l)+  DeltaFromVectorS shm ld -> case V.uncons ld of+    Nothing -> error "evalFwd: empty vector"+    Just (d, _) | FTKS shn x <- ftkDelta d ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shn $+      let (s2, l) = mapAccumL' (evalFwd params) s $ V.toList ld+      in (s2, tsfromVectorN shm $ V.fromListN (V.length ld) l)+  DeltaSum0S @_ @r DeltaZero{} -> (s, toADTensorKindShared (FTKScalar @r) 0)+  DeltaSum0S @_ @r d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKS shm FTKScalar ->+           withKnownShS shm $+           second tssum0 $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaSumS @shm @shn shm d | SNat <- shsRank shm -> case ftkDelta d of+    ftk@(FTKS shmshn x) | Dict0 <- lemTKAllNumAD (ftkToSTK ftk) ->+      gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+      gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS (shsTake @(Rank shm) shmshn) $+      withKnownShS (shsDrop @(Rank shm) shmshn) $+      second (tssumN @_ @shm) $ evalFwd params s d+  DeltaScatterS @shm @shn shm shn shp d f -> case ftkDelta d of+    FTKS _ x | Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      withKnownShS shp $+      let (s2, t) = evalFwd params s d+      in (s2, tsscatter @_ @shm @shn t f)+  DeltaReplicate0NS shm d | FTKScalar <- adFTK $ ftkDelta d ->+    second (tsreplicate0N shm) $ evalFwd params s d+  DeltaReplicateS shm d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (tsreplicateN shm) $ evalFwd params s d+  DeltaGatherS @shm @shn shm shn shp d f -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      withKnownShS shp $+      let (s2, t) = evalFwd params s d+      in (s2, tsgather @_ @shm @shn t f)+  DeltaAppendS d e -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let (s2, t) = evalFwd params s d+          (s3, u) = evalFwd params s2 e+      in (s3, tsappend t u)+  DeltaSliceS i n k d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (tsslice i n k) $ evalFwd params s d+  DeltaReverseS d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second tsreverse $ evalFwd params s d+  DeltaTransposeS perm d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (tstranspose perm) $ evalFwd params s d+  DeltaReshapeS sh2 d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (tsreshape sh2) $ evalFwd params s d++  DeltaCastX d -> second txcast $ evalFwd params s d+  DeltaDot0X @r _ DeltaZero{} -> (s, toADTensorKindShared (FTKScalar @r) 0)+  DeltaDot0X @r v d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKX sh FTKScalar ->+           withKnownShX (ssxFromShX sh) $+           second (txdot0 v) $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaIndex0X @_ @r d ix ->+    ifDifferentiable @r+      (second (`txindex0` ix) $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaIndexX @shm @shn shn d ix -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX shn $+      withKnownShX (ssxTakeIx @shm @shn Proxy ix (ssxFromShX sh)) $+      second (`txindex` ix) $ evalFwd params s d+  DeltaFromVectorLinearX shm ld -> case V.uncons ld of+    Nothing -> error "evalFwd: empty vector"+    Just (d, _) | FTKScalar <- adFTK $ ftkDelta d ->+      let (s2, l) = mapAccumL' (evalFwd params) s $ V.toList ld+      in (s2, txfromVectorLinear shm $ V.fromListN (V.length ld) l)+  DeltaFromVectorX shm ld -> case V.uncons ld of+    Nothing -> error "evalFwd: empty vector"+    Just (d, _) | FTKX shn x <- ftkDelta d ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX (ssxFromShX shn) $+      let (s2, l) = mapAccumL' (evalFwd params) s $ V.toList ld+      in (s2, txfromVectorN shm $ V.fromListN (V.length ld) l)+  DeltaSum0X @_ @r DeltaZero{} -> (s, toADTensorKindShared (FTKScalar @r) 0)+  DeltaSum0X @_ @r d ->+    ifDifferentiable @r+      (case ftkDelta d of+         FTKX shm FTKScalar ->+           withKnownShX (ssxFromShX shm) $+           second txsum0 $ evalFwd params s d)+      (s, tkconcrete Z1)+  DeltaSumX @shm @shn shm d | SNat <- ssxRank shm -> case ftkDelta d of+    ftk@(FTKX shmshn x) | Dict0 <- lemTKAllNumAD (ftkToSTK ftk) ->+      gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+      gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX (ssxFromShX $ shxTake @(Rank shm) shmshn) $+      withKnownShX (ssxFromShX $ shxDrop @(Rank shm) shmshn) $+      second (txsumN @_ @shm) $ evalFwd params s d+  DeltaScatterX @shm @shn shm shn shp sh d f -> case ftkDelta d of+    FTKX _ x | Dict0 <- lemTKAllNumAD (ftkToSTK x) ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX shm $+      withKnownShX shn $+      withKnownShX shp $+      let (s2, t) = evalFwd params s d+      in (s2, txscatter @_ @shm @shn sh t f)+  DeltaReplicate0NX shm d | FTKScalar <- adFTK $ ftkDelta d ->+    second (txreplicate0N shm) $ evalFwd params s d+  DeltaReplicateX shm d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (txreplicateN shm) $ evalFwd params s d+  DeltaGatherX @shm @shn shm shn shp sh d f -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      withKnownShX shm $+      withKnownShX shn $+      withKnownShX shp $+      let (s2, t) = evalFwd params s d+      in (s2, txgather @_ @shm @shn sh t f)+  DeltaAppendX d e -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      let (s2, t) = evalFwd params s d+          (s3, u) = evalFwd params s2 e+      in (s3, txappend t u)+  DeltaSliceX i n k d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (txslice i n k) $ evalFwd params s d+  DeltaReverseX d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second txreverse $ evalFwd params s d+  DeltaTransposeX perm d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (txtranspose perm) $ evalFwd params s d+  DeltaReshapeX sh2 d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (adSTK $ ftkToSTK x) $+      second (txreshape sh2) $ evalFwd params s d++  DeltaConvert c1 d ->+    let z = ftkDelta d+    in case matchingFTK z (adFTK z) of+         Just Refl ->+           -- This follows from the property of conversions reflecting+           -- the underlying scalar types unchanged.+           gcastWith (unsafeCoerceRefl :: ADTensorKind y :~: y) $+           second (tconvert c1 (ftkToSTK (ftkDelta d)))+                  (evalFwd params s d)+         _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)
src/HordeAd/Core/Ops.hs view
@@ -25,1182 +25,1310 @@   , AllTargetShow, CommonTargetEqOrd     -- * Helper functions   , rtr, rflatten, str, sflatten, xtr, xflatten-  , tmapAccumR, tmapAccumL-  , rbuild, sbuild, xbuild-    -- * Helper classes and types-  , IntegralHAndIntElt, RealFloatAndFloatElt-  , TensorSupportsX, TensorSupportsS, TensorSupportsR, TensorSupports-  ) where--import Prelude--import Data.Foldable qualified as Foldable-import Data.Int (Int64)-import Data.Kind (Constraint, Type)-import Data.Maybe (fromMaybe)-import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))-import Data.Vector.Generic qualified as V-import Data.Vector.Strict qualified as Data.Vector-import GHC.Exts (IsList (..))-import GHC.TypeLits (KnownNat, type (+), type (<=), type (<=?))-import Type.Reflection (typeRep)--import Data.Array.Nested (type (++))-import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Convert (withShsFromShX)-import Data.Array.Nested.Lemmas-import Data.Array.Nested.Mixed.Shape-import Data.Array.Nested.Permutation qualified as Permutation-import Data.Array.Nested.Ranked.Shape-import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (Init, unsafeCoerceRefl)--import HordeAd.Core.CarriersConcrete-import HordeAd.Core.ConvertTensor-import HordeAd.Core.TensorKind-import HordeAd.Core.Types---- These user API functions are used in default definitions of methods,--- so they have to be defined already here:--rtr :: forall n x target. (KnownSTK x, BaseTensor target)-    => target (TKR2 (2 + n) x) -> target (TKR2 (2 + n) x)-rtr = trtranspose [1, 0]-rflatten :: forall n x target. (KnownSTK x, BaseTensor target)-         => target (TKR2 n x) -> target (TKR2 1 x)-rflatten u = trreshape (rsize u :$: ZSR) u-str :: forall n m sh x target. (KnownSTK x, BaseTensor target)-    => target (TKS2 (n ': m ': sh) x) -> target (TKS2 (m ': n ': sh) x)-str = gcastWith (unsafeCoerceRefl :: (2 <=? Rank (n ': m ': sh)) :~: True) $-      tstranspose (Permutation.makePerm @'[1, 0])-sflatten :: (KnownShS sh, KnownSTK x, BaseTensor target )-         => target (TKS2 sh x) -> target (TKS2 '[Product sh] x)-sflatten @sh | SNat <- shsProduct (knownShS @sh) = tsreshape knownShS-xtr :: forall n m sh x target. (KnownSTK x, BaseTensor target)-    => target (TKX2 (Just n ': Just m ': sh) x)-    -> target (TKX2 (Just m ': Just n ': sh) x)-xtr = gcastWith (unsafeCoerceRefl-                 :: (2 <=? Rank (Just n ': Just m ': sh)) :~: True) $-      txtranspose (Permutation.makePerm @'[1, 0])-xflatten :: forall sh x target. (KnownSTK x, BaseTensor target)-         => target (TKX2 sh x) -> target (TKX2 '[Nothing] x)-xflatten u = txreshape (Nested.SUnknown (xsize u) :$% ZSX) u--rbuild :: (KnownNat m, KnownNat n, KnownSTK x, BaseTensor target)-       => IShR (m + n)  -- ^ the shape of the resulting tensor-       -> (IxROf target m -> target (TKR2 n x))-            -- ^ the function to build with-       -> target (TKR2 (m + n) x)-rbuild @m @n @x @target sh0 f0 =-  let buildSh :: IShR m1 -> (IxROf target m1 -> target (TKR2 n x))-              -> target (TKR2 (m1 + n) x)-      buildSh ZSR f = f ZIR-      buildSh (k :$: sh) f | SNat <- shrRank sh =-        let g i = buildSh sh (\ix -> f (i :.: ix))-        in trbuild1 k g-  in buildSh (shrTake @m @n sh0) f0-sbuild :: (KnownShS (Take m sh), KnownShS sh, KnownSTK x, BaseTensor target)-       => (IxSOf target (Take m sh) -> target (TKS2 (Drop m sh) x))-            -- ^ the function to build with-       -> target (TKS2 sh x)-sbuild @m @sh @x @target =-  let buildSh-        :: forall sh1.-           ShS sh1 -> ShS (sh1 ++ Drop m sh)-        -> (IxSOf target sh1 -> target (TKS2 (Drop m sh) x))-        -> target (TKS2 (sh1 ++ Drop m sh) x)-      buildSh sh1 sh1m f = case (sh1, sh1m) of-        (ZSS, _) -> f ZIS-        (SNat :$$ sh2, _ :$$ sh2m) ->-          withKnownShS sh2m $-          let g i = buildSh sh2 sh2m (f . (i :.$))-          in tsbuild1 g-  in gcastWith (unsafeCoerceRefl :: sh :~: Take m sh ++ Drop m sh)-     $ buildSh (knownShS @(Take m sh)) (knownShS @sh)-xbuild :: ( KnownShX (Take m sh), KnownSTK x-          , BaseTensor target, ConvertTensor target )-       => IShX sh  -- ^ the shape of the resulting tensor-       -> (IxXOf target (Take m sh) -> target (TKX2 (Drop m sh) x))-            -- ^ the function to build with-       -> target (TKX2 sh x)-xbuild @m @sh @x @target sh0 f0 =-  let buildSh :: IShX sh1 -> IShX (sh1 ++ Drop m sh)-              -> (IxXOf target sh1 -> target (TKX2 (Drop m sh) x))-              -> target (TKX2 (sh1 ++ Drop m sh) x)-      buildSh sh1 sh1m f = case (sh1, sh1m) of-        (ZSX, _) -> f ZIX-        (k :$% sh2, _ :$% sh2m) ->-          withKnownShX (ssxFromShX sh2m) $-          let g i = buildSh sh2 sh2m (f . (i :.%))-          in withSNat (fromSMayNat' k) $ \(SNat @n) ->-               xmcast (ssxFromShX sh1m) $ txbuild1 @_ @n g-  in gcastWith (unsafeCoerceRefl :: sh :~: Take m sh ++ Drop m sh)-     $ buildSh (shxTakeSSX (Proxy @(Drop m sh)) (knownShX @(Take m sh)) sh0) sh0 f0---- | A strict right mapAccum.-tmapAccumR-  :: forall accy by ey k target. BaseTensor target-  => Proxy target-  -> SNat k  -- ^ length of the input-  -> FullShapeTK accy  -- ^ shape of the accumulator-  -> FullShapeTK by  -- ^ shape of the output-  -> FullShapeTK ey  -- ^ shape of an individual input-  -> (forall f. ADReady f-      => f accy -> f ey -> f (TKProduct accy by))-       -- ^ the function to mapAccum with-  -> target accy  -- ^ the initial accumulator-  -> target (BuildTensorKind k ey)  -- ^ the inputs-  -> target (TKProduct accy (BuildTensorKind k by))-{-# INLINE tmapAccumR #-}  -- this doesn't want to specialize-tmapAccumR proxy !k !accftk !bftk !eftk f acc0 es =-  let xftk = FTKProduct accftk eftk-      fl :: forall f. ADReady f-         => f (TKProduct accy ey)-         -> f (TKProduct accy by)-      fl !args = ttlet args $ \ !args1 ->-                   f (tproject1 args1) (tproject2 args1)-  in tmapAccumRDer proxy k accftk bftk eftk-                   (tlambda @target xftk (HFun fl))-                   (tjvp @target xftk $ HFun fl)-                   (tvjp @target xftk $ HFun fl)-                   acc0 es--- | A strict left mapAccum.-tmapAccumL-  :: forall accy by ey k target. BaseTensor target-  => Proxy target-  -> SNat k  -- ^ length of the input-  -> FullShapeTK accy  -- ^ shape of the accumulator-  -> FullShapeTK by  -- ^ shape of the output-  -> FullShapeTK ey  -- ^ shape of an individual input-  -> (forall f. ADReady f-      => f accy -> f ey -> f (TKProduct accy by))-       -- ^ the function to mapAccum with-  -> target accy  -- ^ the initial accumulator-  -> target (BuildTensorKind k ey)  -- ^ the inputs-  -> target (TKProduct accy (BuildTensorKind k by))-{-# INLINE tmapAccumL #-}  -- this doesn't want to specialize-tmapAccumL proxy !k !accftk !bftk !eftk f acc0 es =-  let xftk = FTKProduct accftk eftk-      fl :: forall f. ADReady f-         => f (TKProduct accy ey)-         -> f (TKProduct accy by)-      fl !args = ttlet args $ \ !args1 ->-                   f (tproject1 args1) (tproject2 args1)-  in tmapAccumLDer proxy k accftk bftk eftk-                   (tlambda @target xftk (HFun fl))-                   (tjvp @target xftk $ HFun fl)-                   (tvjp @target xftk $ HFun fl)-                   acc0 es--type TensorSupports :: (Type -> Constraint) -> (Type -> Constraint)-                    -> Target -> Constraint-type TensorSupports c1 c2 f =-  forall r. GoodScalar r-            => c1 r => c2 (f (TKScalar r))--type TensorSupportsR :: (Type -> Constraint) -> (Type -> Constraint)-                     -> Target -> Constraint-type TensorSupportsR c1 c2 f =-  forall r n. GoodScalar r-              => c1 r => c2 (f (TKR n r))--type TensorSupportsS :: (Type -> Constraint) -> (Type -> Constraint)-                     -> Target -> Constraint-type TensorSupportsS c1 c2 f =-  forall r sh. GoodScalar r-               => c1 r => c2 (f (TKS sh r))--type TensorSupportsX :: (Type -> Constraint) -> (Type -> Constraint)-                     -> Target -> Constraint-type TensorSupportsX c1 c2 f =-  forall r sh. GoodScalar r-               => c1 r => c2 (f (TKX sh r))--class (RealFloatH r, Nested.FloatElt r)-      => RealFloatAndFloatElt r-instance (RealFloatH r, Nested.FloatElt r)-         => RealFloatAndFloatElt r--class (IntegralH r, Nested.IntElt r)-      => IntegralHAndIntElt r-instance (IntegralH r, Nested.IntElt r)-      => IntegralHAndIntElt r--class LetTensor (target :: Target) where-  ttlet :: target x -> (target x -> target z) -> target z-  ttletPrimal :: PrimalOf target x -> (PrimalOf target x -> target z)-              -> target z-  toShare :: target y -> ShareOf target y-  tunshare :: ShareOf target y -> target y-  tunshare = error "tunshare: this instance should never be used"-  tappend :: forall m n y. BaseTensor target-          => SNat m -> SNat n -> SingletonTK y-          -> target (BuildTensorKind m y) -> target (BuildTensorKind n y)-          -> target (BuildTensorKind (m + n) y)-  tappend msnat@SNat nsnat@SNat stk a b = case stk of-    STKScalar -> tsappend a b-    STKR _ x | Dict <- lemKnownSTK x -> trappend a b-    STKS _ x | Dict <- lemKnownSTK x -> tsappend a b-    STKX _ x | Dict <- lemKnownSTK x -> txappend a b-    STKProduct stk1 stk2 ->-      ttlet a $ \ !aShared -> ttlet b $ \ !bShared ->-        tpair (tappend msnat nsnat stk1 (tproject1 aShared) (tproject1 bShared))-              (tappend msnat nsnat stk2 (tproject2 aShared) (tproject2 bShared))-  tD :: BaseTensor target-     => SingletonTK y -> PrimalOf target y -> DualOf target y-     -> target y-  tD stk p d =-    -- Lets needed, because taddTarget requires duplicable arguments.-    ttletPrimal p $ \pShared ->-    ttlet (tfromDual d) $ \dShared ->-      taddTarget stk (tfromPrimal stk pShared) dShared-  -- | A strict left fold.-  tfold-    :: forall yn ym k. BaseTensor target-    => SNat k  -- ^ length of the input-    -> SingletonTK yn  -- ^ partial shape of the accumulator-    -> SingletonTK ym  -- ^ partial shape of an individual input-    -> (forall f. ADReady f => f yn -> f ym -> f yn)-         -- ^ the function to fold with-    -> target yn  -- ^ the initial accumulator-    -> target (BuildTensorKind k ym)  -- ^ the inputs-    -> target yn-  {-# INLINE tfold #-}  -- this doesn't want to specialize-  tfold k nstk mstk f acc0 es =-    tproject1-    $ tmapAccumL (Proxy @target)-       k-       (tftk nstk acc0)-       (FTKScalar @Z1)-       (razeFTK k mstk (tftk (buildSTK k mstk) es))-       (let g :: forall f. ADReady f-              => f yn -> f ym -> f (TKProduct yn TKUnit)-            g !acc !e = tpair (f acc e) (tkconcrete Z1)-        in g)-       acc0-       es-  -- | A strict left scan.-  tscan-    :: forall yn ym k. BaseTensor target-    => SNat k  -- ^ length of the input-    -> SingletonTK yn  -- ^ partial shape of the accumulator-    -> SingletonTK ym  -- ^ partial shape of an individual input-    -> (forall f. ADReady f => f yn -> f ym -> f yn)-         -- ^ the function to scan with-    -> target yn  -- ^ the initial accumulator-    -> target (BuildTensorKind k ym)  -- ^ the inputs-    -> target (BuildTensorKind (1 + k) yn)-  {-# INLINE tscan #-}  -- this doesn't want to specialize-  tscan k nstk mstk f acc0 es =-    let bs :: target (BuildTensorKind k yn)-        bs = tproject2-             $ tmapAccumL (Proxy @target)-                k-                (tftk nstk acc0)-                (tftk nstk acc0)-                (razeFTK k mstk (tftk (buildSTK k mstk) es))-              (let g :: forall f. ADReady f-                     => f yn -> f ym -> f (TKProduct yn yn)-                   g !acc !e = ttlet (f acc e) $ \ !res -> tpair res res-               in g)-              acc0-              es-    in tappend (SNat @1) k nstk-               (tfromVector (SNat @1) nstk (V.fromList [acc0])) bs--class ShareTensor (target :: Target) where-  tshare :: target y -> target y-  tunpair :: target (TKProduct x z) -> (target x, target z)-  -- This would suffer from lack of sharing with LetTensor, because-  -- ttlet doesn't work over a list. With sharing it's fine.-  tunravelToListShare :: forall y k. (BaseTensor target, ConvertTensor target)-                      => SNat k -> SingletonTK y-                      -> target (BuildTensorKind k y)-                      -> [target y]-  tunravelToListShare snat@SNat stk u = case stk of-    STKScalar -> let !uShared = tshare u-                 in map kfromS $ tsunravelToList uShared-    STKR SNat x | Dict <- lemKnownSTK x -> let !uShared = tshare u-                                           in trunravelToList uShared-    STKS sh x | Dict <- lemKnownSTK x -> let !uShared = tshare u-                                         in withKnownShS sh-                                            $ tsunravelToList uShared-    STKX sh x | Dict <- lemKnownSTK x -> let !uShared = tshare u-                                         in withKnownShX sh-                                            $ txunravelToList uShared-    STKProduct stk1 stk2 ->-      let (!u1, !u2) = tunpair u-      in zipWith tpair (tunravelToListShare snat stk1 u1)-                       (tunravelToListShare snat stk2 u2)---- | The superclasses indicate that it's not only a container array,--- but also a mathematical tensor, sporting numeric operations.-class ( Num (IntOf target)-      , IntegralH (IntOf target)-      , TensorSupports Num Num target-      , TensorSupports RealFloatAndFloatElt Floating target-      , TensorSupports RealFloatAndFloatElt RealFloatH target-      , TensorSupports IntegralHAndIntElt IntegralH target-      , TensorSupportsR Num Num target-      , TensorSupportsR RealFloatAndFloatElt Floating target-      , TensorSupportsR RealFloatAndFloatElt RealFloatH target-      , TensorSupportsR IntegralHAndIntElt IntegralH target-      , TensorSupportsS Num Num target-      , TensorSupportsS RealFloatAndFloatElt Floating target-      , TensorSupportsS RealFloatAndFloatElt RealFloatH target-      , TensorSupportsS IntegralHAndIntElt IntegralH target-      , TensorSupportsX Num Num target-      , TensorSupportsX RealFloatAndFloatElt Floating target-      , TensorSupportsX RealFloatAndFloatElt RealFloatH target-      , TensorSupportsX IntegralHAndIntElt IntegralH target )-      => BaseTensor (target :: Target) where--  -- First type argument being @target@ is acceptable here, since these-  -- operations are mostly used when the shape is not known at the type level,-  -- so it can't be used as an explicit type argument.-  rshape :: forall n x. KnownSTK x-         => target (TKR2 n x) -> IShR n-  rlength :: forall n x. KnownSTK x-          => target (TKR2 n x) -> Int-  rlength = shrLength . rshape-  rsize :: forall n x. KnownSTK x-        => target (TKR2 n x) -> Int-  rsize = shrSize . rshape-  rwidth :: forall n x. KnownSTK x-          => target (TKR2 (1 + n) x) -> Int-  rwidth a = case rshape a of-    k :$: _ -> k--  sshape :: forall sh x. KnownSTK x-         => target (TKS2 sh x) -> ShS sh-  slength :: forall sh x. KnownSTK x-          => target (TKS2 sh x) -> Int-  slength = shsLength . sshape-  ssize :: forall sh x. KnownSTK x-        => target (TKS2 sh x) -> Int-  ssize = shsSize . sshape-  swidth :: forall n sh x. KnownSTK x-          => target (TKS2 (n ': sh) x) -> Int-  swidth a = case sshape a of-    n :$$ _ -> sNatValue n--  xshape :: forall sh x. KnownSTK x-         => target (TKX2 sh x) -> IShX sh-  xlength :: forall sh x. KnownSTK x-          => target (TKX2 sh x) -> Int-  xlength = shxLength . xshape-  xsize :: forall sh x. KnownSTK x-        => target (TKX2 sh x) -> Int-  xsize = shxSize . xshape-  xwidth :: forall mn sh x. KnownSTK x-          => target (TKX2 (mn ': sh) x) -> Int-  xwidth a = case xshape a of-    mn :$% _ -> fromSMayNat' mn--  tsize :: SingletonTK y -> target y -> Int-  tsize stk a = case stk of-    STKScalar @r -> case testEquality (typeRep @r) (typeRep @Z1) of-      Just Refl -> 0-      _ -> 1-    STKR _ x | Dict <- lemKnownSTK x -> rsize a-    STKS _ x | Dict <- lemKnownSTK x -> ssize a-    STKX _ x | Dict <- lemKnownSTK x -> xsize a-    STKProduct stk1 stk2 ->-      tsize stk1 (tproject1 a) + tsize stk2 (tproject2 a)-  tftk :: SingletonTK y -> target y -> FullShapeTK y--  -- Unlikely to require type applications at all-  tpair :: target x -> target z -> target (TKProduct x z)-  tproject1 :: target (TKProduct x z) -> target x-  tproject2 :: target (TKProduct x z) -> target z--  ------------  -- Everything below is indended to be rarely used and usually there are-  -- more specific and/or more convienient functions that do the same job-  -- in other modules.-  -------------------  -- | The operation is potentially strict in all arguments.-  tcond :: Boolean (BoolOf target)-        => SingletonTK y-        -> BoolOf target -> target y -> target y -> target y--  -- A more precise type would have `PrimalOf target`, but it's require-  -- the user to convert, so we leave that precision to the AST only-  -- which means the AST instance will automatically insert such-  -- conversions as needed. The same holds for trfloor and many others.-  trconcrete :: GoodScalar r-             => Nested.Ranked n r -> target (TKR n r)-  tsconcrete :: GoodScalar r-             => Nested.Shaped sh r -> target (TKS sh r)-  txconcrete :: GoodScalar r-             => Nested.Mixed sh r -> target (TKX sh r)-  tkconcrete :: GoodScalar r => r -> target (TKScalar r)-  tconcrete :: FullShapeTK y -> Concrete y -> target y--  -- These nine methods can't be replaced by tfromVector, because the concrete-  -- instance has much faster implementations.-  ---  -- This is morally non-empty strict vectors:-  trfromVector :: (KnownNat n, KnownSTK x)-               => Data.Vector.Vector (target (TKR2 n x))-               -> target (TKR2 (1 + n) x)-  trfromVector v = withSNat (V.length v) $ \k ->-    tfromVector k (STKR SNat knownSTK) v-  trfromVector0N :: forall n x. KnownSTK x-                 => IShR n -> Data.Vector.Vector (target (TKR2 0 x))-                 -> target (TKR2 n x)-  trfromVector0N sh v | Dict <- eltDictRep (knownSTK @x) =-    if V.null v-    then let arr = Nested.remptyArray-         in trreshape sh $ tconcrete (tftkG knownSTK arr) (Concrete arr)-    else trreshape sh $ trfromVector v-  trunravelToList :: (KnownNat n, KnownSTK x)-                  => target (TKR2 (1 + n) x) -> [target (TKR2 n x)]-  trunravelToList @n @x t =-    let f :: Int -> target (TKR2 n x)-        f i = trindex t (fromIntegral i :.: ZIR)-    in map f [0 .. rwidth t - 1]--  tsfromVector :: (KnownNat n, KnownShS sh, KnownSTK x)-               => Data.Vector.Vector (target (TKS2 sh x))-               -> target (TKS2 (n ': sh) x)-  tsfromVector = tfromVector SNat (STKS knownShS knownSTK)-  tsfromVector0N :: (KnownShS sh, KnownSTK x)-                 => Data.Vector.Vector (target (TKS2 '[] x))-                 -> target (TKS2 sh x)-  tsfromVector0N @sh @x v | Dict <- eltDictRep (knownSTK @x)-                          , SNat <- shsProduct (knownShS @sh) =-    if V.null v-    then gcastWith (unsafeCoerceRefl :: Product sh :~: 0) $-         let arr = Nested.semptyArray ZSS-         in tsreshape knownShS $ tconcrete (tftkG knownSTK arr) (Concrete arr)-    else tsreshape (knownShS @sh) $ tsfromVector v-  tsunravelToList :: (KnownNat n, KnownShS sh, KnownSTK x)-                  => target (TKS2 (n ': sh) x) -> [target (TKS2 sh x)]-  tsunravelToList @_ @sh @x t =-    let f :: Int -> target (TKS2 sh x)-        f i = tsindex t (fromIntegral i :.$ ZIS)-    in map f [0 .. swidth t - 1]--  txfromVector :: (KnownNat n, KnownShX sh, KnownSTK x)-               => Data.Vector.Vector (target (TKX2 sh x))-               -> target (TKX2 (Just n ': sh) x)-  txfromVector = tfromVector SNat (STKX knownShX knownSTK)-  txfromVector0N :: forall sh x. KnownSTK x-                 => IShX sh -> Data.Vector.Vector (target (TKX2 '[] x))-                 -> target (TKX2 sh x)-  txfromVector0N sh v | Dict <- eltDictRep (knownSTK @x) =-    if V.null v-    then let arr = Nested.memptyArray ZSX-         in txreshape sh $ tconcrete (tftkG knownSTK arr) (Concrete arr)-    else withSNat (shxSize sh) $ \(SNat @n) ->-           txreshape @_ @'[Just n] sh $ txfromVector v-  txunravelToList :: (KnownNat n, KnownShX sh, KnownSTK x)-                  => target (TKX2 (Just n ': sh) x) -> [target (TKX2 sh x)]-  txunravelToList @_ @sh @x t =-    let f :: Int -> target (TKX2 sh x)-        f i = txindex t (fromIntegral i :.% ZIX)-    in map f [0 .. xwidth t - 1]--  tfromVector-    :: forall y k.-       SNat k -> SingletonTK y -> Data.Vector.Vector (target y)-    -> target (BuildTensorKind k y)-  tfromListR :: SingletonTK y -> ListR k (target y)-             -> target (BuildTensorKind k y)-  tfromListR stk l =-    tfromVector (listrRank l) stk . V.fromList . Foldable.toList $ l--  -- A number suffix in the name may indicate the rank of the codomain,-  -- if bounded. Suffix 1 may also mean the operations builds up codomain-  -- by 1 dimension.-  trsum :: (KnownNat n, KnownSTK x)-        => target (TKR2 (1 + n) x) -> target (TKR2 n x)-  -- This op (and it's Delta constructor) is worthwhile, because flattening-  -- is O(n) sometimes, unlike transpose, etc.-  trsum0 :: (KnownNat n, KnownSTK x)-         => target (TKR2 n x) -> target (TKR2 0 x)-  trsum0 = trsum . rflatten-  trdot0 :: (KnownNat n, GoodScalar r)-         => target (TKR n r) -> target (TKR n r) -> target (TKR 0 r)-  trdot0 t u = trsum (rflatten (t * u))-  trdot1In :: (KnownNat n, GoodScalar r)-           => target (TKR (1 + n) r) -> target (TKR (1 + n) r)-           -> target (TKR n r)-  trdot1In @n t u = trsum $ trtranspose (permCycle $ 1 + valueOf @n) (t * u)-  trmatvecmul :: GoodScalar r-              => target (TKR 2 r) -> target (TKR 1 r) -> target (TKR 1 r)--- How to generalize (#69)? The few straightforward generalizations--- differ in types but all are far from matmul2.--- rmatvecmul m v = rflatten $ rmap1 (rreplicate 1 . rdot0 v) m-  trmatvecmul m v = trbuild1 (rwidth m) (\i -> trdot0 v (m `trindex` [i]))-  trmatmul2 :: GoodScalar r-            => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)--- How to generalize to tmatmul (#69)?--- Just rmatmul2 the two outermost dimensions?--- rmatmul2 m1 m2 = rmap1 (rmatvecmul (rtr m2)) m1-  trmatmul2 m1 m2 =-    trbuild1 (rwidth m1) (\i -> trmatvecmul (rtr m2) (m1 `trindex` [i]))-  trreplicate :: (KnownNat n, KnownSTK x)-              => Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)-  trreplicate0N :: (KnownNat n, KnownSTK x)-                => IShR n -> target (TKR2 0 x) -> target (TKR2 n x)-  trreplicate0N sh = trreshape sh . trreplicate (shrSize sh)--  tssum :: (KnownNat n, KnownShS sh, KnownSTK x)-        => target (TKS2 (n ': sh) x) -> target (TKS2 sh x)-  tssum0 :: (KnownShS sh, KnownSTK x)-         => target (TKS2 sh x) -> target (TKS2 '[] x)-  tssum0 @sh | SNat <- shsProduct (knownShS @sh) = tssum . sflatten-  tsdot0 :: (KnownShS sh, GoodScalar r)-         => target (TKS sh r) -> target (TKS sh r) -> target (TKS '[] r)-  tsdot0 @sh t u | SNat <- shsProduct (knownShS @sh) = tssum (sflatten (t * u))-  tsdot1In :: (KnownShS sh, GoodScalar r)-           => SNat n -> target (TKS (sh ++ '[n]) r)-           -> target (TKS (sh ++ '[n]) r)-           -> target (TKS sh r)-  tsdot1In @sh (SNat @n) t u =-    let cpermR = permCycle $ 1 + shsLength (knownShS @sh)-    in Permutation.permFromList cpermR $ \(cperm :: Permutation.Perm cperm) ->-         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (sh ++ '[n])) $-         gcastWith (unsafeCoerceRefl-                    :: Permutation.PermutePrefix cperm (sh ++ '[n])-                       :~: n : sh) $-         fromMaybe (error "tsdot1In: impossible non-permutation")-         $ Permutation.permCheckPermutation cperm-         $ tssum $ tstranspose cperm (t * u)-  tsmatvecmul :: (KnownNat m, KnownNat n, GoodScalar r)-              => target (TKS '[m, n] r) -> target (TKS '[n] r)-              -> target (TKS '[m] r)-  tsmatvecmul @m m v = tsbuild1 @_ @m (\i -> tsdot0 v (m `tsindex` (i :.$ ZIS)))-  tsmatmul2 :: (KnownNat m, KnownNat n, KnownNat p, GoodScalar r)-            => target (TKS '[m, n] r) -> target (TKS '[n, p] r)-            -> target (TKS '[m, p] r)-  tsmatmul2 @m m1 m2 =-    tsbuild1 @_ @m (\i -> tsmatvecmul (str m2) (m1 `tsindex` (i :.$ ZIS)))-  tsreplicate :: forall sh k x. KnownSTK x-              => SNat k -> ShS sh -> target (TKS2 sh x)-              -> target (TKS2 (k ': sh) x)-  tsreplicate0N :: forall sh x. KnownSTK x-                => ShS sh -> target (TKS2 '[] x)-                -> target (TKS2 sh x)-  tsreplicate0N sh = tsreshape sh . tsreplicate (shsProduct sh) ZSS--  -- The choice in BuildTensorKind makes it hard to support this one,-  -- due to DeltaSum and AstSum being typed with BuildTensorKind:-  -- xsum :: (KnownShX sh, KnownShX (mn ': sh), KnownSTK x)-  --     => target (TKX2 (mn ': sh) x) -> target (TKX2 sh x)-  txsum :: (KnownNat n, KnownShX sh, KnownSTK x)-        => target (TKX2 (Just n ': sh) x) -> target (TKX2 sh x)-  txsum0 :: (KnownShX sh, KnownSTK x, ConvertTensor target)-         => target (TKX2 sh x) -> target (TKX2 '[] x)-  txsum0 t = withSNat (shxSize $ xshape t) $ \snat ->-    txsum (xmcast (Nested.SKnown snat :!% ZKX) $ xflatten t)-  txdot0 :: (KnownShX sh, GoodScalar r, ConvertTensor target)-         => target (TKX sh r) -> target (TKX sh r) -> target (TKX '[] r)-  txdot0 t u = withSNat (shxSize $ xshape t) $ \snat ->-    txsum (xmcast (Nested.SKnown snat :!% ZKX) $ xflatten (t * u))-  txdot1In :: (KnownShX sh, GoodScalar r)-           => SNat n -> target (TKX (sh ++ '[Just n]) r)-           -> target (TKX (sh ++ '[Just n]) r)-           -> target (TKX sh r)-  txdot1In @sh (SNat @n) t u =-    let cpermR = permCycle $ 1 + sNatValue (ssxRank (knownShX @sh))-    in Permutation.permFromList cpermR $ \(cperm :: Permutation.Perm cperm) ->-         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (sh ++ '[Just n])) $-         gcastWith (unsafeCoerceRefl-                    :: Permutation.PermutePrefix cperm (sh ++ '[Just n])-                       :~: Just n : sh) $-         fromMaybe (error "txdot1In: impossible non-permutation")-         $ Permutation.permCheckPermutation cperm-         $ txsum $ txtranspose cperm (t * u)-  txmatvecmul :: forall mm mn r. (GoodScalar r, ConvertTensor target)-              => Nested.SMayNat Int SNat mm -> Nested.SMayNat Int SNat mn-              -> target (TKX '[mm, mn] r) -> target (TKX '[mn] r)-              -> target (TKX '[mm] r)-  txmatvecmul mm mn m v =-    withKnownShX (ssxFromShX $ mm :$% ZSX) $-    withKnownShX (ssxFromShX $ mn :$% ZSX) $-    withSNat (fromSMayNat' mm) $ \(SNat @k) ->-      xmcast (ssxFromShX $ mm :$% ZSX)-      $ txbuild1 @_ @k (\i -> txdot0 v (m `txindex` (i :.% ZIX)))-  txmatmul2 :: ( KnownNat m, KnownNat n, KnownNat p, GoodScalar r-               , ConvertTensor target )-            => target (TKX '[Just m, Just n] r)-            -> target (TKX '[Just n, Just p] r)-            -> target (TKX '[Just m, Just p] r)-  txmatmul2 @m @n @p m1 m2 =-    txbuild1 @_ @m (\i ->-      txmatvecmul (Nested.SKnown (SNat @p)) (Nested.SKnown (SNat @n))-                  (xtr m2) (m1 `txindex` (i :.% ZIX)))-  txreplicate :: forall sh k x. KnownSTK x-              => SNat k -> StaticShX sh -> target (TKX2 sh x)-              -> target (TKX2 (Just k ': sh) x)-  txreplicate0N :: (KnownShX sh, KnownSTK x)-                => IShX sh -> target (TKX2 '[] x) -> target (TKX2 sh x)-  txreplicate0N sh = withSNat (shxSize sh) $ \snat ->-    txreshape sh . txreplicate snat knownShX--  trindex :: (KnownNat m, KnownNat n, KnownSTK x)-          => target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)-  trindex0 :: (KnownNat m, KnownSTK x)-           => target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)-  trindex0 = trindex-  troneHot :: ( KnownNat m, KnownNat n, KnownSTK x-              , BoolOf (PrimalOf target) ~ BoolOf target-              , EqH (PrimalOf target) (TKScalar Int64))-           => IShR m -> target (TKR2 n x) -> IxROf target m-           -> target (TKR2 (m + n) x)-  {-# INLINE troneHot #-}-  troneHot @_ @_ @x sh v ix = case knownSTK @x of-    STKScalar ->-      trscatter @_ @0 (shrAppend sh (rshape v)) v (const ix)-    _ -> case tftk knownSTK v of-      FTKR _ ftk2 ->-        -- TODO: def at out of bounds-        let f ix2 = tcond knownSTK-                          (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true-                           $ zip (toList ix) (toList ix2))-                          (trindex0 v (ixrDrop ix2))-                          (tdefTarget (FTKR ZSR ftk2))-        in rbuild (shrAppend sh (rshape v)) f-           -- TODO: if this is used often, maybe express this as the gather that-           -- would come out of vectorization, making sure it simplifies well-  trscatter :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x)-            => IShR (p + n) -> target (TKR2 (m + n) x)-            -> (IxROf target m -> IxROf target p)-            -> target (TKR2 (p + n) x)-  trscatter1 :: (KnownNat n, KnownNat p, KnownSTK x)-             => IShR (p + n) -> target (TKR2 (1 + n) x)-             -> (IntOf target -> IxROf target p)-             -> target (TKR2 (p + n) x)-  trscatter1 sh v f = trscatter @target @1 sh v (\(i :.: ZIR) -> f i)-  trgather :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x)-           => IShR (m + n) -> target (TKR2 (p + n) x)-           -> (IxROf target m -> IxROf target p)-           -> target (TKR2 (m + n) x)-  trgather1 :: (KnownNat n, KnownNat p, KnownSTK x)-            => Int -> target (TKR2 (p + n) x)-            -> (IntOf target -> IxROf target p)-            -> target (TKR2 (1 + n) x)-  trgather1 k v f = trgather @target @1-                             (k :$: shrDrop (rshape v)) v-                             (\(i :.: ZIR) -> f i)--  tsindex :: (KnownShS shm, KnownShS shn, KnownSTK x)-          => target (TKS2 (shm ++ shn) x) -> IxSOf target shm-          -> target (TKS2 shn x)-  tsindex0 :: (KnownShS sh1, KnownSTK x)-           => target (TKS2 sh1 x) -> IxSOf target sh1-           -> target (TKS2 '[] x)-  tsindex0 @sh1 | Refl <- lemAppNil @sh1 = tsindex-  tsoneHot :: ( KnownShS sh1, KnownShS sh2, KnownSTK x-              , BoolOf (PrimalOf target) ~ BoolOf target-              , EqH (PrimalOf target) (TKScalar Int64) )-           => target (TKS2 sh2 x) -> IxSOf target sh1-           -> target (TKS2 (sh1 ++ sh2) x)-  {-# INLINE tsoneHot #-}  -- this doesn't want to specialize-  tsoneHot @sh1 @sh2 @x v ix-   | SNat <- shsRank (knownShS @sh1) = case knownSTK @x of-    STKScalar ->-      gcastWith (unsafeCoerceRefl :: Take (Rank sh1) (sh1 ++ sh2) :~: sh1) $-      gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) (sh1 ++ sh2) :~: sh2) $-      tsscatter @_ @'[] @_ @sh1 v (const ix)-    _ -> case tftk knownSTK v of-      FTKS _ ftk2 ->-        -- TODO: def at out of bounds-        gcastWith (unsafeCoerceRefl-                   :: Drop (Rank (sh1 ++ sh2)) (sh1 ++ sh2) :~: '[]) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank (sh1 ++ sh2)) (sh1 ++ sh2) :~: (sh1 ++ sh2)) $-        gcastWith (unsafeCoerceRefl-                   :: Drop (Rank sh1) (sh1 ++ sh2) :~: sh2) $-        withKnownShS (knownShS @sh1 `shsAppend` knownShS @sh2) $-        let f ix2 = tcond knownSTK-                          (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true-                           $ zip (Foldable.toList ix) (Foldable.toList ix2))-                          (tsindex0 v (ixsDrop @(Rank sh1) ix2))-                          (tdefTarget (FTKS ZSS ftk2))-        in sbuild @(Rank (sh1 ++ sh2)) f-  tsscatter-     :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x)-     => target (TKS2 (shm ++ shn) x)-     -> (IxSOf target shm -> IxSOf target shp)-     -> target (TKS2 (shp ++ shn) x)-  tsscatter1-     :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x)-     => target (TKS2 (n2 ': shn) x)-     -> (IntOf target -> IxSOf target shp)-     -> target (TKS2 (shp ++ shn) x)-  tsscatter1 @n2 v f = tsscatter @_ @'[n2] v (\(i :.$ _) -> f i)-  tsgather-     :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x)-     => target (TKS2 (shp ++ shn) x)-     -> (IxSOf target shm -> IxSOf target shp)-     -> target (TKS2 (shm ++ shn) x)-  tsgather1-     :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x)-     => target (TKS2 (shp ++ shn) x)-     -> (IntOf target -> IxSOf target shp)-     -> target (TKS2 (n2 ': shn) x)-  tsgather1 @n2 v f = tsgather @target @'[n2] v (\(i :.$ _) -> f i)--  txindex :: (KnownShX sh1, KnownShX sh2, KnownSTK x)-          => target (TKX2 (sh1 ++ sh2) x) -> IxXOf target sh1-          -> target (TKX2 sh2 x)-  txindex0 :: (KnownShX sh1, KnownSTK x)-           => target (TKX2 sh1 x) -> IxXOf target sh1-           -> target (TKX2 '[] x)-  txindex0 @sh1 | Refl <- lemAppNil @sh1 = txindex-  txoneHot :: ( KnownShX sh1, KnownShX sh2, KnownSTK x-              , BoolOf (PrimalOf target) ~ BoolOf target-              , EqH (PrimalOf target) (TKScalar Int64), ConvertTensor target )-           => IShX sh1 -> target (TKX2 sh2 x) -> IxXOf target sh1-           -> target (TKX2 (sh1 ++ sh2) x)-  {-# INLINE txoneHot #-}-  txoneHot @sh1 @sh2 @x sh1 v ix-   | SNat <- ssxRank (knownShX @sh1) = case knownSTK @x of-    STKScalar ->-      gcastWith (unsafeCoerceRefl :: Take (Rank sh1) (sh1 ++ sh2) :~: sh1) $-      gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) (sh1 ++ sh2) :~: sh2) $-      txscatter @_ @'[] @_ @sh1 (shxAppend sh1 (xshape v)) v (const ix)-    _ -> case tftk knownSTK v of-      FTKX _ ftk2 ->-        -- TODO: def at out of bounds-        gcastWith (unsafeCoerceRefl-                   :: Drop (Rank (sh1 ++ sh2)) (sh1 ++ sh2) :~: '[]) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank (sh1 ++ sh2)) (sh1 ++ sh2) :~: (sh1 ++ sh2)) $-        gcastWith (unsafeCoerceRefl-                   :: Drop (Rank sh1) (sh1 ++ sh2) :~: sh2) $-        withKnownShX (knownShX @sh1 `ssxAppend` knownShX @sh2) $-        let f ix2 = tcond knownSTK-                          (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true-                           $ zip (Foldable.toList ix) (Foldable.toList ix2))-                          (txindex0 v (ixxDrop' @(Rank sh1) ix2))-                          (tdefTarget (FTKX ZSX ftk2))-        in xbuild @(Rank (sh1 ++ sh2)) (shxAppend sh1 (xshape v)) f-  txscatter :: (KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x)-            => IShX (shp ++ shn) -> target (TKX2 (shm ++ shn) x)-            -> (IxXOf target shm -> IxXOf target shp)-            -> target (TKX2 (shp ++ shn) x)-  -- TODO: generalize this to non-Just types.-  txscatter1 :: (KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x)-             => IShX (shp ++ shn) -> target (TKX2 (Just n2 ': shn) x)-             -> (IntOf target -> IxXOf target shp)-             -> target (TKX2 (shp ++ shn) x)-  txscatter1 @n2 @_ @shp @x sh v f = txscatter @_ @'[Just n2] @_ @shp @x sh v-                                               (\(i :.% _) -> f i)-  txgather :: (KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x)-           => IShX (shm ++ shn)-           -> target (TKX2 (shp ++ shn) x)-           -> (IxXOf target shm -> IxXOf target shp)-           -> target (TKX2 (shm ++ shn) x)-  txgather1 :: (KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x)-            => SNat n2 -> target (TKX2 (shp ++ shn) x)-            -> (IntOf target -> IxXOf target shp)-            -> target (TKX2 (Just n2 ': shn) x)-  txgather1 @n2 @_ @shp k v f =-    txgather @target @'[Just n2]-             (Nested.SKnown k :$% shxDropSSX (knownShX @shp) (xshape v)) v-             (\(i :.% ZIX) -> f i)--  trfloor :: (GoodScalar r, RealFrac r, GoodScalar r2, Integral r2)-          => target (TKR n r) -> target (TKR n r2)-  trfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2)-                 => target (TKR n r1) -> target (TKR n r2)-  trcast :: (RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2)-         => target (TKR n r1) -> target (TKR n r2)-  trminIndex, trmaxIndex  -- partial-    :: forall n r r2. (GoodScalar r, GoodScalar r2)-    => target (TKR (1 + n) r) -> target (TKR n r2)-  triota :: GoodScalar r => Int -> target (TKR 1 r)  -- from 0 to n - 1--  tsfloor :: (GoodScalar r, RealFrac r, GoodScalar r2, Integral r2)-          => target (TKS sh r) -> target (TKS sh r2)-  tsfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2)-                 => target (TKS sh r1) -> target (TKS sh r2)-  tscast :: (RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2)-         => target (TKS sh r1) -> target (TKS sh r2)-  tsminIndex, tsmaxIndex  -- partial-    :: forall n sh r r2. (GoodScalar r, GoodScalar r2)-    => target (TKS (n ': sh) r) -> target (TKS (Init (n ': sh)) r2)-  tsiota :: (KnownNat n, GoodScalar r)-         => target (TKS '[n] r)  -- from 0 to n - 1--  txfloor :: (GoodScalar r, RealFrac r, GoodScalar r2, Integral r2)-          => target (TKX sh r) -> target (TKX sh r2)-  txfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2)-                 => target (TKX sh r1) -> target (TKX sh r2)-  txcast :: (RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2)-         => target (TKX sh r1) -> target (TKX sh r2)-  txminIndex, txmaxIndex  -- partial-    :: forall mn sh r r2. (GoodScalar r, GoodScalar r2)-    => target (TKX (mn ': sh) r) -> target (TKX (Init (mn ': sh)) r2)-  txiota :: (KnownNat n, GoodScalar r)-         => target (TKX '[Just n] r)  -- from 0 to n - 1--  tkfloor :: (GoodScalar r, RealFrac r, GoodScalar r2, Integral r2)-          => target (TKScalar r) -> target (TKScalar r2)-  tkfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2)-                 => target (TKScalar r1) -> target (TKScalar r2)-  tkcast :: (RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2)-         => target (TKScalar r1) -> target (TKScalar r2)--  trappend :: forall n x. KnownSTK x-           => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)-           -> target (TKR2 (1 + n) x)-  trslice :: forall n x. KnownSTK x-          => Int -> Int -> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)-  trreverse :: forall n x. KnownSTK x-            => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)-  trtranspose :: forall n x. KnownSTK x-              => Permutation.PermR -> target (TKR2 n x) -> target (TKR2 n x)-  trreshape :: forall n m x. KnownSTK x-            => IShR m -> target (TKR2 n x) -> target (TKR2 m x)--  tsappend :: forall m n sh x. KnownSTK x-           => target (TKS2 (m ': sh) x) -> target (TKS2 (n ': sh) x)-           -> target (TKS2 ((m + n) ': sh) x)-  tsslice :: forall i n k sh x. KnownSTK x-          => SNat i -> SNat n -> SNat k-          -> target (TKS2 (i + n + k ': sh) x) -> target (TKS2 (n ': sh) x)-  tsreverse :: forall n sh x. KnownSTK x-            => target (TKS2 (n ': sh) x) -> target (TKS2 (n ': sh) x)-  tstranspose :: ( Permutation.IsPermutation perm-                 , Rank perm <= Rank sh, KnownSTK x )-              => Permutation.Perm perm -> target (TKS2 sh x)-              -> target (TKS2 (Permutation.PermutePrefix perm sh) x)-  tsreshape :: (Product sh ~ Product sh2, KnownSTK x)-            => ShS sh2 -> target (TKS2 sh x) -> target (TKS2 sh2 x)--  txappend :: forall m n sh x. KnownSTK x-           => target (TKX2 (Just m ': sh) x) -> target (TKX2 (Just n ': sh) x)-           -> target (TKX2 (Just (m + n) ': sh) x)-  txslice :: forall i n k sh x. KnownSTK x-          => SNat i -> SNat n -> SNat k-          -> target (TKX2 (Just (i + n + k) ': sh) x)-          -> target (TKX2 (Just n ': sh) x)-  txreverse :: forall mn sh x. KnownSTK x-            => target (TKX2 (mn ': sh) x) -> target (TKX2 (mn ': sh) x)-  txtranspose :: ( Permutation.IsPermutation perm-                 , Rank perm <= Rank sh, KnownSTK x )-              => Permutation.Perm perm -> target (TKX2 sh x)-              -> target (TKX2 (Permutation.PermutePrefix perm sh) x)-  txreshape :: forall sh sh2 x. KnownSTK x-            => IShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)--  trbuild1 :: (KnownNat n, KnownSTK x)-           => Int -> (IntOf target -> target (TKR2 n x))-           -> target (TKR2 (1 + n) x)-  trmap0N :: (KnownNat n, KnownSTK x, KnownSTK x1)-          => (target (TKR2 0 x1) -> target (TKR2 0 x)) -> target (TKR2 n x1)-          -> target (TKR2 n x)-  trmap0N f v = rbuild (rshape v) (f . trindex0 v)-  trzipWith0N :: (KnownNat n, KnownSTK x, KnownSTK x1, KnownSTK x2)-              => (target (TKR2 0 x1) -> target (TKR2 0 x2) -> target (TKR2 0 x))-              -> target (TKR2 n x1) -> target (TKR2 n x2) -> target (TKR2 n x)-  trzipWith0N f u v =-    rbuild (rshape v) (\ix -> f (trindex0 u ix) (trindex0 v ix))--  tsbuild1 :: (KnownNat k, KnownShS sh, KnownSTK x)-           => (IntOf target -> target (TKS2 sh x))-           -> target (TKS2 (k ': sh) x)-  tsmap0N :: (KnownShS sh, KnownSTK x, KnownSTK x1)-          => (target (TKS2 '[] x1) -> target (TKS2 '[] x))-          -> target (TKS2 sh x1)-          -> target (TKS2 sh x)-  tsmap0N @sh f v =-    gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])-    $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)-    $ sbuild @(Rank sh) (f . tsindex0 v)-  tszipWith0N :: (KnownShS sh, KnownSTK x, KnownSTK x1, KnownSTK x2)-              => (target (TKS2 '[] x1) -> target (TKS2 '[] x2)-                  -> target (TKS2 '[] x))-              -> target (TKS2 sh x1) -> target (TKS2 sh x2)-              -> target (TKS2 sh x)-  tszipWith0N @sh f u v =-    gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])-    $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)-    $ sbuild @(Rank sh) (\ix -> f (tsindex0 u ix) (tsindex0 v ix))--  txbuild1 :: (KnownNat k, KnownShX sh, KnownSTK x)-           => (IntOf target -> target (TKX2 sh x))-           -> target (TKX2 (Just k ': sh) x)--  tbuild1 :: forall y k. ConvertTensor target-               -- y comes first, because k easy to set via SNat-          => SNat k -> SingletonTK y -> (IntOf target -> target y)-          -> target (BuildTensorKind k y)-  tbuild1 snat@SNat stk0 f =-    let replSTK :: SingletonTK z -> (IntOf target -> target z)-                -> target (BuildTensorKind k z)-        replSTK stk g = case stk of-          STKScalar -> tsbuild1 (sfromK . g)-          STKR SNat x | Dict <- lemKnownSTK x -> trbuild1 (sNatValue snat) g-          STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsbuild1 g-          STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txbuild1 g-          STKProduct @z1 @z2 stk1 stk2 ->-              let f1 i = tproject1 @_ @z1 @z2 $ g i-                  f2 i = tproject2 @_ @z1 @z2 $ g i-                    -- TODO: looks expensive, but hard to do better,-                    -- so let's hope g is full of variables-              in tpair (replSTK stk1 f1) (replSTK stk2 f2)-    in replSTK stk0 f--  -- | A strict right mapAccum.-  ---  -- The applications of 'tjvp' and 'tvjp' performed already at this point-  -- ensure that the computation of a derivative is not repeated-  -- and that its result is shared. However, most of the time-  -- the computation is unnneeded, so the AST instance uses a non-strict-  -- constructor 'HordeAd.Core.Ast.AstLambda' for it's instance of 'HFunOf'.-  ---  -- If the same argument functions are passed to many mapAccum calls, as in-  -- > let f = ... in ... (tmapAccumR ... f ...) ... (tmapAccumL ... f ...)-  -- extra care is needed to prevent double derivative computation.-  -- One needs to use tmapAccumRDer manually as in (simplified)-  -- > let f = ...; df = tjvp f; rf = tgrad f-  -- > in ... (tmapAccumRDer f df rf ...) ... (tmapAccumLDer f df rf ...)-  tmapAccumRDer-    :: forall accy by ey k.-       Proxy target-    -> SNat k  -- ^ length of the input-    -> FullShapeTK accy  -- ^ shape of the accumulator-    -> FullShapeTK by  -- ^ shape of the output-    -> FullShapeTK ey  -- ^ shape of an individual input-    -> HFunOf target (TKProduct accy ey) (TKProduct accy by)-         -- ^ the function to mapAccum with-    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))-                                (TKProduct accy ey))-                     (ADTensorKind (TKProduct accy by))-         -- ^ the derivative of the function to mapAccum with-    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))-                                (TKProduct accy ey))-                     (ADTensorKind (TKProduct accy ey))-         -- ^ the reverse derivative of the function to mapAccum with-    -> target accy  -- ^ the initial accumulator-    -> target (BuildTensorKind k ey)  -- ^ the inputs-    -> target (TKProduct accy (BuildTensorKind k by))-  -- | A strict left mapAccum.-  tmapAccumLDer-    :: forall accy by ey k.-       Proxy target-    -> SNat k  -- ^ length of the input-    -> FullShapeTK accy  -- ^ shape of the accumulator-    -> FullShapeTK by  -- ^ shape of the output-    -> FullShapeTK ey  -- ^ shape of an individual input-    -> HFunOf target (TKProduct accy ey) (TKProduct accy by)-         -- ^ the function to mapAccum with-    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))-                                (TKProduct accy ey))-                     (ADTensorKind (TKProduct accy by))-         -- ^ the derivative of the function to mapAccum with-    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))-                                (TKProduct accy ey))-                     (ADTensorKind (TKProduct accy ey))-         -- ^ the reverse derivative of the function to mapAccum with-    -> target accy  -- ^ the initial accumulator-    -> target (BuildTensorKind k ey)  -- ^ the inputs-    -> target (TKProduct accy (BuildTensorKind k by))-  tApply :: HFunOf target x z -> target x -> target z-  tlambda :: FullShapeTK x -> HFun x z -> HFunOf target x z--  -- | Reverse derivative.-  ---  -- The followign methods (and tlambda) are exactly what is needed as arguments-  -- of tmapAccumRDer.-  tgrad-    :: FullShapeTK x  -- ^ shape of x and dx-    -> HFun x (TKScalar r)  -- ^ x |-> TKScalar r-    -> HFunOf target x (ADTensorKind x)  -- ^ x |-> dx-  tvjp-    :: FullShapeTK x  -- ^ shape of x and dx-    -> HFun x z  -- ^ x |-> z-    -> HFunOf target (TKProduct (ADTensorKind z) x) (ADTensorKind x)-         -- ^ (dz, x) |-> dx-  tjvp-    :: FullShapeTK x  -- ^ shape of x and dx-    -> HFun x z  -- ^ x |-> z-    -> HFunOf target (TKProduct (ADTensorKind x) x) (ADTensorKind z)-         -- ^ (dx, x) |-> dz--  tprimalPart :: target y -> PrimalOf target y-  tdualPart :: SingletonTK y -> target y -> DualOf target y-  tfromPrimal :: SingletonTK y -> PrimalOf target y -> target y-  tfromDual :: DualOf target y -> target y-  tScale :: (Num (target y), Num (PrimalOf target y))-         => SingletonTK y -> PrimalOf target y -> DualOf target y-         -> DualOf target y-  tScale stk s t =-    tdualPart stk $ tfromPrimal @target stk s * tfromDual t--  -- General operations that use ShareTensor if available, LetTensor otherwise-  tsum-    :: forall z k. ConvertTensor target-    => SNat k -> SingletonTK z -> target (BuildTensorKind k z)-    -> target z-  default tsum-    :: forall z k. (ShareTensor target, ConvertTensor target)-    => SNat k -> SingletonTK z -> target (BuildTensorKind k z)-    -> target z-  tsum snat@SNat stk u = case stk of-    STKScalar -> kfromS $ tssum u-    STKR SNat x | Dict <- lemKnownSTK x -> trsum u-    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u-    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u-    STKProduct stk1 stk2 ->-      let (u1, u2) = tunpair u-      in tpair (tsum snat stk1 u1)-               (tsum snat stk2 u2)-  treplicate-    :: forall z k. ConvertTensor target-    => SNat k -> SingletonTK z -> target z-    -> target (BuildTensorKind k z)-  default treplicate-    :: forall z k. (ShareTensor target, ConvertTensor target)-    => SNat k -> SingletonTK z -> target z-    -> target (BuildTensorKind k z)-  treplicate snat@SNat stk u = case stk of-    STKScalar -> tsreplicate snat ZSS $ sfromK u-    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (sNatValue snat) u-    STKS sh x | Dict <- lemKnownSTK x -> tsreplicate snat sh u-    STKX sh x | Dict <- lemKnownSTK x -> txreplicate snat sh u-    STKProduct stk1 stk2 ->-      let (u1, u2) = tunpair u-      in tpair (treplicate snat stk1 u1)-               (treplicate snat stk2 u2)-  tindexBuild-    :: forall z k. ConvertTensor target-    => SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> IntOf target-    -> target z-  default tindexBuild-    :: forall z k. (ShareTensor target, ConvertTensor target)-    => SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> IntOf target-    -> target z-  tindexBuild snat@SNat stk u i = case stk of-    STKScalar -> kfromS $ tsindex u (i :.$ ZIS)-    STKR SNat x | Dict <- lemKnownSTK x -> trindex u (i :.: ZIR)-    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsindex u (i :.$ ZIS)-    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txindex u (i :.% ZIX)-    STKProduct stk1 stk2 ->-      let (u1, u2) = tunpair u-      in tpair (tindexBuild snat stk1 u1 i)-               (tindexBuild snat stk2 u2 i)--  -- Unwinding methods, needed mostly to split off the Unwind module.-  -- | Construct tensors with the given constant in each cell.-  treplTarget :: (forall r. GoodScalar r => r) -> FullShapeTK y -> target y-  -- | Construct tensors with @def@ in each cell.-  tdefTarget :: FullShapeTK y -> target y-  -- | Add pointwise all corresponding tensors within nested product, if any.-  ---  -- Requires duplicable arguments or a ShareTensor instance.-  taddTarget :: SingletonTK y -> target y -> target y -> target y-  -- | Multiply pointwise all corresponding tensors within nested products,-  -- if any.-  ---  -- Requires duplicable arguments or a ShareTensor instance.-  tmultTarget :: SingletonTK y -> target y -> target y -> target y-  -- | Sum all dimensions of each component and then sum it all. Ignore all-  -- tensors with non-differentiable elements.-  ---  -- Requires duplicable arguments or a ShareTensor instance.-  tsum0Target :: FullShapeTK y -> target y-              -> target (TKScalar Double)-  -- | Dot product each component and then sum it all. Ignore all-  -- tensors with non-differentiable elements.-  ---  -- Requires duplicable arguments or a ShareTensor instance.-  tdot0Target :: FullShapeTK y -> target y -> target y-              -> target (TKScalar Double)--  -- TODO: express without ConvertTensor or move there-  xmcast :: (KnownSTK x, KnownShX sh, Rank sh ~ Rank sh2, ConvertTensor target)-         => StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)-  xmcast sh2 a = case tftk knownSTK a of-    FTKX sh' _ ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        withKnownShX sh2 $-        withKnownShS sh $-        xfromS $ sfromX @_ @sh a---- These are user-accessible, so the constraint is `ADReady`, which means--- lets, but no shares.-type role HFun nominal nominal-newtype HFun (x :: TK) (z :: TK) =-  HFun {unHFun :: forall f. ADReady f-               => f x -> f z}--instance Show (HFun x y) where-  show _ = "<lambda>"----- * The mega-constraint--type ADReady target =-  ( ADReadyNoLet target-  , LetTensor target--- The following can't be added, because we have instances like ADVal (AstRaw),--- so AstRaw would need to have a LetTensor instance:---  , LetTensor (PrimalOf target)-  )--type ADReadyNoLet target =-  ( ADReadyEqsClasses target-  , ADReadyEqsClasses (ShareOf target)-  , ShareTensor (ShareOf target)-  , ShareTensor (PrimalOf (ShareOf target))-  , ShareOf (ShareOf target) ~ ShareOf target-  )--type ADReadyEqsClasses target =-  ( ADReadyEqs target-  , ADReadyClasses target-  , ADReadyClasses (PrimalOf target)-  )--type ADReadyEqs target =-  ( BoolOf (PrimalOf target) ~ BoolOf target-  )--type ADReadyClasses target =-  ( BaseTensor target-  , ConvertTensor target-  , Boolean (BoolOf target)-  , AllTargetShow target-  , CommonTargetEqOrd target-  )---- This is illegal:--- type AllTargetShow target = forall y. KnownSTK y => Show (target y)--type AllTargetShow :: Target -> Constraint-class (forall y. KnownSTK y => Show (target y))-      => AllTargetShow target where-instance-      (forall y. KnownSTK y => Show (target y))-      => AllTargetShow target where--type CommonTargetEqOrd :: Target -> Constraint-class ( forall r. GoodScalar r => EqH target (TKScalar r)-      , forall r. GoodScalar r => OrdH target (TKScalar r)-      , forall r n. GoodScalar r => EqH target (TKR n r)-      , forall r n. GoodScalar r => OrdH target (TKR n r)-      , forall r sh. GoodScalar r => EqH target (TKS sh r)-      , forall r sh. GoodScalar r => OrdH target (TKS sh r)-      , forall r sh. GoodScalar r => EqH target (TKX sh r)-      , forall r sh. GoodScalar r => OrdH target (TKX sh r) )-      => CommonTargetEqOrd target where-instance-      ( forall r. GoodScalar r => EqH target (TKScalar r)-      , forall r. GoodScalar r => OrdH target (TKScalar r)-      , forall r n. GoodScalar r => EqH target (TKR n r)-      , forall r n. GoodScalar r => OrdH target (TKR n r)-      , forall r sh. GoodScalar r => EqH target (TKS sh r)-      , forall r sh. GoodScalar r => OrdH target (TKS sh r)-      , forall r sh. GoodScalar r => EqH target (TKX sh r)-      , forall r sh. GoodScalar r => OrdH target (TKX sh r) )+  , rrepl, srepl, xrepl, treplTarget, tdefTarget+    -- * Helper classes and types+  , IntegralHAndIntElt, RealFloatAndFloatElt+  , TensorSupportsX, TensorSupportsS, TensorSupportsR, TensorSupports+  ) where++import Prelude++import Data.Foldable qualified as Foldable+import Data.Kind (Constraint, Type)+import Data.List (foldl1')+import Data.List.NonEmpty (NonEmpty)+import Data.List.NonEmpty qualified as NonEmpty+import Data.Maybe (fromMaybe)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Data.Vector.Strict qualified as Data.Vector+import GHC.TypeLits (KnownNat, type (+), type (<=), type (<=?))+import Type.Reflection (typeRep)++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (withShsFromShX)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Init, fromSNat', unsafeCoerceRefl)++import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- These user API functions are used in default definitions of methods,+-- so they have to be defined already here:++rtr :: forall n x target. (KnownSTK x, BaseTensor target)+    => target (TKR2 (2 + n) x) -> target (TKR2 (2 + n) x)+{-# INLINE rtr #-}+rtr = trtranspose [1, 0]+rflatten :: forall n x target. (KnownSTK x, BaseTensor target)+         => target (TKR2 n x) -> target (TKR2 1 x)+{-# INLINE rflatten #-}+rflatten u = trreshape (rsize u :$: ZSR) u+str :: forall n m sh x target. (KnownSTK x, BaseTensor target)+    => target (TKS2 (n ': m ': sh) x) -> target (TKS2 (m ': n ': sh) x)+{-# INLINE str #-}+str = gcastWith (unsafeCoerceRefl :: (2 <=? Rank (n ': m ': sh)) :~: True) $+      tstranspose (Permutation.makePerm @'[1, 0])+sflatten :: (KnownShS sh, KnownSTK x, BaseTensor target )+         => target (TKS2 sh x) -> target (TKS2 '[Product sh] x)+{-# INLINE sflatten #-}+sflatten @sh | SNat <- shsProduct (knownShS @sh) = tsreshape knownShS+xtr :: forall n m sh x target. (KnownSTK x, BaseTensor target)+    => target (TKX2 (Just n ': Just m ': sh) x)+    -> target (TKX2 (Just m ': Just n ': sh) x)+{-# INLINE xtr #-}+xtr = gcastWith (unsafeCoerceRefl+                 :: (2 <=? Rank (Just n ': Just m ': sh)) :~: True) $+      txtranspose (Permutation.makePerm @'[1, 0])+xflatten :: forall sh x target. (KnownSTK x, BaseTensor target)+         => target (TKX2 sh x) -> target (TKX2 '[Nothing] x)+{-# INLINE xflatten #-}+xflatten u = txreshape (SUnknown (xsize u) :$% ZSX) u++rrepl :: forall n r target. (GoodScalar r, BaseTensor target)+      => IShR n -> r -> target (TKR n r)+{-# INLINE rrepl #-}+rrepl sh a = tconcrete (FTKR sh FTKScalar)+                       (Concrete $ Nested.rreplicatePrim sh a)++srepl :: (KnownShS sh, GoodScalar r, BaseTensor target)+      => r -> target (TKS sh r)+{-# INLINE srepl #-}+srepl = tsconcrete . Nested.sreplicatePrim knownShS++xrepl :: forall sh r target. (GoodScalar r, BaseTensor target)+      => IShX sh -> r -> target (TKX sh r)+{-# INLINE xrepl #-}+xrepl sh a = tconcrete (FTKX sh FTKScalar)+                       (Concrete $ Nested.mreplicatePrim sh a)++-- | Construct tensors with the given constant in each cell.+treplTarget :: (TKAllNum y, BaseTensor target)+            => (forall r. NumScalar r => r) -> FullShapeTK y -> target y+treplTarget r ftk = tconcrete ftk $ Concrete $ replTargetRep r ftk++-- | Construct tensors with @def@ in each cell.+tdefTarget :: BaseTensor target+           => FullShapeTK y -> target y+tdefTarget ftk = tconcrete ftk $ Concrete $ defTargetRep ftk++type TensorSupports :: (Type -> Constraint) -> (Type -> Constraint)+                    -> Target -> Constraint+type TensorSupports c1 c2 f =+  forall r. NumScalar r+            => c1 r => c2 (f (TKScalar r))++type TensorSupportsR :: (Type -> Constraint) -> (Type -> Constraint)+                     -> Target -> Constraint+type TensorSupportsR c1 c2 f =+  forall r n. NumScalar r+              => c1 r => c2 (f (TKR n r))++type TensorSupportsS :: (Type -> Constraint) -> (Type -> Constraint)+                     -> Target -> Constraint+type TensorSupportsS c1 c2 f =+  forall r sh. NumScalar r+               => c1 r => c2 (f (TKS sh r))++type TensorSupportsX :: (Type -> Constraint) -> (Type -> Constraint)+                     -> Target -> Constraint+type TensorSupportsX c1 c2 f =+  forall r sh. NumScalar r+               => c1 r => c2 (f (TKX sh r))++class    Differentiable r+         => RealFloatAndFloatElt r+instance Differentiable r+         => RealFloatAndFloatElt r++class    (IntegralH r, Nested.IntElt r)+         => IntegralHAndIntElt r+instance (IntegralH r, Nested.IntElt r)+         => IntegralHAndIntElt r++class LetTensor (target :: Target) where+  ttlet :: target x -> (target x -> target z) -> target z+  ttletPrimal :: PrimalOf target x -> (PrimalOf target x -> target z)+              -> target z+  ttletPlain :: PlainOf target x -> (PlainOf target x -> target z)+             -> target z+  toShare :: target y -> ShareOf target y+  tunshare :: ShareOf target y -> target y+  tunshare = error "tunshare: this instance should never be used"+  tappend :: forall m n y. BaseTensor target+          => SNat m -> SNat n -> SingletonTK y+          -> target (BuildTensorKind m y) -> target (BuildTensorKind n y)+          -> target (BuildTensorKind (m + n) y)+  {-# INLINE tappend #-}+  tappend msnat@SNat nsnat@SNat stk a b = case stk of+    STKScalar -> tsappend a b+    STKR _ x | Dict <- lemKnownSTK x -> trappend a b+    STKS _ x | Dict <- lemKnownSTK x -> tsappend a b+    STKX _ x | Dict <- lemKnownSTK x -> txappend a b+    STKProduct stk1 stk2 ->+      ttlet a $ \ !aShared -> ttlet b $ \ !bShared ->+        tpair (tappend msnat nsnat stk1 (tproject1 aShared) (tproject1 bShared))+              (tappend msnat nsnat stk2 (tproject2 aShared) (tproject2 bShared))+  tD :: (TKAllNum y, BaseTensor target)+     => SingletonTK y -> PrimalOf target y -> DualOf target y+     -> target y+  tD stk p d =+    -- Lets needed, because taddTarget requires duplicable arguments.+    ttletPrimal p $ \pShared ->+    ttlet (tfromDual d) $ taddTarget stk (tfromPrimal stk pShared)+  -- | A strict left fold.+  tfold+    :: forall yn ym k. BaseTensor target+    => SNat k  -- ^ length of the input+    -> SingletonTK yn  -- ^ partial shape of the accumulator+    -> SingletonTK ym  -- ^ partial shape of an individual input+    -> (forall f. ADReady f => f yn -> f ym -> f yn)+         -- ^ the function to fold with+    -> target yn  -- ^ the initial accumulator+    -> target (BuildTensorKind k ym)  -- ^ the inputs+    -> target yn+  {-# INLINE tfold #-}  -- this doesn't want to specialize+  tfold k nstk mstk f acc0 es =+    tproject1+    $ tmapAccumL (Proxy @target)+       k+       (tftk nstk acc0)+       (FTKScalar @Z1)+       (razeFTK k mstk (tftk (buildSTK k mstk) es))+       (let g :: forall f. ADReady f+              => f yn -> f ym -> f (TKProduct yn TKUnit)+            g !acc !e = tpair (f acc e) (tkconcrete Z1)+        in g)+       acc0+       es+  -- | A strict left scan.+  tscan+    :: forall yn ym k. (BaseTensor target, ConvertTensor target)+    => SNat k  -- ^ length of the input+    -> SingletonTK yn  -- ^ partial shape of the accumulator+    -> SingletonTK ym  -- ^ partial shape of an individual input+    -> (forall f. ADReady f => f yn -> f ym -> f yn)+         -- ^ the function to scan with+    -> target yn  -- ^ the initial accumulator+    -> target (BuildTensorKind k ym)  -- ^ the inputs+    -> target (BuildTensorKind (1 + k) yn)+  {-# INLINE tscan #-}  -- this doesn't want to specialize+  tscan k nstk mstk f acc0' es = ttlet acc0' $ \acc0 ->  -- sharing just in case+    let bs :: target (BuildTensorKind k yn)+        bs = tproject2+             $ tmapAccumL (Proxy @target)+                k+                (tftk nstk acc0)+                (tftk nstk acc0)+                (razeFTK k mstk (tftk (buildSTK k mstk) es))+              (let g :: forall f. ADReady f+                     => f yn -> f ym -> f (TKProduct yn yn)+                   g !acc !e = ttlet (f acc e) $ \ !res -> tpair res res+               in g)+              acc0+              es+    in tappend (SNat @1) k nstk (treplicate (SNat @1) nstk acc0) bs++class ShareTensor (target :: Target) where+  tshare :: target y -> target y+  tunpair :: target (TKProduct x z) -> (target x, target z)+  -- This would suffer from lack of sharing if in LetTensor, because+  -- ttlet doesn't work over a list. With sharing it's fine.+  tunravelToListShare :: forall y k. BaseTensor target+                      => SNat k -> SingletonTK y+                      -> target (BuildTensorKind k y)+                      -> [target y]+  {-# INLINE tunravelToListShare #-}+  tunravelToListShare snat@SNat stk u = case stk of+    STKScalar -> let !uShared = tshare u+                 in tstoListLinear uShared+    STKR SNat x | Dict <- lemKnownSTK x -> let !uShared = tshare u+                                           in trunravelToList uShared+    STKS sh x | Dict <- lemKnownSTK x -> let !uShared = tshare u+                                         in withKnownShS sh+                                            $ tsunravelToList uShared+    STKX sh x | Dict <- lemKnownSTK x -> let !uShared = tshare u+                                         in withKnownShX sh+                                            $ txunravelToList uShared+    STKProduct stk1 stk2 ->+      -- TODO: prevent the proliferation of sharing, maybe via+      -- the unnest (unzip) trick, similarly as in vectorization,+      -- though ox-arrays has two recurseive calls as well+      -- (but not sharing):+      -- mtoListOuter (M_Tup2 a b) =+      --   zipWith M_Tup2 (mtoListOuter a) (mtoListOuter b)+      let (!u1, !u2) = tunpair u+      in zipWith tpair (tunravelToListShare snat stk1 u1)+                       (tunravelToListShare snat stk2 u2)++-- | The superclasses indicate that it's not only a container array,+-- but also a mathematical tensor, sporting numeric operations.+class ( Num (IntOf target)+          -- This is not redundant, because it constrains @PrimaOf f@,+          -- not @f@, as below, and so the user doesn't need+          -- to require @ADReady@ instead of @BaseTensor@ as often.+          -- Also, it leads to a choice of the more specific+          -- @Num (ADVal f (TKScalar r))@ instance (in GHC 9.12 at least),+          -- making some more cases of @fromInteger@ usage accepted.+      , TensorSupports Num Num target+      , TensorSupports RealFloatAndFloatElt Floating target+      , TensorSupports RealFloatAndFloatElt RealFloatH target+      , TensorSupports IntegralHAndIntElt IntegralH target+      , TensorSupportsR Num Num target+      , TensorSupportsR RealFloatAndFloatElt Floating target+      , TensorSupportsR RealFloatAndFloatElt RealFloatH target+      , TensorSupportsR IntegralHAndIntElt IntegralH target+      , TensorSupportsS Num Num target+      , TensorSupportsS RealFloatAndFloatElt Floating target+      , TensorSupportsS RealFloatAndFloatElt RealFloatH target+      , TensorSupportsS IntegralHAndIntElt IntegralH target+      , TensorSupportsX Num Num target+      , TensorSupportsX RealFloatAndFloatElt Floating target+      , TensorSupportsX RealFloatAndFloatElt RealFloatH target+      , TensorSupportsX IntegralHAndIntElt IntegralH target )+      => BaseTensor (target :: Target) where++  isConcreteInstance :: Bool+  -- First type argument being @target@ is acceptable here, since these+  -- operations are mostly used when the shape is not known at the type level,+  -- so it can't be used as an explicit type argument.+  rshape :: forall n x. KnownSTK x+         => target (TKR2 n x) -> IShR n+  rlength :: forall n x. KnownSTK x+          => target (TKR2 n x) -> Int+  {-# INLINE rlength #-}+  rlength = shrLength . rshape+  rsize :: forall n x. KnownSTK x+        => target (TKR2 n x) -> Int+  {-# INLINE rsize #-}+  rsize = shrSize . rshape+  rwidth :: forall n x. KnownSTK x+         => target (TKR2 (1 + n) x) -> Int+  {-# INLINE rwidth #-}+  rwidth a = case rshape a of+    k :$: _ -> k++  sshape :: forall sh x. KnownSTK x+         => target (TKS2 sh x) -> ShS sh+  slength :: forall sh x. KnownSTK x+          => target (TKS2 sh x) -> Int+  {-# INLINE slength #-}+  slength = shsLength . sshape+  ssize :: forall sh x. KnownSTK x+        => target (TKS2 sh x) -> Int+  {-# INLINE ssize #-}+  ssize = shsSize . sshape+  swidth :: forall n sh x. KnownSTK x+         => target (TKS2 (n ': sh) x) -> Int+  {-# INLINE swidth #-}+  swidth a = case sshape a of+    n :$$ _ -> fromSNat' n++  xshape :: forall sh x. KnownSTK x+         => target (TKX2 sh x) -> IShX sh+  xlength :: forall sh x. KnownSTK x+          => target (TKX2 sh x) -> Int+  {-# INLINE xlength #-}+  xlength = shxLength . xshape+  xsize :: forall sh x. KnownSTK x+        => target (TKX2 sh x) -> Int+  {-# INLINE xsize #-}+  xsize = shxSize . xshape+  xwidth :: forall mn sh x. KnownSTK x+         => target (TKX2 (mn ': sh) x) -> Int+  {-# INLINE xwidth #-}+  xwidth a = case xshape a of+    mn :$% _ -> fromSMayNat' mn++  tsize :: SingletonTK y -> target y -> Int+  {-# INLINE tsize #-}+  tsize stk a = case stk of+    STKScalar @r -> case testEquality (typeRep @r) (typeRep @Z1) of+      Just Refl -> 0+      _ -> 1+    STKR _ x | Dict <- lemKnownSTK x -> rsize a+    STKS _ x | Dict <- lemKnownSTK x -> ssize a+    STKX _ x | Dict <- lemKnownSTK x -> xsize a+    STKProduct stk1 stk2 ->+      tsize stk1 (tproject1 a) + tsize stk2 (tproject2 a)+  tftk :: SingletonTK y -> target y -> FullShapeTK y++  -- Unlikely to require type applications at all+  tpair :: target x -> target z -> target (TKProduct x z)+  tproject1 :: target (TKProduct x z) -> target x+  tproject2 :: target (TKProduct x z) -> target z++  -----------+  -- Everything below is indended to be rarely used and usually there are+  -- more specific and/or more convienient functions that do the same job+  -- in other modules.+  -----------------++  -- | These operations are potentially strict in all arguments.+  kcond :: (Boolean (BoolOf target), GoodScalar r)+        => BoolOf target -> target (TKScalar r) -> target (TKScalar r)+        -> target (TKScalar r)+  scond :: (Boolean (BoolOf target), KnownShS sh, KnownSTK x)+        => BoolOf target -> target (TKS2 sh x) -> target (TKS2 sh x)+        -> target (TKS2 sh x)+  tcond :: Boolean (BoolOf target)+        => SingletonTK y+        -> BoolOf target -> target y -> target y -> target y++  -- A more precise type would have `PrimalOf target`, but it's require+  -- the user to convert, so we leave that precision to the AST only+  -- which means the AST instance will automatically insert such+  -- conversions as needed. The same holds for trfloor and many others.+  tkconcrete :: GoodScalar r+             => r -> target (TKScalar r)+  trconcrete :: GoodScalar r+             => Nested.Ranked n r -> target (TKR n r)+  tsconcrete :: GoodScalar r+             => Nested.Shaped sh r -> target (TKS sh r)+  txconcrete :: GoodScalar r+             => Nested.Mixed sh r -> target (TKX sh r)+  tconcrete :: FullShapeTK y -> Concrete y -> target y++  -- The argument is assumed to be a non-empty strict vector:+  trfromVector :: (KnownNat n, KnownSTK x)+               => Data.Vector.Vector (target (TKR2 n x))+               -> target (TKR2 (1 + n) x)+  trfromVector t = trfromVectorN (V.length t :$: ZSR) t+  trfromVectorN :: forall m n x. (KnownNat n, KnownSTK x)+                => IShR m -> Data.Vector.Vector (target (TKR2 n x))+                -> target (TKR2 (m + n) x)+  trfromVectorLinear :: forall m r. GoodScalar r+                     => IShR m -> Data.Vector.Vector (target (TKScalar r))+                     -> target (TKR m r)+  trunravelToList :: (KnownNat n, KnownSTK x)+                  => target (TKR2 (1 + n) x) -> [target (TKR2 n x)]+  trunravelToList t = trunravelToListN (rwidth t :$: ZSR) t+  trunravelToListN :: forall m n x. (KnownNat n, KnownSTK x)+                   => IShR m -> target (TKR2 (m + n) x)+                   -> [target (TKR2 n x)]+  trunravelToListN shm t = map (trindex t) (shrEnum' shm)+  trtoListLinear :: forall m r. GoodScalar r+                 => target (TKR m r) -> [target (TKScalar r)]+  trtoListLinear t = map (trindex0 t) (shrEnum' (rshape t))++  tsfromVector :: (KnownNat k, KnownShS shn, KnownSTK x)+               => Data.Vector.Vector (target (TKS2 shn x))+               -> target (TKS2 (k ': shn) x)+  tsfromVector = tsfromVectorN (SNat :$$ ZSS)+  tsfromVectorN :: forall shm shn x. (KnownShS shn, KnownSTK x)+                => ShS shm -> Data.Vector.Vector (target (TKS2 shn x))+                -> target (TKS2 (shm ++ shn) x)+  tsfromVectorLinear :: forall shm r. GoodScalar r+                     => ShS shm -> Data.Vector.Vector (target (TKScalar r))+                     -> target (TKS shm r)+  tsunravelToList :: (KnownNat k, KnownShS shn, KnownSTK x)+                  => target (TKS2 (k ': shn) x) -> [target (TKS2 shn x)]+  tsunravelToList = tsunravelToListN (SNat :$$ ZSS)+  tsunravelToListN :: forall shm shn x. (KnownShS shn, KnownSTK x)+                   => ShS shm -> target (TKS2 (shm ++ shn) x)+                   -> [target (TKS2 shn x)]+  tsunravelToListN shm t = map (tsindex t) (shsEnum' shm)+  tstoListLinear :: forall shm r. GoodScalar r+                 => target (TKS shm r) -> [target (TKScalar r)]+  tstoListLinear t = map (tsindex0 t) (shsEnum' (sshape t))+  tsfromIxR :: forall k shn x. (1 <= k, KnownShS shn, KnownSTK x)+            => IxR k (target (TKS2 shn x))+            -> target (TKS2 (k ': shn) x)+  tsfromIxR l | SNat <- ixrRank l =+    tsfromVector $ V.fromList $ Foldable.toList l++  txfromVector :: (KnownNat k, KnownShX shn, KnownSTK x)+               => Data.Vector.Vector (target (TKX2 shn x))+               -> target (TKX2 (Just k ': shn) x)+  txfromVector = txfromVectorN (SKnown SNat :$% ZSX)+  txfromVectorN :: forall shm shn x. (KnownShX shn, KnownSTK x)+                => IShX shm -> Data.Vector.Vector (target (TKX2 shn x))+                -> target (TKX2 (shm ++ shn) x)+  txfromVectorLinear :: forall shm r. GoodScalar r+                     => IShX shm -> Data.Vector.Vector (target (TKScalar r))+                     -> target (TKX shm r)+  txunravelToList :: (KnownNat k, KnownShX shn, KnownSTK x)+                  => target (TKX2 (Just k ': shn) x) -> [target (TKX2 shn x)]+  txunravelToList = txunravelToListN (SKnown SNat :$% ZSX)+  txunravelToListN :: forall shm shn x. (KnownShX shn, KnownSTK x)+                   => IShX shm -> target (TKX2 (shm ++ shn) x)+                   -> [target (TKX2 shn x)]+  txunravelToListN shm t = map (txindex t) (shxEnum' shm)+  txtoListLinear :: forall shm r. GoodScalar r+                 => target (TKX shm r) -> [target (TKScalar r)]+  txtoListLinear t = map (txindex0 t) (shxEnum' (xshape t))++  -- A number suffix in the name may indicate the rank of the codomain,+  -- if bounded. Suffix 1 may also mean the operations builds up codomain+  -- by 1 dimension.+  trsum :: forall n x. (KnownNat n, TKAllNum x, KnownSTK x)+        => target (TKR2 (1 + n) x) -> target (TKR2 n x)+  trsum = trsumN+  trsumN :: forall m n x. (KnownNat m, KnownNat n, TKAllNum x, KnownSTK x)+         => target (TKR2 (m + n) x) -> target (TKR2 n x)+  -- This op (and it's Delta constructor) is worthwhile, because flattening+  -- is O(n) sometimes, unlike transpose, etc.+  trsum0 :: forall m r. NumScalar r+         => target (TKR m r) -> target (TKScalar r)+  trdot0 :: forall n r. (KnownNat n, NumScalar r)+         => target (TKR n r) -> target (TKR n r) -> target (TKScalar r)+  trdot0 t u = trsum0 (t * u)+  trdot1In :: forall n r. (KnownNat n, NumScalar r)+           => target (TKR (1 + n) r)+           -> target (TKR (1 + n) r)+           -> target (TKR n r)+  trdot1In t u = trsum $ trtranspose (permCycle $ 1 + valueOf @n) (t * u)+  trmatvecmul :: (NumScalar r, ConvertTensor target)+              => target (TKR 2 r) -> target (TKR 1 r) -> target (TKR 1 r)+-- How to generalize (#69)? The few straightforward generalizations+-- differ in types but all are far from matmul2.+-- rmatvecmul m v = rflatten $ rmap1 (rreplicate 1 . rdot0 v) m+  trmatvecmul m v =+    trbuild1 (rwidth m) (\i -> rfromK $ trdot0 v (m `trindex` [i]))+  trmatmul2 :: (NumScalar r, ConvertTensor target)+            => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)+-- How to generalize to tmatmul (#69)?+-- Just rmatmul2 the two outermost dimensions?+-- rmatmul2 m1 m2 = rmap1 (rmatvecmul (rtr m2)) m1+  trmatmul2 m1 m2 =+    trbuild1 (rwidth m1) (\i -> trmatvecmul (rtr m2) (m1 `trindex` [i]))+  trreplicate :: forall n x. KnownSTK x+              => Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)+  trreplicate k = trreplicateN (k :$: ZSR)+  trreplicateN :: forall m n x. KnownSTK x+               => IShR m -> target (TKR2 n x) -> target (TKR2 (m + n) x)+  trreplicate0N :: forall m r. GoodScalar r+                => IShR m -> target (TKScalar r) -> target (TKR m r)++  tssum :: forall n sh x. (KnownNat n, KnownShS sh, TKAllNum x, KnownSTK x)+        => target (TKS2 (n ': sh) x) -> target (TKS2 sh x)+  tssum = tssumN @_ @(n : '[])+  tssumN :: forall shm shn x.+            (KnownShS shm, KnownShS shn, TKAllNum x, KnownSTK x)+         => target (TKS2 (shm ++ shn) x) -> target (TKS2 shn x)+  tssum0 :: forall shm r. NumScalar r+         => target (TKS shm r) -> target (TKScalar r)+  tsdot0 :: forall sh r. (KnownShS sh, NumScalar r)+         => target (TKS sh r) -> target (TKS sh r) -> target (TKScalar r)+  tsdot0 t u = tssum0 (t * u)+  tsdot1In :: forall sh n r. (KnownShS sh, NumScalar r)+           => SNat n+           -> target (TKS (sh ++ '[n]) r)+           -> target (TKS (sh ++ '[n]) r)+           -> target (TKS sh r)+  tsdot1In SNat t u =+    let cpermR = permCycle $ 1 + shsLength (knownShS @sh)+    in Permutation.permFromListCont cpermR+       $ \(cperm :: Permutation.Perm cperm) ->+         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (sh ++ '[n])) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix cperm (sh ++ '[n])+                       :~: n : sh) $+         fromMaybe (error "tsdot1In: impossible non-permutation")+         $ Permutation.permCheckPermutation cperm+         $ tssum $ tstranspose cperm (t * u)+  tsmatvecmul :: (KnownNat m, KnownNat n, NumScalar r)+              => target (TKS '[m, n] r) -> target (TKS '[n] r)+              -> target (TKS '[m] r)+  tsmatvecmul @m m v =+    tkbuild1 @_ @m (\i -> tsdot0 v (m `tsindex` (i :.$ ZIS)))+  tsmatmul2 :: (KnownNat m, KnownNat n, KnownNat p, NumScalar r)+            => target (TKS '[m, n] r) -> target (TKS '[n, p] r)+            -> target (TKS '[m, p] r)+  tsmatmul2 @m m1 m2 =+    tsbuild1 @_ @m (\i -> tsmatvecmul (str m2) (m1 `tsindex` (i :.$ ZIS)))+  tsreplicate :: forall k sh x. KnownSTK x+              => SNat k -> target (TKS2 sh x)+              -> target (TKS2 (k ': sh) x)+  tsreplicate snat = tsreplicateN (snat :$$ ZSS)+  tsreplicateN :: forall shm shn x. KnownSTK x+               => ShS shm -> target (TKS2 shn x) -> target (TKS2 (shm ++ shn) x)+  tsreplicate0N :: forall shm r. GoodScalar r+                => ShS shm -> target (TKScalar r) -> target (TKS shm r)++  -- The choice in BuildTensorKind makes it hard to support this one,+  -- due to DeltaSum and AstSum being typed with BuildTensorKind:+  -- xsum :: (KnownShX sh, KnownShX (mn ': sh), KnownSTK x)+  --     => target (TKX2 (mn ': sh) x) -> target (TKX2 sh x)+  txsum :: forall n sh x. (KnownNat n, KnownShX sh, TKAllNum x, KnownSTK x)+        => target (TKX2 (Just n ': sh) x) -> target (TKX2 sh x)+  txsum = txsumN @_ @(Just n : '[])+  txsumN :: forall shm shn x.+            (KnownShX shm, KnownShX shn, TKAllNum x, KnownSTK x)+         => target (TKX2 (shm ++ shn) x) -> target (TKX2 shn x)+  txsum0 :: forall shm r. NumScalar r+         => target (TKX shm r) -> target (TKScalar r)+  txdot0 :: forall sh r. (KnownShX sh, NumScalar r)+         => target (TKX sh r) -> target (TKX sh r) -> target (TKScalar r)+  txdot0 t u = txsum0 (t * u)+  txdot1In :: forall sh n r. (KnownShX sh, NumScalar r)+           => SNat n+           -> target (TKX (sh ++ '[Just n]) r)+           -> target (TKX (sh ++ '[Just n]) r)+           -> target (TKX sh r)+  txdot1In SNat t u =+    let cpermR = permCycle $ 1 + fromSNat' (ssxRank (knownShX @sh))+    in Permutation.permFromListCont cpermR+       $ \(cperm :: Permutation.Perm cperm) ->+         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (sh ++ '[Just n])) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix cperm (sh ++ '[Just n])+                       :~: Just n : sh) $+         fromMaybe (error "txdot1In: impossible non-permutation")+         $ Permutation.permCheckPermutation cperm+         $ txsum $ txtranspose cperm (t * u)+  txmatvecmul :: forall mm mn r. (NumScalar r, ConvertTensor target)+              => Nested.SMayNat Int mm -> Nested.SMayNat Int mn+              -> target (TKX '[mm, mn] r) -> target (TKX '[mn] r)+              -> target (TKX '[mm] r)+  txmatvecmul mm mn m v =+    withKnownShX (ssxFromShX $ mm :$% ZSX) $+    withKnownShX (ssxFromShX $ mn :$% ZSX) $+    withSNat (fromSMayNat' mm) $ \(SNat @k) ->+      xmcast (ssxFromShX $ mm :$% ZSX)+      $ txbuild1 @_ @k (\i -> xfromK $ txdot0 v (m `txindex` (i :.% ZIX)))+  txmatmul2 :: ( KnownNat m, KnownNat n, KnownNat p, NumScalar r+               , ConvertTensor target )+            => target (TKX '[Just m, Just n] r)+            -> target (TKX '[Just n, Just p] r)+            -> target (TKX '[Just m, Just p] r)+  txmatmul2 @m @n @p m1 m2 =+    txbuild1 @_ @m (\i ->+      txmatvecmul (SKnown (SNat @p)) (SKnown (SNat @n))+                  (xtr m2) (m1 `txindex` (i :.% ZIX)))+  txreplicate :: forall k sh x. KnownSTK x+              => SNat k -> target (TKX2 sh x)+              -> target (TKX2 (Just k ': sh) x)+  txreplicate snat = txreplicateN (SKnown snat :$% ZSX)+  txreplicateN :: forall shm shn x. KnownSTK x+               => IShX shm -> target (TKX2 shn x)+               -> target (TKX2 (shm ++ shn) x)+  txreplicate0N :: forall shm r. GoodScalar r+                => IShX shm -> target (TKScalar r) -> target (TKX shm r)++  trindex :: forall m n x. (KnownNat n, KnownSTK x)+          => target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)+  trindex0 :: forall m r. GoodScalar r+           => target (TKR m r) -> IxROf target m -> target (TKScalar r)+  troneHot :: ( KnownNat m, KnownNat n, TKAllNum x, KnownSTK x+              , PlainOf (PlainOf target) ~ PlainOf target+              , EqH (PlainOf target) (TKScalar Int))+           => IShR m -> target (TKR2 n x) -> IxROf target m+           -> target (TKR2 (m + n) x)+  {-# INLINE troneHot #-}+  troneHot sh v ix = trscatter @_ @0 sh v (const ix)+      -- this code is often better for differentiable contexts, because+      -- a gather results, though this code is problematic if vectorization+      -- blows up the dimensions+    {- _ ->+      -- TODO: def at out of bounds and handle empty arrays+      let f ix2 = tcond knownSTK+                        (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true+                         $ zip (Foldable.toList ix) (Foldable.toList ix2))+                        v+                        (tdefTarget (tftk knownSTK v))+      in trbuild (shrAppend sh (rshape v)) f+        -- this code is probably better for non-differentiable contexts+        -- (even though it seems to do more work than the scatter above),+        -- because this code vectorizes better, often to the same size gather+        -- TODO: so maybe check if all scalars in v are non-differentiable+        -- and choose this if so?+        -- TODO: if this is used often, maybe express this as the gather that+        -- would come out of vectorization, to help it simplify well -}+  trscatter :: (KnownNat m, KnownNat n, KnownNat p, TKAllNum x, KnownSTK x)+            => IShR p -> target (TKR2 (m + n) x)+            -> (IxROf target m -> IxROf target p)+            -> target (TKR2 (p + n) x)+  trscatter1 :: (KnownNat n, KnownNat p, TKAllNum x, KnownSTK x)+             => IShR p -> target (TKR2 (1 + n) x)+             -> (IntOf target -> IxROf target p)+             -> target (TKR2 (p + n) x)+  {-# INLINE trscatter1 #-}+  trscatter1 sh v f = trscatter @target @1 sh v (\(i :.: ZIR) -> f i)+  trgather :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x)+           => IShR m -> target (TKR2 (p + n) x)+           -> (IxROf target m -> IxROf target p)+           -> target (TKR2 (m + n) x)+  trgather1 :: (KnownNat n, KnownNat p, KnownSTK x)+            => Int -> target (TKR2 (p + n) x)+            -> (IntOf target -> IxROf target p)+            -> target (TKR2 (1 + n) x)+  {-# INLINE trgather1 #-}+  trgather1 k v f = trgather @target @1 (k :$: ZSR) v (\(i :.: ZIR) -> f i)++  tsindex :: forall shm shn x. (KnownShS shn, KnownSTK x)+          => target (TKS2 (shm ++ shn) x) -> IxSOf target shm+          -> target (TKS2 shn x)+  tsindex0 :: forall shm r. GoodScalar r+           => target (TKS shm r) -> IxSOf target shm -> target (TKScalar r)+  tsoneHot :: ( KnownShS sh1, KnownShS sh2, TKAllNum x, KnownSTK x+              , PlainOf (PlainOf target) ~ PlainOf target+              , EqH (PlainOf target) (TKScalar Int) )+           => target (TKS2 sh2 x) -> IxSOf target sh1+           -> target (TKS2 (sh1 ++ sh2) x)+  {-# INLINE tsoneHot #-}+  tsoneHot v ix = tsscatter @_ @'[] v (const ix)+    {- _ | SNat <- shsRank (knownShS @sh1)+         , Refl <- lemAppNil @sh2 ->+      -- TODO: def at out of bounds and handle empty arrays+      withKnownShS (knownShS @sh1 `shsAppend` knownShS @sh2) $+          let f ix2 = tcond knownSTK+                            (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true+                             $ zip (Foldable.toList ix) (Foldable.toList ix2))+                            v+                            (tdefTarget (tftk knownSTK v))+          in tsbuild @_ @(Rank sh1) SNat f -}+  tsscatter+     :: (KnownShS shm, KnownShS shn, KnownShS shp, TKAllNum x, KnownSTK x)+     => target (TKS2 (shm ++ shn) x)+     -> (IxSOf target shm -> IxSOf target shp)+     -> target (TKS2 (shp ++ shn) x)+  tsscatter1+     :: (KnownNat n2, KnownShS shn, KnownShS shp, TKAllNum x, KnownSTK x)+     => target (TKS2 (n2 ': shn) x)+     -> (IntOf target -> IxSOf target shp)+     -> target (TKS2 (shp ++ shn) x)+  {-# INLINE tsscatter1 #-}+  tsscatter1 @n2 v f = tsscatter @_ @'[n2] v (\(i :.$ _) -> f i)+  tsgather+     :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x)+     => target (TKS2 (shp ++ shn) x)+     -> (IxSOf target shm -> IxSOf target shp)+     -> target (TKS2 (shm ++ shn) x)+  tsgather1+     :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x)+     => target (TKS2 (shp ++ shn) x)+     -> (IntOf target -> IxSOf target shp)+     -> target (TKS2 (n2 ': shn) x)+  {-# INLINE tsgather1 #-}+  tsgather1 @n2 v f = tsgather @target @'[n2] v (\(i :.$ _) -> f i)++  txindex :: forall shm shn x. (KnownShX shn, KnownSTK x)+          => target (TKX2 (shm ++ shn) x) -> IxXOf target shm+          -> target (TKX2 shn x)+  txindex0 :: forall shm r. GoodScalar r+           => target (TKX shm r) -> IxXOf target shm -> target (TKScalar r)+  txoneHot :: ( KnownShX sh1, KnownShX sh2, TKAllNum x, KnownSTK x+              , PlainOf (PlainOf target) ~ PlainOf target+              , EqH (PlainOf target) (TKScalar Int) )+           => IShX sh1 -> target (TKX2 sh2 x) -> IxXOf target sh1+           -> target (TKX2 (sh1 ++ sh2) x)+  {-# INLINE txoneHot #-}+  txoneHot sh1 v ix = txscatter @_ @'[] sh1 v (const ix)+    {- _ | SNat <- ssxRank (knownShX @sh1)+         , Refl <- lemAppNil @sh2 ->+      -- TODO: def at out of bounds and handle empty arrays+      withKnownShX (knownShX @sh1 `ssxAppend` knownShX @sh2) $+          let f ix2 = tcond knownSTK+                            (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true+                             $ zip (Foldable.toList ix) (Foldable.toList ix2))+                            v+                            (tdefTarget (tftk knownSTK v))+          in txbuild @_ @(Rank sh1) SNat (shxAppend sh1 (xshape v)) f -}+  txscatter :: ( KnownShX shm, KnownShX shn, KnownShX shp+               , TKAllNum x, KnownSTK x )+            => IShX shp -> target (TKX2 (shm ++ shn) x)+            -> (IxXOf target shm -> IxXOf target shp)+            -> target (TKX2 (shp ++ shn) x)+  -- TODO: generalize this to non-Just types.+  txscatter1 :: ( KnownNat n2, KnownShX shn, KnownShX shp+                , TKAllNum x, KnownSTK x )+             => IShX shp -> target (TKX2 (Just n2 ': shn) x)+             -> (IntOf target -> IxXOf target shp)+             -> target (TKX2 (shp ++ shn) x)+  {-# INLINE txscatter1 #-}+  txscatter1 @n2 @_ @shp @x sh v f = txscatter @_ @'[Just n2] @_ @shp @x sh v+                                               (\(i :.% _) -> f i)+  txgather :: (KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x)+           => IShX shm+           -> target (TKX2 (shp ++ shn) x)+           -> (IxXOf target shm -> IxXOf target shp)+           -> target (TKX2 (shm ++ shn) x)+  txgather1 :: (KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x)+            => SNat n2 -> target (TKX2 (shp ++ shn) x)+            -> (IntOf target -> IxXOf target shp)+            -> target (TKX2 (Just n2 ': shn) x)+  {-# INLINE txgather1 #-}+  txgather1 @n2 k v f =+    txgather @target @'[Just n2] (SKnown k :$% ZSX) v (\(i :.% ZIX) -> f i)++  tkfloor :: (NumScalar r, Differentiable r, NumScalar r2, Integral r2)+          => target (TKScalar r) -> target (TKScalar r2)+  -- The NumScalar r1 is only needed to rewrite well, thanks to the conversion+  -- respecting the NumScalar property.+  tkfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2)+                 => target (TKScalar r1) -> target (TKScalar r2)+  tkcast :: (Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2)+         => target (TKScalar r1) -> target (TKScalar r2)+  tkargMin, tkargMax  -- partial+    :: forall n r. NumScalar r+    => target (TKS '[n] r) -> target (TKScalar Int)++  trfloor :: (NumScalar r, Differentiable r, NumScalar r2, Integral r2)+          => target (TKR n r) -> target (TKR n r2)+  trfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2)+                 => target (TKR n r1) -> target (TKR n r2)+  trcast :: (Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2)+         => target (TKR n r1) -> target (TKR n r2)+  trargMin, trargMax  -- partial+    :: forall n r. NumScalar r+    => target (TKR (1 + n) r) -> target (TKR n Int)+  triota :: NumScalar r => Int -> target (TKR 1 r)  -- from 0 to n - 1++  tsfloor :: (NumScalar r, Differentiable r, NumScalar r2, Integral r2)+          => target (TKS sh r) -> target (TKS sh r2)+  tsfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2)+                 => target (TKS sh r1) -> target (TKS sh r2)+  tscast :: (Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2)+         => target (TKS sh r1) -> target (TKS sh r2)+  tsargMin, tsargMax  -- partial+    :: forall n sh r. NumScalar r+    => target (TKS (n ': sh) r) -> target (TKS (Init (n ': sh)) Int)+  tsiota :: (KnownNat n, NumScalar r)+         => target (TKS '[n] r)  -- from 0 to n - 1++  txfloor :: (NumScalar r, Differentiable r, NumScalar r2, Integral r2)+          => target (TKX sh r) -> target (TKX sh r2)+  txfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2)+                 => target (TKX sh r1) -> target (TKX sh r2)+  txcast :: (Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2)+         => target (TKX sh r1) -> target (TKX sh r2)+  txargMin, txargMax  -- partial+    :: forall mn sh r. NumScalar r+    => target (TKX (mn ': sh) r) -> target (TKX (Init (mn ': sh)) Int)+  txiota :: (KnownNat n, NumScalar r)+         => target (TKX '[Just n] r)  -- from 0 to n - 1++  trappend :: forall n x. KnownSTK x+           => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+           -> target (TKR2 (1 + n) x)+  trconcat :: forall n x. KnownSTK x+           => NonEmpty (target (TKR2 (1 + n) x)) -> target (TKR2 (1 + n) x)+  trconcat = foldl1' trappend . NonEmpty.toList+  trslice :: forall n x. KnownSTK x+          => Int -> Int -> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+  trreverse :: forall n x. KnownSTK x+            => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+  trtranspose :: forall n x. KnownSTK x+              => Permutation.PermR -> target (TKR2 n x) -> target (TKR2 n x)+  trreshape :: forall n m x. KnownSTK x+            => IShR m -> target (TKR2 n x) -> target (TKR2 m x)++  tsappend :: forall m n sh x. KnownSTK x+           => target (TKS2 (m ': sh) x) -> target (TKS2 (n ': sh) x)+           -> target (TKS2 ((m + n) ': sh) x)+  tsslice :: forall i n k sh x. KnownSTK x+          => SNat i -> SNat n -> SNat k+          -> target (TKS2 (i + n + k ': sh) x) -> target (TKS2 (n ': sh) x)+  tsreverse :: forall n sh x. KnownSTK x+            => target (TKS2 (n ': sh) x) -> target (TKS2 (n ': sh) x)+  tstranspose :: ( Permutation.IsPermutation perm+                 , Rank perm <= Rank sh, KnownSTK x )+              => Permutation.Perm perm -> target (TKS2 sh x)+              -> target (TKS2 (Permutation.PermutePrefix perm sh) x)+  tsreshape :: (Product sh ~ Product sh2, KnownSTK x)+            => ShS sh2 -> target (TKS2 sh x) -> target (TKS2 sh2 x)++  txappend :: forall m n sh x. KnownSTK x+           => target (TKX2 (m ': sh) x) -> target (TKX2 (n ': sh) x)+           -> target (TKX2 (AddMaybe m n ': sh) x)+  txconcat :: forall sh x. KnownSTK x+           => NonEmpty (target (TKX2 (Nothing ': sh) x))+           -> target (TKX2 (Nothing ': sh) x)+  txconcat = foldl1' txappend . NonEmpty.toList+  txslice :: forall i n k sh x. KnownSTK x+          => SMayNat Int i -> SMayNat Int n -> SMayNat Int k+          -> target (TKX2 (AddMaybe (AddMaybe i n) k ': sh) x)+          -> target (TKX2 (n ': sh) x)+  txreverse :: forall mn sh x. KnownSTK x+            => target (TKX2 (mn ': sh) x) -> target (TKX2 (mn ': sh) x)+  txtranspose :: ( Permutation.IsPermutation perm+                 , Rank perm <= Rank sh, KnownSTK x )+              => Permutation.Perm perm -> target (TKX2 sh x)+              -> target (TKX2 (Permutation.PermutePrefix perm sh) x)+  txreshape :: forall sh sh2 x. KnownSTK x+            => IShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)++  tkbuild1 :: (KnownNat k, GoodScalar r)+           => (IntOf target -> target (TKScalar r))+           -> target (TKS '[k] r)+  tkbuild :: (KnownShS sh, GoodScalar r, ConvertTensor target)+          => (IxSOf target sh -> target (TKScalar r))+          -> target (TKS sh r)+  {-# INLINE tkbuild #-}+  tkbuild @sh @r =+    let buildSh+          :: forall sh1.+             ShS sh1+          -> (IxSOf target sh1 -> target (TKScalar r))+          -> target (TKS sh1 r)+        buildSh sh1 f = case sh1 of+          ZSS -> sfromK $ f ZIS+          SNat :$$ ZSS ->+            tkbuild1 (\i -> f (i :.$ ZIS))+          SNat :$$ sh2 ->+            withKnownShS sh2 $+            let g i = buildSh sh2 (f . (i :.$))+            in tsbuild1 g+    in buildSh (knownShS @sh)+  trbuild1 :: (KnownNat n, KnownSTK x)+           => Int -> (IntOf target -> target (TKR2 n x))+           -> target (TKR2 (1 + n) x)+  trbuild :: (KnownNat m, KnownNat n, KnownSTK x)+          => IShR m+          -> (IxROf target m -> target (TKR2 n x))+          -> target (TKR2 (m + n) x)+  {-# INLINE trbuild #-}+  trbuild @_ @n @x shm f0 =+    let buildSh :: IShR m1 -> (IxROf target m1 -> target (TKR2 n x))+                -> target (TKR2 (m1 + n) x)+        buildSh ZSR f = f ZIR+        buildSh (k :$: sh) f | SNat <- shrRank sh =+          let g i = buildSh sh (\ix -> f (i :.: ix))+          in trbuild1 k g+    in buildSh shm f0+  trmap0N :: (KnownNat n, GoodScalar r1, GoodScalar r, ConvertTensor target)+          => (target (TKScalar r1) -> target (TKScalar r)) -> target (TKR n r1)+          -> target (TKR n r)+  trmap0N f v = trbuild (rshape v) (rfromK . f . trindex0 v)+  trzipWith0N :: ( KnownNat n, GoodScalar r, GoodScalar r1, GoodScalar r2+                 , ConvertTensor target )+              => (target (TKScalar r1) -> target (TKScalar r2)+                  -> target (TKScalar r))+              -> target (TKR n r1) -> target (TKR n r2) -> target (TKR n r)+  trzipWith0N f u v =+    trbuild (rshape v) (\ix -> rfromK $ f (trindex0 u ix) (trindex0 v ix))++  tsbuild1 :: (KnownNat k, KnownShS sh, KnownSTK x)+           => (IntOf target -> target (TKS2 sh x))+           -> target (TKS2 (k ': sh) x)+  tsbuild :: (KnownShS shm, KnownShS shn, KnownSTK x)+          => (IxSOf target shm -> target (TKS2 shn x))+          -> target (TKS2 (shm ++ shn) x)+  {-# INLINE tsbuild #-}+  tsbuild @shm @shn @x =+    let buildSh+          :: ShS shm1+          -> (IxSOf target shm1 -> target (TKS2 shn x))+          -> target (TKS2 (shm1 ++ shn) x)+        buildSh shm1 f = case shm1 of+          ZSS -> f ZIS+          SNat :$$ shm2 ->+            let g i = buildSh shm2 (f . (i :.$))+            in withKnownShS (shm2 `shsAppend` knownShS @shn) $+               tsbuild1 g+    in buildSh (knownShS @shm)+  tsmap0N :: (KnownShS sh, GoodScalar r1, GoodScalar r, ConvertTensor target)+          => (target (TKScalar r1) -> target (TKScalar r))+          -> target (TKS sh r1)+          -> target (TKS sh r)+  tsmap0N @sh f v | Refl <- lemAppNil @sh = tkbuild (f . tsindex0 v)+  tszipWith0N :: ( KnownShS sh, GoodScalar r, GoodScalar r1, GoodScalar r2+                 , ConvertTensor target )+              => (target (TKScalar r1) -> target (TKScalar r2)+                  -> target (TKScalar r))+              -> target (TKS sh r1) -> target (TKS sh r2)+              -> target (TKS sh r)+  tszipWith0N @sh f u v | Refl <- lemAppNil @sh =+    tkbuild (\ix -> f (tsindex0 u ix) (tsindex0 v ix))++  txbuild1 :: (KnownNat k, KnownShX sh, KnownSTK x)+           => (IntOf target -> target (TKX2 sh x))+           -> target (TKX2 (Just k ': sh) x)++  txbuild :: forall shm shn x. (KnownShX shn, KnownSTK x)+          => IShX shm+          -> (IxXOf target shm -> target (TKX2 shn x))+          -> target (TKX2 (shm ++ shn) x)+  {-# INLINE txbuild #-}+  txbuild shm f0 =+    let buildSh :: IShX shm1+                -> (IxXOf target shm1 -> target (TKX2 shn x))+                -> target (TKX2 (shm1 ++ shn) x)+        buildSh shm1 f = case shm1 of+          ZSX -> f ZIX+          k :$% shm2 ->+            let g i = buildSh shm2 (f . (i :.%))+            in withKnownShX (ssxFromShX shm2 `ssxAppend` knownShX @shn) $+               withSNat (fromSMayNat' k) $ \(SNat @n) ->+                 xmcast (ssxFromShX shm1 `ssxAppend` knownShX @shn)+                 $ txbuild1 @_ @n g+    in buildSh shm f0+  tbuild1 :: forall y k. ConvertTensor target+               -- y comes first, because k easy to set via SNat+          => SNat k -> SingletonTK y -> (IntOf target -> target y)+          -> target (BuildTensorKind k y)+  {-# INLINE tbuild1 #-}+  tbuild1 snat@SNat stk0 f =+    let replSTK :: SingletonTK z -> (IntOf target -> target z)+                -> target (BuildTensorKind k z)+        replSTK stk g = case stk of+          STKScalar -> tkbuild1 g+          STKR SNat x | Dict <- lemKnownSTK x -> trbuild1 (fromSNat' snat) g+          STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsbuild1 g+          STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txbuild1 g+          STKProduct @z1 @z2 stk1 stk2 ->+              let f1 i = tproject1 @_ @z1 @z2 $ g i+                  f2 i = tproject2 @_ @z1 @z2 $ g i+                    -- TODO: looks expensive, but hard to do better,+                    -- so let's hope g is full of variables+              in tpair (replSTK stk1 f1) (replSTK stk2 f2)+    in replSTK stk0 f++  -- | A strict right mapAccum.+  tmapAccumR+    :: forall accy by ey k.+       Proxy target+    -> SNat k  -- ^ length of the input+    -> FullShapeTK accy  -- ^ shape of the accumulator+    -> FullShapeTK by  -- ^ shape of the output+    -> FullShapeTK ey  -- ^ shape of an individual input+    -> (forall f. ADReady f+        => f accy -> f ey -> f (TKProduct accy by))+         -- ^ the function to mapAccum with+    -> target accy  -- ^ the initial accumulator+    -> target (BuildTensorKind k ey)  -- ^ the inputs+    -> target (TKProduct accy (BuildTensorKind k by))+  default tmapAccumR+    :: forall accy by ey k. ShareTensor target+    => Proxy target+    -> SNat k  -- ^ length of the input+    -> FullShapeTK accy  -- ^ shape of the accumulator+    -> FullShapeTK by  -- ^ shape of the output+    -> FullShapeTK ey  -- ^ shape of an individual input+    -> (forall f. ADReady f+        => f accy -> f ey -> f (TKProduct accy by))+         -- ^ the function to mapAccum with+    -> target accy  -- ^ the initial accumulator+    -> target (BuildTensorKind k ey)  -- ^ the inputs+    -> target (TKProduct accy (BuildTensorKind k by))+  {-# INLINE tmapAccumR #-}  -- this doesn't want to specialize+  tmapAccumR proxy !k !accftk !bftk !eftk f acc0 es =+    let (acc, l) =+          tunpair $ tmapAccumL proxy k accftk bftk eftk f acc0+                               (treverse k (ftkToSTK eftk) es)+    in tpair acc (treverse k (ftkToSTK bftk) l)+  -- | A strict left mapAccum.+  tmapAccumL+    :: forall accy by ey k.+       Proxy target+    -> SNat k  -- ^ length of the input+    -> FullShapeTK accy  -- ^ shape of the accumulator+    -> FullShapeTK by  -- ^ shape of the output+    -> FullShapeTK ey  -- ^ shape of an individual input+    -> (forall f. ADReady f+        => f accy -> f ey -> f (TKProduct accy by))+         -- ^ the function to mapAccum with+    -> target accy  -- ^ the initial accumulator+    -> target (BuildTensorKind k ey)  -- ^ the inputs+    -> target (TKProduct accy (BuildTensorKind k by))+  {-# INLINE tmapAccumL #-}  -- this doesn't want to specialize+  tmapAccumL proxy !k !accftk !bftk !eftk f acc0 es =+    let xftk = FTKProduct accftk eftk+        fl :: forall f. ADReady f+           => f (TKProduct accy ey)+           -> f (TKProduct accy by)+        fl !args = ttlet args $ \ !args1 ->  -- sharing just in case+                     f (tproject1 args1) (tproject2 args1)+    in tmapAccumLDer proxy k accftk bftk eftk+                     (tlambda @target xftk (HFun fl))+                     (tjvp @target xftk $ HFun fl)+                     (tvjp @target xftk $ HFun fl)+                     acc0 es+  -- | An strict right mapAccum with explicitly provided derivatives.+  tmapAccumRDer+    :: forall accy by ey k.+       Proxy target+    -> SNat k  -- ^ length of the input+    -> FullShapeTK accy  -- ^ shape of the accumulator+    -> FullShapeTK by  -- ^ shape of the output+    -> FullShapeTK ey  -- ^ shape of an individual input+    -> HFunOf target (TKProduct accy ey) (TKProduct accy by)+         -- ^ the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy by))+         -- ^ the derivative of the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy ey))+         -- ^ the reverse derivative of the function to mapAccum with+    -> target accy  -- ^ the initial accumulator+    -> target (BuildTensorKind k ey)  -- ^ the inputs+    -> target (TKProduct accy (BuildTensorKind k by))+  default tmapAccumRDer+    :: forall accy by ey k. ShareTensor target+    => Proxy target+    -> SNat k  -- ^ length of the input+    -> FullShapeTK accy  -- ^ shape of the accumulator+    -> FullShapeTK by  -- ^ shape of the output+    -> FullShapeTK ey  -- ^ shape of an individual input+    -> HFunOf target (TKProduct accy ey) (TKProduct accy by)+         -- ^ the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy by))+         -- ^ the derivative of the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy ey))+         -- ^ the reverse derivative of the function to mapAccum with+    -> target accy  -- ^ the initial accumulator+    -> target (BuildTensorKind k ey)  -- ^ the inputs+    -> target (TKProduct accy (BuildTensorKind k by))+  tmapAccumRDer proxy !k !accftk !bftk !eftk f df rf acc0 es =+    let (acc, l) =+          tunpair $ tmapAccumLDer proxy k accftk bftk eftk f df rf acc0+                                  (treverse k (ftkToSTK eftk) es)+    in tpair acc (treverse k (ftkToSTK bftk) l)+  -- | A strict left mapAccum with explicitly provided derivatives.+  --+  -- The applications of 'tjvp' and 'tvjp' performed already at this point+  -- ensure that the computation of a derivative is not repeated+  -- and that its result is shared. However, most of the time+  -- the computation is unnneeded, so the AST instance uses a non-strict+  -- constructor 'HordeAd.Core.Ast.AstLambda' for it's instance of 'HFunOf'.+  --+  -- If the same argument functions are passed to many mapAccum calls, as in+  -- > let f = ... in ... (tmapAccumR ... f ...) ... (tmapAccumL ... f ...)+  -- extra care is needed to prevent double derivative computation.+  -- One needs to use tmapAccumRDer manually as in (simplified)+  -- > let f = ...; df = tjvp f; rf = tgrad f+  -- > in ... (tmapAccumRDer f df rf ...) ... (tmapAccumLDer f df rf ...)+  tmapAccumLDer+    :: forall accy by ey k.+       Proxy target+    -> SNat k  -- ^ length of the input+    -> FullShapeTK accy  -- ^ shape of the accumulator+    -> FullShapeTK by  -- ^ shape of the output+    -> FullShapeTK ey  -- ^ shape of an individual input+    -> HFunOf target (TKProduct accy ey) (TKProduct accy by)+         -- ^ the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy by))+         -- ^ the derivative of the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy ey))+         -- ^ the reverse derivative of the function to mapAccum with+    -> target accy  -- ^ the initial accumulator+    -> target (BuildTensorKind k ey)  -- ^ the inputs+    -> target (TKProduct accy (BuildTensorKind k by))+  tapply :: HFunOf target x z -> target x -> target z+  tlambda :: FullShapeTK x -> HFun x z -> HFunOf target x z++  -- | Reverse derivative.+  --+  -- The following methods (and tlambda) are exactly what is needed as arguments+  -- of tmapAccumRDer.+  tgrad+    :: forall x r. GoodScalar r+    => FullShapeTK x  -- ^ shape of x and dx+    -> HFun x (TKScalar r)  -- ^ x |-> TKScalar r+    -> HFunOf target x (ADTensorKind x)  -- ^ x |-> dx+  tvjp+    :: FullShapeTK x  -- ^ shape of x and dx+    -> HFun x z  -- ^ x |-> z+    -> HFunOf target (TKProduct (ADTensorKind z) x) (ADTensorKind x)+         -- ^ (dz, x) |-> dx+  tjvp+    :: FullShapeTK x  -- ^ shape of x and dx+    -> HFun x z  -- ^ x |-> z+    -> HFunOf target (TKProduct (ADTensorKind x) x) (ADTensorKind z)+         -- ^ (dx, x) |-> dz++  tprimalPart :: target y -> PrimalOf target y+  tdualPart :: SingletonTK y -> target y -> DualOf target y+  tplainPart :: target y -> PlainOf target y+  tfromPrimal :: SingletonTK y -> PrimalOf target y -> target y+  tfromDual :: DualOf target y -> target y+  tfromPlain :: SingletonTK y -> PlainOf target y -> target y+  tScale :: (Num (target y), Num (PrimalOf target y))+         => SingletonTK y -> PrimalOf target y -> DualOf target y+         -> DualOf target y+  tScale stk s t =+    tdualPart stk $ tfromPrimal @target stk s * tfromDual t++  -- General operations that use ShareTensor if available, LetTensor otherwise+  tsum+    :: forall z k. TKAllNum z+    => SNat k -> SingletonTK z -> target (BuildTensorKind k z)+    -> target z+  default tsum+    :: forall z k. (ShareTensor target, ConvertTensor target, TKAllNum z)+    => SNat k -> SingletonTK z -> target (BuildTensorKind k z)+    -> target z+  {-# INLINE tsum #-}+  tsum snat@SNat stk u = case stk of+    STKScalar -> kfromS $ tssum u+    STKR SNat x | Dict <- lemKnownSTK x -> trsum u+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u+    STKProduct stk1 stk2 ->+      let (u1, u2) = tunpair u+      in tpair (tsum snat stk1 u1)+               (tsum snat stk2 u2)+  treplicate+    :: forall z k.+       SNat k -> SingletonTK z -> target z+    -> target (BuildTensorKind k z)+  default treplicate+    :: forall z k. (ShareTensor target, ConvertTensor target)+    => SNat k -> SingletonTK z -> target z+    -> target (BuildTensorKind k z)+  {-# INLINE treplicate #-}+  treplicate snat@SNat stk u = case stk of+    STKScalar -> tsreplicate snat $ sfromK u+    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (fromSNat' snat) u+    STKS _ x | Dict <- lemKnownSTK x -> tsreplicate snat u+    STKX _ x | Dict <- lemKnownSTK x -> txreplicate snat u+    STKProduct stk1 stk2 ->+      let (u1, u2) = tunpair u+      in tpair (treplicate snat stk1 u1)+               (treplicate snat stk2 u2)+  treverse+    :: forall z k.+       SNat k -> SingletonTK z -> target (BuildTensorKind k z)+    -> target (BuildTensorKind k z)+  default treverse+    :: forall z k. ShareTensor target+    => SNat k -> SingletonTK z -> target (BuildTensorKind k z)+    -> target (BuildTensorKind k z)+  {-# INLINE treverse #-}+  treverse snat stk u = case stk of+    STKScalar -> tsreverse u+    STKR _ x | Dict <- lemKnownSTK x -> trreverse u+    STKS _ x | Dict <- lemKnownSTK x -> tsreverse u+    STKX _ x | Dict <- lemKnownSTK x -> txreverse u+    STKProduct stk1 stk2 ->+      let (u1, u2) = tunpair u+      in tpair (treverse snat stk1 u1)+               (treverse snat stk2 u2)++  -- Unwinding methods, needed mostly to split off the Unwind module.+  -- | Add pointwise all corresponding tensors within nested product, if any.+  --+  -- Requires duplicable arguments or a 'ShareTensor' instance.+  taddTarget :: TKAllNum y+             => SingletonTK y -> target y -> target y -> target y+  -- | Multiply pointwise all corresponding tensors within nested products,+  -- if any.+  --+  -- Requires duplicable arguments or a 'ShareTensor' instance.+  tmultTarget :: TKAllNum y+              => SingletonTK y -> target y -> target y -> target y+  -- | Sum all dimensions of each component and then sum it all. Ignore all+  -- subtensors with non-differentiable elements.+  --+  -- Requires duplicable arguments or a 'ShareTensor' instance.+  tsum0Target :: TKAllNum y+              => FullShapeTK y -> target y+              -> target (TKScalar Double)+  -- | Dot product each component and then sum it all. Ignore all+  -- tensors with non-differentiable elements.+  --+  -- Requires duplicable arguments or a 'ShareTensor' instance.+  tdot0Target :: TKAllNum y+              => FullShapeTK y -> target y -> target y+              -> target (TKScalar Double)++  xmcast+    :: (KnownSTK x, KnownShX sh, Rank sh ~ Rank sh2)+    => StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)+  default xmcast+    :: (KnownSTK x, KnownShX sh, Rank sh ~ Rank sh2, ConvertTensor target)+    => StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)+  xmcast sh2 a = case tftk knownSTK a of+    FTKX sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        withKnownShX sh2 $+        withKnownShS sh $+        xfromS $ sfromX @_ @sh a++-- These are user-accessible, so the constraint is `ADReady`, which means+-- lets, but no shares.+type role HFun nominal nominal+newtype HFun (x :: TK) (z :: TK) =+  HFun {unHFun :: forall f. ADReady f+               => f x -> f z}++instance Show (HFun x y) where+  show _ = "<lambda>"+++-- * The mega-constraint++type ADReady target =+  ( ADReadyNoLet target+  , LetTensor target+  )++type ADReadyNoLet target =+  ( ADReadyEqsClasses target+  , ADReadyEqsClasses (ShareOf target)+  , ShareTensor (ShareOf target)+  , ShareTensor (PrimalOf (ShareOf target))+  , ShareTensor (PlainOf (ShareOf target))+  , ShareOf (ShareOf target) ~ ShareOf target+  )++type ADReadyEqsClasses f =+  ( ADReadyEqs f+  , LetTensor (PlainOf f)+  , ADReadyClasses f+  , ADReadyClasses (PrimalOf f)+  , ADReadyClasses (PlainOf f)+  )++type ADReadyEqs f =+  ( PlainOf (PlainOf f) ~ PlainOf f+  , PlainOf (PrimalOf f) ~ PlainOf f+  )++type ADReadyClasses f =+  ( BaseTensor f+  , ConvertTensor f+  , Boolean (BoolOf f)+  , AllTargetShow f+  , CommonTargetEqOrd f+  )++-- This is illegal:+-- type AllTargetShow target = forall y. KnownSTK y => Show (target y)++type AllTargetShow :: Target -> Constraint+class (forall y. KnownSTK y => Show (target y))+      => AllTargetShow target where+instance+      (forall y. KnownSTK y => Show (target y))+      => AllTargetShow target where++type CommonTargetEqOrd :: Target -> Constraint+class ( forall r. NumScalar r => EqH target (TKScalar r)+      , forall r. NumScalar r => OrdH target (TKScalar r)+      , forall r n. NumScalar r => EqH target (TKR n r)+      , forall r n. NumScalar r => OrdH target (TKR n r)+      , forall r sh. NumScalar r => EqH target (TKS sh r)+      , forall r sh. NumScalar r => OrdH target (TKS sh r)+      , forall r sh. NumScalar r => EqH target (TKX sh r)+      , forall r sh. NumScalar r => OrdH target (TKX sh r) )+      => CommonTargetEqOrd target where+instance+      ( forall r. NumScalar r => EqH target (TKScalar r)+      , forall r. NumScalar r => OrdH target (TKScalar r)+      , forall r n. NumScalar r => EqH target (TKR n r)+      , forall r n. NumScalar r => OrdH target (TKR n r)+      , forall r sh. NumScalar r => EqH target (TKS sh r)+      , forall r sh. NumScalar r => OrdH target (TKS sh r)+      , forall r sh. NumScalar r => EqH target (TKX sh r)+      , forall r sh. NumScalar r => OrdH target (TKX sh r) )       => CommonTargetEqOrd target where
src/HordeAd/Core/OpsADVal.hs view
@@ -4,16 +4,17 @@ -- are generic over whether the dual numbers are built from concrete arrays -- of floats or from AST terms or anything else (e.g., nested 'ADVal'). module HordeAd.Core.OpsADVal-  ( crevOnADInputs, crevOnParams, cfwdOnParams+  ( crevOnADInputs, crevOnParams, crevOnParamsDt, cfwdOnParams   ) where -import Prelude hiding (foldl')+import Prelude  import Data.Maybe (fromMaybe) import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Type.Equality (testEquality, (:~:) (Refl)) import Data.Vector.Generic qualified as V import GHC.TypeLits (sameNat)+import Type.Reflection (typeRep)  import Data.Array.Nested (Replicate, type (++)) import Data.Array.Nested qualified as Nested@@ -23,28 +24,30 @@ import Data.Array.Nested.Permutation qualified as Permutation import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (unsafeCoerceRefl)+import Data.Array.Nested.Types (pattern SZ)  import HordeAd.Core.CarriersADVal import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Conversion import HordeAd.Core.ConvertTensor import HordeAd.Core.Delta import HordeAd.Core.DeltaEval import HordeAd.Core.Ops import HordeAd.Core.TensorKind import HordeAd.Core.Types-import HordeAd.Core.Unwind+import HordeAd.Core.UnwindNum  -- * Non-symbolic (or at least non-sharing) reverse and forward derivative computation  -- The user-written function f can do anything, so the inputs -- argument has to be duplicable. crevOnADInputs-  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  :: forall x z target.+     (TKAllNum (ADTensorKind z), ADReadyNoLet target, ShareTensor target)   => Maybe (target (ADTensorKind z))   -> (ADVal target x -> ADVal target z)   -> FullShapeTK x -> ADVal target x-  -> (target (ADTensorKind x), target z)+  -> (target z, target (ADTensorKind x)) -- Break the inline chain to prevent false positives in inspection testing. -- {-# INLINE crevOnADInputs #-} crevOnADInputs mdt f xftk inputs =@@ -53,40 +56,72 @@       !(D v delta) = f inputs in   let zftk = ftkDelta delta       dt = fromMaybe (treplTarget 1 $ adFTK zftk) mdt-      !gradient = gradientFromDelta xftk zftk dt delta-  in (gradient, v)+      !gradient = gradientFromDelta xftk dt delta+  in (v, gradient)  crevOnParams-  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  :: forall x z target.+     (TKAllNum (ADTensorKind z), ADReadyNoLet target, ShareTensor target)   => Maybe (target (ADTensorKind z))   -> (ADVal target x -> ADVal target z)   -> FullShapeTK x -> target x-  -> (target (ADTensorKind x), target z)+  -> (target z, target (ADTensorKind x)) {-# INLINE crevOnParams #-} crevOnParams edt f xftk parameters =   let deltaInputs = generateDeltaInputs xftk       inputs = dDnotShared parameters deltaInputs   in crevOnADInputs edt f xftk inputs +-- These two functions are as above, but the dt must be provided and so,+-- due to technical reasons, the type is less constrained.+-- The user-written function f can do anything, so the inputs+-- argument has to be duplicable.+crevOnADInputsDt+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => target (ADTensorKind z)+  -> (ADVal target x -> ADVal target z)+  -> FullShapeTK x -> ADVal target x+  -> (target z, target (ADTensorKind x))+-- Break the inline chain to prevent false positives in inspection testing.+-- {-# INLINE crevOnADInputsDt #-}+crevOnADInputsDt dt f xftk inputs =+  let -- Evaluate completely after terms constructed, to free memory+      -- before evaluation allocates new memory and new FFI is started.+      !(D v delta) = f inputs in+  let gradient = gradientFromDelta xftk dt delta+  in (v, gradient)++crevOnParamsDt+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => target (ADTensorKind z)+  -> (ADVal target x -> ADVal target z)+  -> FullShapeTK x -> target x+  -> (target z, target (ADTensorKind x))+{-# INLINE crevOnParamsDt #-}+crevOnParamsDt dt f xftk parameters =+  let deltaInputs = generateDeltaInputs xftk+      inputs = dDnotShared parameters deltaInputs+  in crevOnADInputsDt dt f xftk inputs+ cfwdOnADInputs   :: forall x z target. (ADReadyNoLet target, ShareTensor target)   => FullShapeTK x -> ADVal target x   -> (ADVal target x -> ADVal target z)   -> target (ADTensorKind x)-  -> (target (ADTensorKind z), target z)+  -> (target z, target (ADTensorKind z)) -- Break the inline chain to prevent false positives in inspection testing. -- {-# INLINE cfwdOnADInputs #-} cfwdOnADInputs xftk inputs f ds =   let !(D v delta) = f inputs in-  let !derivative = derivativeFromDelta @x delta (adFTK xftk) ds-  in (derivative, v)+  let !derivative = derivativeFromDelta (Proxy @x) delta (adFTK xftk) ds+  in (v, derivative)  cfwdOnParams   :: forall x z target. (ADReadyNoLet target, ShareTensor target)   => FullShapeTK x -> target x   -> (ADVal target x -> ADVal target z)   -> target (ADTensorKind x)-  -> (target (ADTensorKind z), target z)+  -> (target z, target (ADTensorKind z)) {-# INLINE cfwdOnParams #-} cfwdOnParams xftk parameters f ds =   let deltaInputs = generateDeltaInputs xftk@@ -96,11 +131,9 @@  -- * Instances -fromPrimalFTK :: FullShapeTK z -> f z -> ADVal f z-fromPrimalFTK ftk a = dDnotShared a (DeltaZero ftk)--instance ( ADReadyNoLet target, ShareTensor target-         , ShareTensor (PrimalOf target) )+-- TODO: inline all ops that take functions here, in OpsAst and OpsConcrete+-- and benchmark to see if that's a good idea+instance (ADReadyNoLet target, ShareTensor target, ShareTensor (PlainOf target))          => LetTensor (ADVal target) where   ttlet (D u u') f =     let !var2 = tshare u@@ -108,6 +141,9 @@   ttletPrimal u f =     let !var2 = tshare u     in f var2+  ttletPlain u f =+    let !var2 = tshare u+    in f var2   toShare = id   tunshare = id   tD _stk = dD@@ -126,12 +162,74 @@ -- needed for the interpretation of Ast in ADVal. -- The ADVal Concrete instantiation is used in other pipelines and tests. instance ( ADReadyNoLet target, ShareTensor target-         , ShareTensor (PrimalOf target) )+         , ShareTensor (PrimalOf target), ShareTensor (PlainOf target) )          => BaseTensor (ADVal target) where-  -- Ranked ops+  isConcreteInstance = False   rshape (D u _) = rshape u-  trsum (D u u') = withSNat (rwidth u) $ \snat ->-    dD (trsum u) (DeltaSum snat knownSTK u')+  sshape (D u _) = sshape u+  xshape (D u _) = xshape u+  tftk _stk (D _ u') = ftkDelta u'+  tpair (D u u') (D v v') = dDnotShared (tpair u v) (DeltaPair u' v')+  tproject1 (D u u') = dDnotShared (tproject1 u) (fst $ unDeltaPairUnshared u')+  tproject2 (D u u') = dDnotShared (tproject2 u) (snd $ unDeltaPairUnshared u')+  -- Bangs are for the proper order of sharing stamps.+  kcond !b !u !v =+    let uv = tsfromVectorLinear (SNat @2 :$$ ZSS) (V.fromList [u, v])+    in tsindex0 uv (kcond b 0 1 :.$ ZIS)+  scond !b !u !v =+    let uv = tsfromVector @_ @2 (V.fromList [u, v])+    in tsindex uv (kcond b 0 1 :.$ ZIS)+  tcond stk !b !u !v = case stk of+    STKScalar -> kcond b u v+    STKR _ x | FTKR sh0 _ <- tftk stk u ->+      withShsFromShR sh0 $ \(sh :: ShS sh) ->+        withKnownShS sh $+        withKnownSTK x $+        rfromS $ scond b (sfromR @_ @sh u) (sfromR v)+    STKS sh x -> withKnownShS sh $ withKnownSTK x $ scond b u v+    STKX ssx0 x | FTKX sh0 _ <- tftk stk u ->+      withShsFromShX sh0 $ \(sh :: ShS sh) ->+        withKnownShS sh $+        withKnownShX ssx0 $+        withKnownSTK x $+        xfromS $ scond b (sfromX @_ @sh u) (sfromX v)+    STKProduct stk1 stk2 ->+      let (u1, u2) = tunpair u+          (v1, v2) = tunpair v+      in tpair (tcond stk1 b u1 v1) (tcond stk2 b u2 v2)+  tkconcrete a =+    let v = tkconcrete a+    in fromPrimalFTK FTKScalar v+  trconcrete a =+    let v = trconcrete a+    in fromPrimalFTK (FTKR (Nested.rshape a) FTKScalar) v+  tsconcrete a =+    let v = tsconcrete a+    in fromPrimalFTK (FTKS (Nested.sshape a) FTKScalar) v+  txconcrete a =+    let v = txconcrete a+    in fromPrimalFTK (FTKX (Nested.mshape a) FTKScalar) v+  tconcrete ftk t | Dict <- lemKnownSTK (ftkToSTK ftk) =+    fromPrimalFTK ftk $ tconcrete ftk t+  trfromVectorN shm lu =+    dD (trfromVectorN shm $ V.map (\(D u _) -> u) lu)+       (DeltaFromVectorR shm $ V.map (\(D _ u') -> u') lu)+  trfromVectorLinear shm lu =+    dD (trfromVectorLinear shm $ V.map (\(D u _) -> u) lu)+       (DeltaFromVectorLinearR shm $ V.map (\(D _ u') -> u') lu)+  tsfromVectorN shm lu =+    dD (tsfromVectorN shm $ V.map (\(D u _) -> u) lu)+       (DeltaFromVectorS shm $ V.map (\(D _ u') -> u') lu)+  tsfromVectorLinear shm lu =+    dD (tsfromVectorLinear shm $ V.map (\(D u _) -> u) lu)+       (DeltaFromVectorLinearS shm $ V.map (\(D _ u') -> u') lu)+  txfromVectorN shm lu =+    dD (txfromVectorN shm $ V.map (\(D u _) -> u) lu)+       (DeltaFromVectorX shm $ V.map (\(D _ u') -> u') lu)+  txfromVectorLinear shm lu =+    dD (txfromVectorLinear shm $ V.map (\(D u _) -> u) lu)+       (DeltaFromVectorLinearX shm $ V.map (\(D _ u') -> u') lu)+  trsumN (D u u') = dD (trsumN u) (DeltaSumR SNat u')   trsum0 (D u u') = dD (trsum0 u) (DeltaSum0R u')   trdot0 (D ue u') (D ve v') =     -- The bangs below are neccessary for GHC 9.2.7 test results to match 9.4.@@ -145,58 +243,11 @@     _ :$: width2 :$: ZSR ->       trsum (trtranspose [2,1,0] (trreplicate width2 m1)              * trtranspose [1,0] (trreplicate (rwidth m1) m2))-  trreplicate k (D u u') = withSNat k $ \snat ->-    dD (trreplicate k u) (DeltaReplicate snat knownSTK u')-  trindex (D u u') i =-    let !ix = tshare . tprimalPart <$> i-    in dD (trindex u ix) (DeltaIndexR SNat u' ix)-  trscatter sh (D u u') f =-    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)-    in dD (trscatter sh u g) (DeltaScatterR SNat SNat SNat sh u' g)-  trgather sh (D u u') f =-    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)-    in dD (trgather sh u g) (DeltaGatherR SNat SNat SNat sh u' g)-      -- Note how f is not interpreted as a function on dual numbers-      -- but just on integers and so no cotangents for results of application-      -- of f have to be computed and stored in contangent maps later on.-      -- Note also how g is duplicated and this leads to loss of sharing-      -- of indexes in AST instances.-  trconcrete a =-    let v = trconcrete a-    in fromPrimalFTK (FTKR (Nested.rshape a) FTKScalar) v-  trfloor (D u _) =-    let v = trfloor u-    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v-  trfromIntegral (D u _) =-    let v = trfromIntegral u-    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v-  trcast (D u u') = dD (trcast u) (DeltaCastR u')-  trminIndex (D u _) =-    let v = trminIndex u-    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v-  trmaxIndex (D u _) =-    let v = trmaxIndex u-    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v-  triota n = fromPrimalFTK (FTKR (n :$: ZSR) FTKScalar) $ triota n-  trappend (D u u') (D v v') = dD (trappend u v) (DeltaAppendR u' v')-  trslice i n (D u u') = dD (trslice i n u) (DeltaSliceR i n u')-  trreverse (D u u') = dD (trreverse u) (DeltaReverseR u')-  trtranspose perm (D u u') = dD (trtranspose perm u) (DeltaTransposeR perm u')-  trreshape sh (D u u') = dD (trreshape sh u) (DeltaReshapeR sh u')-  trbuild1 @n @x k f =-    let l = [0 .. fromIntegral k - 1]-    in if null l-       then case sameNat (Proxy @n) (Proxy @0) of-         Just Refl | Dict <- eltDictRep (knownSTK @x) ->-           let arr = Nested.remptyArray-           in tconcrete (tftkG knownSTK arr) (Concrete arr)-         Nothing -> error "rbuild1: shape ambiguity"-       else trfromVector $ V.fromList $ map (f . fromInteger) l-              -- hope this fuses--  -- Shaped ops-  sshape (D u _) = sshape u-  tssum (D u u') = dD (tssum u) (DeltaSum SNat knownSTK u')+  trreplicateN shm (D u u') =+    dD (trreplicateN shm u) (DeltaReplicateR shm u')+  trreplicate0N shm (D u u') =+    dD (trreplicate0N shm u) (DeltaReplicate0NR shm u')+  tssumN @shm (D u u') = dD (tssumN @_ @shm u) (DeltaSumS (knownShS @shm) u')   tssum0 (D u u') = dD (tssum0 u) (DeltaSum0S u')   tsdot0 (D ue u') (D ve v') =     -- The bangs below are neccessary for GHC 9.2.7 test results to match 9.4.@@ -205,55 +256,17 @@     in dD (tsdot0 u v) (dAdd (DeltaDot0S v u') (DeltaDot0S u v'))   -- These two are manually vectorized to avoid delta blowup when run   -- via primitive pipelines.-  tsmatvecmul m v = tssum (str (tsreplicate SNat knownShS v * m))+  tsmatvecmul m v = tssum (str (tsreplicate SNat v * m))   tsmatmul2 m1 m2 =     tssum (tstranspose (Permutation.makePerm @'[2, 1, 0])-                       (tsreplicate SNat knownShS m1)+                       (tsreplicate SNat m1)            * tstranspose (Permutation.makePerm @'[1, 0])-                         (tsreplicate SNat knownShS m2))-  tsindex (D u u') i =-    let !ix = tshare . tprimalPart <$> i-    in dD (tsindex u ix) (DeltaIndexS knownShS u' ix)-  tsscatter @shm @shn @shp (D u u') f =-    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)-    in dD (tsscatter @_ @shm @shn @shp u g)-          (DeltaScatterS @shm @shn @shp knownShS knownShS knownShS u' g)-  tsgather @shm @shn @shp (D u u') f =-    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)-    in dD (tsgather @_ @shm @shn @shp u g)-          (DeltaGatherS @shm @shn @shp knownShS knownShS knownShS u' g)-  tsconcrete a =-    let v = tsconcrete a-    in fromPrimalFTK (FTKS (Nested.sshape a) FTKScalar) v-  tsfloor (D u _) =-    let v = tsfloor u-    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v-  tsfromIntegral (D u _) =-    let v = tsfromIntegral u-    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v-  tscast (D u u') = dD (tscast u) (DeltaCastS u')-  tsminIndex (D u _) =-    let v = tsminIndex u-    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v-  tsmaxIndex (D u _) =-    let v = tsmaxIndex u-    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v-  tsiota = fromPrimalFTK (FTKS (SNat :$$ ZSS) FTKScalar) tsiota-  tsappend (D u u') (D v v') = dD (tsappend u v) (DeltaAppendS u' v')-  tsslice i n k (D u u') = dD (tsslice i n k u) (DeltaSliceS i n k u')-  tsreverse (D u u') = dD (tsreverse u) (DeltaReverseS u')-  tsbuild1 @k @sh @r f | Dict <- eltDictRep (knownSTK @r) =-    let l = [0 .. valueOf @k - 1]-    in if null l-       then let arr = Nested.semptyArray @(RepConcrete r) (knownShS @sh)-            in gcastWith (unsafeCoerceRefl :: k :~: 0) $-               tconcrete (tftkG knownSTK arr) (Concrete arr)-       else tsfromVector $ V.fromList $ map (f . fromInteger) l-              -- hope this fuses--  -- Mixed ops-  xshape (D u _) = xshape u-  txsum (D u u') = dD (txsum u) (DeltaSum SNat knownSTK u')+                         (tsreplicate SNat m2))+  tsreplicateN shm (D u u') =+    dD (tsreplicateN shm u) (DeltaReplicateS shm u')+  tsreplicate0N shm (D u u') =+    dD (tsreplicate0N shm u) (DeltaReplicate0NS shm u')+  txsumN @shm (D u u') = dD (txsumN @_ @shm u) (DeltaSumX (knownShX @shm) u')   txsum0 (D u u') = dD (txsum0 u) (DeltaSum0X u')   txdot0 (D ue u') (D ve v') =     -- The bangs below are neccessary for GHC 9.2.7 test results to match 9.4.@@ -268,33 +281,106 @@     withSNat (fromSMayNat' mm) $ \(SNat @m) ->     withSNat (fromSMayNat' mn) $ \(SNat @n) ->       xmcast (ssxFromShX (mm :$% ZSX))-      $ txsum (xtr (txreplicate (SNat @m) knownShX-                      (xmcast (ssxFromShX (Nested.SKnown (SNat @n)-                                             :$% ZSX)) v)-                    * xmcast (ssxFromShX (Nested.SKnown (SNat @m)-                                            :$% Nested.SKnown (SNat @n)-                                            :$% ZSX)) m))+      $ txsum (xtr (txreplicate (SNat @m)+                      (xmcast (ssxFromShX (SKnown (SNat @n) :$% ZSX)) v)+                    * xmcast (ssxFromShX (SKnown (SNat @m)+                                          :$% SKnown (SNat @n)+                                          :$% ZSX)) m))   txmatmul2 m1 m2 =     txsum (txtranspose (Permutation.makePerm @'[2, 1, 0])-                       (txreplicate SNat knownShX m1)+                       (txreplicate SNat m1)            * txtranspose (Permutation.makePerm @'[1, 0])-                         (txreplicate SNat knownShX m2))-  txreplicate snat sh (D u u') =-    dD (txreplicate snat sh u) (DeltaReplicate snat (STKX sh knownSTK) u')+                         (txreplicate SNat m2))+  txreplicateN shm (D u u') =+    dD (txreplicateN shm u) (DeltaReplicateX shm u')+  txreplicate0N shm (D u u') =+    dD (txreplicate0N shm u) (DeltaReplicate0NX shm u')+  trindex (D u u') i =+    let !ix = tshare <$> i+    in dD (trindex u ix) (DeltaIndexR SNat u' ix)+  trindex0 (D u u') i =+    let !ix = tshare <$> i+    in dD (trindex0 u ix) (DeltaIndex0R u' ix)+  trscatter sh (D u u') f =+    dD (trscatter sh u f) (DeltaScatterR SNat SNat SNat sh u' f)+  trgather sh (D u u') f =+    dD (trgather sh u f) (DeltaGatherR SNat SNat SNat sh u' f)+      -- Note how f is not interpreted as a function on dual numbers+      -- but just on integers and so no cotangents for results of application+      -- of f have to be computed and stored in contangent maps later on.+      -- Oh, and also a fun little detail: note how f is duplicated and this+      -- leads to loss of sharing of indexes in AST instances. If the index+      -- expressions are large, e.g., they perform a lookup in a non-constant+      -- big table, the user needs to work around this loss of sharing by, e.g.,+      -- adding a let expression with the said lookup table. Inlining should+      -- not inline such tables if the interation is non-trivial.+      -- Another workaround is to make the tables constant, precompute them+      -- and then all their copies are shared as Haskell values.+  tsindex (D u u') i =+    let !ix = tshare <$> i+    in dD (tsindex u ix) (DeltaIndexS knownShS u' ix)+  tsindex0 (D u u') i =+    let !ix = tshare <$> i+    in dD (tsindex0 u ix) (DeltaIndex0S u' ix)+  tsscatter @shm @shn @shp (D u u') f =+    dD (tsscatter @_ @shm @shn @shp u f)+       (DeltaScatterS @shm @shn @shp knownShS knownShS knownShS u' f)+  tsgather @shm @shn @shp (D u u') f =+    dD (tsgather @_ @shm @shn @shp u f)+       (DeltaGatherS @shm @shn @shp knownShS knownShS knownShS u' f)   txindex (D u u') i =-    let !ix = tshare . tprimalPart <$> i+    let !ix = tshare <$> i     in dD (txindex u ix) (DeltaIndexX knownShX u' ix)+  txindex0 (D u u') i =+    let !ix = tshare <$> i+    in dD (txindex0 u ix) (DeltaIndex0X u' ix)   txscatter @shm @shn @shp sh (D u u') f =-    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)-    in dD (txscatter @_ @shm @shn @shp sh u g)-          (DeltaScatterX @shm @shn @shp knownShX knownShX knownShX sh u' g)+    dD (txscatter @_ @shm @shn @shp sh u f)+       (DeltaScatterX @shm @shn @shp knownShX knownShX knownShX sh u' f)   txgather @shm @shn @shp sh (D u u') f =-    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)-    in dD (txgather @_ @shm @shn @shp sh u g)-          (DeltaGatherX @shm @shn @shp knownShX knownShX knownShX sh u' g)-  txconcrete a =-    let v = txconcrete a-    in fromPrimalFTK (FTKX (Nested.mshape a) FTKScalar) v+    dD (txgather @_ @shm @shn @shp sh u f)+       (DeltaGatherX @shm @shn @shp knownShX knownShX knownShX sh u' f)+  tkfloor (D u _) =+    let v = tkfloor u+    in fromPrimalFTK FTKScalar v+  tkfromIntegral (D u _) =+    let v = tkfromIntegral u+    in fromPrimalFTK FTKScalar v+  tkcast (D u u') = dD (tkcast u) (DeltaCastK u')+  tkargMin (D u _) =+    let v = tkargMin u+    in fromPrimalFTK FTKScalar v+  tkargMax (D u _) =+    let v = tkargMax u+    in fromPrimalFTK FTKScalar v+  trfloor (D u _) =+    let v = trfloor u+    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v+  trfromIntegral (D u _) =+    let v = trfromIntegral u+    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v+  trcast (D u u') = dD (trcast u) (DeltaCastR u')+  trargMin (D u _) =+    let v = trargMin u+    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v+  trargMax (D u _) =+    let v = trargMax u+    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v+  triota n = fromPrimalFTK (FTKR (n :$: ZSR) FTKScalar) $ triota n+  tsfloor (D u _) =+    let v = tsfloor u+    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v+  tsfromIntegral (D u _) =+    let v = tsfromIntegral u+    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v+  tscast (D u u') = dD (tscast u) (DeltaCastS u')+  tsargMin (D u _) =+    let v = tsargMin u+    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v+  tsargMax (D u _) =+    let v = tsargMax u+    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v+  tsiota = fromPrimalFTK (FTKS (SNat :$$ ZSS) FTKScalar) tsiota   txfloor (D u _) =     let v = txfloor u     in fromPrimalFTK (FTKX (xshape v) FTKScalar) v@@ -302,117 +388,93 @@     let v = txfromIntegral u     in fromPrimalFTK (FTKX (xshape v) FTKScalar) v   txcast (D u u') = dD (txcast u) (DeltaCastX u')-  txminIndex (D u _) =-    let v = txminIndex u+  txargMin (D u _) =+    let v = txargMin u     in fromPrimalFTK (FTKX (xshape v) FTKScalar) v-  txmaxIndex (D u _) =-    let v = txmaxIndex u+  txargMax (D u _) =+    let v = txargMax u     in fromPrimalFTK (FTKX (xshape v) FTKScalar) v-  txiota = fromPrimalFTK (FTKX (Nested.SKnown SNat :$% ZSX) FTKScalar) txiota+  txiota = fromPrimalFTK (FTKX (SKnown SNat :$% ZSX) FTKScalar) txiota+  trappend (D u u') (D v v') = dD (trappend u v) (DeltaAppendR u' v')+  trslice i n (D u u') = dD (trslice i n u) (DeltaSliceR i n u')+  trreverse (D u u') = dD (trreverse u) (DeltaReverseR u')+  trtranspose perm (D u u') = dD (trtranspose perm u) (DeltaTransposeR perm u')+  trreshape sh (D u u') = dD (trreshape sh u) (DeltaReshapeR sh u')+  tsappend (D u u') (D v v') = dD (tsappend u v) (DeltaAppendS u' v')+  tsslice i n k (D u u') = dD (tsslice i n k u) (DeltaSliceS i n k u')+  tsreverse (D u u') = dD (tsreverse u) (DeltaReverseS u')+  tstranspose perm (D u u') =+    dD (tstranspose perm u) (DeltaTransposeS @_ @_ @_ @target perm u')+  tsreshape sh (D u u') = dD (tsreshape sh u) (DeltaReshapeS sh u')   txappend (D u u') (D v v') = dD (txappend u v) (DeltaAppendX u' v')   txslice i n k (D u u') = dD (txslice i n k u) (DeltaSliceX i n k u')   txreverse (D u u') = dD (txreverse u) (DeltaReverseX u')   txtranspose perm (D u u') =     dD (txtranspose perm u) (DeltaTransposeX @_ @_ @_ @target perm u')   txreshape sh (D u u') = dD (txreshape sh u) (DeltaReshapeX sh u')-  txbuild1 @k @sh @r f =-    let l = [0 .. valueOf @k - 1]-    in if null l-       then case testEquality (knownShX @sh) ZKX of-         Just Refl | Dict <- eltDictRep (knownSTK @r) ->-           let arr = Nested.memptyArray @(RepConcrete r) ZSX-           in gcastWith (unsafeCoerceRefl :: k :~: 0) $-              tconcrete (tftkG knownSTK arr) (Concrete arr)-         Nothing -> error "xbuild1: shape ambiguity"-       else txfromVector $ V.fromList $ map (f . fromInteger) l+  tkbuild1 @k f =+    case SNat @k of+      SZ ->+        let arr = Nested.semptyArray ZSS+        in tconcrete (FTKS (SNat @0 :$$ ZSS) FTKScalar) (Concrete arr)+      _ ->+        let l = [0 .. valueOf @k - 1 :: Int]+        in tsfromVectorLinear (SNat :$$ ZSS)+           $ V.fromListN (valueOf @k) $ map (f . fromIntegral) l+             -- hope this fuses+  trbuild1 @n @x k f =+    if k == 0+    then case sameNat (Proxy @n) (Proxy @0) of+      Just Refl | Dict <- eltDictRep (knownSTK @x) ->+        let arr = Nested.remptyArray+        in tconcrete (tftkG knownSTK arr) (Concrete arr)+      Nothing -> error "rbuild1: shape ambiguity"+    else let l = [0 .. k - 1]+         in trfromVector $ V.fromListN k $ map (f . fromIntegral) l               -- hope this fuses--  -- Scalar ops-  tkconcrete a =-    let v = tkconcrete a-    in fromPrimalFTK FTKScalar v-  tkfloor (D u _) =-    let v = tkfloor u-    in fromPrimalFTK FTKScalar v-  tkfromIntegral (D u _) =-    let v = tkfromIntegral u-    in fromPrimalFTK FTKScalar v-  tkcast (D u u') = dD (tkcast u) (DeltaCastK u')--  -- General operations that don't require LetTensor nor ShareTensor-  tftk _stk (D _ u') = ftkDelta u'-  tconcrete ftk t | Dict <- lemKnownSTK (ftkToSTK ftk) =-    fromPrimalFTK ftk $ tconcrete ftk t-  tpair (D u u') (D v v') = dDnotShared (tpair u v) (DeltaPair u' v')-  tproject1 (D u u') = dDnotShared (tproject1 u) (fst $ unDeltaPairUnshared u')-  tproject2 (D u u') = dDnotShared (tproject2 u) (snd $ unDeltaPairUnshared u')-  tsreplicate snat sh (D u u') =-    dD (tsreplicate snat sh u) (DeltaReplicate snat (STKS sh knownSTK) u')-  tstranspose perm (D u u') =-    dD (tstranspose perm u) (DeltaTransposeS @_ @_ @_ @target perm u')-  tsreshape sh (D u u') = dD (tsreshape sh u) (DeltaReshapeS sh u')-  tmapAccumRDer @accy @by @ey _ !k accftk bftk eftk f df rf acc0D esD-   | Dict <- lemKnownSTKOfBuild k (ftkToSTK accftk)+  tsbuild1 @k @sh @r f | Dict <- eltDictRep (knownSTK @r) =+    case SNat @k of+      SZ ->+        let arr = Nested.semptyArray @_ @(RepConcrete r) (knownShS @sh)+        in tconcrete (tftkG knownSTK arr) (Concrete arr)+      _ ->+        let l = [0 .. valueOf @k - 1 :: Int]+        in tsfromVector $ V.fromListN (valueOf @k) $ map (f . fromIntegral) l+             -- hope this fuses+  txbuild1 @k @sh @r f =+    case SNat @k of+      SZ -> case testEquality (knownShX @sh) ZKX of+        Just Refl | Dict <- eltDictRep (knownSTK @r) ->+          let arr = Nested.memptyArray @_ @(RepConcrete r) ZSX+          in tconcrete (tftkG knownSTK arr) (Concrete arr)+        _ -> error "xbuild1: shape ambiguity"+      _ ->+        let l = [0 .. valueOf @k - 1 :: Int]+        in txfromVector $ V.fromListN (valueOf @k) $ map (f . fromIntegral) l+             -- hope this fuses+  tmapAccumLDer @accy @_ @ey  -- special case to speed up folds+                _ k accftk bftk@(FTKScalar @z1) eftk f df rf acc0D esD+   | Just Refl <- testEquality (typeRep @z1) (typeRep @Z1)+   , Dict <- lemKnownSTKOfBuild k (ftkToSTK accftk)    , Dict <- lemKnownSTKOfBuild k (ftkToSTK eftk) =     let !(D acc0 acc0') = acc0D in     let !(D esNotShared es') = esD in     let !es = tshare esNotShared-        codomainShs = FTKProduct accftk bftk+        codomainShs = accftk         g :: forall f. ADReady f-          => f (TKProduct accy ey)-          -> f (TKProduct accy (TKProduct accy by))-        g !acc_e =-          ttlet acc_e $ \ !acc_e1 ->-          ttlet (unHFun f acc_e) $ \ !accRes_bRes ->-            tpair (tproject1 accRes_bRes)-                  (tpair (tproject1 acc_e1) (tproject2 accRes_bRes))-        dg :: forall f. ADReady f-           => f (TKProduct (ADTensorKind (TKProduct accy ey))-                           (TKProduct accy ey))-           -> f (ADTensorKind (TKProduct accy (TKProduct accy by)))-        dg !dacc_de_acc_e =-          ttlet dacc_de_acc_e $ \ !dacc_de_acc_e1 ->-            let (!dacc_de, !_acc_e) =-                  (tproject1 dacc_de_acc_e1, tproject2 dacc_de_acc_e1)-                !dacc1 = tproject1 dacc_de-            in ttlet (unHFun df dacc_de_acc_e) $ \ !accRes_bRes ->-                 tpair (tproject1 accRes_bRes)-                       (tpair dacc1 (tproject2 accRes_bRes))-        rg :: forall f. ADReady f-           => f (TKProduct (ADTensorKind (TKProduct accy-                                         (TKProduct accy by)))-                           (TKProduct accy ey))-           -> f (ADTensorKind (TKProduct accy ey))-        rg !args =-          ttlet args $ \ args1 ->-            let (!dx_db, !acc_e) = (tproject1 args1, tproject2 args1)-            in ttlet dx_db $ \ !dx_db1 ->-              let (!dx, !db) = (tproject1 dx_db1, tproject2 dx_db1)-              in ttlet db $ \ !db1 ->-                let dx_dbRes = tpair dx (tproject2 db1)-                in ttlet (unHFun rf (tpair dx_dbRes acc_e))-                   $ \ !daccRes_deRes ->-                  let added = taddTarget (adSTK $ ftkToSTK accftk)-                                         (tproject1 daccRes_deRes)-                                         (tproject1 db1)-                  in tpair added (tproject2 daccRes_deRes)-        p = tmapAccumRDer (Proxy @target)-                          k accftk codomainShs eftk-                          (tlambda @target (FTKProduct accftk eftk)-                           $ HFun g)-                          (tlambda @target-                             (FTKProduct (adFTK (FTKProduct accftk eftk))-                                         (FTKProduct accftk eftk))-                           $ HFun dg)-                          (tlambda @target-                             (FTKProduct (adFTK (FTKProduct accftk codomainShs))-                                         (FTKProduct accftk eftk))-                           $ HFun rg)-                          acc0 es-        (accFin, qbs) = tunpair p-        (q, bs) = tunpair qbs-        dual = DeltaMapAccumR k bftk eftk q es df rf acc0' es'-    in dD (tpair accFin bs) dual+          => f accy -> f ey+          -> f (TKProduct accy accy)+        g !acc !e =+          let accRes_bRes = unHFun f (tpair acc e)+          in tpair (tproject1 accRes_bRes) acc+        p = tmapAccumL (Proxy @target)+                       k accftk codomainShs eftk+                       g+                       acc0 es+        (accFin, as) = tunpair p+        dual = DeltaMapAccumL k bftk eftk as es df rf acc0' es'+        lout2 = tsconcrete $ Nested.sreplicatePrim (k :$$ ZSS) Z1+    in dD (tpair accFin lout2) dual   tmapAccumLDer @accy @by @ey _ !k accftk bftk eftk f df rf acc0D esD    | Dict <- lemKnownSTKOfBuild k (ftkToSTK accftk)    , Dict <- lemKnownSTKOfBuild k (ftkToSTK eftk) =@@ -421,78 +483,29 @@     let !es = tshare esNotShared         codomainShs = FTKProduct accftk bftk         g :: forall f. ADReady f-          => f (TKProduct accy ey)+          => f accy -> f ey           -> f (TKProduct accy (TKProduct accy by))-        g !acc_e =-          ttlet acc_e $ \ !acc_e1 ->-          ttlet (unHFun f acc_e) $ \ !accRes_bRes ->+        g !acc !e =+          ttlet (unHFun f (tpair acc e)) $ \ !accRes_bRes ->             tpair (tproject1 accRes_bRes)-                  (tpair (tproject1 acc_e1) (tproject2 accRes_bRes))-        dg :: forall f. ADReady f-           => f (TKProduct (ADTensorKind (TKProduct accy ey))-                           (TKProduct accy ey))-           -> f (ADTensorKind (TKProduct accy (TKProduct accy by)))-        dg !dacc_de_acc_e =-          ttlet dacc_de_acc_e $ \ !dacc_de_acc_e1 ->-            let (!dacc_de, !_acc_e) =-                  (tproject1 dacc_de_acc_e1, tproject2 dacc_de_acc_e1)-                !dacc1 = tproject1 dacc_de-            in ttlet (unHFun df dacc_de_acc_e) $ \ !accRes_bRes ->-                 tpair (tproject1 accRes_bRes)-                       (tpair dacc1 (tproject2 accRes_bRes))-        rg :: forall f. ADReady f-           => f (TKProduct (ADTensorKind (TKProduct accy-                                         (TKProduct accy by)))-                           (TKProduct accy ey))-           -> f (ADTensorKind (TKProduct accy ey))-        rg !args =-          ttlet args $ \ args1 ->-            let (!dx_db, !acc_e) = (tproject1 args1, tproject2 args1)-            in ttlet dx_db $ \ !dx_db1 ->-              let (!dx, !db) = (tproject1 dx_db1, tproject2 dx_db1)-              in ttlet db $ \ !db1 ->-                let dx_dbRes = tpair dx (tproject2 db1)-                in ttlet (unHFun rf (tpair dx_dbRes acc_e))-                   $ \ !daccRes_deRes ->-                  let added = taddTarget (adSTK $ ftkToSTK accftk)-                                         (tproject1 daccRes_deRes)-                                         (tproject1 db1)-                  in tpair added (tproject2 daccRes_deRes)-        p = tmapAccumLDer (Proxy @target)-                          k accftk codomainShs eftk-                          (tlambda @target (FTKProduct accftk eftk)-                           $ HFun g)-                          (tlambda @target-                             (FTKProduct (adFTK (FTKProduct accftk eftk))-                                         (FTKProduct accftk eftk))-                           $ HFun dg)-                          (tlambda @target-                             (FTKProduct (adFTK (FTKProduct accftk codomainShs))-                                         (FTKProduct accftk eftk))-                           $ HFun rg)-                          acc0 es-        (accFin, qbs) = tunpair p-        (q, bs) = tunpair qbs-        dual = DeltaMapAccumL k bftk eftk q es df rf acc0' es'+                  (tpair acc (tproject2 accRes_bRes))+        p = tmapAccumL (Proxy @target)+                       k accftk codomainShs eftk+                       g+                       acc0 es+        (accFin, asbs) = tunpair p+        (as, bs) = tunpair asbs+        dual = DeltaMapAccumL k bftk eftk as es df rf acc0' es'     in dD (tpair accFin bs) dual-  tApply (HFun f) = f+  tapply (HFun f) = f   tlambda _ = id-  -- Bangs are for the proper order of sharing stamps.-  tcond !stk !b !u !v =-    let uv = tfromVector (SNat @2) stk (V.fromList [u, v])-    in tindexBuild (SNat @2) stk uv (tcond knownSTK b 0 1)-  tprimalPart (D u _) = u-  tdualPart _stk (D _ u') = u'-  tfromPrimal stk t = fromPrimalFTK (tftk stk t) t-  tfromDual t = dDnotShared (tdefTarget (ftkDelta t)) t-  tScale _stk = dScale-  tgrad @x xftk h =+  tgrad @x @r xftk h | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =     let rf :: forall f. ADReady f            => f x            -> f (ADTensorKind x)         -- This computes the derivative of g again for each new a.-        rf !a = ttlet a $ \ !aShared ->-          tunshare $ fst $ crevOnParams+        rf !a = ttlet a $ \ !aShared ->  -- sharing just in case+          tunshare $ snd $ crevOnParams                              Nothing                              (unHFun h @(ADVal (ShareOf f)))                              xftk@@ -504,8 +517,8 @@            -> f (ADTensorKind x)         -- This computes the derivative of g again for each new db and a.         rf !db_a = ttlet db_a $ \ !db_aShared ->-          tunshare $ fst $ crevOnParams-                             (Just $ toShare $ tproject1 db_aShared)+          tunshare $ snd $ crevOnParamsDt+                             (toShare $ tproject1 db_aShared)                              (unHFun h @(ADVal (ShareOf f)))                              xftk                              (toShare $ tproject2 db_aShared)@@ -516,26 +529,26 @@            -> f (ADTensorKind z)         -- This computes the derivative of g again for each new da and a.         df !da_a = ttlet da_a $ \ !da_aShared ->-          tunshare $ fst $ cfwdOnParams+          tunshare $ snd $ cfwdOnParams                              xftk                              (toShare $ tproject2 da_aShared)                              (unHFun h @(ADVal (ShareOf f)))                              (toShare $ tproject1 da_aShared)     in HFun df--  tfromVector snat stk lu =-    dD (tfromVector snat stk $ V.map (\(D u _) -> u) lu)-       (DeltaFromVector snat stk $ V.map (\(D _ u') -> u') lu)--  treplTarget r ftk = dDnotShared (treplTarget r ftk) (DeltaZero ftk)-  tdefTarget ftk = dDnotShared (tdefTarget ftk) (DeltaZero ftk)+  tprimalPart (D u _) = u+  tdualPart _stk (D _ u') = u'+  tplainPart (D u _) = tplainPart u+  tfromPrimal stk t = fromPrimalFTK (tftk stk t) t+  tfromDual t = dDnotShared (tdefTarget (ftkDelta t)) t+  tfromPlain stk t = fromPrimalFTK (tftk stk t) (tfromPlain stk t)+  tScale _stk = dScale   taddTarget = addTarget   tmultTarget = multTarget   tsum0Target = sum0Target   tdot0Target = dot0Target  instance ( ADReadyNoLet target, ShareTensor target-         , ShareTensor (PrimalOf target) )+         , ShareTensor (PrimalOf target), ShareTensor (PlainOf target) )          => ConvertTensor (ADVal target) where   tconvert c astk (D u u') =     dDnotShared (tconvert c astk u)@@ -559,9 +572,9 @@   rzip @_ @_ @n (D u u')    | Refl <- lemRankReplicate (Proxy @n) = case ftkDelta u' of     ftk@(FTKProduct (FTKR _sh y) (FTKR _ z)) ->-      let c = ConvCmp+      let c = convCmp                 (ConvXR (ftkToSTK (FTKProduct y z)))-                (ConvCmp+                (convCmp                    (ConvZip (ftkToSTK y) (ftkToSTK z))                    (ConvT2 ConvRX ConvRX))       in dD (tconvert c (ftkToSTK ftk) u)@@ -569,27 +582,27 @@   runzip @_ @_ @n (D u u')    | Refl <- lemRankReplicate (Proxy @n) = case ftkDelta u' of     ftk@(FTKR _sh (FTKProduct y z)) ->-      let c = ConvCmp+      let c = convCmp                 (ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))-                (ConvCmp+                (convCmp                    (ConvUnzip (ftkToSTK y) (ftkToSTK z))                    ConvRX)       in dD (tconvert c (ftkToSTK ftk) u)             (DeltaConvert c u')   szip (D u u') = case ftkDelta u' of     ftk@(FTKProduct (FTKS _sh y) (FTKS _ z)) ->-      let c = ConvCmp+      let c = convCmp                 ConvXS-                (ConvCmp+                (convCmp                    (ConvZip (ftkToSTK y) (ftkToSTK z))                    (ConvT2 ConvSX ConvSX))       in dD (tconvert c (ftkToSTK ftk) u)             (DeltaConvert c u')   sunzip (D u u') = case ftkDelta u' of     ftk@(FTKS _sh (FTKProduct y z)) ->-      let c = ConvCmp+      let c = convCmp                 (ConvT2 ConvXS ConvXS)-                (ConvCmp+                (convCmp                    (ConvUnzip (ftkToSTK y) (ftkToSTK z))                    ConvSX)       in dD (tconvert c (ftkToSTK ftk) u)@@ -609,14 +622,14 @@     | Refl <- lemRankReplicate (Proxy @m) =       let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)                             (TKX2 sh1 (TKR2 m x))-          c = ConvCmp+          c = convCmp                 (ConvXX (ConvXR (knownSTK @x)))                 (ConvNest @_ @_ @(Replicate m Nothing)                           (STKX sh1 (knownSTK @x)))       in dD (tconvert c (ftkToSTK $ ftkDelta u') u)             (DeltaConvert c u')   xnestS @_ @_ @x sh1 (D u u') =-    let c = ConvCmp (ConvXX ConvXS)+    let c = convCmp (ConvXX ConvXS)                     (ConvNest (STKX sh1 (knownSTK @x)))     in dD (tconvert c (ftkToSTK $ ftkDelta u') u)           (DeltaConvert c u')@@ -625,12 +638,12 @@     in dD (tconvert c (ftkToSTK $ ftkDelta u') u)           (DeltaConvert c u')   xunNestR (D u u') =-    let c = ConvCmp ConvUnnest+    let c = convCmp ConvUnnest                     (ConvXX ConvRX)     in dD (tconvert c (ftkToSTK $ ftkDelta u') u)           (DeltaConvert c u')   xunNestS (D u u') =-    let c = ConvCmp ConvUnnest+    let c = convCmp ConvUnnest                    (ConvXX ConvSX)     in dD (tconvert c (ftkToSTK $ ftkDelta u') u)           (DeltaConvert c u')
src/HordeAd/Core/OpsAst.hs view
@@ -1,1608 +1,2070 @@-{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}-{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}-{-# OPTIONS_GHC -Wno-orphans #-}--- | Tensor class instances for AST terms. Most of these instances--- vectorize any term with the build constructor in the root.--- The AST term instances can be used as building blocks for ADVal(AST)--- instances defined in "HordeAd.Core.OpsADVal" but may also be used standalone.-module HordeAd.Core.OpsAst-  ( IncomingCotangentHandling(..)-  , forwardPassByInterpretation-  , revArtifactFromForwardPass, revProduceArtifact-  , fwdArtifactFromForwardPass, fwdProduceArtifact-  ) where--import Prelude--import Data.Maybe (fromMaybe)-import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))-import Data.Vector.Generic qualified as V-import GHC.Exts (inline)-import GHC.TypeLits (OrderingI (..), cmpNat, type (+), type (-), type (<=?))-import System.IO.Unsafe (unsafePerformIO)-import Unsafe.Coerce (unsafeCoerce)--import Data.Array.Nested (Replicate, type (++))-import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Convert-  ( ixrFromIxS-  , ixsFromIxR-  , ixsFromIxR'-  , ixsFromIxX'-  , ixxFromIxS-  , withShsFromShR-  , withShsFromShX-  )-import Data.Array.Nested.Lemmas-import Data.Array.Nested.Mixed.Shape-import Data.Array.Nested.Permutation qualified as Permutation-import Data.Array.Nested.Ranked.Shape-import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (Init, snatPlus, unsafeCoerceRefl)--import HordeAd.AstEngine-import HordeAd.Core.Ast-import HordeAd.Core.AstEnv-import HordeAd.Core.AstFreshId-import HordeAd.Core.AstInline-import HordeAd.Core.AstInterpret-import HordeAd.Core.AstSimplify-import HordeAd.Core.AstTools-import HordeAd.Core.AstVectorize-import HordeAd.Core.CarriersADVal-import HordeAd.Core.CarriersAst-import HordeAd.Core.CarriersConcrete-import HordeAd.Core.ConvertTensor-import HordeAd.Core.Delta-import HordeAd.Core.DeltaEval-import HordeAd.Core.Ops-import HordeAd.Core.OpsConcrete ()-import HordeAd.Core.TensorKind-import HordeAd.Core.Types-import HordeAd.Core.Unwind---- * Symbolic reverse and forward derivative computation--data IncomingCotangentHandling = UseIncomingCotangent | IgnoreIncomingCotangent---- Here a choice is made that derivatives are PrimalSpan terms.--- This makes them easier to simplify and expresses via type that they--- don't introduce tangents nor cotangents, but are purely primal functions.--- They can still be liften to dual number functions via interpretations,--- as is done, e.g., in tgrad below.-forwardPassByInterpretation-  :: forall x z.-     (AstTensor AstMethodLet FullSpan x-      -> AstTensor AstMethodLet FullSpan z)-  -> AstEnv (ADVal (AstRaw PrimalSpan))-  -> AstTensor AstMethodShare PrimalSpan x-  -> AstVarName FullSpan x-  -> AstTensor AstMethodLet FullSpan x-  -> ADVal (AstRaw PrimalSpan) z-{-# INLINE forwardPassByInterpretation #-}-forwardPassByInterpretation g envInit astVarPrimal var astVar0 =-  let deltaInputs = generateDeltaInputs $ varNameToFTK var-      varInputs = dDnotShared (AstRaw astVarPrimal) deltaInputs-      ast = g astVar0-      env = extendEnv var varInputs envInit-  in interpretAstFull env ast--revArtifactFromForwardPass-  :: forall x z.-     IncomingCotangentHandling-  -> (AstTensor AstMethodShare PrimalSpan x-      -> AstVarName FullSpan x-      -> AstTensor AstMethodLet FullSpan x-      -> ADVal (AstRaw PrimalSpan) z)-  -> FullShapeTK x-  -> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)--- Break the inline chain to prevent false positives in inspection testing--- and protect the unsafePerformIO.-{-# NOINLINE revArtifactFromForwardPass #-}-revArtifactFromForwardPass cotangentHandling-                           forwardPass xftk = unsafePerformIO $ do-  -- IO and bangs and the compound function to fix the numbering of variables-  -- for pretty-printing and prevent sharing the impure values-  -- in tests that reset the impure counters.-  (!varPrimal, astVarPrimal, var, astVar0) <- funToAstRevIO xftk-  -- Evaluate completely after terms constructed, to free memory-  -- before gradientFromDelta allocates new memory and new FFI is started.-  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0-  let zftk = ftkAst $ unAstRaw primalBody-      (!varDt, astDt) = funToAst (adFTK zftk) Nothing id-  let oneAtF = treplTarget 1 $ adFTK zftk-      !dt = case cotangentHandling of-        UseIncomingCotangent -> AstRaw astDt-        IgnoreIncomingCotangent -> oneAtF-  let !gradient = gradientFromDelta xftk zftk dt delta-      !unGradient = unshareAstTensor $ unAstRaw gradient-      !unPrimal = unshareAstTensor $ unAstRaw primalBody-  return (AstArtifactRev varDt varPrimal unGradient unPrimal, delta)--revProduceArtifact-  :: forall x z.-     IncomingCotangentHandling-  -> (AstTensor AstMethodLet FullSpan x-      -> AstTensor AstMethodLet FullSpan z)-  -> AstEnv (ADVal (AstRaw PrimalSpan))-  -> FullShapeTK x-  -> AstArtifactRev x z-revProduceArtifact cotangentHandling g envInit xftk =-  fst $ inline revArtifactFromForwardPass-          cotangentHandling (forwardPassByInterpretation g envInit) xftk--fwdArtifactFromForwardPass-  :: forall x z.-     (AstTensor AstMethodShare PrimalSpan x-      -> AstVarName FullSpan x-      -> AstTensor AstMethodLet FullSpan x-      -> ADVal (AstRaw PrimalSpan) z)-  -> FullShapeTK x-  -> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)--- Break the inline chain to prevent false positives in inspection testing--- and protect the unsafePerformIO.-{-# NOINLINE fwdArtifactFromForwardPass #-}-fwdArtifactFromForwardPass forwardPass xftk = unsafePerformIO $ do-  (!varPrimalD, astVarD, varPrimal, astVarPrimal, var, astVar0)-    <- funToAstFwdIO xftk-  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0-  let !derivative = derivativeFromDelta @x delta (adFTK xftk) (AstRaw astVarD)-      !unDerivative = unshareAstTensor $ unAstRaw derivative-      !unPrimal = unshareAstTensor $ unAstRaw primalBody-  return (AstArtifactFwd varPrimalD varPrimal unDerivative unPrimal, delta)--fwdProduceArtifact-  :: forall x z.-     (AstTensor AstMethodLet FullSpan x-      -> AstTensor AstMethodLet FullSpan z)-  -> AstEnv (ADVal (AstRaw PrimalSpan))-  -> FullShapeTK x-  -> AstArtifactFwd x z-{-# INLINE fwdProduceArtifact #-}-fwdProduceArtifact f envInit xftk =-  fst $ inline fwdArtifactFromForwardPass-          (forwardPassByInterpretation f envInit) xftk----- * AstTensor instances---- | This is a vectorizing combinator that also simplifies--- the terms touched during vectorization, but not any others.--- Due to how the Ast tensor instances are defined, vectorization--- works bottom-up, which removes the need to backtrack in the vectorization--- pass or repeat until a fixed point is reached.--- This combinator also introduces new variable names.-astBuild1Vectorize-  :: AstSpan s-  => SNat k -> SingletonTK y-  -> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)-  -> AstTensor AstMethodLet s (BuildTensorKind k y)-astBuild1Vectorize k@(SNat @k) stk f =-  build1Vectorize k stk $ funToAstI (Just (0, valueOf @k - 1)) f--instance AstSpan s => LetTensor (AstTensor AstMethodLet s) where-  ttlet = astLetFun-  ttletPrimal = astLetFun-  toShare t = AstRaw $ AstToShare t-  -- For convenience and simplicity we define this for all spans,-  -- but it can only ever be used for PrimalSpan.-  tunshare =-    case sameAstSpan @s @PrimalSpan of-      Just Refl -> unshareAstTensor . unAstRaw-      _ -> error "tunshare: used not at PrimalSpan"---- | The checks and error messages in these functions result in complete--- shape-checking of the ranked and mixed user code (shaped is already--- fully checked by the Haskell type system).-instance AstSpan s => BaseTensor (AstTensor AstMethodLet s) where-  -- Ranked ops-  rshape t = case ftkAst t of-    FTKR sh _ -> sh-  trsum v = withSNat (rwidth v) $ \snat -> astSum snat knownSTK v-  trreplicate k = withSNat k $ \snat -> astReplicate snat knownSTK-  trindex @m @n a ix = case ftkAst a of-    FTKR @_ @x shmshn x ->-      withShsFromShR shmshn $ \(sh :: ShS sh) ->-        withKnownShS sh $-        gcastWith (unsafeCoerceRefl :: Rank (Take m sh) :~: m) $-        gcastWith (unsafeCoerceRefl :: Rank (Drop m sh) :~: n) $-        gcastWith (unsafeCoerceRefl :: Take m sh ++ Drop m sh :~: sh) $-        withKnownShS (shsTake @m sh) $-        astFromS' @(TKS2 (Drop m sh) x) (FTKR (shrDrop @m shmshn) x)-        $ astIndexS @(Take m sh) @(Drop m sh)-                    (shsDrop @m sh) (astSFromR' @sh sh a)-                    (ixsFromIxR' knownShS ix)-  trscatter @m @_ @p shpshn0 t f = case ftkAst t of-    FTKR @_ @x shmshn0 x ->-      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->-      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->-        withKnownShS shmshn $-        withKnownShS shpshn $-        withKnownShS (shsTake @m shmshn) $-        withKnownShS (shsDrop @m shmshn) $-        withKnownShS (shsTake @p shpshn) $-        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $-        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $-        gcastWith (unsafeCoerceRefl-                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $-        case testEquality (shsDrop @p shpshn) (shsDrop @m shmshn) of-          Just Refl ->-            astFromS' @(TKS2 shpshn x) (FTKR shpshn0 x)-            $ astScatterS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)-                          knownShS (astSFromR' shmshn t)-            $ funToAstIxS knownShS (ixsFromIxR' knownShS . f . ixrFromIxS)-                -- this introduces new variable names-          _ -> error $ "rscatter: shapes don't match: "-                       ++ show (shsDrop @p shpshn, shsDrop @m shmshn)-  trgather @m @_ @p shmshn0 t f = case ftkAst t of-    FTKR shpshn0 x ->-      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->-      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->-        withKnownShS shmshn $-        withKnownShS shpshn $-        withKnownShS (shsTake @m shmshn) $-        withKnownShS (shsDrop @m shmshn) $-        withKnownShS (shsTake @p shpshn) $-        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $-        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $-        gcastWith (unsafeCoerceRefl-                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $-        case testEquality (shsDrop @p shpshn) (shsDrop @m shmshn) of-          Just Refl ->-            astFromS' (FTKR shmshn0 x)-            $ astGatherS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)-                         knownShS (astSFromR' shpshn t)-            $ funToAstIxS knownShS (ixsFromIxR' knownShS . f . ixrFromIxS)-                -- this introduces new variable names-          _ -> error $ "rgather: shapes don't match: "-                       ++ show (shsDrop @p shpshn, shsDrop @m shmshn)-  trconcrete a = tconcrete (FTKR (Nested.rshape a) FTKScalar) (Concrete a)-  trfloor @_ @r2 a = case ftkAst a of-    FTKR sh' FTKScalar ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        astFromS' @(TKS sh r2) (FTKR sh' FTKScalar)-        . fromPrimal . astFloorS . primalPart . astSFromR' @sh sh $ a-  trfromIntegral @_ @r2 a = case ftkAst a of-    FTKR sh' FTKScalar ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        astFromS' @(TKS sh r2) (FTKR sh' FTKScalar)-        . fromPrimal . astFromIntegralS . primalPart . astSFromR' @sh sh $ a-  trcast @_ @r2 a = case ftkAst a of-    FTKR sh' FTKScalar ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        astFromS' @(TKS sh r2) (FTKR sh' FTKScalar)-        . astCastS . astSFromR' sh $ a-  trminIndex @_ @_ @r2 a = case ftkAst a of-    FTKR sh' _ ->-      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of-        (:$$) @_ @rest _ _ ->-          -- unfortunately, this is not enough:-          -- gcastWith (unsafeCoerceRefl :: Rank sh :~: 1 + Rank (Init sh)) $-          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $-          astFromS' @(TKS (Init sh) r2) (FTKR (shrInit sh') FTKScalar)-          . fromPrimal . AstMinIndexS . primalPart . astSFromR' @sh sh $ a-        ZSS -> error "rminIndex: impossible empty shape"-  trmaxIndex @_ @_ @r2 a = case ftkAst a of-    FTKR sh' _ ->-      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of-        (:$$) @_ @rest _ _ ->-          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $-          astFromS' @(TKS (Init sh) r2) (FTKR (shrInit sh') FTKScalar)-          . fromPrimal . AstMaxIndexS . primalPart . astSFromR' @sh sh $ a-        ZSS -> error "rmaxIndex: impossible empty shape"-  triota @r n =-    withSNat n $ \(SNat @n) ->-      astFromS' (FTKR (n :$: ZSR) FTKScalar)-      $ fromPrimal $ AstIotaS @n @r SNat-  trappend u v = case ftkAst u of-    FTKR shu' x -> case ftkAst v of-      FTKR shv' _ ->-        withShsFromShR shu' $ \shu -> case shu of-          _ :$$ restu ->-            withShsFromShR shv' $ \shv -> case shv of-              _ :$$ restv ->-                case testEquality restu restv of-                  Just Refl ->-                    astFromS' (FTKR shu' x)-                    $ astAppendS (astSFromR' shu u) (astSFromR' shv v)-                  _ -> error $ "rappend: shapes don't match: "-                               ++ show (restu, restv)-              ZSS -> error "rappend: impossible shape"-          ZSS -> error "rappend: impossible shape"-  trslice i n a = case ftkAst a of-    FTKR sh' x ->-      withShsFromShR sh' $ \sh -> case sh of-        msnat@(SNat @m) :$$ _ ->-          withSNat i $ \isnat@(SNat @i) -> withSNat n $ \nsnat@(SNat @n) ->-            case cmpNat (snatPlus isnat nsnat) msnat of-              GTI -> error $ "rslice: argument tensor too narrow: "-                             ++ show (i, n, sNatValue msnat)-              EQI ->-                astFromS' (FTKR sh' x)-                . astSliceS isnat nsnat (SNat @(m - (i + n)))-                . astSFromR' sh $ a-              LTI ->-                astFromS' (FTKR sh' x)-                . astSliceS isnat nsnat (SNat @(m - (i + n)))-                . astSFromR' sh $ a-        ZSS -> error "xslice: impossible shape"-  trreverse a = case ftkAst a of-    FTKR sh' x ->-      withShsFromShR sh' $ \sh -> case sh of-        _ :$$ _ ->-          astFromS' (FTKR sh' x)-          . astReverseS . astSFromR' sh $ a-        ZSS -> error "xreverse: impossible shape"-  trtranspose @n @r permr a = case ftkAst a of-    FTKR sh' x ->-      withShsFromShR sh' $ \(sh :: ShS sh)  ->-        Permutation.permFromList permr $ \(perm :: Permutation.Perm perm) ->-          let result :: AstTensor AstMethodLet s (TKR2 n r)-              result =-                -- A noble lie, verified down below.-                gcastWith (unsafeCoerceRefl-                           :: (Rank perm <=? Rank sh) :~: True) $-                fromMaybe (error "rtranspose: impossible non-permutation")-                $ Permutation.permCheckPermutation perm-                $ astFromS' (FTKR (shrPermutePrefix permr sh') x)-                  . astTransposeS perm . astSFromR' sh $ a-          in case (Permutation.permRank perm, shsRank sh) of-            (psnat@SNat, shsnat@SNat) ->-              case cmpNat psnat shsnat of-                GTI -> error $ "rtranspose: rank mismatch: "-                               ++ show (sNatValue psnat, sNatValue shsnat)-                EQI -> result-                LTI -> result-  trreshape sh2' a = case ftkAst a of-    FTKR sh' x ->-      withShsFromShR sh' $ \sh ->-      withShsFromShR sh2' $ \sh2 ->-        case testEquality (shsProduct sh) (shsProduct sh2) of-          Just Refl -> astFromS' (FTKR sh2' x)-                       . astReshapeS sh2 . astSFromR' sh $ a-          _ -> error $ "rreshape: tensor size mismatch: "-                       ++ show ( sNatValue (shsProduct sh)-                               , sNatValue (shsProduct sh2) )-  trbuild1 @n @x k f = withSNat k $ \snat ->-    astBuild1Vectorize snat (STKR (SNat @n) (knownSTK @x)) f--  -- Shaped ops-  sshape t = case ftkAst t of-    FTKS sh _ -> sh-  tssum = astSum SNat knownSTK-  tsindex = astIndexS knownShS-  tsscatter @shm @shn @shp t f =-    astScatterS @shm @shn @shp knownShS t-    $ funToAstIxS knownShS f  -- this introduces new variable names-  tsgather @shm @shn @shp t f =-    astGatherS @shm @shn @shp knownShS t-    $ funToAstIxS knownShS f  -- this introduces new variable names-  tsconcrete = fromPrimal . AstConcreteS-  tsfloor = fromPrimal . astFloorS . primalPart-  tsfromIntegral = fromPrimal . astFromIntegralS . primalPart-  tscast = astCastS-  tsminIndex = fromPrimal . AstMinIndexS . primalPart-  tsmaxIndex = fromPrimal . AstMaxIndexS . primalPart-  tsiota = fromPrimal $ AstIotaS SNat-  tsappend = astAppendS-  tsslice = astSliceS-  tsreverse = astReverseS-  tsbuild1 @k @sh @x =-    astBuild1Vectorize (SNat @k) (STKS (knownShS @sh) (knownSTK @x))--  -- Mixed ops-  xshape t = case ftkAst t of-    FTKX sh _ -> sh-  txsum = astSum SNat knownSTK-  txreplicate snat sh = astReplicate snat (STKX sh knownSTK)-  txindex @sh1 @sh2 a ix = case ftkAst a of-    FTKX @sh1sh2 @x sh1sh2 x | SNat <- ssxRank (knownShX @sh1) ->-      withShsFromShX sh1sh2 $ \(sh :: ShS sh) ->-        withKnownShX (ssxFromShX sh1sh2) $-        withKnownShS sh $-        gcastWith (unsafeCoerceRefl :: Rank (Drop (Rank sh1) sh) :~: Rank sh2) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank sh1) sh ++ Drop (Rank sh1) sh :~: sh) $-        gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) sh1sh2 :~: sh2) $-        withKnownShS (shsTake @(Rank sh1) sh) $-        astFromS' @(TKS2 (Drop (Rank sh1) sh) x)-                  (FTKX (shxDrop @(Rank sh1) sh1sh2) x)-        $ astIndexS @(Take (Rank sh1) sh) @(Drop (Rank sh1) sh)-                    (shsDrop @(Rank sh1) sh) (astSFromX' @sh @sh1sh2 sh a)-                    (ixsFromIxX' knownShS ix)-  txscatter @shm @_ @shp shpshn0 t f = case ftkAst t of-    FTKX shmshn0 x | SNat <- ssxRank (knownShX @shm)-                   , SNat <- ssxRank (knownShX @shp) ->-      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->-      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->-        withKnownShS shmshn $-        withKnownShS shpshn $-        withKnownShS (shsTake @(Rank shm) shmshn) $-        withKnownShS (shsDrop @(Rank shm) shmshn) $-        withKnownShS (shsTake @(Rank shp) shpshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn-                      :~: shmshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn-                      :~: shpshn) $-        case testEquality (shsDrop @(Rank shp) shpshn)-                          (shsDrop @(Rank shm) shmshn) of-          Just Refl ->-            astFromS' (FTKX shpshn0 x)-            $ astScatterS @(Take (Rank shm) shmshn)-                          @(Drop (Rank shm) shmshn)-                          @(Take (Rank shp) shpshn)-                          knownShS (astSFromX' shmshn t)-            $ funToAstIxS knownShS (ixsFromIxX' knownShS-                                    . f . ixxCast knownShX . ixxFromIxS)-                -- this introduces new variable names-          _ -> error $ "xscatter: shapes don't match: "-                       ++ show ( shsDrop @(Rank shp) shpshn-                               , shsDrop @(Rank shm) shmshn )-  txgather @shm @_ @shp shmshn0 t f = case ftkAst t of-    FTKX shpshn0 x | SNat <- ssxRank (knownShX @shm)-                   , SNat <- ssxRank (knownShX @shp) ->-      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->-      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->-        withKnownShS shmshn $-        withKnownShS shpshn $-        withKnownShS (shsTake @(Rank shm) shmshn) $-        withKnownShS (shsDrop @(Rank shm) shmshn) $-        withKnownShS (shsTake @(Rank shp) shpshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn-                      :~: shmshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn-                      :~: shpshn) $-        case testEquality (shsDrop @(Rank shp) shpshn)-                          (shsDrop @(Rank shm) shmshn) of-          Just Refl ->-            astFromS' (FTKX shmshn0 x)-            $ astGatherS @(Take (Rank shm) shmshn)-                         @(Drop (Rank shm) shmshn)-                         @(Take (Rank shp) shpshn)-                         knownShS (astSFromX' shpshn t)-            $ funToAstIxS knownShS (ixsFromIxX' knownShS-                                    . f . ixxCast knownShX . ixxFromIxS)-                -- this introduces new variable names-          _ -> error $ "xgather: shapes don't match: "-                       ++ show ( shsDrop @(Rank shp) shpshn-                               , shsDrop @(Rank shm) shmshn )-  txconcrete a = tconcrete (FTKX (Nested.mshape a) FTKScalar) (Concrete a)-  txfloor @_ @r2 @sh' a = case ftkAst a of-    FTKX sh' FTKScalar ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        astFromS' @(TKS sh r2) (FTKX sh' FTKScalar)-        . fromPrimal . astFloorS . primalPart . astSFromX' @sh @sh' sh $ a-  txfromIntegral @_ @r2 @sh' a = case ftkAst a of-    FTKX sh' FTKScalar ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        astFromS' @(TKS sh r2) (FTKX sh' FTKScalar)-        . fromPrimal . astFromIntegralS-        . primalPart . astSFromX' @sh @sh' sh $ a-  txcast @_ @r2 a = case ftkAst a of-    FTKX sh' FTKScalar ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        astFromS' @(TKS sh r2) (FTKX sh' FTKScalar)-        . astCastS . astSFromX' sh $ a-  txminIndex @_ @_ @_ @r2 a = case ftkAst a of-    FTKX @sh' sh' _ ->-      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of-        (:$$) @n @rest _ _ ->-          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $-          astFromS' @(TKS (Init sh) r2)-                    (FTKX (shxInit sh') FTKScalar)-          . fromPrimal . AstMinIndexS @n @rest-          . primalPart . astSFromX' @sh @sh' sh $ a-  txmaxIndex @_ @_ @_ @r2 a = case ftkAst a of-    FTKX @sh' sh' _ ->-      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of-        (:$$) @n @rest _ _ ->-          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $-          astFromS' @(TKS (Init sh) r2)-                   (FTKX (shxInit sh') FTKScalar)-          . fromPrimal . AstMaxIndexS @n @rest-          . primalPart . astSFromX' @sh @sh' sh $ a-  txiota @n @r = astFromS' (FTKX (SKnown (SNat @n) :$% ZSX) FTKScalar)-                 $ fromPrimal $ AstIotaS @n @r SNat-  txappend u v = case ftkAst u of-    FTKX (Nested.SKnown m@SNat :$% shu') x -> case ftkAst v of-      FTKX (Nested.SKnown n@SNat :$% shv') _ ->-        withShsFromShX shu' $ \(shu :: ShS shu) ->-          withShsFromShX shv' $ \(shv :: ShS shv) ->-            case shxEqual shu' shv' of-              Just Refl ->-                gcastWith (unsafeCoerceRefl :: shu :~: shv) $-                astFromS' (FTKX (Nested.SKnown (snatPlus m n) :$% shu') x)-                $ astAppendS (astSFromX' (m :$$ shu) u)-                             (astSFromX' (n :$$ shv) v)-              _ -> error $ "xappend: shapes don't match: "-                           ++ show (shu', shv')-  txslice i n@SNat k a = case ftkAst a of-    FTKX sh'@(_ :$% sh2') x ->-      withShsFromShX sh' $ \sh@(msnat :$$ _) ->-        case testEquality (snatPlus i (snatPlus n k)) msnat of-          Just Refl ->-            astFromS' (FTKX (SKnown n :$% sh2') x)-            . astSliceS i n k . astSFromX' sh $ a-          _ -> error $ "xslice: argument tensor too narrow: "-                       ++ show ( sNatValue i, sNatValue n, sNatValue k-                               , sNatValue msnat )-  txreverse a = case ftkAst a of-    FTKX sh' x ->-      withShsFromShX sh' $ \(sh@(_ :$$ _) :: ShS sh) ->-        astFromS' (FTKX sh' x)-        . astReverseS . astSFromX' @sh sh $ a-  txtranspose perm a = case ftkAst a of-    FTKX sh' x ->-      let sh2' = shxPermutePrefix perm sh'-      in withShsFromShX sh' $ \sh ->-           astFromS' (FTKX sh2' x)-           . astTransposeS perm . astSFromX' sh $ a-  txreshape sh2' a = case ftkAst a of-    FTKX sh' x ->-      withShsFromShX sh' $ \sh ->-      withShsFromShX sh2' $ \sh2 ->-        case testEquality (shsProduct sh) (shsProduct sh2) of-          Just Refl ->-            astFromS' (FTKX sh2' x)-            . astReshapeS sh2 . astSFromX' sh $ a-          _ -> error $ "xreshape: tensor size mismatch: "-                       ++ show ( sNatValue (shsProduct sh)-                               , sNatValue (shsProduct sh2) )-  txbuild1 @k @sh @x =-    astBuild1Vectorize (SNat @k) (STKX (knownShX @sh) (knownSTK @x))--  -- Scalar ops-  tkconcrete = fromPrimal . AstConcreteK-  tkfloor = fromPrimal . astFloorK . primalPart-  tkfromIntegral = fromPrimal . astFromIntegralK . primalPart-  tkcast = astCastK--  -- General operations that don't require LetTensor nor ShareTensor-  tftk _stk = ftkAst-  tconcrete ftk a = fromPrimal $ astConcrete ftk a-  tpair = astPair-  tproject1 = astProject1-  tproject2 = astProject2-  tsreplicate snat sh = astReplicate snat (STKS sh knownSTK)-  tstranspose = astTransposeS-  tsreshape = astReshapeS-  tmapAccumRDer _ !k _ !bftk !eftk f df rf acc0 es =-    astMapAccumRDer k bftk eftk f df rf acc0 es-  tmapAccumLDer _ !k _ !bftk !eftk f df rf acc0 es =-    astMapAccumLDer k bftk eftk f df rf acc0 es-  tApply = astApply-  tlambda ftk f =-    let (var, ast) = funToAst ftk Nothing $ unHFun f-    in AstLambda var ast-  tcond _ !b !u !v = astCond b u v-  tprimalPart = primalPart-  tdualPart _ = dualPart-  tfromPrimal _ = fromPrimal-  tfromDual = fromDual-  tgrad xftk f =-    -- We don't have an AST constructor to hold it, so we compute outright.-    ---    -- This computes the (AST of) derivative of f once and interprets it again-    -- for each new tensor of arguments, which is better than computing it anew.-    let -- No bangs here, because this goes under lambda and should not be-        -- evaluated too early (which at some point was even incorrect-        -- and triggered error "tunshare: used not at PrimalSpan"; maybe this-        -- is related to terms getting spans converted when interpreted)-        AstArtifactRev{..} =-          revProduceArtifact-            IgnoreIncomingCotangent (simplifyInline . unHFun f) emptyEnv xftk-        -- A new variable is created to give it the right span as opposed-        -- to the fixed PrimalSpan that artVarDomainRev has.-        (varP, ast) = funToAst xftk Nothing $ \ !astP ->-          let env = extendEnv artVarDomainRev astP emptyEnv-          in interpretAst env $ simplifyInline artDerivativeRev-    in AstLambda varP ast-  tvjp ftkx f =-    -- This computes the (AST of) derivative of f once and interprets it again-    -- for each new tensor of arguments, which is better than computing it anew.-    let AstArtifactRev{..} =-          revProduceArtifact-            UseIncomingCotangent (simplifyInline . unHFun f) emptyEnv ftkx-        ftkz = varNameToFTK artVarDtRev-        ftk2 = FTKProduct ftkz ftkx-        (varP, ast) = funToAst ftk2 Nothing $ \ !astP ->-          let env = extendEnv artVarDtRev (astProject1 astP)-                    $ extendEnv artVarDomainRev (astProject2 astP) emptyEnv-          in interpretAst env $ simplifyInline artDerivativeRev-    in AstLambda varP ast-  tjvp ftkx f =-    -- This computes the (AST of) derivative of f once and interprets it again-    -- for each new tensor of arguments, which is better than computing it anew.-    let AstArtifactFwd{..} =-          fwdProduceArtifact (simplifyInline . unHFun f) emptyEnv ftkx-        ftk2 = FTKProduct (adFTK ftkx) ftkx-        (varP, ast) = funToAst ftk2 Nothing $ \ !astP ->-          let env = extendEnv artVarDsFwd (astProject1 astP)-                    $ extendEnv artVarDomainFwd (astProject2 astP) emptyEnv-          in interpretAst env $ simplifyInline artDerivativeFwd-    in AstLambda varP ast--  tfromVector = astFromVector-  tsum snat@SNat stk u = case stk of-    STKScalar -> kfromS $ tssum u-    STKR SNat x | Dict <- lemKnownSTK x -> trsum u-    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u-    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u-    STKProduct stk1 stk2 ->-      ttlet u $ \ !u3 ->-        tpair (tsum snat stk1 (tproject1 u3))-              (tsum snat stk2 (tproject2 u3))-  treplicate snat@SNat stk u = case stk of-    STKScalar -> tsreplicate snat ZSS $ sfromK u-    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (sNatValue snat) u-    STKS sh x | Dict <- lemKnownSTK x -> tsreplicate snat sh u-    STKX sh x | Dict <- lemKnownSTK x -> txreplicate snat sh u-    STKProduct stk1 stk2 ->-      ttlet u $ \ !u3 ->-        tpair (treplicate snat stk1 (tproject1 u3))-              (treplicate snat stk2 (tproject2 u3))-  tindexBuild snat@SNat stk u i = case stk of-    STKScalar -> kfromS $ tsindex u (i :.$ ZIS)-    STKR SNat x | Dict <- lemKnownSTK x -> trindex u (i :.: ZIR)-    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsindex u (i :.$ ZIS)-    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txindex u (i :.% ZIX)-    STKProduct stk1 stk2 ->-      ttlet u $ \ !u3 ->-        tpair (tindexBuild snat stk1 (tproject1 u3) i)-              (tindexBuild snat stk2 (tproject2 u3) i)--  treplTarget = replTarget-  tdefTarget = defTarget-  taddTarget = addTarget-  tmultTarget = multTarget-  tsum0Target = sum0Target-  tdot0Target = dot0Target----- * AstRaw instances--instance AstSpan s => ShareTensor (AstRaw s) where-  -- For convenience and simplicity we define this for all spans,-  -- but it can only ever be used for PrimalSpan.-  tshare t =-    let u = unAstRaw t-    in if astIsSmall True u-       then t-       else AstRaw $ fun1ToAst (ftkAst u) $ \ !var -> AstShare var u-  tunpair (AstRaw (AstPair t1 t2)) = (AstRaw t1, AstRaw t2)-  tunpair t = let tShared = tshare t-              in (tproject1 tShared, tproject2 tShared)--instance AstSpan s => BaseTensor (AstRaw s) where-  -- Ranked ops-  rshape t = case ftkAst $ unAstRaw t of-    FTKR sh _ -> sh-  trsum v = withSNat (rwidth v) $ \snat ->-              AstRaw . AstSum snat knownSTK . unAstRaw $ v-  trreplicate k = withSNat k $ \snat ->-    AstRaw . AstReplicate snat knownSTK . unAstRaw-  trindex @m @n (AstRaw a) ix = AstRaw $ case ftkAst a of-    FTKR @_ @x shmshn x ->-      withShsFromShR shmshn $ \(sh :: ShS sh) ->-        withKnownShS sh $-        gcastWith (unsafeCoerceRefl :: Rank (Take m sh) :~: m) $-        gcastWith (unsafeCoerceRefl :: Rank (Drop m sh) :~: n) $-        gcastWith (unsafeCoerceRefl :: Take m sh ++ Drop m sh :~: sh) $-        withKnownShS (shsTake @m sh) $-        cAstFromS @(TKS2 (Drop m sh) x) (FTKR (shrDrop @m shmshn) x)-        $ AstIndexS @(Take m sh) @(Drop m sh)-                    (shsDrop @m sh) (cAstSFromR @sh sh a)-                    (ixsFromIxR knownShS (unAstRaw <$> ix))-  trscatter @m @_ @p shpshn0 (AstRaw t) f = AstRaw $ case ftkAst t of-    FTKR @_ @x shmshn0 x ->-      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->-      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->-        withKnownShS shmshn $-        withKnownShS shpshn $-        withKnownShS (shsTake @m shmshn) $-        withKnownShS (shsDrop @m shmshn) $-        withKnownShS (shsTake @p shpshn) $-        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $-        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $-        gcastWith (unsafeCoerceRefl-                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $-        case testEquality (shsDrop @p shpshn) (shsDrop @m shmshn) of-          Just Refl ->-            cAstFromS @(TKS2 shpshn x) (FTKR shpshn0 x)-            $ AstScatterS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)-                          knownShS (cAstSFromR shmshn t)-            $ funToAstIxS knownShS (fmap unAstRaw . ixsFromIxR knownShS-                                    . f . ixrFromIxS-                                    . fmap AstRaw)-                -- this introduces new variable names-          _ -> error $ "rscatter: shapes don't match: "-                       ++ show (shsDrop @p shpshn, shsDrop @m shmshn)-  trgather @m @_ @p shmshn0 (AstRaw t) f = AstRaw $ case ftkAst t of-    FTKR shpshn0 x ->-      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->-      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->-        withKnownShS shmshn $-        withKnownShS shpshn $-        withKnownShS (shsTake @m shmshn) $-        withKnownShS (shsDrop @m shmshn) $-        withKnownShS (shsTake @p shpshn) $-        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $-        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $-        gcastWith (unsafeCoerceRefl-                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $-        case testEquality (shsDrop @p shpshn) (shsDrop @m shmshn) of-          Just Refl ->-            cAstFromS (FTKR shmshn0 x)-            $ AstGatherS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)-                         knownShS (cAstSFromR shpshn t)-            $ funToAstIxS knownShS (fmap unAstRaw . ixsFromIxR knownShS-                                    . f . ixrFromIxS-                           . fmap AstRaw)-                -- this introduces new variable names-          _ -> error $ "rgather: shapes don't match: "-                       ++ show (shsDrop @p shpshn, shsDrop @m shmshn)-  trconcrete a = tconcrete (FTKR (Nested.rshape a) FTKScalar) (Concrete a)-  trfloor @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of-    FTKR sh' FTKScalar ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        cAstFromS @(TKS sh r2) (FTKR sh' FTKScalar)-        . fromPrimal . AstFloorS . primalPart . cAstSFromR @sh sh $ a-  trfromIntegral @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of-    FTKR sh' FTKScalar ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        cAstFromS @(TKS sh r2) (FTKR sh' FTKScalar)-        . fromPrimal . AstFromIntegralS . primalPart . cAstSFromR @sh sh $ a-  trcast @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of-    FTKR sh' FTKScalar ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        cAstFromS @(TKS sh r2) (FTKR sh' FTKScalar)-        . AstCastS . cAstSFromR sh $ a-  trminIndex @_ @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of-    FTKR sh' _ ->-      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of-        (:$$) @_ @rest _ _ ->-          -- unfortunately, this is not enough:-          -- gcastWith (unsafeCoerceRefl :: Rank sh :~: 1 + Rank (Init sh)) $-          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $-          cAstFromS @(TKS (Init sh) r2) (FTKR (shrInit sh') FTKScalar)-          . fromPrimal . AstMinIndexS . primalPart . cAstSFromR @sh sh $ a-        ZSS -> error "rminIndex: impossible shape"-  trmaxIndex @_ @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of-    FTKR sh' _ ->-      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of-        (:$$) @_ @rest _ _ ->-          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $-          cAstFromS @(TKS (Init sh) r2) (FTKR (shrInit sh') FTKScalar)-          . fromPrimal . AstMaxIndexS . primalPart . cAstSFromR @sh sh $ a-        ZSS -> error "rmaxIndex: impossible shape"-  triota @r n =-    AstRaw-    $ withSNat n $ \(SNat @n) ->-        cAstFromS (FTKR (n :$: ZSR) FTKScalar)-        $ fromPrimal $ AstIotaS @n @r SNat-  trappend (AstRaw u) (AstRaw v) = AstRaw $ case ftkAst u of-    ftk@(FTKR shu' _) -> case ftkAst v of-      FTKR shv' _ ->-        withShsFromShR shu' $ \shu -> case shu of-          _ :$$ restu ->-            withShsFromShR shv' $ \shv -> case shv of-              _ :$$ restv ->-                case testEquality restu restv of-                  Just Refl ->-                    cAstFromS ftk-                    $ AstAppendS (cAstSFromR shu u) (cAstSFromR shv v)-                  _ -> error $ "rappend: shapes don't match: "-                               ++ show (restu, restv)-              ZSS -> error "rappend: impossible shape"-          ZSS -> error "rappend: impossible shape"-  trslice i n (AstRaw a) = AstRaw $ case ftkAst a of-    ftk@(FTKR sh' _) ->-      withShsFromShR sh' $ \sh -> case sh of-        msnat@(SNat @m) :$$ _ ->-          withSNat i $ \isnat@(SNat @i) -> withSNat n $ \nsnat@(SNat @n) ->-            case cmpNat (snatPlus isnat nsnat) msnat of-              GTI -> error $ "rslice: argument tensor too narrow: "-                             ++ show (i, n, sNatValue msnat)-              EQI ->-                cAstFromS ftk-                . AstSliceS isnat nsnat (SNat @(m - (i + n)))-                . cAstSFromR sh $ a-              LTI ->-                cAstFromS ftk-                . AstSliceS isnat nsnat (SNat @(m - (i + n)))-                . cAstSFromR sh $ a-        ZSS -> error "xslice: impossible shape"-  trreverse (AstRaw a) = AstRaw $ case ftkAst a of-    ftk@(FTKR sh' _) ->-      withShsFromShR sh' $ \sh -> case sh of-        _ :$$ _ ->-          cAstFromS ftk-          . AstReverseS . cAstSFromR sh $ a-        ZSS -> error "xreverse: impossible shape"-  trtranspose @n @r permr (AstRaw a) = AstRaw $ case ftkAst a of-    ftk@(FTKR sh' _) ->-      withShsFromShR sh' $ \(sh :: ShS sh) ->-        Permutation.permFromList permr $ \(perm :: Permutation.Perm perm) ->-          let result :: AstTensor AstMethodShare s (TKR2 n r)-              result =-                -- A noble lie, verified down below.-                gcastWith (unsafeCoerceRefl-                           :: (Rank perm <=? Rank sh) :~: True) $-                fromMaybe (error "rtranspose: impossible non-permutation")-                $ Permutation.permCheckPermutation perm-                $ cAstFromS ftk-                  . AstTransposeS perm . cAstSFromR sh $ a-          in case (Permutation.permRank perm, shsRank sh) of-            (psnat@SNat, shsnat@SNat) ->-              case cmpNat psnat shsnat of-                GTI -> error $ "rtranspose: rank mismatch: "-                               ++ show (sNatValue psnat, sNatValue shsnat)-                EQI -> result-                LTI -> result-  trreshape sh2' (AstRaw a) = AstRaw $ case ftkAst a of-    FTKR sh' x ->-      withShsFromShR sh' $ \sh ->-      withShsFromShR sh2' $ \sh2 ->-        case testEquality (shsProduct sh) (shsProduct sh2) of-          Just Refl -> cAstFromS (FTKR sh2' x)-                       . AstReshapeS sh2 . cAstSFromR sh $ a-          _ -> error $ "rreshape: tensor size mismatch: "-                       ++ show ( sNatValue (shsProduct sh)-                               , sNatValue (shsProduct sh2) )-  trbuild1 k f = withSNat k $ \snat ->-    AstRaw $ AstBuild1 snat knownSTK-    $ funToAstI (Just (0, fromIntegral k - 1))-        -- this introduces new variable names-    $ unAstRaw . f . AstRaw--  -- Shaped ops-  sshape t = case ftkAst $ unAstRaw t of-    FTKS sh _ -> sh-  tssum = AstRaw . AstSum SNat knownSTK . unAstRaw-  tsindex v ix = AstRaw $ AstIndexS knownShS (unAstRaw v) (unAstRaw <$> ix)-  tsscatter @shm @shn @shp t f =-    AstRaw $ AstScatterS @shm @shn @shp knownShS (unAstRaw t)-           $ funToAstIxS knownShS (fmap unAstRaw . f . fmap AstRaw)-               -- this introduces new variable names-  tsgather @shm @shn @shp t f =-    AstRaw $ AstGatherS @shm @shn @shp knownShS (unAstRaw t)-           $ funToAstIxS knownShS (fmap unAstRaw . f . fmap AstRaw)-               -- this introduces new variable names-  tsconcrete = AstRaw . fromPrimal . AstConcreteS-  tsfloor = AstRaw . fromPrimal . AstFloorS . primalPart . unAstRaw-  tsfromIntegral =-    AstRaw . fromPrimal . AstFromIntegralS . primalPart . unAstRaw-  tscast = AstRaw . AstCastS . unAstRaw-  tsminIndex = AstRaw . fromPrimal . AstMinIndexS . primalPart . unAstRaw-  tsmaxIndex = AstRaw . fromPrimal . AstMaxIndexS . primalPart . unAstRaw-  tsiota = AstRaw . fromPrimal $ AstIotaS SNat-  tsappend u v = AstRaw $ AstAppendS (unAstRaw u) (unAstRaw v)-  tsslice i n k = AstRaw . AstSliceS i n k . unAstRaw-  tsreverse = AstRaw . AstReverseS . unAstRaw-  tsbuild1 @k f = AstRaw $ AstBuild1 (SNat @k) knownSTK-                  $ funToAstI (Just (0, valueOf @k - 1))-                      -- this introduces new variable names-                  $ unAstRaw . f . AstRaw--  -- Mixed ops-  xshape t = case ftkAst $ unAstRaw t of-    FTKX sh _ -> sh-  txsum = AstRaw . AstSum SNat knownSTK . unAstRaw-  txreplicate snat sh = AstRaw . AstReplicate snat (STKX sh knownSTK) . unAstRaw-  txindex @sh1 @sh2 (AstRaw a) ix = case ftkAst a of-    FTKX @sh1sh2 @x sh1sh2 x | SNat <- ssxRank (knownShX @sh1) ->-      withShsFromShX sh1sh2 $ \(sh :: ShS sh) ->-        withKnownShX (ssxFromShX sh1sh2) $-        withKnownShS sh $-        gcastWith (unsafeCoerceRefl :: Rank (Drop (Rank sh1) sh) :~: Rank sh2) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank sh1) sh ++ Drop (Rank sh1) sh :~: sh) $-        gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) sh1sh2 :~: sh2) $-        withKnownShS (shsTake @(Rank sh1) sh) $-        AstRaw $ cAstFromS @(TKS2 (Drop (Rank sh1) sh) x)-                           (FTKX (shxDrop @(Rank sh1) sh1sh2) x)-        $ AstIndexS @(Take (Rank sh1) sh) @(Drop (Rank sh1) sh)-                    (shsDrop @(Rank sh1) sh) (cAstSFromX @sh @sh1sh2 sh a)-                    (ixsFromIxX' knownShS (unAstRaw <$> ix))-  txscatter @shm @_ @shp shpshn0 (AstRaw t) f = AstRaw $ case ftkAst t of-    FTKX shmshn0 x | SNat <- ssxRank (knownShX @shm)-                   , SNat <- ssxRank (knownShX @shp) ->-      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->-      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->-        withKnownShS shmshn $-        withKnownShS shpshn $-        withKnownShS (shsTake @(Rank shm) shmshn) $-        withKnownShS (shsDrop @(Rank shm) shmshn) $-        withKnownShS (shsTake @(Rank shp) shpshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn-                      :~: shmshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn-                      :~: shpshn) $-        case testEquality (shsDrop @(Rank shp) shpshn)-                          (shsDrop @(Rank shm) shmshn) of-          Just Refl ->-            cAstFromS (FTKX shpshn0 x)-            $ AstScatterS @(Take (Rank shm) shmshn)-                          @(Drop (Rank shm) shmshn)-                          @(Take (Rank shp) shpshn)-                          knownShS (cAstSFromX shmshn t)-            $ funToAstIxS knownShS (fmap unAstRaw-                                    . ixsFromIxX' knownShS-                                    . f . ixxCast knownShX . ixxFromIxS-                                    . fmap AstRaw)-                -- this introduces new variable names-          _ -> error $ "xscatter: shapes don't match: "-                       ++ show ( shsDrop @(Rank shp) shpshn-                               , shsDrop @(Rank shm) shmshn )-  txgather @shm @_ @shp shmshn0 (AstRaw t) f = AstRaw $ case ftkAst t of-    FTKX shpshn0 x | SNat <- ssxRank (knownShX @shm)-                   , SNat <- ssxRank (knownShX @shp) ->-      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->-      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->-        withKnownShS shmshn $-        withKnownShS shpshn $-        withKnownShS (shsTake @(Rank shm) shmshn) $-        withKnownShS (shsDrop @(Rank shm) shmshn) $-        withKnownShS (shsTake @(Rank shp) shpshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn-                      :~: shmshn) $-        gcastWith (unsafeCoerceRefl-                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn-                      :~: shpshn) $-        case testEquality (shsDrop @(Rank shp) shpshn)-                          (shsDrop @(Rank shm) shmshn) of-          Just Refl ->-            cAstFromS (FTKX shmshn0 x)-            $ AstGatherS @(Take (Rank shm) shmshn)-                         @(Drop (Rank shm) shmshn)-                         @(Take (Rank shp) shpshn)-                         knownShS (cAstSFromX shpshn t)-            $ funToAstIxS knownShS (fmap unAstRaw-                                    . ixsFromIxX' knownShS-                                    . f . ixxCast knownShX . ixxFromIxS-                                    . fmap AstRaw)-                -- this introduces new variable names-          _ -> error $ "xgather: shapes don't match: "-                       ++ show ( shsDrop @(Rank shp) shpshn-                               , shsDrop @(Rank shm) shmshn )-  txconcrete a = tconcrete (FTKX (Nested.mshape a) FTKScalar) (Concrete a)-  txfloor @_ @r2 @sh' (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX sh' FTKScalar ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        cAstFromS @(TKS sh r2) (FTKX sh' FTKScalar)-        . fromPrimal . AstFloorS . primalPart . cAstSFromX @sh @sh' sh $ a-  txfromIntegral @_ @r2 @sh' (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX sh' FTKScalar ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        cAstFromS @(TKS sh r2) (FTKX sh' FTKScalar)-        . fromPrimal . AstFromIntegralS-        . primalPart . cAstSFromX @sh @sh' sh $ a-  txcast @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX sh' FTKScalar ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        cAstFromS @(TKS sh r2) (FTKX sh' FTKScalar)-        . AstCastS . cAstSFromX sh $ a-  txminIndex @_ @_ @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX @sh' sh' _ ->-      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of-        (:$$) @n @rest _ _ ->-          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $-          cAstFromS @(TKS (Init sh) r2)-                    (FTKX (shxInit sh') FTKScalar)-          . fromPrimal . AstMinIndexS @n @rest-          . primalPart . cAstSFromX @sh @sh' sh $ a-  txmaxIndex @_ @_ @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX @sh' sh' _ ->-      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of-        (:$$) @n @rest _ _ ->-          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $-          cAstFromS @(TKS (Init sh) r2)-                   (FTKX (shxInit sh') FTKScalar)-          . fromPrimal . AstMaxIndexS @n @rest-          . primalPart . cAstSFromX @sh @sh' sh $ a-  txiota @n @r = AstRaw $ cAstFromS (FTKX (SKnown (SNat @n) :$% ZSX) FTKScalar)-                 $ fromPrimal $ AstIotaS @n @r SNat-  txappend (AstRaw u) (AstRaw v) = AstRaw $ case ftkAst u of-    FTKX (Nested.SKnown m@SNat :$% shu') x -> case ftkAst v of-      FTKX (Nested.SKnown n@SNat :$% shv') _ ->-        withShsFromShX shu' $ \(shu :: ShS shu) ->-          withShsFromShX shv' $ \(shv :: ShS shv) ->-            case shxEqual shu' shv' of-              Just Refl ->-                gcastWith (unsafeCoerceRefl :: shu :~: shv) $-                cAstFromS (FTKX (SKnown (snatPlus m n) :$% shu') x)-                $ AstAppendS (cAstSFromX (m :$$ shu) u)-                             (cAstSFromX (n :$$ shv) v)-              _ -> error $ "xappend: shapes don't match: "-                           ++ show (shu', shv')-  txslice i n@SNat k (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX sh'@(_ :$% sh2') x ->-      withShsFromShX sh' $ \sh@(msnat :$$ _) ->-        case testEquality (snatPlus i (snatPlus n k)) msnat of-          Just Refl ->-            cAstFromS (FTKX (SKnown n :$% sh2') x)-            . AstSliceS i n k . cAstSFromX sh $ a-          _ -> error $ "xslice: argument tensor too narrow: "-                       ++ show ( sNatValue i, sNatValue n, sNatValue k-                               , sNatValue msnat )-  txreverse (AstRaw a) = AstRaw $ case ftkAst a of-    ftk@(FTKX sh' _) ->-      withShsFromShX sh' $ \(sh@(_ :$$ _) :: ShS sh) ->-        cAstFromS ftk-        . AstReverseS . cAstSFromX @sh sh $ a-  txtranspose perm (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX sh' x ->-      let sh2' = shxPermutePrefix perm sh'-      in withShsFromShX sh' $ \sh ->-           cAstFromS (FTKX sh2' x)-           . AstTransposeS perm-           . cAstSFromX sh $ a-  txreshape sh2' (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX sh' x ->-      withShsFromShX sh' $ \sh ->-      withShsFromShX sh2' $ \sh2 ->-        case testEquality (shsProduct sh) (shsProduct sh2) of-          Just Refl ->-            cAstFromS (FTKX sh2' x)-            . AstReshapeS sh2 . cAstSFromX sh $ a-          _ -> error $ "xreshape: tensor size mismatch: "-                       ++ show ( sNatValue (shsProduct sh)-                               , sNatValue (shsProduct sh2) )-  txbuild1 @k f = AstRaw $ AstBuild1 (SNat @k) knownSTK-                  $ funToAstI (Just (0, valueOf @k - 1))-                      -- this introduces new variable names-                  $ unAstRaw . f . AstRaw--  -- Scalar ops-  tkconcrete = AstRaw . fromPrimal . AstConcreteK-  tkfloor = AstRaw . fromPrimal . AstFloorK . primalPart . unAstRaw-  tkfromIntegral = AstRaw . fromPrimal . AstFromIntegralK-                   . primalPart . unAstRaw-  tkcast = AstRaw . AstCastK . unAstRaw--  -- General operations that don't require LetTensor nor ShareTensor-  tftk _stk = ftkAst . unAstRaw-  tconcrete ftk a = AstRaw $ fromPrimal $ unAstRaw $ astConcreteRaw ftk a-  tpair t1 t2 = AstRaw $ AstPair (unAstRaw t1) (unAstRaw t2)-  tproject1 t = AstRaw $ AstProject1 $ unAstRaw t-  tproject2 t = AstRaw $ AstProject2 $ unAstRaw t-  tsreplicate snat sh = AstRaw . AstReplicate snat (STKS sh knownSTK) . unAstRaw-  tstranspose perm = AstRaw . AstTransposeS perm . unAstRaw-  tsreshape sh = AstRaw . AstReshapeS sh . unAstRaw-  tmapAccumRDer _ !k _ !bftk !eftk f df rf acc0 es =-    AstRaw $ AstMapAccumRDer k bftk eftk f df rf (unAstRaw acc0) (unAstRaw es)-  tmapAccumLDer _ !k _ !bftk !eftk f df rf acc0 es =-    AstRaw $ AstMapAccumLDer k bftk eftk f df rf (unAstRaw acc0) (unAstRaw es)-  tApply t ll = AstRaw $ AstApply t (unAstRaw ll)-  tlambda = tlambda @(AstTensor AstMethodLet s)-  tcond _ !b !u !v = AstRaw $ AstCond b (unAstRaw u) (unAstRaw v)-  tprimalPart t = AstRaw $ primalPart $ unAstRaw t-  tdualPart _ t = dualPart $ unAstRaw t-  tfromPrimal _ t = AstRaw $ fromPrimal $ unAstRaw t-  tfromDual t = AstRaw $ fromDual t-  -- These three methods are called at this type in delta evaluation via-  -- tmapAccumR and tmapAccumL, so they have to work. We could refrain from-  -- simplifying the resulting terms, but it's not clear that's more consistent.-  tgrad = tgrad @(AstTensor AstMethodLet s)-  tvjp = tvjp @(AstTensor AstMethodLet s)-  tjvp = tjvp @(AstTensor AstMethodLet s)--  tfromVector k stk =-    AstRaw . AstFromVector k stk . V.map unAstRaw--  treplTarget = replTarget-  tdefTarget = defTarget-  taddTarget = addTarget-  tmultTarget = multTarget-  tsum0Target = sum0Target-  tdot0Target = dot0Target--instance AstSpan s => ConvertTensor (AstRaw s) where-  tconvert c _astk = AstRaw . cAstConvert c . unAstRaw--  rfromX a = case ftkAst $ unAstRaw a of-    FTKX sh' _ ->-      withShsFromShX sh' $ \(sh :: ShS sh) ->-        withKnownShS sh $-        rfromS $ sfromX @_ @sh a-  xfromR a = case ftkAst $ unAstRaw a of-    FTKR shr _ ->-      withShsFromShR shr $ \(sh :: ShS sh) ->-        withKnownShS sh $-        xfromS @_ @sh $ sfromR a--  sfromR = AstRaw . cAstSFromR knownShS . unAstRaw-  sfromX = AstRaw . cAstSFromX knownShS . unAstRaw-  xfromS @_ @sh' = AstRaw . cAstXFromS (knownShX @sh') . unAstRaw--  rzip @_ @_ @n (AstRaw a)-   | Refl <- lemRankReplicate (Proxy @n) = AstRaw $ case ftkAst a of-    FTKProduct (FTKR _sh y) (FTKR _ z) ->-      let c = ConvCmp-                (ConvXR (ftkToSTK (FTKProduct y z)))-                (ConvCmp-                   (ConvZip (ftkToSTK y) (ftkToSTK z))-                   (ConvT2 ConvRX ConvRX))-      in cAstConvert c a-  runzip @_ @_ @n (AstRaw a)-   | Refl <- lemRankReplicate (Proxy @n) = AstRaw $ case ftkAst a of-    FTKR _sh (FTKProduct y z) ->-      let c = ConvCmp-                (ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))-                (ConvCmp-                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))-                   ConvRX)-      in cAstConvert c a-  szip (AstRaw a) = AstRaw $ case ftkAst a of-    FTKProduct (FTKS _sh y) (FTKS _ z) ->-      let c = ConvCmp-                ConvXS-                (ConvCmp-                   (ConvZip (ftkToSTK y) (ftkToSTK z))-                   (ConvT2 ConvSX ConvSX))-      in cAstConvert c a-  sunzip (AstRaw a) = AstRaw $ case ftkAst a of-    FTKS _sh (FTKProduct y z) ->-      let c = ConvCmp-                (ConvT2 ConvXS ConvXS)-                (ConvCmp-                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))-                   ConvSX)-      in cAstConvert c a-  xzip (AstRaw a) = AstRaw $ case ftkAst a of-    FTKProduct (FTKX _sh y) (FTKX _ z) ->-      let c = ConvZip (ftkToSTK y) (ftkToSTK z)-      in cAstConvert c a-  xunzip (AstRaw a) = AstRaw $ case ftkAst a of-    FTKX _sh (FTKProduct y z) ->-      let c = ConvUnzip (ftkToSTK y) (ftkToSTK z)-      in cAstConvert c a--  xnestR @sh1 @m @x sh1 (AstRaw a)-    | Refl <- lemRankReplicate (Proxy @m) = AstRaw $-      let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)-                            (TKX2 sh1 (TKR2 m x))-          c = ConvCmp-                (ConvXX (ConvXR (knownSTK @x)))-                (ConvNest @_ @_ @(Replicate m Nothing)-                          (STKX sh1 (knownSTK @x)))-      in cAstConvert c a-  xnestS @_ @_ @x sh1 (AstRaw a) = AstRaw $-    let c = ConvCmp (ConvXX ConvXS)-                    (ConvNest (STKX sh1 (knownSTK @x)))-    in cAstConvert c a-  xnest @_ @_ @x sh1 (AstRaw a) = AstRaw $-    let c = ConvNest (STKX sh1 (knownSTK @x))-    in cAstConvert c a-  xunNestR (AstRaw a) = AstRaw $-    let c = ConvCmp ConvUnnest-                    (ConvXX ConvRX)-    in cAstConvert c a-  xunNestS (AstRaw a) = AstRaw $-    let c = ConvCmp ConvUnnest-                    (ConvXX ConvSX)-    in cAstConvert c a-  xunNest (AstRaw a) = AstRaw $-    let c = ConvUnnest-    in cAstConvert c a--  tpairConv = tpair-  tunpairConv = tunpair---- All but the last case are shortcuts for common forms.-astConcreteRaw :: FullShapeTK y -> Concrete y-               -> AstRaw PrimalSpan y-astConcreteRaw ftk v = case ftk of-  FTKScalar -> AstRaw $ AstConcreteK $ unConcrete v-  FTKR sh' FTKScalar -> AstRaw $-    withShsFromShR sh' $ \(sh :: ShS sh) ->-      withKnownShS sh $-      cAstFromS ftk $ AstConcreteS (unConcrete $ sfromR @_ @sh v)-  FTKS _ FTKScalar -> AstRaw $ AstConcreteS $ unConcrete v-  FTKX sh' FTKScalar -> AstRaw $-    withShsFromShX sh' $ \(sh :: ShS sh) ->-      withKnownShS sh $-      cAstFromS ftk $ AstConcreteS (unConcrete $ sfromX @_ @sh v)-  FTKProduct ftk1 ftk2 -> AstRaw $-    AstPair (unAstRaw $ astConcreteRaw ftk1 (tproject1 v))-            (unAstRaw $ astConcreteRaw ftk2 (tproject2 v))-  _ -> concreteTarget (tkconcrete . unConcrete) (tsconcrete . unConcrete)-                      (\ftk2 a -> AstRaw $ cAstFromS ftk2 $ unAstRaw a)-                      (ftkToSTK ftk) v----- * AstNoVectorize instances--instance AstSpan s => LetTensor (AstNoVectorize s) where-  ttlet u f = AstNoVectorize-              $ ttlet (unAstNoVectorize u)-                      (unAstNoVectorize . f . AstNoVectorize)-  ttletPrimal u f = AstNoVectorize-                    $ ttletPrimal (unAstNoVectorize u)-                                  (unAstNoVectorize . f . AstNoVectorize)-  toShare t = toShare $ unAstNoVectorize t--instance AstSpan s => BaseTensor (AstNoVectorize s) where-  -- Ranked ops-  rshape = rshape . unAstNoVectorize-  trsum = AstNoVectorize . trsum . unAstNoVectorize-  trreplicate k = AstNoVectorize . trreplicate k . unAstNoVectorize-  trindex v ix =-    AstNoVectorize $ trindex (unAstNoVectorize v) (unAstNoVectorize <$> ix)-  trscatter sh t f =-    AstNoVectorize $ trscatter sh (unAstNoVectorize t)-                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize-  trgather sh t f =-    AstNoVectorize $ trgather sh (unAstNoVectorize t)-                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize-  trconcrete = AstNoVectorize . trconcrete-  trfloor = AstNoVectorize . trfloor . unAstNoVectorize-  trfromIntegral = AstNoVectorize . trfromIntegral . unAstNoVectorize-  trcast = AstNoVectorize . trcast . unAstNoVectorize-  trminIndex = AstNoVectorize . trminIndex . unAstNoVectorize-  trmaxIndex = AstNoVectorize . trmaxIndex . unAstNoVectorize-  triota = AstNoVectorize . triota-  trappend u v =-    AstNoVectorize $ trappend (unAstNoVectorize u) (unAstNoVectorize v)-  trslice i n = AstNoVectorize . trslice i n . unAstNoVectorize-  trreverse = AstNoVectorize . trreverse . unAstNoVectorize-  trtranspose perm = AstNoVectorize . trtranspose perm . unAstNoVectorize-  trreshape sh = AstNoVectorize . trreshape sh . unAstNoVectorize-  trbuild1 k f = withSNat k $ \snat ->-    AstNoVectorize $ AstBuild1 snat knownSTK-    $ funToAstI (Just (0, fromIntegral k - 1))-        -- this introduces new variable names-    $ unAstNoVectorize . f . AstNoVectorize--  -- Shaped ops-  sshape = sshape . unAstNoVectorize-  tssum = AstNoVectorize . tssum . unAstNoVectorize-  tsindex v ix =-    AstNoVectorize $ tsindex (unAstNoVectorize v) (unAstNoVectorize <$> ix)-  tsscatter @_ @shm @shn @shp t f =-    AstNoVectorize $ tsscatter @_ @_ @shm @shn @shp (unAstNoVectorize t)-                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize-  tsgather @_ @shm @shn @shp t f =-    AstNoVectorize $ tsgather @_ @_ @shm @shn @shp (unAstNoVectorize t)-                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize-  tsconcrete = AstNoVectorize . tsconcrete-  tsfloor = AstNoVectorize . tsfloor . unAstNoVectorize-  tsfromIntegral = AstNoVectorize . tsfromIntegral . unAstNoVectorize-  tscast = AstNoVectorize . tscast . unAstNoVectorize-  tsminIndex = AstNoVectorize . tsminIndex . unAstNoVectorize-  tsmaxIndex = AstNoVectorize . tsmaxIndex . unAstNoVectorize-  tsiota = AstNoVectorize tsiota-  tsappend u v =-    AstNoVectorize $ tsappend (unAstNoVectorize u) (unAstNoVectorize v)-  tsslice i n k = AstNoVectorize . tsslice i n k . unAstNoVectorize-  tsreverse = AstNoVectorize . tsreverse . unAstNoVectorize-  tsbuild1 @k f = AstNoVectorize $ AstBuild1 (SNat @k) knownSTK-                  $ funToAstI (Just (0, valueOf @k - 1))-                      -- this introduces new variable names-                  $ unAstNoVectorize . f . AstNoVectorize--  -- Mixed ops-  xshape = xshape . unAstNoVectorize-  txsum = AstNoVectorize . txsum . unAstNoVectorize-  txreplicate snat sh = AstNoVectorize . txreplicate snat sh . unAstNoVectorize-  txindex v ix =-    AstNoVectorize $ txindex (unAstNoVectorize v) (unAstNoVectorize <$> ix)-  txscatter @_ @shm @shn @shp sh t f =-    AstNoVectorize $ txscatter @_ @_ @shm @shn @shp sh (unAstNoVectorize t)-                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize-  txgather @_ @shm @shn @shp sh t f =-    AstNoVectorize $ txgather @_ @_ @shm @shn @shp sh (unAstNoVectorize t)-                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize-  txconcrete = AstNoVectorize . txconcrete-  txfloor = AstNoVectorize . txfloor . unAstNoVectorize-  txfromIntegral = AstNoVectorize . txfromIntegral . unAstNoVectorize-  txcast = AstNoVectorize . txcast . unAstNoVectorize-  txminIndex = AstNoVectorize . txminIndex . unAstNoVectorize-  txmaxIndex = AstNoVectorize . txmaxIndex . unAstNoVectorize-  txiota @n = AstNoVectorize $ txiota @_ @n-  txappend u v =-    AstNoVectorize $ txappend (unAstNoVectorize u) (unAstNoVectorize v)-  txslice i n k = AstNoVectorize . txslice i n k . unAstNoVectorize-  txreverse = AstNoVectorize . txreverse . unAstNoVectorize-  txtranspose perm = AstNoVectorize . txtranspose perm . unAstNoVectorize-  txreshape sh = AstNoVectorize . txreshape sh . unAstNoVectorize-  txbuild1 @k f = AstNoVectorize $ AstBuild1 (SNat @k) knownSTK-                  $ funToAstI (Just (0, valueOf @k - 1))-                      -- this introduces new variable names-                  $ unAstNoVectorize . f . AstNoVectorize--  -- Scalar ops-  tkconcrete = AstNoVectorize . tkconcrete-  tkfloor = AstNoVectorize . tkfloor . unAstNoVectorize-  tkfromIntegral = AstNoVectorize . tkfromIntegral . unAstNoVectorize-  tkcast = AstNoVectorize . tkcast . unAstNoVectorize--  -- General operations that don't require LetTensor nor ShareTensor-  tftk stk = tftk stk . unAstNoVectorize-  tconcrete ftk a = AstNoVectorize $ tconcrete ftk a-  tpair t1 t2 =-    AstNoVectorize $ tpair (unAstNoVectorize t1) (unAstNoVectorize t2)-  tproject1 t = AstNoVectorize $ tproject1 $ unAstNoVectorize t-  tproject2 t = AstNoVectorize $ tproject2 $ unAstNoVectorize t-  tsreplicate snat sh = AstNoVectorize . tsreplicate snat sh. unAstNoVectorize-  tstranspose perm =-    AstNoVectorize . tstranspose perm . unAstNoVectorize-  tsreshape sh = AstNoVectorize . tsreshape sh . unAstNoVectorize-  tmapAccumRDer _ !k !accftk !bftk !eftk f df rf acc0 es =-    AstNoVectorize $ tmapAccumRDer Proxy k accftk bftk eftk f df rf-                       (unAstNoVectorize acc0) (unAstNoVectorize es)-  tmapAccumLDer _ !k !accftk !bftk !eftk f df rf acc0 es =-    AstNoVectorize $ tmapAccumLDer Proxy k accftk bftk eftk f df rf-                       (unAstNoVectorize acc0) (unAstNoVectorize es)-  tApply t ll = AstNoVectorize $ tApply t (unAstNoVectorize ll)-  tlambda = tlambda @(AstTensor AstMethodLet s)-  tcond !stk !b !u !v =-    AstNoVectorize $ tcond stk b (unAstNoVectorize u) (unAstNoVectorize v)-  tprimalPart t = AstNoVectorize $ tprimalPart $ unAstNoVectorize t-  tdualPart stk t = tdualPart stk $ unAstNoVectorize t-  tfromPrimal stk t = AstNoVectorize $ tfromPrimal stk $ unAstNoVectorize t-  tfromDual t = AstNoVectorize $ tfromDual t-  tgrad = tgrad @(AstTensor AstMethodLet s)-  tvjp = tvjp @(AstTensor AstMethodLet s)-  tjvp = tjvp @(AstTensor AstMethodLet s)--  tfromVector k stk =-    AstNoVectorize . tfromVector k stk . V.map unAstNoVectorize-  tsum k stk =-    AstNoVectorize . tsum k stk . unAstNoVectorize-  treplicate k stk =-    AstNoVectorize . treplicate k stk . unAstNoVectorize-  tindexBuild k stk u i =-    AstNoVectorize $ tindexBuild k stk (unAstNoVectorize u) (unAstNoVectorize i)--  treplTarget r ftk = AstNoVectorize $ treplTarget r ftk-  tdefTarget = AstNoVectorize . tdefTarget-  taddTarget stk a b = AstNoVectorize $ taddTarget stk (unAstNoVectorize a)-                                                       (unAstNoVectorize b)-  tmultTarget stk a b = AstNoVectorize $ tmultTarget stk (unAstNoVectorize a)-                                                         (unAstNoVectorize b)-  tsum0Target stk a = AstNoVectorize $ tsum0Target stk (unAstNoVectorize a)-  tdot0Target stk a b = AstNoVectorize $ tdot0Target stk (unAstNoVectorize a)-                                                         (unAstNoVectorize b)--instance AstSpan s => ConvertTensor (AstNoVectorize s) where-  tconvert c _astk = AstNoVectorize . astConvert c . unAstNoVectorize--  rfromX = AstNoVectorize . rfromX . unAstNoVectorize-  xfromR = AstNoVectorize . xfromR . unAstNoVectorize--  sfromR = AstNoVectorize . sfromR . unAstNoVectorize-  sfromX = AstNoVectorize . sfromX . unAstNoVectorize-  xfromS = AstNoVectorize . xfromS . unAstNoVectorize--  rzip = AstNoVectorize . rzip . unAstNoVectorize-  runzip = AstNoVectorize . runzip . unAstNoVectorize-  szip = AstNoVectorize . szip . unAstNoVectorize-  sunzip = AstNoVectorize . sunzip . unAstNoVectorize-  xzip = AstNoVectorize . xzip . unAstNoVectorize-  xunzip = AstNoVectorize . xunzip . unAstNoVectorize--  xnestR sh = AstNoVectorize . xnestR sh . unAstNoVectorize-  xnestS sh = AstNoVectorize . xnestS sh . unAstNoVectorize-  xnest sh = AstNoVectorize . xnest sh . unAstNoVectorize-  xunNestR = AstNoVectorize . xunNestR . unAstNoVectorize-  xunNestS = AstNoVectorize . xunNestS . unAstNoVectorize-  xunNest = AstNoVectorize . xunNest . unAstNoVectorize--  tpairConv = tpair-  tunpairConv a = let (b, c) = tunpairConv $ unAstNoVectorize a-                  in (AstNoVectorize b, AstNoVectorize c)----- * AstNoSimplify instances--instance AstSpan s => LetTensor (AstNoSimplify s) where-  ttlet u f = AstNoSimplify-              $ astLetFunNoSimplify (unAstNoSimplify u)-                                    (unAstNoSimplify . f . AstNoSimplify)-  ttletPrimal u f = AstNoSimplify-                    $ astLetFunNoSimplify (unAstNoSimplify u)-                                          (unAstNoSimplify . f . AstNoSimplify)-  toShare t = AstRaw $ AstToShare $ unAstNoSimplify t--wAstNoSimplify :: AstRaw s y -> AstNoSimplify s y-wAstNoSimplify =-  AstNoSimplify-  . (unsafeCoerce :: AstTensor AstMethodShare s y-                  -> AstTensor AstMethodLet s y)-  . unAstRaw--wunAstNoSimplify :: AstNoSimplify s y -> AstRaw s y-wunAstNoSimplify =-  AstRaw-  . (unsafeCoerce :: AstTensor AstMethodLet s y-                  -> AstTensor AstMethodShare s y)-  . unAstNoSimplify--instance AstSpan s => BaseTensor (AstNoSimplify s) where-  -- The implementation of these methods differs from the AstRaw instance:-  trbuild1 @n @x k f = withSNat k $ \snat ->-    AstNoSimplify-    $ astBuild1Vectorize snat (STKR (SNat @n) (knownSTK @x))-                         (unAstNoSimplify . f . AstNoSimplify)-  tsbuild1 @k @sh @x f =-    AstNoSimplify-    $ astBuild1Vectorize (SNat @k) (STKS (knownShS @sh) (knownSTK @x))-                         (unAstNoSimplify . f . AstNoSimplify)-  txbuild1 @k @sh @x f =-    AstNoSimplify-    $ astBuild1Vectorize (SNat @k) (STKX (knownShX @sh) (knownSTK @x))-                         (unAstNoSimplify . f . AstNoSimplify)-  -- These three have tricky types, so we repaat the AstRaw definitions:-  tcond _ !b !u !v =-    AstNoSimplify $ AstCond b (unAstNoSimplify u) (unAstNoSimplify v)-  tdualPart _ t = dualPart $ unAstNoSimplify t-  tfromDual t = AstNoSimplify $ fromDual t--  -- All the following implementations piggy-back on AstRaw implementations.-  -- Ranked ops-  rshape = rshape . wunAstNoSimplify-  trsum = wAstNoSimplify . trsum . wunAstNoSimplify-  trreplicate k = wAstNoSimplify . trreplicate k . wunAstNoSimplify-  trindex v ix =-    wAstNoSimplify $ trindex (wunAstNoSimplify v) (wunAstNoSimplify <$> ix)-  trscatter sh t f =-    wAstNoSimplify $ trscatter sh (wunAstNoSimplify t)-                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify-  trgather sh t f =-    wAstNoSimplify $ trgather sh (wunAstNoSimplify t)-                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify-  trconcrete = wAstNoSimplify . trconcrete-  trfloor = wAstNoSimplify . trfloor . wunAstNoSimplify-  trfromIntegral = wAstNoSimplify . trfromIntegral . wunAstNoSimplify-  trcast = wAstNoSimplify . trcast . wunAstNoSimplify-  trminIndex = wAstNoSimplify . trminIndex . wunAstNoSimplify-  trmaxIndex = wAstNoSimplify . trmaxIndex . wunAstNoSimplify-  triota = wAstNoSimplify . triota-  trappend u v =-    wAstNoSimplify $ trappend (wunAstNoSimplify u) (wunAstNoSimplify v)-  trslice i n = wAstNoSimplify . trslice i n . wunAstNoSimplify-  trreverse = wAstNoSimplify . trreverse . wunAstNoSimplify-  trtranspose perm = wAstNoSimplify . trtranspose perm . wunAstNoSimplify-  trreshape sh = wAstNoSimplify . trreshape sh . wunAstNoSimplify--  -- Shaped ops-  sshape = sshape . wunAstNoSimplify-  tssum = wAstNoSimplify . tssum . wunAstNoSimplify-  tsindex v ix =-    wAstNoSimplify $ tsindex (wunAstNoSimplify v) (wunAstNoSimplify <$> ix)-  tsscatter @_ @shm @shn @shp t f =-    wAstNoSimplify $ tsscatter @_ @_ @shm @shn @shp (wunAstNoSimplify t)-                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify-  tsgather @_ @shm @shn @shp t f =-    wAstNoSimplify $ tsgather @_ @_ @shm @shn @shp (wunAstNoSimplify t)-                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify-  tsconcrete = wAstNoSimplify . tsconcrete-  tsfloor = wAstNoSimplify . tsfloor . wunAstNoSimplify-  tsfromIntegral = wAstNoSimplify . tsfromIntegral . wunAstNoSimplify-  tscast = wAstNoSimplify . tscast . wunAstNoSimplify-  tsminIndex = wAstNoSimplify . tsminIndex . wunAstNoSimplify-  tsmaxIndex = wAstNoSimplify . tsmaxIndex . wunAstNoSimplify-  tsiota = wAstNoSimplify tsiota-  tsappend u v =-    wAstNoSimplify $ tsappend (wunAstNoSimplify u) (wunAstNoSimplify v)-  tsslice i n k = wAstNoSimplify . tsslice i n k . wunAstNoSimplify-  tsreverse = wAstNoSimplify . tsreverse . wunAstNoSimplify--  -- Mixed ops-  xshape = xshape . wunAstNoSimplify-  txsum = wAstNoSimplify . txsum . wunAstNoSimplify-  txreplicate snat sh = wAstNoSimplify . txreplicate snat sh . wunAstNoSimplify-  txindex v ix =-    wAstNoSimplify $ txindex (wunAstNoSimplify v) (wunAstNoSimplify <$> ix)-  txscatter @_ @shm @shn @shp sh t f =-    wAstNoSimplify $ txscatter @_ @_ @shm @shn @shp sh (wunAstNoSimplify t)-                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify-  txgather @_ @shm @shn @shp sh t f =-    wAstNoSimplify $ txgather @_ @_ @shm @shn @shp sh (wunAstNoSimplify t)-                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify-  txconcrete = wAstNoSimplify . txconcrete-  txfloor = wAstNoSimplify . txfloor . wunAstNoSimplify-  txfromIntegral = wAstNoSimplify . txfromIntegral . wunAstNoSimplify-  txcast = wAstNoSimplify . txcast . wunAstNoSimplify-  txminIndex = wAstNoSimplify . txminIndex . wunAstNoSimplify-  txmaxIndex = wAstNoSimplify . txmaxIndex . wunAstNoSimplify-  txiota @n = wAstNoSimplify $ txiota @_ @n-  txappend u v =-    wAstNoSimplify $ txappend (wunAstNoSimplify u) (wunAstNoSimplify v)-  txslice i n k = wAstNoSimplify . txslice i n k . wunAstNoSimplify-  txreverse = wAstNoSimplify . txreverse . wunAstNoSimplify-  txtranspose perm = wAstNoSimplify . txtranspose perm . wunAstNoSimplify-  txreshape sh = wAstNoSimplify . txreshape sh . wunAstNoSimplify--  -- Scalar ops-  tkconcrete = wAstNoSimplify . tkconcrete-  tkfloor = wAstNoSimplify . tkfloor . wunAstNoSimplify-  tkfromIntegral = wAstNoSimplify . tkfromIntegral . wunAstNoSimplify-  tkcast = wAstNoSimplify . tkcast . wunAstNoSimplify--  -- General operations that don't require LetTensor nor ShareTensor-  tftk stk = tftk stk . wunAstNoSimplify-  tconcrete ftk a = wAstNoSimplify $ tconcrete ftk a-  tpair t1 t2 =-    wAstNoSimplify $ tpair (wunAstNoSimplify t1) (wunAstNoSimplify t2)-  tproject1 t = wAstNoSimplify $ tproject1 $ wunAstNoSimplify t-  tproject2 t = wAstNoSimplify $ tproject2 $ wunAstNoSimplify t-  tsreplicate snat sh = wAstNoSimplify . tsreplicate snat sh . wunAstNoSimplify-  tstranspose perm =-    wAstNoSimplify . tstranspose perm . wunAstNoSimplify-  tsreshape sh = wAstNoSimplify . tsreshape sh . wunAstNoSimplify-  tmapAccumRDer _ !k !accftk !bftk !eftk f df rf acc0 es =-    wAstNoSimplify $ tmapAccumRDer Proxy k accftk bftk eftk f df rf-                       (wunAstNoSimplify acc0) (wunAstNoSimplify es)-  tmapAccumLDer _ !k !accftk !bftk !eftk f df rf acc0 es =-    wAstNoSimplify $ tmapAccumLDer Proxy k accftk bftk eftk f df rf-                       (wunAstNoSimplify acc0) (wunAstNoSimplify es)-  tApply t ll = wAstNoSimplify $ tApply t (wunAstNoSimplify ll)-  tlambda = tlambda @(AstRaw s)-  tprimalPart t = wAstNoSimplify $ tprimalPart $ wunAstNoSimplify t-  tfromPrimal stk t = wAstNoSimplify $ tfromPrimal stk $ wunAstNoSimplify t-  tgrad = tgrad @(AstRaw s)-  tvjp = tvjp @(AstRaw s)-  tjvp = tjvp @(AstRaw s)--  tfromVector k stk =-    wAstNoSimplify . tfromVector k stk . V.map wunAstNoSimplify-  tsum snat@SNat stk u = case stk of-    STKScalar -> kfromS $ tssum u-    STKR SNat x | Dict <- lemKnownSTK x -> trsum u-    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u-    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u-    STKProduct stk1 stk2 ->-      ttlet u $ \ !u3 ->-        tpair (tsum snat stk1 (tproject1 u3))-              (tsum snat stk2 (tproject2 u3))-  treplicate snat@SNat stk u = case stk of-    STKScalar -> tsreplicate snat ZSS $ sfromK u-    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (sNatValue snat) u-    STKS sh x | Dict <- lemKnownSTK x -> tsreplicate snat sh u-    STKX sh x | Dict <- lemKnownSTK x -> txreplicate snat sh u-    STKProduct stk1 stk2 ->-      ttlet u $ \ !u3 ->-        tpair (treplicate snat stk1 (tproject1 u3))-              (treplicate snat stk2 (tproject2 u3))-  tindexBuild snat@SNat stk u i = case stk of-    STKScalar -> kfromS $ tsindex u (i :.$ ZIS)-    STKR SNat x | Dict <- lemKnownSTK x -> trindex u (i :.: ZIR)-    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsindex u (i :.$ ZIS)-    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txindex u (i :.% ZIX)-    STKProduct stk1 stk2 ->-      ttlet u $ \ !u3 ->-        tpair (tindexBuild snat stk1 (tproject1 u3) i)-              (tindexBuild snat stk2 (tproject2 u3) i)--  treplTarget r ftk = wAstNoSimplify $ treplTarget r ftk-  tdefTarget = wAstNoSimplify . tdefTarget-  taddTarget stk a b = wAstNoSimplify $ taddTarget stk (wunAstNoSimplify a)-                                                       (wunAstNoSimplify b)-  tmultTarget stk a b = wAstNoSimplify $ tmultTarget stk (wunAstNoSimplify a)-                                                         (wunAstNoSimplify b)-  tsum0Target stk a = wAstNoSimplify $ tsum0Target stk (wunAstNoSimplify a)-  tdot0Target stk a b = wAstNoSimplify $ tdot0Target stk (wunAstNoSimplify a)-                                                         (wunAstNoSimplify b)--instance AstSpan s => ConvertTensor (AstNoSimplify s) where+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Tensor class instances for AST terms. Most of these instances+-- vectorize any term with the build constructor in the root.+-- The AST term instances can be used as building blocks for ADVal(AST)+-- instances defined in "HordeAd.Core.OpsADVal" but may also be used standalone.+module HordeAd.Core.OpsAst+  ( IncomingCotangentHandling(..)+  , forwardPassByInterpretation+  , revArtifactFromForwardPass, revProduceArtifact, revProduceArtifactDt+  , fwdArtifactFromForwardPass, fwdProduceArtifact+  ) where++import Prelude++import Data.Coerce (Coercible, coerce)+import Data.Maybe (fromMaybe)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import GHC.TypeLits (OrderingI (..), cmpNat, type (+), type (-), type (<=?))+import System.IO.Unsafe (unsafePerformIO)+import Unsafe.Coerce (unsafeCoerce)++import Data.Array.Nested (Replicate, type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert+  ( ixrFromIxS'+  , ixsFromIxR'+  , ixsFromIxX'+  , ixxFromIxS'+  , withShsFromShR+  , withShsFromShX+  )+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation (DropLen, TakeLen)+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Init, fromSNat', snatPlus, unsafeCoerceRefl)++import HordeAd.Core.Ast+import HordeAd.Core.AstEngine+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInline+import HordeAd.Core.AstInterpret+import HordeAd.Core.AstMethodLet ()+import HordeAd.Core.AstMethodShare+import HordeAd.Core.AstSimplify+import HordeAd.Core.AstTools+import HordeAd.Core.AstVectorize+import HordeAd.Core.CarriersADVal+import HordeAd.Core.CarriersAst+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Conversion+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Delta+import HordeAd.Core.DeltaEval+import HordeAd.Core.Ops+import HordeAd.Core.OpsConcrete ()+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind+import HordeAd.Core.UnwindNum++-- * Symbolic reverse and forward derivative computation++data IncomingCotangentHandling = UseIncomingCotangent | IgnoreIncomingCotangent++forwardPassByInterpretation+  :: forall x z.+     (AstTensor AstMethodLet FullSpan x+      -> AstTensor AstMethodLet FullSpan z)+  -> AstEnv (ADVal (AstRaw FullSpan))+  -> AstTensor AstMethodShare FullSpan x+  -> AstVarName '(FullSpan, x)+  -> AstTensor AstMethodLet FullSpan x+  -> ADVal (AstRaw FullSpan) z+{-# INLINE forwardPassByInterpretation #-}+forwardPassByInterpretation g envInit astVarPrimal var astVar0 =+  let deltaInputs = generateDeltaInputs $ varNameToFTK var+      varInputs = dDnotShared (AstRaw astVarPrimal) deltaInputs+      ast = g astVar0+      env = extendEnv var varInputs envInit+  in interpretAstFull env ast++revArtifactFromForwardPass+  :: forall x z. TKAllNum (ADTensorKind z)+  => IncomingCotangentHandling+  -> (AstTensor AstMethodShare FullSpan x+      -> AstVarName '(FullSpan, x)+      -> AstTensor AstMethodLet FullSpan x+      -> ADVal (AstRaw FullSpan) z)+  -> FullShapeTK x+  -> (AstArtifactRev x z, Delta (AstRaw FullSpan) z)+-- Break the inline chain to prevent false positives in inspection testing+-- and protect the unsafePerformIO.+{-# NOINLINE revArtifactFromForwardPass #-}+revArtifactFromForwardPass cotangentHandling+                           forwardPass xftk = unsafePerformIO $ do+  -- IO and bangs and the compound function to fix the numbering of variables+  -- for pretty-printing and prevent sharing the impure values+  -- in tests that reset the impure counters.+  (!astVarPrimal, var, astVar0) <- funToAstRevIO xftk+  -- Evaluate completely after terms constructed, to free memory+  -- before gradientFromDelta allocates new memory and new FFI is started.+  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0+  let zftk = ftkAst $ unAstRaw primalBody+      (!varDt, astDt) = funToAst (adFTK zftk) id+  let oneAtF = treplTarget 1 $ adFTK zftk+      !dt = case cotangentHandling of+        UseIncomingCotangent -> AstRaw astDt+        IgnoreIncomingCotangent -> oneAtF+  let !gradient = gradientFromDelta xftk dt delta+      !unGradient = unshareAstTensor $ unAstRaw gradient+      unPrimal = unshareAstTensor $ unAstRaw primalBody+  return (AstArtifactRev varDt var unGradient unPrimal, delta)++revProduceArtifact+  :: forall x z. TKAllNum (ADTensorKind z)+  => IncomingCotangentHandling+  -> (AstTensor AstMethodLet FullSpan x+      -> AstTensor AstMethodLet FullSpan z)+  -> AstEnv (ADVal (AstRaw FullSpan))+  -> FullShapeTK x+  -> AstArtifactRev x z+{-# INLINE revProduceArtifact #-}+revProduceArtifact cotangentHandling g envInit xftk =+  fst $ revArtifactFromForwardPass+          cotangentHandling (forwardPassByInterpretation g envInit) xftk++-- These two functions are as above, but the dt must be provided and so,+-- due to technical reasons, the type is less constrained.+revArtifactFromForwardPassDt+  :: forall x z.+     (AstTensor AstMethodShare FullSpan x+      -> AstVarName '(FullSpan, x)+      -> AstTensor AstMethodLet FullSpan x+      -> ADVal (AstRaw FullSpan) z)+  -> FullShapeTK x+  -> (AstArtifactRev x z, Delta (AstRaw FullSpan) z)+-- Break the inline chain to prevent false positives in inspection testing+-- and protect the unsafePerformIO.+{-# NOINLINE revArtifactFromForwardPassDt #-}+revArtifactFromForwardPassDt forwardPass xftk = unsafePerformIO $ do+  -- IO and bangs and the compound function to fix the numbering of variables+  -- for pretty-printing and prevent sharing the impure values+  -- in tests that reset the impure counters.+  (!astVarPrimal, var, astVar0) <- funToAstRevIO xftk+  -- Evaluate completely after terms constructed, to free memory+  -- before gradientFromDelta allocates new memory and new FFI is started.+  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0+  let zftk = ftkAst $ unAstRaw primalBody+      (!varDt, !dt) = funToAst (adFTK zftk) id+  let !gradient = gradientFromDelta xftk (AstRaw dt) delta+      !unGradient = unshareAstTensor $ unAstRaw gradient+      unPrimal = unshareAstTensor $ unAstRaw primalBody+  return (AstArtifactRev varDt var unGradient unPrimal, delta)++revProduceArtifactDt+  :: forall x z.+     (AstTensor AstMethodLet FullSpan x+      -> AstTensor AstMethodLet FullSpan z)+  -> AstEnv (ADVal (AstRaw FullSpan))+  -> FullShapeTK x+  -> AstArtifactRev x z+{-# INLINE revProduceArtifactDt #-}+revProduceArtifactDt g envInit xftk =+  fst $ revArtifactFromForwardPassDt+          (forwardPassByInterpretation g envInit) xftk++fwdArtifactFromForwardPass+  :: forall x z.+     (AstTensor AstMethodShare FullSpan x+      -> AstVarName '(FullSpan, x)+      -> AstTensor AstMethodLet FullSpan x+      -> ADVal (AstRaw FullSpan) z)+  -> FullShapeTK x+  -> (AstArtifactFwd x z, Delta (AstRaw FullSpan) z)+-- Break the inline chain to prevent false positives in inspection testing+-- and protect the unsafePerformIO.+{-# NOINLINE fwdArtifactFromForwardPass #-}+fwdArtifactFromForwardPass forwardPass xftk = unsafePerformIO $ do+  (!varPrimalD, astVarD, astVarPrimal, var, astVar0) <- funToAstFwdIO xftk+  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0+  let !derivative =+        derivativeFromDelta (Proxy @x) delta (adFTK xftk) (AstRaw astVarD)+      !unDerivative = unshareAstTensor $ unAstRaw derivative+      unPrimal = unshareAstTensor $ unAstRaw primalBody+  return (AstArtifactFwd varPrimalD var unDerivative unPrimal, delta)++fwdProduceArtifact+  :: forall x z.+     (AstTensor AstMethodLet FullSpan x+      -> AstTensor AstMethodLet FullSpan z)+  -> AstEnv (ADVal (AstRaw FullSpan))+  -> FullShapeTK x+  -> AstArtifactFwd x z+{-# INLINE fwdProduceArtifact #-}+fwdProduceArtifact f envInit xftk =+  fst $ fwdArtifactFromForwardPass+          (forwardPassByInterpretation f envInit) xftk+++-- * Main AstTensor instances++-- | This is a vectorizing combinator that also simplifies+-- the terms touched during vectorization, but not any others.+-- Due to how the Ast tensor instances are defined, vectorization+-- works bottom-up, which removes the need to backtrack in the vectorization+-- pass or repeat until a fixed point is reached.+-- This combinator also introduces new variable names.+astBuild1Vectorize+  :: KnownSpan s+  => SNat k -> SingletonTK y+  -> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)+  -> AstTensor AstMethodLet s (BuildTensorKind k y)+{-# NOINLINE astBuild1Vectorize #-}+astBuild1Vectorize k stk f = unsafePerformIO $ do+  varx <- funToAstIntIO (0, fromSNat' k - 1) f+  build1Vectorize k stk varx++instance KnownSpan s => LetTensor (AstTensor AstMethodLet s) where+  ttlet = astLetFun+  ttletPrimal = astLetFun+  ttletPlain = astLetFun+  toShare t = AstRaw $ AstToShare t+  -- For convenience and simplicity we define this for all spans,+  -- but it can only ever be used for FullSpan.+  tunshare =+    case knownSpan @s of+      SFullSpan -> unshareAstTensor . unAstRaw+      _ -> error "tunshare: used not at FullSpan"++-- | The checks and error messages in these functions result in complete+-- shape-checking of the ranked and mixed user code (shaped is already+-- fully checked by the Haskell type system).+--+-- Them methods are listed in ranked, shaped, mixed order to keep+-- similar code transformations together.+instance KnownSpan s => BaseTensor (AstTensor AstMethodLet s) where+  isConcreteInstance = False+  -- Ranked ops+  rshape t = case ftkAst t of+    FTKR sh _ -> sh+  trfromVectorN shm0 l = case V.uncons l of+    Just (v, _) | FTKR shn0 x <- ftkAst v ->+      withShsFromShR shn0 $ \(shn :: ShS shn) ->+      withShsFromShR shm0 $ \(shm :: ShS shm) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ shn) :~: Rank shm + Rank shn) $+        astConvUpRFromS (shm `shsAppend` shn) x+        . astFromVectorS shm . V.map (astConvDownSFromR shn x) $ l+    Nothing -> error "trfromVectorN: empty vector"+  trfromVectorLinear shm0 l =+    withShsFromShR shm0 $ \(shm :: ShS shm) ->+      astConvUpRFromS shm FTKScalar . astFromVectorK shm $ l+  trsumN @m @n t = case ftkAst t of+    FTKR shmshn0 x ->+      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->+        gcastWith (unsafeCoerceRefl :: Rank (Drop m shmshn) :~: n) $+        gcastWith (unsafeCoerceRefl+                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $+        astConvUpRFromS (shsDrop @m shmshn) x . astSumS (shsTake @m shmshn)+        . astConvDownSFromR shmshn x $ t+  trsum0 t = case ftkAst t of+    FTKR shm x ->+      withShsFromShR shm $ \(sh :: ShS sh) ->+        astSumK (astConvDownSFromR sh x t)+  trreplicateN shm0 t = case ftkAst t of+    FTKR shn0 x ->+      withShsFromShR shn0 $ \(shn :: ShS shn) ->+      withShsFromShR shm0 $ \(shm :: ShS shm) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ shn) :~: Rank shm + Rank shn) $+        astConvUpRFromS (shm `shsAppend` shn) x . astReplicateS shm+        . astConvDownSFromR shn x $ t+  trreplicate0N shm0 t =+    withShsFromShR shm0 $ \(shm :: ShS shm) ->+      astConvUpRFromS shm FTKScalar . astReplicateK shm $ t+  trconcrete a = tconcrete (FTKR (Nested.rshape a) FTKScalar) (Concrete a)+  trfloor a = case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        astConvUpRFromS sh FTKScalar+        . fromPlain . astFloorS . astPlainPart+        . astConvDownSFromR sh FTKScalar $ a+  trfromIntegral a = case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        astConvUpRFromS sh FTKScalar+        . fromPlain . astFromIntegralS . astPlainPart+        . astConvDownSFromR sh FTKScalar $ a+  trcast a = case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        astConvUpRFromS sh FTKScalar+        . astCastS . astConvDownSFromR sh FTKScalar $ a+  trindex @m @n a ix | SNat <- ixrRank ix = case ftkAst a of+    FTKR shmshn x ->+      withShsFromShR shmshn $ \(sh :: ShS sh) ->+        gcastWith (unsafeCoerceRefl :: Rank (Take m sh) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Drop m sh) :~: n) $+        gcastWith (unsafeCoerceRefl :: Take m sh ++ Drop m sh :~: sh) $+        astConvUpRFromS (shsDrop @m sh) x+        $ astIndexS @(Take m sh) @(Drop m sh)+                    (shsDrop @m sh) (astConvDownSFromR sh x a)+                    (ixsFromIxR' (shsTake @m sh) ix)+  trindex0 a ix = case ftkAst a of+    FTKR shm x ->+      withShsFromShR shm $ \(sh :: ShS sh) ->+        astIndexK (astConvDownSFromR sh x a) (ixsFromIxR' sh ix)+  trscatter @m @n shp0 t f = case ftkAst t of+    FTKR shmshn0 x ->+      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShR shp0 $ \(shp :: ShS shp) ->+        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Drop m shmshn) :~: n) $+        gcastWith (unsafeCoerceRefl+                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (shp ++ Drop m shmshn)+                      :~: Rank shp + Rank (Drop m shmshn)) $+        astConvUpRFromS (shp `shsAppend` shsDrop @m shmshn) x+        $ funToVarsIxS (shsTake @m shmshn) $ \vars ix ->+            let !ix2 = ixsFromIxR' shp . f . ixrFromIxS' (SNat @m) $ ix+            in astScatterS (shsTake @m shmshn)+                           (shsDrop @m shmshn)+                           shp+                           (astConvDownSFromR shmshn x t) (vars, ix2)+            -- this introduces new variable names+  trgather @m @n @p shm0 t f = case ftkAst t of+    FTKR shpshn0 x ->+      withShsFromShR shm0 $ \(shm :: ShS shm) ->+      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->+        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $+        gcastWith (unsafeCoerceRefl :: Rank (Drop p shpshn) :~: n) $+        gcastWith (unsafeCoerceRefl+                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ Drop p shpshn)+                      :~: Rank shm + Rank (Drop p shpshn)) $+        astConvUpRFromS (shm `shsAppend` shsDrop @p shpshn) x+        $ funToVarsIxS shm $ \vars ix ->+            let !ix2 = ixsFromIxR' (shsTake @p shpshn)+                       . f . ixrFromIxS' (SNat @m) $ ix+            in astGatherS shm+                          (shsDrop @p shpshn)+                          (shsTake @p shpshn)+                          (astConvDownSFromR shpshn x t) (vars, ix2)+            -- this introduces new variable names+  -- Depite the warning, the pattern match is exhaustive and if a dummy+  -- pattern is added, GHC 9.14.1 complains about that, in turn.+  trargMin a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @_ @rest _ _ ->+          -- unfortunately, this is not enough:+          -- gcastWith (unsafeCoerceRefl :: Rank sh :~: 1 + Rank (Init sh)) $+          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $+          astConvUpRFromS (shsInit sh) FTKScalar+          . fromPlain . AstArgMinS . astPlainPart . astConvDownSFromR sh x $ a+  trargMax a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @_ @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $+          astConvUpRFromS (shsInit sh) FTKScalar+          . fromPlain . AstArgMaxS . astPlainPart . astConvDownSFromR sh x $ a+  triota @r n =+    withSNat n $ \snat@(SNat @n) ->+      astConvUpRFromS (snat :$$ ZSS) FTKScalar+      $ fromPlain $ AstIotaS @n @r SNat+  trappend u v = case ftkAst u of+    FTKR shu' x -> case ftkAst v of+      FTKR shv' _ ->+        withShsFromShR shu' $ \shu -> case shu of+          usnat :$$ restu ->+            withShsFromShR shv' $ \shv -> case shv of+              vsnat :$$ restv ->+                case testEquality restu restv of+                  Just Refl ->+                    astConvUpRFromS (snatPlus usnat vsnat :$$ restu) x+                    $ astAppendS (astConvDownSFromR shu x u)+                                 (astConvDownSFromR shv x v)+                  _ -> error $ "rappend: shapes don't match: "+                               ++ show (restu, restv)+  trslice i n a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh@(_ :$$ shRest) -> case sh of+        msnat@(SNat @m) :$$ _ ->+          withSNat i $ \isnat@(SNat @i) -> withSNat n $ \nsnat@(SNat @n) ->+            case cmpNat (snatPlus isnat nsnat) msnat of+              GTI -> error $ "rslice: argument tensor too narrow: "+                             ++ show (i, n, fromSNat' msnat)+              EQI ->+                astConvUpRFromS (nsnat :$$ shRest) x+                . astSliceS isnat nsnat (SNat @(m - (i + n)))+                . astConvDownSFromR sh x $ a+              LTI ->+                astConvUpRFromS (nsnat :$$ shRest) x+                . astSliceS isnat nsnat (SNat @(m - (i + n)))+                . astConvDownSFromR sh x $ a+  trreverse a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh -> case sh of+        _ :$$ _ ->+          astConvUpRFromS sh x+          . astReverseS . astConvDownSFromR sh x $ a+  trtranspose @n @r permr a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh)  ->+        Permutation.permFromListCont permr $ \(perm :: Permutation.Perm perm) ->+          let result :: AstTensor AstMethodLet s (TKR2 n r)+              result =+                -- A noble lie, verified down below.+                gcastWith (unsafeCoerceRefl+                           :: (Rank perm <=? Rank sh) :~: True) $+                gcastWith (unsafeCoerceRefl+                           :: Rank (Permutation.PermutePrefix perm sh)+                              :~: Rank sh) $+                gcastWith (unsafeCoerceRefl+                           :: Permutation.Permute perm (TakeLen perm sh)+                              ++ DropLen perm sh+                              :~: sh) $+                fromMaybe (error "rtranspose: impossible non-permutation")+                $ Permutation.permCheckPermutation perm+                $ astConvUpRFromS sh x+                  . astTransposeS perm . astConvDownSFromR sh x $ a+          in case (Permutation.permRank perm, shsRank sh) of+            (psnat@SNat, shsnat@SNat) ->+              case cmpNat psnat shsnat of+                GTI -> error $ "rtranspose: rank mismatch: "+                               ++ show (fromSNat' psnat, fromSNat' shsnat)+                EQI -> result+                LTI -> result+  trreshape sh2' a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh ->+      withShsFromShR sh2' $ \sh2 ->+        case testEquality (shsProduct sh) (shsProduct sh2) of+          Just Refl -> astConvUpRFromS sh2 x+                       . astReshapeS sh2 . astConvDownSFromR sh x $ a+          _ -> error $ "rreshape: tensor size mismatch: "+                       ++ show ( fromSNat' (shsProduct sh)+                               , fromSNat' (shsProduct sh2) )+  trbuild1 @n @x k f = withSNat k $ \snat ->+    astBuild1Vectorize snat (STKR (SNat @n) (knownSTK @x)) f++  -- Shaped ops+  sshape t = case ftkAst t of+    FTKS sh _ -> sh+  tsfromVectorN = astFromVectorS+  tsfromVectorLinear = astFromVectorK+  tssumN @shm @shn t | SNat <- shsRank (knownShS @shm) =+    gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+    astSumS (shsTake @(Rank shm) (sshape t)) t+  tssum0 = astSumK+  tsreplicateN = astReplicateS+  tsreplicate0N = astReplicateK+  tsconcrete = fromPlain . AstConcreteS+  tsfloor = fromPlain . astFloorS . astPlainPart+  tsfromIntegral = fromPlain . astFromIntegralS . astPlainPart+  tscast = astCastS+  tsindex = astIndexS knownShS+  tsindex0 = astIndexK+  tsscatter @shm @shn @shp t f =+    funToVarsIxS knownShS $ \vars ix ->+      let !ix2 = f ix+      in astScatterS (knownShS @shm) (knownShS @shn) (knownShS @shp)+                     t (vars, ix2)+      -- this introduces new variable names+  tsgather @shm @shn @shp t f =+    funToVarsIxS knownShS $ \vars ix ->+      let !ix2 = f ix+      in astGatherS (knownShS @shm) (knownShS @shn) (knownShS @shp)+                    t (vars, ix2)+      -- this introduces new variable names+  tsargMin = fromPlain . AstArgMinS . astPlainPart+  tsargMax = fromPlain . AstArgMaxS . astPlainPart+  tsiota = fromPlain $ AstIotaS SNat+  tsappend = astAppendS+  tsslice = astSliceS+  tsreverse = astReverseS+  tstranspose = astTransposeS+  tsreshape = astReshapeS+  tsbuild1 @k @sh @x =+    astBuild1Vectorize (SNat @k) (STKS (knownShS @sh) (knownSTK @x))++  -- Mixed ops+  xshape t = case ftkAst t of+    FTKX sh _ -> sh+  txfromVectorN @shm0 @shn0 shm0 l = case V.uncons l of+    Just (v, _) | FTKX shn0 x <- ftkAst v ->+      withShsFromShX shn0 $ \(shn :: ShS shn) ->+      withShsFromShX shm0 $ \(shm :: ShS shm) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm0 ++ shn0) :~: Rank shm0 + Rank shn0) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ shn) :~: Rank shm + Rank shn) $+        astConvUpXFromS (shm0 `shxAppend` shn0) x+        . astFromVectorS shm . V.map (astConvDownSFromX shn x) $ l+    Nothing -> error "trfromVectorN: empty vector"+  txfromVectorLinear shm0 l =+    withShsFromShX shm0 $ \(shm :: ShS shm) ->+      astConvUpXFromS shm0 FTKScalar . astFromVectorK shm $ l+  txsumN @shm0 @shn0 t+    | SNat <- ssxRank (knownShX @shm0) = case ftkAst t of+      FTKX shmshn0 x ->+        withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->+          gcastWith (unsafeCoerceRefl+                     :: Rank (Drop (Rank shm0) shmshn) :~: Rank shn0) $+          gcastWith (unsafeCoerceRefl+                     :: Take (Rank shm0) shmshn ++ Drop (Rank shm0) shmshn+                        :~: shmshn) $+          astConvUpXFromS @(Drop (Rank shm0) shmshn)+                           (shxDropSSX (knownShX @shm0) shmshn0) x+          $ astSumS (shsTake @(Rank shm0) shmshn)+          $ astConvDownSFromX shmshn x t+  txsum0 t = case ftkAst t of+    FTKX shm x ->+      withShsFromShX shm $ \(sh :: ShS sh) ->+        astSumK (astConvDownSFromX sh x t)+  txreplicateN @shm0 @shn0 shm0 t = case ftkAst t of+    FTKX shn0 x ->+      withShsFromShX shn0 $ \(shn :: ShS shn) ->+      withShsFromShX shm0 $ \(shm :: ShS shm) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm0 ++ shn0) :~: Rank shm0 + Rank shn0) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ shn) :~: Rank shm + Rank shn) $+        astConvUpXFromS @(shm ++ shn) (shm0 `shxAppend` shn0) x+        . astReplicateS shm . astConvDownSFromX shn x $ t+  txreplicate0N shm0 t =+    withShsFromShX shm0 $ \(shm :: ShS shm) ->+      astConvUpXFromS shm0 FTKScalar . astReplicateK shm $ t+  txconcrete a = tconcrete (FTKX (Nested.mshape a) FTKScalar) (Concrete a)+  txfloor a = case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        astConvUpXFromS sh' FTKScalar+        . fromPlain . astFloorS . astPlainPart+        . astConvDownSFromX sh FTKScalar $ a+  txfromIntegral a = case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        astConvUpXFromS sh' FTKScalar+        . fromPlain . astFromIntegralS+        . astPlainPart . astConvDownSFromX sh FTKScalar $ a+  txcast a = case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        astConvUpXFromS sh' FTKScalar+        . astCastS . astConvDownSFromX sh FTKScalar $ a+  txindex @sh1 @sh2 a ix = case ftkAst a of+    FTKX @sh1sh2 sh1sh2 x | SNat <- ixxRank ix ->+      withShsFromShX sh1sh2 $ \(sh :: ShS sh) ->+        gcastWith (unsafeCoerceRefl :: Rank (Drop (Rank sh1) sh) :~: Rank sh2) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank sh1) sh ++ Drop (Rank sh1) sh :~: sh) $+        gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) sh1sh2 :~: sh2) $+        astConvUpXFromS (shxDropIx ix sh1sh2) x+        $ astIndexS @(Take (Rank sh1) sh) @(Drop (Rank sh1) sh)+                    (shsDrop @(Rank sh1) sh)+                    (astConvDownSFromX sh x a)+                    (ixsFromIxX' (shsTake @(Rank sh1) sh) ix)+  txindex0 a ix = case ftkAst a of+    FTKX shm x ->+      withShsFromShX shm $ \(sh :: ShS sh) ->+        astIndexK (astConvDownSFromX sh x a) (ixsFromIxX' sh ix)+  txscatter @shm @shn @shp shp0 t f = case ftkAst t of+    FTKX shmshn0 x | SNat <- ssxRank (knownShX @shm)+                   , SNat <- ssxRank (knownShX @shp) ->+      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShX shp0 $ \(shp :: ShS shp2) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shp2 ++ Drop (Rank shm) shmshn)+                      :~: Rank (shp ++ shn)) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn+                      :~: shmshn) $+        astConvUpXFromS (shp0 `shxAppend` shxDropSSX @_ @shn+                                            (knownShX @shm) shmshn0) x+        $ funToVarsIxS (shsTake @(Rank shm) shmshn) $ \vars ix ->+            let !ix2 = ixsFromIxX' shp . f . ixxFromIxS' knownShX $ ix+            in astScatterS (shsTake @(Rank shm) shmshn)+                           (shsDrop @(Rank shm) shmshn)+                           shp+                           (astConvDownSFromX shmshn x t) (vars, ix2)+            -- this introduces new variable names+  txgather @shm @shn @shp shm0 t f = case ftkAst t of+    FTKX shpshn0 x | SNat <- ssxRank (knownShX @shm)+                   , SNat <- ssxRank (knownShX @shp) ->+      withShsFromShX shm0 $ \(shm :: ShS shm2) ->+      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm2 ++ Drop (Rank shp) shpshn)+                      :~: Rank (shm ++ shn)) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn+                      :~: shpshn) $+        astConvUpXFromS (shm0 `shxAppend` shxDropSSX @_ @shn+                                            (knownShX @shp) shpshn0) x+        $ funToVarsIxS shm $ \vars ix ->+            let !ix2 = ixsFromIxX' (shsTake @(Rank shp) shpshn)+                       . f . ixxFromIxS' (ssxFromShX shm0) $ ix+            in astGatherS shm+                          (shsDrop @(Rank shp) shpshn)+                          (shsTake @(Rank shp) shpshn)+                          (astConvDownSFromX shpshn x t) (vars, ix2)+            -- this introduces new variable names+  txargMin a = case ftkAst a of+    FTKX @sh' sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @n @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $+          astConvUpXFromS (shxInit sh') FTKScalar+          . fromPlain . AstArgMinS @n @rest+          . astPlainPart . astConvDownSFromX sh x $ a+  txargMax a = case ftkAst a of+    FTKX @sh' sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @n @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $+          astConvUpXFromS (shxInit sh') FTKScalar+          . fromPlain . AstArgMaxS @n @rest+          . astPlainPart . astConvDownSFromX sh x $ a+  txiota @n @r = astConvUpXFromS (SKnown (SNat @n) :$% ZSX) FTKScalar+                 $ fromPlain $ AstIotaS @n @r SNat+  txappend u v = case ftkAst u of+    FTKX (m' :$% shu') x -> case ftkAst v of+      FTKX (n' :$% shv') _ ->+        withSNat (fromSMayNat' m') $ \m ->+        withSNat (fromSMayNat' n') $ \n ->+        withShsFromShX shu' $ \(shu :: ShS shu) ->+        withShsFromShX shv' $ \(shv :: ShS shv) ->+          case shxEqual shu' shv' of+            Just Refl ->+              gcastWith (unsafeCoerceRefl :: shu :~: shv) $+              astConvUpXFromS (smnAddMaybe m' n' :$% shu') x+              $ astAppendS (astConvDownSFromX (m :$$ shu) x u)+                           (astConvDownSFromX (n :$$ shv) x v)+            _ -> error $ "xappend: shapes don't match: "+                         ++ show (shu', shv')+  txslice i' n' k' a = case ftkAst a of+    FTKX sh'@(_ :$% sh2') x ->+      withSNat (fromSMayNat' i') $ \i ->+      withSNat (fromSMayNat' n') $ \n@(SNat @n) ->+      withSNat (fromSMayNat' k') $ \k ->+      withShsFromShX sh' $ \sh@(msnat :$$ (_ :: ShS sh2)) ->+        case testEquality (snatPlus (snatPlus i n) k) msnat of+          Just Refl ->+            astConvUpXFromS @(n ': sh2) (n' :$% sh2') x+            . astSliceS i n k . astConvDownSFromX sh x $ a+          _ -> error $ "xslice: argument tensor has a wrong width: "+                       ++ show ( fromSNat' i, fromSNat' n, fromSNat' k+                               , fromSNat' msnat )+  txreverse a = case ftkAst a of+    FTKX sh' x ->+      withShsFromShX sh' $ \(sh@(_ :$$ _) :: ShS sh) ->+        astConvUpXFromS @sh sh' x+        . astReverseS . astConvDownSFromX sh x $ a+  txtranspose @perm perm a = case ftkAst a of+    FTKX @sh' sh' x ->+      let sh2' = shxPermutePrefix perm sh'+      in withShsFromShX sh' $ \(sh :: ShS sh) ->+           gcastWith (unsafeCoerceRefl+                      :: Rank (Permutation.PermutePrefix perm sh')+                         :~: Rank sh') $+           gcastWith (unsafeCoerceRefl+                      :: Rank (Permutation.PermutePrefix perm sh)+                         :~: Rank sh) $+           astConvUpXFromS sh2' x+           . astTransposeS perm+           . astConvDownSFromX sh x $ a+  txreshape sh2' a = case ftkAst a of+    FTKX sh' x ->+      withShsFromShX sh' $ \sh ->+      withShsFromShX sh2' $ \sh2 ->+        case testEquality (shsProduct sh) (shsProduct sh2) of+          Just Refl ->+            astConvUpXFromS sh2' x+            . astReshapeS sh2 . astConvDownSFromX sh x $ a+          _ -> error $ "xreshape: tensor size mismatch: "+                       ++ show ( fromSNat' (shsProduct sh)+                               , fromSNat' (shsProduct sh2) )+  txbuild1 @k @sh @x =+    astBuild1Vectorize (SNat @k) (STKX (knownShX @sh) (knownSTK @x))++  -- Scalar ops+  tkconcrete = fromPlain . AstConcreteK+  tkfloor = fromPlain . astFloorK . astPlainPart+  tkfromIntegral = fromPlain . astFromIntegralK . astPlainPart+  tkcast = astCastK+  tkargMin = fromPlain . AstArgMinK . astPlainPart+  tkargMax = fromPlain . AstArgMaxK . astPlainPart+  tkbuild1 @k = astBuild1Vectorize (SNat @k) STKScalar++  -- General operations that don't require LetTensor nor ShareTensor+  tftk _stk = ftkAst+  tpair = astPair+  tproject1 = astProject1+  tproject2 = astProject2+  kcond !b !u !v = astCondKInitial b u v+  scond !b !u !v = astCondSInitial b u v+  tcond stk !b !u !v = case stk of+    STKScalar -> astCondKInitial b u v+    STKR{} | FTKR sh0 x <- ftkAst u ->+      withShsFromShR sh0 $ \(sh :: ShS sh) ->+        astConvUpRFromS sh x+        $ astCondSInitial b (astConvDownSFromR sh x u)+                            (astConvDownSFromR sh x v)+    STKS{} -> astCondSInitial b u v+    STKX{} | FTKX sh0 x <- ftkAst u ->+      withShsFromShX sh0 $ \(sh :: ShS sh) ->+        astConvUpXFromS sh0 x+        $ astCondSInitial b (astConvDownSFromX sh x u)+                            (astConvDownSFromX sh x v)+    STKProduct stk1 stk2 -> ttlet u $ \uShared -> ttlet v $ \vShared ->+      tpair (tcond stk1 b (tproject1 uShared) (tproject1 vShared))+            (tcond stk2 b (tproject2 uShared) (tproject2 vShared))+  tconcrete ftk a = fromPlain $ astConcrete ftk a+  tmapAccumR proxy !k !accftk !bftk !eftk f acc0 es =+    ttlet (tmapAccumL proxy k accftk bftk eftk f acc0+                      (treverse k (ftkToSTK eftk) es)) $ \ !res ->+      tpair (tproject1 res) (treverse k (ftkToSTK bftk) $ tproject2 res)+  tmapAccumRDer proxy !k !accftk !bftk !eftk f df rf acc0 es =+    ttlet (tmapAccumLDer proxy k accftk bftk eftk f df rf acc0+                         (treverse k (ftkToSTK eftk) es)) $ \ !res ->+      tpair (tproject1 res) (treverse k (ftkToSTK bftk) $ tproject2 res)+  tmapAccumLDer _ !k _ !bftk !eftk f df rf acc0 es =+    astMapAccumLDer k bftk eftk f df rf acc0 es+  tapply = astApply+  tlambda ftk f =+    let (var, ast) = funToAst ftk $ unHFun f+    in AstLambda var ast+  tgrad @_ @r xftk f | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =+    -- We don't have an AST constructor to hold it, so we compute outright.+    --+    -- This computes the (AST of) derivative of f once and interprets it again+    -- for each new tensor of arguments, which is better than computing it anew.+    gcastWith (unsafeCoerceRefl+               :: SpanTargetFam (AstTensor AstMethodLet FullSpan) s+                  :~: AstTensor AstMethodLet s) $+    let -- No bangs here, because this goes under lambda and should not be+        -- evaluated too early (which at some point was even incorrect+        -- and triggered error "tunshare: used not at FullSpan"; maybe this+        -- is related to terms getting spans converted when interpreted)+        AstArtifactRev{..} =+          revProduceArtifact+            IgnoreIncomingCotangent (simplifyUserCode . unHFun f) emptyEnv xftk+        -- A new variable is created to give it the right span as opposed+        -- to the fixed FullSpan that artVarDomainRev has.+        (varP, ast) = funToAst xftk $ \ !astP ->+          simplifyUserCode+          $ fromFullSpan (knownSpan @s)+          $ substituteAst (toFullSpan (ftkToSTK xftk) (knownSpan @s) astP)+                          artVarDomainRev+                          artDerivativeRev+    in AstLambda varP ast+  tvjp ftkx f =+    -- This computes the (AST of) derivative of f once and interprets it again+    -- for each new tensor of arguments, which is better than computing it anew.+    gcastWith (unsafeCoerceRefl+               :: SpanTargetFam (AstTensor AstMethodLet FullSpan) s+                  :~: AstTensor AstMethodLet s) $+    let AstArtifactRev{..} =+          revProduceArtifactDt+            (simplifyUserCode . unHFun f) emptyEnv ftkx+        ftkz = varNameToFTK artVarDtRev+        ftk2 = FTKProduct ftkz ftkx+        (varP, ast) = funToAst ftk2 $ \ !astP ->+          simplifyUserCode+          $ fromFullSpan (knownSpan @s)+          $ substituteAst+              (toFullSpan (ftkToSTK ftkx) (knownSpan @s) (astProject2 astP))+              artVarDomainRev+          $ substituteAst+              (toFullSpan (ftkToSTK ftkz) (knownSpan @s) (astProject1 astP))+              artVarDtRev+              artDerivativeRev+    in AstLambda varP ast+  tjvp ftkx f =+    -- This computes the (AST of) derivative of f once and interprets it again+    -- for each new tensor of arguments, which is better than computing it anew.+    gcastWith (unsafeCoerceRefl+               :: SpanTargetFam (AstTensor AstMethodLet FullSpan) s+                  :~: AstTensor AstMethodLet s) $+    let AstArtifactFwd{..} =+          fwdProduceArtifact (simplifyUserCode . unHFun f) emptyEnv ftkx+        ftk2 = FTKProduct (adFTK ftkx) ftkx+        (varP, ast) = funToAst ftk2 $ \ !astP ->+          simplifyUserCode+          $ fromFullSpan (knownSpan @s)+          $ substituteAst+              (toFullSpan (ftkToSTK ftkx) (knownSpan @s) (astProject2 astP))+              artVarDomainFwd+          $ substituteAst+              (toFullSpan (ftkToSTK (adFTK ftkx)) (knownSpan @s)+                          (astProject1 astP))+              artVarDsFwd+              artDerivativeFwd+    in AstLambda varP ast++  {-# INLINE tsum #-}+  tsum snat@SNat stk u = case stk of+    STKScalar -> kfromS $ tssum u+    STKR SNat x | Dict <- lemKnownSTK x -> trsum u+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (tsum snat stk1 (tproject1 u3))+              (tsum snat stk2 (tproject2 u3))+  {-# INLINE treplicate #-}+  treplicate snat@SNat stk u = case stk of+    STKScalar -> tsreplicate snat $ sfromK u+    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (fromSNat' snat) u+    STKS _ x | Dict <- lemKnownSTK x -> tsreplicate snat u+    STKX _ x | Dict <- lemKnownSTK x -> txreplicate snat u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (treplicate snat stk1 (tproject1 u3))+              (treplicate snat stk2 (tproject2 u3))+  {-# INLINE treverse #-}+  treverse snat stk u = case stk of+    STKScalar -> tsreverse u+    STKR _ x | Dict <- lemKnownSTK x -> trreverse u+    STKS _ x | Dict <- lemKnownSTK x -> tsreverse u+    STKX _ x | Dict <- lemKnownSTK x -> txreverse u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (treverse snat stk1 (tproject1 u3))+              (treverse snat stk2 (tproject2 u3))++  tprimalPart = astPrimalPart+  tdualPart _ = dualPart+  tplainPart = astPlainPart+  tfromPrimal _ = fromPrimal+  tfromDual = fromDual+  tfromPlain _ = fromPlain++  taddTarget = addTarget+  tmultTarget = multTarget+  tsum0Target = sum0Target+  tdot0Target = dot0Target+++-- * AstRaw instances++fmapAstRaw :: Coercible (f (AstTensor AstMethodShare s y)) (f (AstRaw s y))+           => f (AstTensor AstMethodShare s y) -> f (AstRaw s y)+fmapAstRaw = coerce++fmapUnAstRaw :: Coercible (f (AstRaw s y)) (f (AstTensor AstMethodShare s y))+             => f (AstRaw s y) -> f (AstTensor AstMethodShare s y)+fmapUnAstRaw = coerce++instance KnownSpan s => LetTensor (AstRaw s) where+  ttlet u f =+    let !var2 = tshare u+    in f var2+  ttletPrimal u f =+    let !var2 = tshare u+    in f var2+  ttletPlain u f =+    let !var2 = tshare u+    in f var2+  toShare = id+  tunshare = id++instance KnownSpan s => ShareTensor (AstRaw s) where+  tshare t = AstRaw $ astShareNoSimplify $ unAstRaw t+  tunpair (AstRaw (AstPair t1 t2)) = (AstRaw t1, AstRaw t2)+  tunpair t = let tShared = tshare t+              in (tproject1 tShared, tproject2 tShared)++instance KnownSpan s => BaseTensor (AstRaw s) where+  isConcreteInstance = False+  -- Ranked ops+  rshape t = case ftkAst $ unAstRaw t of+    FTKR sh _ -> sh+  trfromVectorN shm0 l = AstRaw $ case V.uncons l of+    Just (AstRaw v, _) | FTKR shn0 x <- ftkAst v ->+      withShsFromShR shn0 $ \(shn :: ShS shn) ->+      withShsFromShR shm0 $ \(shm :: ShS shm) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ shn) :~: Rank shm + Rank shn) $+        cAstConvUpRFromS (shm `shsAppend` shn) x+        . AstFromVectorS shm . V.map (cAstConvDownSFromR shn x) $ fmapUnAstRaw l+    Nothing -> error "trfromVectorN: empty vector"+  trfromVectorLinear shm0 l = AstRaw $+    withShsFromShR shm0 $ \(shm :: ShS shm) ->+      cAstConvUpRFromS shm FTKScalar . AstFromVectorK shm $ fmapUnAstRaw l+  trsumN @m @n (AstRaw t) = AstRaw $ case ftkAst t of+    FTKR shmshn0 x ->+      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->+        gcastWith (unsafeCoerceRefl :: Rank (Drop m shmshn) :~: n) $+        gcastWith (unsafeCoerceRefl+                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $+        cAstConvUpRFromS (shsDrop @m shmshn) x . AstSumS (shsTake @m shmshn)+        . cAstConvDownSFromR shmshn x $ t+  trsum0  (AstRaw t) = AstRaw $ case ftkAst t of+    FTKR shm x ->+      withShsFromShR shm $ \(sh :: ShS sh) ->+        AstSumK (cAstConvDownSFromR sh x t)+  trreplicateN shm0 (AstRaw t) = AstRaw $ case ftkAst t of+    FTKR shn0 x ->+      withShsFromShR shn0 $ \(shn :: ShS shn) ->+      withShsFromShR shm0 $ \(shm :: ShS shm) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ shn) :~: Rank shm + Rank shn) $+        cAstConvUpRFromS (shm `shsAppend` shn) x . AstReplicateS shm+        . cAstConvDownSFromR shn x $ t+  trreplicate0N shm0  (AstRaw t) = AstRaw $+    withShsFromShR shm0 $ \(shm :: ShS shm) ->+      cAstConvUpRFromS shm FTKScalar . AstReplicateK shm $ t+  trconcrete a = tconcrete (FTKR (Nested.rshape a) FTKScalar) (Concrete a)+  trfloor (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        cAstConvUpRFromS sh FTKScalar+        . fromPlain . AstFloorS . plainPart+        . cAstConvDownSFromR sh FTKScalar $ a+  trfromIntegral (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        cAstConvUpRFromS sh FTKScalar+        . fromPlain . AstFromIntegralS . plainPart+        . cAstConvDownSFromR sh FTKScalar $ a+  trcast (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        cAstConvUpRFromS sh FTKScalar+        . AstCastS . cAstConvDownSFromR sh FTKScalar $ a+  trindex @m @n (AstRaw a) ix | SNat <- ixrRank ix = AstRaw $ case ftkAst a of+    FTKR shmshn x ->+      withShsFromShR shmshn $ \(sh :: ShS sh) ->+        gcastWith (unsafeCoerceRefl :: Rank (Take m sh) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Drop m sh) :~: n) $+        gcastWith (unsafeCoerceRefl :: Take m sh ++ Drop m sh :~: sh) $+        cAstConvUpRFromS (shsDrop @m sh) x+        $ AstIndexS @(Take m sh) @(Drop m sh)+                    (shsDrop @m sh) (cAstConvDownSFromR sh x a)+                    (ixsFromIxR' (shsTake @m sh) (fmapUnAstRaw ix))+  trindex0 (AstRaw a) ix = AstRaw $ case ftkAst a of+    FTKR shm x ->+      withShsFromShR shm $ \(sh :: ShS sh) ->+        AstIndexK (cAstConvDownSFromR sh x a) (ixsFromIxR' sh (fmapUnAstRaw ix))+  trscatter @m @n shp0 (AstRaw t) f = AstRaw $ case ftkAst t of+    FTKR shmshn0 x ->+      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShR shp0 $ \(shp :: ShS shp) ->+        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Drop m shmshn) :~: n) $+        gcastWith (unsafeCoerceRefl+                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (shp ++ Drop m shmshn)+                      :~: Rank shp + Rank (Drop m shmshn)) $+        cAstConvUpRFromS (shp `shsAppend` shsDrop @m shmshn) x+        $ funToVarsIxS (shsTake @m shmshn) $ \vars ix ->+            let !ix2 = fmapUnAstRaw . ixsFromIxR' shp+                       . f . ixrFromIxS' (SNat @m) . fmapAstRaw $ ix+            in AstScatterS (shsTake @m shmshn)+                           (shsDrop @m shmshn)+                           shp+                           (cAstConvDownSFromR shmshn x t) (vars, ix2)+            -- this introduces new variable names+  trgather @m @n @p shm0 (AstRaw t) f = AstRaw $ case ftkAst t of+    FTKR shpshn0 x ->+      withShsFromShR shm0 $ \(shm :: ShS shm) ->+      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->+        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $+        gcastWith (unsafeCoerceRefl :: Rank (Drop p shpshn) :~: n) $+        gcastWith (unsafeCoerceRefl+                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ Drop p shpshn)+                      :~: Rank shm + Rank (Drop p shpshn)) $+        cAstConvUpRFromS (shm `shsAppend` shsDrop @p shpshn) x+        $ funToVarsIxS shm $ \vars ix ->+            let !ix2 = fmapUnAstRaw . ixsFromIxR' (shsTake @p shpshn)+                       . f . ixrFromIxS' (SNat @m) . fmapAstRaw $ ix+            in AstGatherS shm+                          (shsDrop @p shpshn)+                          (shsTake @p shpshn)+                          (cAstConvDownSFromR shpshn x t) (vars, ix2)+            -- this introduces new variable names+  trargMin (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @_ @rest _ _ ->+          -- unfortunately, this is not enough:+          -- gcastWith (unsafeCoerceRefl :: Rank sh :~: 1 + Rank (Init sh)) $+          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $+          cAstConvUpRFromS (shsInit sh) FTKScalar+          . fromPlain . AstArgMinS . plainPart . cAstConvDownSFromR sh x $ a+  trargMax (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @_ @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $+          cAstConvUpRFromS (shsInit sh) FTKScalar+          . fromPlain . AstArgMaxS . plainPart . cAstConvDownSFromR sh x $ a+  triota @r n =+    AstRaw+    $ withSNat n $ \snat@(SNat @n) ->+        cAstConvUpRFromS (snat :$$ ZSS) FTKScalar+        $ fromPlain $ AstIotaS @n @r SNat+  trappend (AstRaw u) (AstRaw v) = AstRaw $ case ftkAst u of+    FTKR shu' x -> case ftkAst v of+      FTKR shv' _ ->+        withShsFromShR shu' $ \shu -> case shu of+          usnat :$$ restu ->+            withShsFromShR shv' $ \shv -> case shv of+              vsnat :$$ restv ->+                case testEquality restu restv of+                  Just Refl ->+                    cAstConvUpRFromS (snatPlus usnat vsnat :$$ restu) x+                    $ AstAppendS (cAstConvDownSFromR shu x u)+                                 (cAstConvDownSFromR shv x v)+                  _ -> error $ "rappend: shapes don't match: "+                               ++ show (restu, restv)+  trslice i n (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh@(_ :$$ shRest) -> case sh of+        msnat@(SNat @m) :$$ _ ->+          withSNat i $ \isnat@(SNat @i) -> withSNat n $ \nsnat@(SNat @n) ->+            case cmpNat (snatPlus isnat nsnat) msnat of+              GTI -> error $ "rslice: argument tensor too narrow: "+                             ++ show (i, n, fromSNat' msnat)+              EQI ->+                cAstConvUpRFromS (nsnat :$$ shRest) x+                . AstSliceS isnat nsnat (SNat @(m - (i + n)))+                . cAstConvDownSFromR sh x $ a+              LTI ->+                cAstConvUpRFromS (nsnat :$$ shRest) x+                . AstSliceS isnat nsnat (SNat @(m - (i + n)))+                . cAstConvDownSFromR sh x $ a+  trreverse (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh -> case sh of+        _ :$$ _ ->+          cAstConvUpRFromS sh x+          . AstReverseS . cAstConvDownSFromR sh x $ a+  trtranspose @n @r permr (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        Permutation.permFromListCont permr $ \(perm :: Permutation.Perm perm) ->+          let result :: AstTensor AstMethodShare s (TKR2 n r)+              result =+                -- A noble lie, verified down below.+                gcastWith (unsafeCoerceRefl+                           :: (Rank perm <=? Rank sh) :~: True) $+                gcastWith (unsafeCoerceRefl+                           :: Rank (Permutation.PermutePrefix perm sh)+                              :~: Rank sh) $+                gcastWith (unsafeCoerceRefl+                           :: Permutation.Permute perm (TakeLen perm sh)+                              ++ DropLen perm sh+                              :~: sh) $+                fromMaybe (error "rtranspose: impossible non-permutation")+                $ Permutation.permCheckPermutation perm+                $ cAstConvUpRFromS sh x+                  . AstTransposeS perm . cAstConvDownSFromR sh x $ a+          in case (Permutation.permRank perm, shsRank sh) of+            (psnat@SNat, shsnat@SNat) ->+              case cmpNat psnat shsnat of+                GTI -> error $ "rtranspose: rank mismatch: "+                               ++ show (fromSNat' psnat, fromSNat' shsnat)+                EQI -> result+                LTI -> result+  trreshape sh2' (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh ->+      withShsFromShR sh2' $ \sh2 ->+        case testEquality (shsProduct sh) (shsProduct sh2) of+          Just Refl -> cAstConvUpRFromS sh2 x+                       . AstReshapeS sh2 . cAstConvDownSFromR sh x $ a+          _ -> error $ "rreshape: tensor size mismatch: "+                       ++ show ( fromSNat' (shsProduct sh)+                               , fromSNat' (shsProduct sh2) )+  trbuild1 k f = withSNat k $ \snat ->+    AstRaw $ AstBuild1 snat knownSTK+    $ funToAstInt (0, k - 1)+        -- this introduces new variable names+    $ unAstRaw . f . AstRaw++  -- Shaped ops+  sshape t = case ftkAst $ unAstRaw t of+    FTKS sh _ -> sh+  tsfromVectorN shm = AstRaw . AstFromVectorS shm . fmapUnAstRaw+  tsfromVectorLinear shm = AstRaw . AstFromVectorK shm . fmapUnAstRaw+  tssumN @shm @shn t | SNat <- shsRank (knownShS @shm) =+    gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+    AstRaw . AstSumS (shsTake @(Rank shm) (sshape t)) . unAstRaw $ t+  tssum0 = AstRaw . AstSumK . unAstRaw+  tsreplicateN shm = AstRaw . AstReplicateS shm . unAstRaw+  tsreplicate0N shm = AstRaw . AstReplicateK shm . unAstRaw+  tsconcrete = AstRaw . fromPlain . AstConcreteS+  tsfloor = AstRaw . fromPlain . AstFloorS . plainPart . unAstRaw+  tsfromIntegral =+    AstRaw . fromPlain . AstFromIntegralS . plainPart . unAstRaw+  tscast = AstRaw . AstCastS . unAstRaw+  tsindex v ix = AstRaw $ AstIndexS knownShS (unAstRaw v) (fmapUnAstRaw ix)+  tsindex0 v ix = AstRaw $ AstIndexK (unAstRaw v) (fmapUnAstRaw ix)+  tsscatter @shm @shn @shp t f = AstRaw $+    funToVarsIxS knownShS $ \vars ix ->+      let !ix2 = fmapUnAstRaw . f . fmapAstRaw $ ix+      in AstScatterS (knownShS @shm) (knownShS @shn) (knownShS @shp)+                     (unAstRaw t) (vars, ix2)+      -- this introduces new variable names+  tsgather @shm @shn @shp t f = AstRaw $+    funToVarsIxS knownShS $ \vars ix ->+      let !ix2 = fmapUnAstRaw . f . fmapAstRaw $ ix+      in AstGatherS (knownShS @shm) (knownShS @shn) (knownShS @shp)+                    (unAstRaw t) (vars, ix2)+  tsargMin = AstRaw . fromPlain . AstArgMinS . plainPart . unAstRaw+  tsargMax = AstRaw . fromPlain . AstArgMaxS . plainPart . unAstRaw+  tsiota = AstRaw . fromPlain $ AstIotaS SNat+  tsappend u v = AstRaw $ AstAppendS (unAstRaw u) (unAstRaw v)+  tsslice i n k = AstRaw . AstSliceS i n k . unAstRaw+  tsreverse = AstRaw . AstReverseS . unAstRaw+  tstranspose perm = AstRaw . AstTransposeS perm . unAstRaw+  tsreshape sh = AstRaw . AstReshapeS sh . unAstRaw+  tsbuild1 @k f = AstRaw $ AstBuild1 (SNat @k) knownSTK+                  $ funToAstInt (0, valueOf @k - 1)+                      -- this introduces new variable names+                  $ unAstRaw . f . AstRaw++  -- Mixed ops+  xshape t = case ftkAst $ unAstRaw t of+    FTKX sh _ -> sh+  txfromVectorN @shm0 @shn0 shm0 l = AstRaw $ case V.uncons l of+    Just (AstRaw v, _) | FTKX shn0 x <- ftkAst v ->+      withShsFromShX shn0 $ \(shn :: ShS shn) ->+      withShsFromShX shm0 $ \(shm :: ShS shm) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm0 ++ shn0) :~: Rank shm0 + Rank shn0) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ shn) :~: Rank shm + Rank shn) $+        cAstConvUpXFromS (shm0 `shxAppend` shn0) x+        . AstFromVectorS shm . V.map (cAstConvDownSFromX shn x) $ fmapUnAstRaw l+    Nothing -> error "trfromVectorN: empty vector"+  txfromVectorLinear shm0 l = AstRaw $+    withShsFromShX shm0 $ \(shm :: ShS shm) ->+      cAstConvUpXFromS shm0 FTKScalar . AstFromVectorK shm $ fmapUnAstRaw l+  txsumN @shm0 @shn0 (AstRaw t)+    | SNat <- ssxRank (knownShX @shm0) = AstRaw $ case ftkAst t of+      FTKX shmshn0 x ->+        withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->+          gcastWith (unsafeCoerceRefl+                     :: Rank (Drop (Rank shm0) shmshn) :~: Rank shn0) $+          gcastWith (unsafeCoerceRefl+                     :: Take (Rank shm0) shmshn ++ Drop (Rank shm0) shmshn+                        :~: shmshn) $+          cAstConvUpXFromS @(Drop (Rank shm0) shmshn)+                           (shxDropSSX (knownShX @shm0) shmshn0) x+          $ AstSumS (shsTake @(Rank shm0) shmshn)+          $ cAstConvDownSFromX shmshn x t+  txsum0 (AstRaw t) = AstRaw $ case ftkAst t of+    FTKX shm x ->+      withShsFromShX shm $ \(sh :: ShS sh) ->+        AstSumK (cAstConvDownSFromX sh x t)+  txreplicateN @shm0 @shn0 shm0 (AstRaw t) = AstRaw $ case ftkAst t of+    FTKX shn0 x ->+      withShsFromShX shn0 $ \(shn :: ShS shn) ->+      withShsFromShX shm0 $ \(shm :: ShS shm) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm0 ++ shn0) :~: Rank shm0 + Rank shn0) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm ++ shn) :~: Rank shm + Rank shn) $+        cAstConvUpXFromS @(shm ++ shn) (shm0 `shxAppend` shn0) x+        . AstReplicateS shm . cAstConvDownSFromX shn x $ t+  txreplicate0N shm0 (AstRaw t) = AstRaw $+    withShsFromShX shm0 $ \(shm :: ShS shm) ->+      cAstConvUpXFromS shm0 FTKScalar . AstReplicateK shm $ t+  txconcrete a = tconcrete (FTKX (Nested.mshape a) FTKScalar) (Concrete a)+  txfloor (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        cAstConvUpXFromS sh' FTKScalar+        . fromPlain . AstFloorS . plainPart+        . cAstConvDownSFromX sh FTKScalar $ a+  txfromIntegral (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        cAstConvUpXFromS sh' FTKScalar+        . fromPlain . AstFromIntegralS+        . plainPart . cAstConvDownSFromX sh FTKScalar $ a+  txcast (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        cAstConvUpXFromS sh' FTKScalar+        . AstCastS . cAstConvDownSFromX sh FTKScalar $ a+  txindex @sh1 @sh2 (AstRaw a) ix = AstRaw $ case ftkAst a of+    FTKX @sh1sh2 sh1sh2 x | SNat <- ixxRank ix ->+      withShsFromShX sh1sh2 $ \(sh :: ShS sh) ->+        gcastWith (unsafeCoerceRefl :: Rank (Drop (Rank sh1) sh) :~: Rank sh2) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank sh1) sh ++ Drop (Rank sh1) sh :~: sh) $+        gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) sh1sh2 :~: sh2) $+        cAstConvUpXFromS (shxDropIx ix sh1sh2) x+        $ AstIndexS @(Take (Rank sh1) sh) @(Drop (Rank sh1) sh)+                    (shsDrop @(Rank sh1) sh)+                    (cAstConvDownSFromX sh x a)+                    (ixsFromIxX' (shsTake @(Rank sh1) sh) (fmapUnAstRaw ix))+  txindex0 (AstRaw a) ix = AstRaw $ case ftkAst a of+    FTKX shm x ->+      withShsFromShX shm $ \(sh :: ShS sh) ->+        AstIndexK (cAstConvDownSFromX sh x a) (ixsFromIxX' sh (fmapUnAstRaw ix))+  txscatter @shm @shn @shp shp0 (AstRaw t) f = AstRaw $ case ftkAst t of+    FTKX shmshn0 x | SNat <- ssxRank (knownShX @shm)+                   , SNat <- ssxRank (knownShX @shp) ->+      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShX shp0 $ \(shp :: ShS shp2) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shp2 ++ Drop (Rank shm) shmshn)+                      :~: Rank (shp ++ shn)) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn+                      :~: shmshn) $+        cAstConvUpXFromS (shp0 `shxAppend` shxDropSSX @_ @shn+                                             (knownShX @shm) shmshn0) x+        $ funToVarsIxS (shsTake @(Rank shm) shmshn) $ \vars ix ->+            let !ix2 = fmapUnAstRaw . ixsFromIxX' shp+                       . f . ixxFromIxS' knownShX . fmapAstRaw $ ix+            in AstScatterS (shsTake @(Rank shm) shmshn)+                           (shsDrop @(Rank shm) shmshn)+                           shp+                           (cAstConvDownSFromX shmshn x t) (vars, ix2)+            -- this introduces new variable names+  txgather @shm @shn @shp shm0 (AstRaw t) f = AstRaw $ case ftkAst t of+    FTKX shpshn0 x | SNat <- ssxRank (knownShX @shm)+                   , SNat <- ssxRank (knownShX @shp) ->+      withShsFromShX shm0 $ \(shm :: ShS shm2) ->+      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->+        gcastWith (unsafeCoerceRefl+                   :: Rank (shm2 ++ Drop (Rank shp) shpshn)+                      :~: Rank (shm ++ shn)) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn+                      :~: shpshn) $+        cAstConvUpXFromS (shm0 `shxAppend` shxDropSSX @_ @shn+                                             (knownShX @shp) shpshn0) x+        $ funToVarsIxS shm $ \vars ix ->+            let !ix2 = fmapUnAstRaw . ixsFromIxX' (shsTake @(Rank shp) shpshn)+                       . f . ixxFromIxS' (ssxFromShX shm0) . fmapAstRaw $ ix+            in AstGatherS shm+                          (shsDrop @(Rank shp) shpshn)+                          (shsTake @(Rank shp) shpshn)+                          (cAstConvDownSFromX shpshn x t) (vars, ix2)+            -- this introduces new variable names+  txargMin (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX @sh' sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @n @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $+          cAstConvUpXFromS (shxInit sh') FTKScalar+          . fromPlain . AstArgMinS @n @rest+          . plainPart . cAstConvDownSFromX sh x $ a+  txargMax (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX @sh' sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @n @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $+          cAstConvUpXFromS (shxInit sh') FTKScalar+          . fromPlain . AstArgMaxS @n @rest+          . plainPart . cAstConvDownSFromX sh x $ a+  txiota @n @r = AstRaw $ cAstConvUpXFromS (SKnown (SNat @n) :$% ZSX) FTKScalar+                 $ fromPlain $ AstIotaS @n @r SNat+  txappend (AstRaw u) (AstRaw v) = AstRaw $ case ftkAst u of+    FTKX (m' :$% shu') x -> case ftkAst v of+      FTKX (n' :$% shv') _ ->+        withSNat (fromSMayNat' m') $ \m ->+        withSNat (fromSMayNat' n') $ \n ->+        withShsFromShX shu' $ \(shu :: ShS shu) ->+        withShsFromShX shv' $ \(shv :: ShS shv) ->+          case shxEqual shu' shv' of+            Just Refl ->+              gcastWith (unsafeCoerceRefl :: shu :~: shv) $+              cAstConvUpXFromS (smnAddMaybe m' n' :$% shu') x+              $ AstAppendS (cAstConvDownSFromX (m :$$ shu) x u)+                           (cAstConvDownSFromX (n :$$ shv) x v)+            _ -> error $ "xappend: shapes don't match: "+                         ++ show (shu', shv')+  txslice i' n' k' (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh'@(_ :$% sh2') x ->+      withSNat (fromSMayNat' i') $ \i ->+      withSNat (fromSMayNat' n') $ \n@(SNat @n) ->+      withSNat (fromSMayNat' k') $ \k ->+      withShsFromShX sh' $ \sh@(msnat :$$ (_ :: ShS sh2)) ->+        case testEquality (snatPlus (snatPlus i n) k) msnat of+          Just Refl ->+            cAstConvUpXFromS @(n ': sh2) (n' :$% sh2') x+            . AstSliceS i n k . cAstConvDownSFromX sh x $ a+          _ -> error $ "xslice: argument tensor has a wrong width: "+                       ++ show ( fromSNat' i, fromSNat' n, fromSNat' k+                               , fromSNat' msnat )+  txreverse (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' x ->+      withShsFromShX sh' $ \(sh@(_ :$$ _) :: ShS sh) ->+        cAstConvUpXFromS @sh sh' x  -- the @sh needed for GHC 9.10 only+        . AstReverseS . cAstConvDownSFromX sh x $ a+  txtranspose @perm perm (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX @sh' sh' x ->+      let sh2' = shxPermutePrefix perm sh'+      in withShsFromShX sh' $ \(sh :: ShS sh) ->+           gcastWith (unsafeCoerceRefl+                      :: Rank (Permutation.PermutePrefix perm sh')+                         :~: Rank sh') $+           gcastWith (unsafeCoerceRefl+                      :: Rank (Permutation.PermutePrefix perm sh)+                         :~: Rank sh) $+           cAstConvUpXFromS sh2' x+           . AstTransposeS perm+           . cAstConvDownSFromX sh x $ a+  txreshape sh2' (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' x ->+      withShsFromShX sh' $ \sh ->+      withShsFromShX sh2' $ \sh2 ->+        case testEquality (shsProduct sh) (shsProduct sh2) of+          Just Refl ->+            cAstConvUpXFromS sh2' x+            . AstReshapeS sh2 . cAstConvDownSFromX sh x $ a+          _ -> error $ "xreshape: tensor size mismatch: "+                       ++ show ( fromSNat' (shsProduct sh)+                               , fromSNat' (shsProduct sh2) )+  txbuild1 @k f = AstRaw $ AstBuild1 (SNat @k) knownSTK+                  $ funToAstInt (0, valueOf @k - 1)+                      -- this introduces new variable names+                  $ unAstRaw . f . AstRaw++  -- Scalar ops+  tkconcrete = AstRaw . fromPlain . AstConcreteK+  tkfloor = AstRaw . fromPlain . AstFloorK . plainPart . unAstRaw+  tkfromIntegral = AstRaw . fromPlain . AstFromIntegralK+                   . plainPart . unAstRaw+  tkcast = AstRaw . AstCastK . unAstRaw+  tkargMin = AstRaw . fromPlain . AstArgMinK . plainPart . unAstRaw+  tkargMax = AstRaw . fromPlain . AstArgMaxK . plainPart . unAstRaw+  tkbuild1 @k f = AstRaw $ AstBuild1 (SNat @k) STKScalar+                  $ funToAstInt (0, valueOf @k - 1)+                      -- this introduces new variable names+                  $ unAstRaw . f . AstRaw++  -- General operations that don't require LetTensor nor ShareTensor+  tftk _stk = ftkAst . unAstRaw+  tpair t1 t2 = AstRaw $ AstPair (unAstRaw t1) (unAstRaw t2)+  tproject1 t = AstRaw $ AstProject1 $ unAstRaw t+  tproject2 t = AstRaw $ AstProject2 $ unAstRaw t+  kcond (AstRaw b) (AstRaw u) (AstRaw v) = AstRaw $ AstCondK b u v+  scond (AstRaw b) (AstRaw u) (AstRaw v) = AstRaw $ AstCondS b u v+  tcond stk b0@(AstRaw b) u0@(AstRaw u) v0@(AstRaw v) = case stk of+    STKScalar -> AstRaw $ AstCondK b u v+    STKR{} | FTKR sh0 x <- ftkAst u -> AstRaw $+      withShsFromShR sh0 $ \(sh :: ShS sh) ->+        cAstConvUpRFromS sh x+        $ AstCondS b (cAstConvDownSFromR sh x u)+                     (cAstConvDownSFromR sh x v)+    STKS{} -> AstRaw $ AstCondS b u v+    STKX{} | FTKX sh0 x <- ftkAst u -> AstRaw $+      withShsFromShX sh0 $ \(sh :: ShS sh) ->+        cAstConvUpXFromS sh0 x+        $ AstCondS b (cAstConvDownSFromX sh x u)+                     (cAstConvDownSFromX sh x v)+    STKProduct stk1 stk2 ->+      let (u1, u2) = tunpair u0+          (v1, v2) = tunpair v0+      in tpair (tcond stk1 b0 u1 v1) (tcond stk2 b0 u2 v2)+  tconcrete ftk a = AstRaw $ fromPlain $ unAstRaw $ astConcreteRaw ftk a+  tmapAccumLDer _ !k _ !bftk !eftk f df rf acc0 es =+    AstRaw $ AstMapAccumLDer k bftk eftk f df rf (unAstRaw acc0) (unAstRaw es)+  tapply f t = AstRaw $ AstApply f (unAstRaw t)+  tlambda = tlambda @(AstTensor AstMethodLet s)+  -- These three methods are called at this type in delta evaluation via+  -- tmapAccumR and tmapAccumL, so they have to work. We could refrain from+  -- simplifying the resulting terms, but it's not clear that's more consistent.+  tgrad = tgrad @(AstTensor AstMethodLet s)+  tvjp = tvjp @(AstTensor AstMethodLet s)+  tjvp = tjvp @(AstTensor AstMethodLet s)++  tprimalPart t = AstRaw $ primalPart $ unAstRaw t+  tdualPart _ t = dualPart $ unAstRaw t+  tplainPart t = AstRaw $ plainPart $ unAstRaw t+  tfromPrimal _ t = AstRaw $ fromPrimal $ unAstRaw t+  tfromDual t = AstRaw $ fromDual t+  tfromPlain _ t = AstRaw $ fromPlain $ unAstRaw t++  taddTarget = addTarget+  tmultTarget = multTarget+  tsum0Target = sum0Target+  tdot0Target = dot0Target++instance KnownSpan s => ConvertTensor (AstRaw s) where+  tconvert c _astk = AstRaw . AstConvert c . unAstRaw++  -- These are somewhat faster than their default implementations.+  kfromR = AstRaw . AstConvert (ConvCmp ConvX0 ConvRX) . unAstRaw+  kfromS = AstRaw . cAstConvDownKFromS . unAstRaw+  kfromX = AstRaw . AstConvert ConvX0 . unAstRaw+  sfromK = AstRaw . cAstConvUpSFromK . unAstRaw++  rfromK = AstRaw+           . AstConvert (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))+           . unAstRaw+  rfromS (AstRaw t) = AstRaw $ case ftkAst t of+    FTKS sh x -> cAstConvUpRFromS sh x t+  rfromX (AstRaw t) = AstRaw $ case ftkAst t of+    FTKX sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        cAstConvUpRFromS sh x $ cAstConvDownSFromX sh x t+  xfromR (AstRaw t) = AstRaw $ case ftkAst t of+    FTKR shr x ->+      withShsFromShR shr $ \(sh :: ShS sh) ->+        cAstConvUpXFromS (shCastSX knownShX sh) x+        $ cAstConvDownSFromR sh x t+  sfromR (AstRaw t) = AstRaw $ case ftkAst t of+    FTKR _ x -> cAstConvDownSFromR knownShS x t+  sfromX (AstRaw t) = AstRaw $ case ftkAst t of+    FTKX _ x -> cAstConvDownSFromX knownShS x t+  xfromK = AstRaw . AstConvert (Conv0X STKScalar) . unAstRaw+  xfromS (AstRaw t) = AstRaw $ case ftkAst t of+    FTKS sh x -> cAstConvUpXFromS (shCastSX knownShX sh) x t++  rzip @_ @_ @n (AstRaw a)+   | Refl <- lemRankReplicate (Proxy @n) = AstRaw $ case ftkAst a of+    FTKProduct (FTKR _sh y) (FTKR _ z) ->+      let c = convCmp+                (ConvXR (ftkToSTK (FTKProduct y z)))+                (convCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvRX ConvRX))+      in AstConvert c a+  runzip @_ @_ @n (AstRaw a)+   | Refl <- lemRankReplicate (Proxy @n) = AstRaw $ case ftkAst a of+    FTKR _sh (FTKProduct y z) ->+      let c = convCmp+                (ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))+                (convCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvRX)+      in AstConvert c a+  szip (AstRaw a) = AstRaw $ case ftkAst a of+    FTKProduct (FTKS _sh y) (FTKS _ z) ->+      let c = convCmp+                ConvXS+                (convCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvSX ConvSX))+      in AstConvert c a+  sunzip (AstRaw a) = AstRaw $ case ftkAst a of+    FTKS _sh (FTKProduct y z) ->+      let c = convCmp+                (ConvT2 ConvXS ConvXS)+                (convCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvSX)+      in AstConvert c a+  xzip (AstRaw a) = AstRaw $ case ftkAst a of+    FTKProduct (FTKX _sh y) (FTKX _ z) ->+      let c = ConvZip (ftkToSTK y) (ftkToSTK z)+      in AstConvert c a+  xunzip (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX _sh (FTKProduct y z) ->+      let c = ConvUnzip (ftkToSTK y) (ftkToSTK z)+      in AstConvert c a++  xnestR @sh1 @m @x sh1 (AstRaw a)+    | Refl <- lemRankReplicate (Proxy @m) = AstRaw $+      let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)+                            (TKX2 sh1 (TKR2 m x))+          c = convCmp+                (ConvXX (ConvXR (knownSTK @x)))+                (ConvNest @_ @_ @(Replicate m Nothing)+                          (STKX sh1 (knownSTK @x)))+      in AstConvert c a+  xnestS @_ @_ @x sh1 (AstRaw a) = AstRaw $+    let c = convCmp (ConvXX ConvXS)+                    (ConvNest (STKX sh1 (knownSTK @x)))+    in AstConvert c a+  xnest @_ @_ @x sh1 (AstRaw a) = AstRaw $+    let c = ConvNest (STKX sh1 (knownSTK @x))+    in AstConvert c a+  xunNestR (AstRaw a) = AstRaw $+    let c = convCmp ConvUnnest+                    (ConvXX ConvRX)+    in AstConvert c a+  xunNestS (AstRaw a) = AstRaw $+    let c = convCmp ConvUnnest+                    (ConvXX ConvSX)+    in AstConvert c a+  xunNest (AstRaw a) = AstRaw $+    let c = ConvUnnest+    in AstConvert c a++  tpairConv = tpair+  tunpairConv = tunpair++-- All but the last case are shortcuts for common forms.+astConcreteRaw :: FullShapeTK y -> Concrete y -> AstRaw PlainSpan y+astConcreteRaw ftk v = case ftk of+  FTKScalar -> AstRaw $ AstConcreteK $ unConcrete v+  FTKR ZSR FTKScalar ->+    AstRaw $ AstConvert (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))+    $ AstConcreteK $ Nested.runScalar $ unConcrete v+  FTKR sh' FTKScalar -> AstRaw $+    withShsFromShR sh' $ \sh ->+      withKnownShS sh $+      cAstConvUpRFromS sh FTKScalar $ AstConcreteS $ unConcrete $ sfromR v+  FTKS ZSS FTKScalar ->+    sfromK $ AstRaw $ AstConcreteK $ Nested.sunScalar $ unConcrete v+  FTKS _ FTKScalar -> AstRaw $ AstConcreteS $ unConcrete v+  FTKX ZSX FTKScalar ->+    AstRaw $ AstConvert (Conv0X STKScalar)+    $ AstConcreteK $ Nested.munScalar $ unConcrete v+  FTKX sh' FTKScalar -> AstRaw $+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      withKnownShS sh $+      cAstConvUpXFromS sh' FTKScalar+      $ AstConcreteS $ unConcrete $ sfromX @_ @sh v+  FTKProduct ftk1 ftk2 -> AstRaw $+    AstPair (unAstRaw $ astConcreteRaw ftk1 (tproject1 v))+            (unAstRaw $ astConcreteRaw ftk2 (tproject2 v))+  _ -> concreteTarget+         (tkconcrete . unConcrete) (tsconcrete . unConcrete)+         (\sh (AstRaw t) -> AstRaw $ cAstConvUpRFromS sh FTKScalar t)+         (\sh' (AstRaw t) -> AstRaw $ cAstConvUpXFromS sh' FTKScalar t)+         (ftkToSTK ftk) v+++-- * AstNoVectorize instances++fmapAstNoVectorize+  :: Coercible (f (AstTensor AstMethodLet s y)) (f (AstNoVectorize s y))+  => f (AstTensor AstMethodLet s y) -> f (AstNoVectorize s y)+fmapAstNoVectorize = coerce++fmapUnAstNoVectorize+  :: Coercible (f (AstNoVectorize s y)) (f (AstTensor AstMethodLet s y))+  => f (AstNoVectorize s y) -> f (AstTensor AstMethodLet s y)+fmapUnAstNoVectorize = coerce++instance KnownSpan s => LetTensor (AstNoVectorize s) where+  ttlet u f = AstNoVectorize+              $ ttlet (unAstNoVectorize u)+                      (unAstNoVectorize . f . AstNoVectorize)+  ttletPrimal u f = AstNoVectorize+                    $ ttletPrimal (unAstNoVectorize u)+                                  (unAstNoVectorize . f . AstNoVectorize)+  ttletPlain u f = AstNoVectorize+                   $ ttletPlain (unAstNoVectorize u)+                                (unAstNoVectorize . f . AstNoVectorize)+  toShare t = toShare $ unAstNoVectorize t++instance KnownSpan s => BaseTensor (AstNoVectorize s) where+  isConcreteInstance = False+  -- Ranked ops+  rshape = rshape . unAstNoVectorize+  trfromVectorN shm =+    AstNoVectorize . trfromVectorN shm . fmapUnAstNoVectorize+  trfromVectorLinear shm =+    AstNoVectorize . trfromVectorLinear shm . fmapUnAstNoVectorize+  trsumN = AstNoVectorize . trsumN . unAstNoVectorize+  trsum0 = AstNoVectorize . trsum0 . unAstNoVectorize+  trreplicateN shm = AstNoVectorize . trreplicateN shm . unAstNoVectorize+  trreplicate0N shm = AstNoVectorize . trreplicate0N shm . unAstNoVectorize+  trconcrete = AstNoVectorize . trconcrete+  trfloor = AstNoVectorize . trfloor . unAstNoVectorize+  trfromIntegral = AstNoVectorize . trfromIntegral . unAstNoVectorize+  trcast = AstNoVectorize . trcast . unAstNoVectorize+  trindex v ix =+    AstNoVectorize $ trindex (unAstNoVectorize v) (fmapUnAstNoVectorize ix)+  trindex0 v ix =+    AstNoVectorize $ trindex0 (unAstNoVectorize v) (fmapUnAstNoVectorize ix)+  trscatter sh t f =+    AstNoVectorize $ trscatter sh (unAstNoVectorize t)+                   $ fmapUnAstNoVectorize . f . fmapAstNoVectorize+  trgather sh t f =+    AstNoVectorize $ trgather sh (unAstNoVectorize t)+                   $ fmapUnAstNoVectorize . f . fmapAstNoVectorize+  trargMin = AstNoVectorize . trargMin . unAstNoVectorize+  trargMax = AstNoVectorize . trargMax . unAstNoVectorize+  triota = AstNoVectorize . triota+  trappend u v =+    AstNoVectorize $ trappend (unAstNoVectorize u) (unAstNoVectorize v)+  trslice i n = AstNoVectorize . trslice i n . unAstNoVectorize+  trreverse = AstNoVectorize . trreverse . unAstNoVectorize+  trtranspose perm = AstNoVectorize . trtranspose perm . unAstNoVectorize+  trreshape sh = AstNoVectorize . trreshape sh . unAstNoVectorize+  trbuild1 k f = withSNat k $ \snat ->+    AstNoVectorize $ AstBuild1 snat knownSTK+    $ funToAstInt (0, k - 1)+        -- this introduces new variable names+    $ unAstNoVectorize . f . AstNoVectorize++  -- Shaped ops+  sshape = sshape . unAstNoVectorize+  tsfromVectorN shm =+    AstNoVectorize . tsfromVectorN shm . fmapUnAstNoVectorize+  tsfromVectorLinear shm =+    AstNoVectorize . tsfromVectorLinear shm . fmapUnAstNoVectorize+  tssumN @shm = AstNoVectorize . tssumN @_ @shm . unAstNoVectorize+  tssum0 = AstNoVectorize . tssum0 . unAstNoVectorize+  tsreplicateN shm = AstNoVectorize . tsreplicateN shm . unAstNoVectorize+  tsreplicate0N shm = AstNoVectorize . tsreplicate0N shm . unAstNoVectorize+  tsconcrete = AstNoVectorize . tsconcrete+  tsfloor = AstNoVectorize . tsfloor . unAstNoVectorize+  tsfromIntegral = AstNoVectorize . tsfromIntegral . unAstNoVectorize+  tscast = AstNoVectorize . tscast . unAstNoVectorize+  tsindex v ix =+    AstNoVectorize $ tsindex (unAstNoVectorize v) (fmapUnAstNoVectorize ix)+  tsindex0 v ix =+    AstNoVectorize $ tsindex0 (unAstNoVectorize v) (fmapUnAstNoVectorize ix)+  tsscatter @_ @shm @shn @shp t f =+    AstNoVectorize $ tsscatter @_ @_ @shm @shn @shp (unAstNoVectorize t)+                   $ fmapUnAstNoVectorize . f . fmapAstNoVectorize+  tsgather @_ @shm @shn @shp t f =+    AstNoVectorize $ tsgather @_ @_ @shm @shn @shp (unAstNoVectorize t)+                   $ fmapUnAstNoVectorize . f . fmapAstNoVectorize+  tsargMin = AstNoVectorize . tsargMin . unAstNoVectorize+  tsargMax = AstNoVectorize . tsargMax . unAstNoVectorize+  tsiota = AstNoVectorize tsiota+  tsappend u v =+    AstNoVectorize $ tsappend (unAstNoVectorize u) (unAstNoVectorize v)+  tsslice i n k = AstNoVectorize . tsslice i n k . unAstNoVectorize+  tsreverse = AstNoVectorize . tsreverse . unAstNoVectorize+  tstranspose perm =+    AstNoVectorize . tstranspose perm . unAstNoVectorize+  tsreshape sh = AstNoVectorize . tsreshape sh . unAstNoVectorize+  tsbuild1 @k f = AstNoVectorize $ AstBuild1 (SNat @k) knownSTK+                  $ funToAstInt (0, valueOf @k - 1)+                      -- this introduces new variable names+                  $ unAstNoVectorize . f . AstNoVectorize++  -- Mixed ops+  xshape = xshape . unAstNoVectorize+  txfromVectorN shm =+    AstNoVectorize . txfromVectorN shm . fmapUnAstNoVectorize+  txfromVectorLinear shm =+    AstNoVectorize . txfromVectorLinear shm . fmapUnAstNoVectorize+  txsumN @shm = AstNoVectorize . txsumN @_ @shm . unAstNoVectorize+  txsum0 = AstNoVectorize . txsum0 . unAstNoVectorize+  txreplicateN shm = AstNoVectorize . txreplicateN shm . unAstNoVectorize+  txreplicate0N shm = AstNoVectorize . txreplicate0N shm . unAstNoVectorize+  txconcrete = AstNoVectorize . txconcrete+  txfloor = AstNoVectorize . txfloor . unAstNoVectorize+  txfromIntegral = AstNoVectorize . txfromIntegral . unAstNoVectorize+  txcast = AstNoVectorize . txcast . unAstNoVectorize+  txindex v ix =+    AstNoVectorize $ txindex (unAstNoVectorize v) (fmapUnAstNoVectorize ix)+  txindex0 v ix =+    AstNoVectorize $ txindex0 (unAstNoVectorize v) (fmapUnAstNoVectorize ix)+  txscatter @_ @shm @shn @shp sh t f =+    AstNoVectorize $ txscatter @_ @_ @shm @shn @shp sh (unAstNoVectorize t)+                   $ fmapUnAstNoVectorize . f . fmapAstNoVectorize+  txgather @_ @shm @shn @shp sh t f =+    AstNoVectorize $ txgather @_ @_ @shm @shn @shp sh (unAstNoVectorize t)+                   $ fmapUnAstNoVectorize . f . fmapAstNoVectorize+  txargMin = AstNoVectorize . txargMin . unAstNoVectorize+  txargMax = AstNoVectorize . txargMax . unAstNoVectorize+  txiota @n = AstNoVectorize $ txiota @_ @n+  txappend u v =+    AstNoVectorize $ txappend (unAstNoVectorize u) (unAstNoVectorize v)+  txslice i n k = AstNoVectorize . txslice i n k . unAstNoVectorize+  txreverse = AstNoVectorize . txreverse . unAstNoVectorize+  txtranspose perm = AstNoVectorize . txtranspose perm . unAstNoVectorize+  txreshape sh = AstNoVectorize . txreshape sh . unAstNoVectorize+  txbuild1 @k f = AstNoVectorize $ AstBuild1 (SNat @k) knownSTK+                  $ funToAstInt (0, valueOf @k - 1)+                      -- this introduces new variable names+                  $ unAstNoVectorize . f . AstNoVectorize++  -- Scalar ops+  tkconcrete = AstNoVectorize . tkconcrete+  tkfloor = AstNoVectorize . tkfloor . unAstNoVectorize+  tkfromIntegral = AstNoVectorize . tkfromIntegral . unAstNoVectorize+  tkcast = AstNoVectorize . tkcast . unAstNoVectorize+  tkargMin = AstNoVectorize . tkargMin . unAstNoVectorize+  tkargMax = AstNoVectorize . tkargMax . unAstNoVectorize+  tkbuild1 @k f = AstNoVectorize $ AstBuild1 (SNat @k) STKScalar+                  $ funToAstInt (0, valueOf @k - 1)+                      -- this introduces new variable names+                  $ unAstNoVectorize . f . AstNoVectorize++  -- General operations that don't require LetTensor nor ShareTensor+  tftk stk = tftk stk . unAstNoVectorize+  tpair t1 t2 =+    AstNoVectorize $ tpair (unAstNoVectorize t1) (unAstNoVectorize t2)+  tproject1 t = AstNoVectorize $ tproject1 $ unAstNoVectorize t+  tproject2 t = AstNoVectorize $ tproject2 $ unAstNoVectorize t+  kcond !b !u !v =+    AstNoVectorize $ kcond (unAstNoVectorize b)+                           (unAstNoVectorize u) (unAstNoVectorize v)+  scond !b !u !v =+    AstNoVectorize $ scond (unAstNoVectorize b)+                           (unAstNoVectorize u) (unAstNoVectorize v)+  tcond !stk !b !u !v =+    AstNoVectorize $ tcond stk (unAstNoVectorize b)+                               (unAstNoVectorize u) (unAstNoVectorize v)+  tconcrete ftk a = AstNoVectorize $ tconcrete ftk a+  tmapAccumR _ !k !accftk !bftk !eftk f acc0 es =+    AstNoVectorize $ tmapAccumR Proxy k accftk bftk eftk f+                       (unAstNoVectorize acc0) (unAstNoVectorize es)+  tmapAccumRDer _ !k !accftk !bftk !eftk f df rf acc0 es =+    AstNoVectorize $ tmapAccumRDer Proxy k accftk bftk eftk f df rf+                       (unAstNoVectorize acc0) (unAstNoVectorize es)+  tmapAccumLDer _ !k !accftk !bftk !eftk f df rf acc0 es =+    AstNoVectorize $ tmapAccumLDer Proxy k accftk bftk eftk f df rf+                       (unAstNoVectorize acc0) (unAstNoVectorize es)+  tapply f t = AstNoVectorize $ tapply f (unAstNoVectorize t)+  tlambda = tlambda @(AstTensor AstMethodLet s)+  tgrad = tgrad @(AstTensor AstMethodLet s)+  tvjp = tvjp @(AstTensor AstMethodLet s)+  tjvp = tjvp @(AstTensor AstMethodLet s)++  tsum k stk =+    AstNoVectorize . tsum k stk . unAstNoVectorize+  treplicate k stk =+    AstNoVectorize . treplicate k stk . unAstNoVectorize+  treverse k stk =+    AstNoVectorize . treverse k stk . unAstNoVectorize++  tprimalPart t = AstNoVectorize $ tprimalPart $ unAstNoVectorize t+  tdualPart stk t = tdualPart stk $ unAstNoVectorize t+  tplainPart t = AstNoVectorize $ tplainPart $ unAstNoVectorize t+  tfromPrimal stk t = AstNoVectorize $ tfromPrimal stk $ unAstNoVectorize t+  tfromDual t = AstNoVectorize $ tfromDual t+  tfromPlain stk t = AstNoVectorize $ tfromPlain stk $ unAstNoVectorize t++  taddTarget stk a b = AstNoVectorize $ taddTarget stk (unAstNoVectorize a)+                                                       (unAstNoVectorize b)+  tmultTarget stk a b = AstNoVectorize $ tmultTarget stk (unAstNoVectorize a)+                                                         (unAstNoVectorize b)+  tsum0Target stk a = AstNoVectorize $ tsum0Target stk (unAstNoVectorize a)+  tdot0Target stk a b = AstNoVectorize $ tdot0Target stk (unAstNoVectorize a)+                                                         (unAstNoVectorize b)++instance KnownSpan s => ConvertTensor (AstNoVectorize s) where+  tconvert c _astk = AstNoVectorize . astConvert c . unAstNoVectorize++  rfromX = AstNoVectorize . rfromX . unAstNoVectorize+  xfromR = AstNoVectorize . xfromR . unAstNoVectorize++  sfromR = AstNoVectorize . sfromR . unAstNoVectorize+  sfromX = AstNoVectorize . sfromX . unAstNoVectorize+  xfromS = AstNoVectorize . xfromS . unAstNoVectorize++  rzip = AstNoVectorize . rzip . unAstNoVectorize+  runzip = AstNoVectorize . runzip . unAstNoVectorize+  szip = AstNoVectorize . szip . unAstNoVectorize+  sunzip = AstNoVectorize . sunzip . unAstNoVectorize+  xzip = AstNoVectorize . xzip . unAstNoVectorize+  xunzip = AstNoVectorize . xunzip . unAstNoVectorize++  xnestR sh = AstNoVectorize . xnestR sh . unAstNoVectorize+  xnestS sh = AstNoVectorize . xnestS sh . unAstNoVectorize+  xnest sh = AstNoVectorize . xnest sh . unAstNoVectorize+  xunNestR = AstNoVectorize . xunNestR . unAstNoVectorize+  xunNestS = AstNoVectorize . xunNestS . unAstNoVectorize+  xunNest = AstNoVectorize . xunNest . unAstNoVectorize++  tpairConv = tpair+  tunpairConv a = let (b, c) = tunpairConv $ unAstNoVectorize a+                  in (AstNoVectorize b, AstNoVectorize c)+++-- * AstNoSimplify instances++instance KnownSpan s => LetTensor (AstNoSimplify s) where+  ttlet u f = AstNoSimplify+              $ astLetFunNoSimplify (unAstNoSimplify u)+                                    (unAstNoSimplify . f . AstNoSimplify)+  ttletPrimal u f = AstNoSimplify+                    $ astLetFunNoSimplify (unAstNoSimplify u)+                                          (unAstNoSimplify . f . AstNoSimplify)+  ttletPlain u f = AstNoSimplify+                   $ astLetFunNoSimplify (unAstNoSimplify u)+                                         (unAstNoSimplify . f . AstNoSimplify)+  toShare t = AstRaw $ AstToShare $ unAstNoSimplify t++-- INLINE here would bloat the binary a lot, probably negating any+-- gains from directly calling the function. Also, this is not a bottleneck.+astLetFunNoSimplify+  :: forall y z s s2. KnownSpan s+  => AstTensor AstMethodLet s y+  -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)+  -> AstTensor AstMethodLet s2 z+{-# NOINLINE astLetFunNoSimplify #-}+astLetFunNoSimplify a f | astIsSmall True a = f a+                            -- too important an optimization to skip+astLetFunNoSimplify a f = case a of+  AstFromPrimal v -> astLetFunNoSimplify v (f . fromPrimal)+  AstFromDual v -> astLetFunNoSimplify v (f . fromDual)+  AstFromPlain v -> astLetFunNoSimplify v (f . fromPlain)+  _ -> unsafePerformIO $ case ftkAst a of+    ftk@FTKScalar -> do+        var <- funToAstAutoBoundsIO ftk a+        pure $! AstLet var a (f $ astVar var)+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+      case (lemRankReplicate (Proxy @(Rank sh)), lemRankMapJust sh) of+        (Refl, Refl) -> do+          let v = AstConvert (ConvCmp (ConvXS' (FTKS sh x)) ConvRX) a+          var <- funToAstNoBoundsIO (FTKS sh x)+          pure $! AstLet var v+                         (f $ AstConvert (ConvCmp (ConvXR (ftkToSTK x)) ConvSX)+                          $ astVar var)+            -- safe, because subsitution ruled out above+    FTKX sh' x ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+      case lemRankMapJust sh of+        Refl -> do+          let v = AstConvert (ConvXS' (FTKS sh x)) a+          var <- funToAstNoBoundsIO (FTKS sh x)+          pure $! AstLet var v+                         (f $ AstConvert (ConvCmp (ConvXX' (FTKX sh' x)) ConvSX)+                          $ astVar var)+    FTKS ZSS x@FTKScalar -> do+        let v = AstConvert (ConvCmp ConvX0 ConvSX) a+        var <- funToAstAutoBoundsIO x v+        pure $! AstLet var v+                       (f $ AstConvert (ConvCmp ConvXS (Conv0X STKScalar))+                        $ astVar var)+    -- calling recursively for product may be not worth it+    ftk -> do+        var <- funToAstNoBoundsIO ftk+        pure $! AstLet var a (f $ astVar var)++wAstNoSimplify :: AstRaw s y -> AstNoSimplify s y+wAstNoSimplify =+  AstNoSimplify+  . (unsafeCoerce :: AstTensor AstMethodShare s y+                  -> AstTensor AstMethodLet s y)+  . unAstRaw++wunAstNoSimplify :: AstNoSimplify s y -> AstRaw s y+wunAstNoSimplify =+  AstRaw+  . (unsafeCoerce :: AstTensor AstMethodLet s y+                  -> AstTensor AstMethodShare s y)+  . unAstNoSimplify++fmapwAstNoSimplify+  :: f (AstRaw s y) -> f (AstNoSimplify s y)+fmapwAstNoSimplify = unsafeCoerce++fmapwUnAstNoSimplify+  :: f (AstNoSimplify s y) -> f (AstRaw s y)+fmapwUnAstNoSimplify = unsafeCoerce++instance KnownSpan s => BaseTensor (AstNoSimplify s) where+  isConcreteInstance = False+  -- The implementation of these methods differs from the AstRaw instance:+  tkbuild1 @k f =+    AstNoSimplify+    $ astBuild1Vectorize (SNat @k) STKScalar+                         (unAstNoSimplify . f . AstNoSimplify)+  trbuild1 @n @x k f = withSNat k $ \snat ->+    AstNoSimplify+    $ astBuild1Vectorize snat (STKR (SNat @n) (knownSTK @x))+                         (unAstNoSimplify . f . AstNoSimplify)+  tsbuild1 @k @sh @x f =+    AstNoSimplify+    $ astBuild1Vectorize (SNat @k) (STKS (knownShS @sh) (knownSTK @x))+                         (unAstNoSimplify . f . AstNoSimplify)+  txbuild1 @k @sh @x f =+    AstNoSimplify+    $ astBuild1Vectorize (SNat @k) (STKX (knownShX @sh) (knownSTK @x))+                         (unAstNoSimplify . f . AstNoSimplify)+  tmapAccumR _ !k !accftk !bftk !eftk f acc0 es =+    AstNoSimplify $ tmapAccumR Proxy k accftk bftk eftk f+                      (unAstNoSimplify acc0) (unAstNoSimplify es)+  tmapAccumRDer _ !k !accftk !bftk !eftk f df rf acc0 es =+    AstNoSimplify $ tmapAccumRDer Proxy k accftk bftk eftk f df rf+                      (unAstNoSimplify acc0) (unAstNoSimplify es)+  -- These two have tricky types, so we repaat the AstRaw definitions:+  tdualPart _ t = dualPart $ unAstNoSimplify t+  tfromDual t = AstNoSimplify $ fromDual t++  -- All the following implementations piggy-back on AstRaw implementations.+  -- Ranked ops+  rshape = rshape . wunAstNoSimplify+  trfromVectorN shm =+    wAstNoSimplify . trfromVectorN shm . fmapwUnAstNoSimplify+  trfromVectorLinear shm =+    wAstNoSimplify . trfromVectorLinear shm . fmapwUnAstNoSimplify+  trsumN = wAstNoSimplify . trsumN . wunAstNoSimplify+  trsum0 = wAstNoSimplify . trsum0 . wunAstNoSimplify+  trreplicateN shm = wAstNoSimplify . trreplicateN shm . wunAstNoSimplify+  trreplicate0N shm = wAstNoSimplify . trreplicate0N shm . wunAstNoSimplify+  trconcrete = wAstNoSimplify . trconcrete+  trfloor = wAstNoSimplify . trfloor . wunAstNoSimplify+  trfromIntegral = wAstNoSimplify . trfromIntegral . wunAstNoSimplify+  trcast = wAstNoSimplify . trcast . wunAstNoSimplify+  trindex v ix =+    wAstNoSimplify $ trindex (wunAstNoSimplify v) (fmapwUnAstNoSimplify ix)+  trindex0 v ix =+    wAstNoSimplify $ trindex0 (wunAstNoSimplify v) (fmapwUnAstNoSimplify ix)+  trscatter sh t f =+    wAstNoSimplify $ trscatter sh (wunAstNoSimplify t)+                   $ fmapwUnAstNoSimplify . f . fmapwAstNoSimplify+  trgather sh t f =+    wAstNoSimplify $ trgather sh (wunAstNoSimplify t)+                   $ fmapwUnAstNoSimplify . f . fmapwAstNoSimplify+  trargMin = wAstNoSimplify . trargMin . wunAstNoSimplify+  trargMax = wAstNoSimplify . trargMax . wunAstNoSimplify+  triota = wAstNoSimplify . triota+  trappend u v =+    wAstNoSimplify $ trappend (wunAstNoSimplify u) (wunAstNoSimplify v)+  trslice i n = wAstNoSimplify . trslice i n . wunAstNoSimplify+  trreverse = wAstNoSimplify . trreverse . wunAstNoSimplify+  trtranspose perm = wAstNoSimplify . trtranspose perm . wunAstNoSimplify+  trreshape sh = wAstNoSimplify . trreshape sh . wunAstNoSimplify++  -- Shaped ops+  sshape = sshape . wunAstNoSimplify+  tsfromVectorN shm =+    wAstNoSimplify . tsfromVectorN shm . fmapwUnAstNoSimplify+  tsfromVectorLinear shm =+    wAstNoSimplify . tsfromVectorLinear shm . fmapwUnAstNoSimplify+  tssumN @shm = wAstNoSimplify . tssumN @_ @shm . wunAstNoSimplify+  tssum0 = wAstNoSimplify . tssum0 . wunAstNoSimplify+  tsreplicateN shm = wAstNoSimplify . tsreplicateN shm . wunAstNoSimplify+  tsreplicate0N shm = wAstNoSimplify . tsreplicate0N shm . wunAstNoSimplify+  tsconcrete = wAstNoSimplify . tsconcrete+  tsfloor = wAstNoSimplify . tsfloor . wunAstNoSimplify+  tsfromIntegral = wAstNoSimplify . tsfromIntegral . wunAstNoSimplify+  tscast = wAstNoSimplify . tscast . wunAstNoSimplify+  tsindex v ix =+    wAstNoSimplify $ tsindex (wunAstNoSimplify v) (fmapwUnAstNoSimplify ix)+  tsindex0 v ix =+    wAstNoSimplify $ tsindex0 (wunAstNoSimplify v) (fmapwUnAstNoSimplify ix)+  tsscatter @_ @shm @shn @shp t f =+    wAstNoSimplify $ tsscatter @_ @_ @shm @shn @shp (wunAstNoSimplify t)+                   $ fmapwUnAstNoSimplify . f . fmapwAstNoSimplify+  tsgather @_ @shm @shn @shp t f =+    wAstNoSimplify $ tsgather @_ @_ @shm @shn @shp (wunAstNoSimplify t)+                   $ fmapwUnAstNoSimplify . f . fmapwAstNoSimplify+  tsargMin = wAstNoSimplify . tsargMin . wunAstNoSimplify+  tsargMax = wAstNoSimplify . tsargMax . wunAstNoSimplify+  tsiota = wAstNoSimplify tsiota+  tsappend u v =+    wAstNoSimplify $ tsappend (wunAstNoSimplify u) (wunAstNoSimplify v)+  tsslice i n k = wAstNoSimplify . tsslice i n k . wunAstNoSimplify+  tsreverse = wAstNoSimplify . tsreverse . wunAstNoSimplify+  tstranspose perm =+    wAstNoSimplify . tstranspose perm . wunAstNoSimplify+  tsreshape sh = wAstNoSimplify . tsreshape sh . wunAstNoSimplify++  -- Mixed ops+  xshape = xshape . wunAstNoSimplify+  txfromVectorN shm =+    wAstNoSimplify . txfromVectorN shm . fmapwUnAstNoSimplify+  txfromVectorLinear shm =+    wAstNoSimplify . txfromVectorLinear shm . fmapwUnAstNoSimplify+  txsumN @shm = wAstNoSimplify . txsumN @_ @shm . wunAstNoSimplify+  txsum0 = wAstNoSimplify . txsum0 . wunAstNoSimplify+  txreplicateN shm = wAstNoSimplify . txreplicateN shm . wunAstNoSimplify+  txreplicate0N shm = wAstNoSimplify . txreplicate0N shm . wunAstNoSimplify+  txconcrete = wAstNoSimplify . txconcrete+  txfloor = wAstNoSimplify . txfloor . wunAstNoSimplify+  txfromIntegral = wAstNoSimplify . txfromIntegral . wunAstNoSimplify+  txcast = wAstNoSimplify . txcast . wunAstNoSimplify+  txindex v ix =+    wAstNoSimplify $ txindex (wunAstNoSimplify v) (fmapwUnAstNoSimplify ix)+  txindex0 v ix =+    wAstNoSimplify $ txindex0 (wunAstNoSimplify v) (fmapwUnAstNoSimplify ix)+  txscatter @_ @shm @shn @shp sh t f =+    wAstNoSimplify $ txscatter @_ @_ @shm @shn @shp sh (wunAstNoSimplify t)+                   $ fmapwUnAstNoSimplify . f . fmapwAstNoSimplify+  txgather @_ @shm @shn @shp sh t f =+    wAstNoSimplify $ txgather @_ @_ @shm @shn @shp sh (wunAstNoSimplify t)+                   $ fmapwUnAstNoSimplify . f . fmapwAstNoSimplify+  txargMin = wAstNoSimplify . txargMin . wunAstNoSimplify+  txargMax = wAstNoSimplify . txargMax . wunAstNoSimplify+  txiota @n = wAstNoSimplify $ txiota @_ @n+  txappend u v =+    wAstNoSimplify $ txappend (wunAstNoSimplify u) (wunAstNoSimplify v)+  txslice i n k = wAstNoSimplify . txslice i n k . wunAstNoSimplify+  txreverse = wAstNoSimplify . txreverse . wunAstNoSimplify+  txtranspose perm = wAstNoSimplify . txtranspose perm . wunAstNoSimplify+  txreshape sh = wAstNoSimplify . txreshape sh . wunAstNoSimplify++  -- Scalar ops+  tkconcrete = wAstNoSimplify . tkconcrete+  tkfloor = wAstNoSimplify . tkfloor . wunAstNoSimplify+  tkfromIntegral = wAstNoSimplify . tkfromIntegral . wunAstNoSimplify+  tkcast = wAstNoSimplify . tkcast . wunAstNoSimplify+  tkargMin = wAstNoSimplify . tkargMin . wunAstNoSimplify+  tkargMax = wAstNoSimplify . tkargMax . wunAstNoSimplify++  -- General operations that don't require LetTensor nor ShareTensor+  tftk stk = tftk stk . wunAstNoSimplify+  tpair t1 t2 =+    wAstNoSimplify $ tpair (wunAstNoSimplify t1) (wunAstNoSimplify t2)+  tproject1 t = wAstNoSimplify $ tproject1 $ wunAstNoSimplify t+  tproject2 t = wAstNoSimplify $ tproject2 $ wunAstNoSimplify t+  kcond !b !u !v =+    wAstNoSimplify $ kcond (wunAstNoSimplify b)+                           (wunAstNoSimplify u) (wunAstNoSimplify v)+  scond !b !u !v =+    wAstNoSimplify $ scond (wunAstNoSimplify b)+                           (wunAstNoSimplify u) (wunAstNoSimplify v)+  tcond stk !b !u !v =+    wAstNoSimplify $ tcond stk (wunAstNoSimplify b)+                           (wunAstNoSimplify u) (wunAstNoSimplify v)+  tconcrete ftk a = wAstNoSimplify $ tconcrete ftk a+  tmapAccumLDer _ !k !accftk !bftk !eftk f df rf acc0 es =+    wAstNoSimplify $ tmapAccumLDer Proxy k accftk bftk eftk f df rf+                       (wunAstNoSimplify acc0) (wunAstNoSimplify es)+  tapply f t = wAstNoSimplify $ tapply f (wunAstNoSimplify t)+  tlambda = tlambda @(AstRaw s)+  tgrad = tgrad @(AstRaw s)+  tvjp = tvjp @(AstRaw s)+  tjvp = tjvp @(AstRaw s)++  tsum snat@SNat stk u = case stk of+    STKScalar -> kfromS $ tssum u+    STKR SNat x | Dict <- lemKnownSTK x -> trsum u+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (tsum snat stk1 (tproject1 u3))+              (tsum snat stk2 (tproject2 u3))+  treplicate snat@SNat stk u = case stk of+    STKScalar -> tsreplicate snat $ sfromK u+    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (fromSNat' snat) u+    STKS _ x | Dict <- lemKnownSTK x -> tsreplicate snat u+    STKX _ x | Dict <- lemKnownSTK x -> txreplicate snat u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (treplicate snat stk1 (tproject1 u3))+              (treplicate snat stk2 (tproject2 u3))+  treverse snat stk u = case stk of+    STKScalar -> tsreverse u+    STKR _ x | Dict <- lemKnownSTK x -> trreverse u+    STKS _ x | Dict <- lemKnownSTK x -> tsreverse u+    STKX _ x | Dict <- lemKnownSTK x -> txreverse u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (treverse snat stk1 (tproject1 u3))+              (treverse snat stk2 (tproject2 u3))++  tprimalPart t = wAstNoSimplify $ tprimalPart $ wunAstNoSimplify t+  tplainPart t = wAstNoSimplify $ tplainPart $ wunAstNoSimplify t+  tfromPrimal stk t = wAstNoSimplify $ tfromPrimal stk $ wunAstNoSimplify t+  tfromPlain stk t = wAstNoSimplify $ tfromPlain stk $ wunAstNoSimplify t++  taddTarget stk a b = wAstNoSimplify $ taddTarget stk (wunAstNoSimplify a)+                                                       (wunAstNoSimplify b)+  tmultTarget stk a b = wAstNoSimplify $ tmultTarget stk (wunAstNoSimplify a)+                                                         (wunAstNoSimplify b)+  tsum0Target stk a = wAstNoSimplify $ tsum0Target stk (wunAstNoSimplify a)+  tdot0Target stk a b = wAstNoSimplify $ tdot0Target stk (wunAstNoSimplify a)+                                                         (wunAstNoSimplify b)++instance KnownSpan s => ConvertTensor (AstNoSimplify s) where   tconvert c astk = wAstNoSimplify . tconvert c astk . wunAstNoSimplify    rfromX = wAstNoSimplify . rfromX . wunAstNoSimplify
src/HordeAd/Core/OpsConcrete.hs view
@@ -1,1315 +1,1972 @@-{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}-{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}-{-# OPTIONS_GHC -Wno-orphans #-}--- | Tensor class instances for concrete arrays backed--- by 'Data.Vector.Storable.Vector'--- and defined in @ox-arrays@ on the basis of @orthotope@.-module HordeAd.Core.OpsConcrete-  () where--import Prelude hiding (foldl')--import Control.Arrow (second)-import Control.Exception.Assert.Sugar-import Data.Coerce (Coercible, coerce)-import Data.Foldable qualified as Foldable-import Data.Function ((&))-import Data.Functor.WithIndex (imap)-import Data.Int (Int64)-import Data.List (foldl', mapAccumL, mapAccumR)-import Data.List.NonEmpty qualified as NonEmpty-import Data.Map.Strict qualified as M-import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))-import Data.Vector.Generic qualified as V-import Data.Vector.Storable qualified as VS-import Data.Vector.Strict qualified as Data.Vector-import GHC.Exts (IsList (..))-import GHC.TypeLits (KnownNat, sameNat, type (+))--import Data.Array.Nested (MapJust, Replicate, type (++))-import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Lemmas-import Data.Array.Nested.Mixed qualified as Mixed-import Data.Array.Nested.Mixed.Shape-import Data.Array.Nested.Permutation qualified as Permutation-import Data.Array.Nested.Ranked qualified as Ranked-import Data.Array.Nested.Ranked.Shape-import Data.Array.Nested.Shaped qualified as Shaped-import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (Init, unsafeCoerceRefl)-import Data.Array.Strided.Orthotope (liftVEltwise1)--import HordeAd.Core.CarriersConcrete-import HordeAd.Core.ConvertTensor-import HordeAd.Core.Ops-import HordeAd.Core.OpsADVal-import HordeAd.Core.TensorKind-import HordeAd.Core.Types-import HordeAd.Core.Unwind---- * Tensor classes instance--instance LetTensor Concrete where-  ttlet = (&)-  ttletPrimal = (&)-  toShare = id-  tunshare = id-  tD _stk t DummyDualTarget{} = t-  tfold k _ stk f x0 as = foldl' f x0 (tunravelToListShare k stk as)-  tscan k@(SNat @k) nstk stk f x0 as =-    tfromVector (SNat @(1 + k)) nstk-    $ V.scanl' f x0 (V.fromList $ tunravelToListShare k stk as)--instance ShareTensor Concrete where-  tshare = id-  tunpair (Concrete (t1, t2)) = (Concrete t1, Concrete t2)--instance BaseTensor Concrete where-  -- Ranked ops-  rshape @_ @r | Dict <- eltDictRep (knownSTK @r) = Nested.rshape . unConcrete-  trfromVector @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.rfromListOuter . NonEmpty.fromList . V.toList-    . fmapUnConcrete-  trfromVector0N @_ @r sh | Dict <- eltDictRep (knownSTK @r) =-    Concrete . tfromVector0NR sh . fmapUnConcrete-  trunravelToList @_ @r | Dict <- eltDictRep (knownSTK @r) =-    fmapConcrete . Nested.rtoListOuter . unConcrete-  trsum t = case tftk knownSTK t of-    FTKR _ FTKScalar ->  -- optimized-      Concrete . Nested.rsumOuter1 . unConcrete $ t-    FTKR _ x ->-      let l = trunravelToList t-          sh = shrTail $ rshape t-      in foldr (taddTarget knownSTK) (tdefTarget (FTKR sh x)) l-        -- Concrete has a ShareTensor instance, so taddTarget arguments-        -- don't need to be duplicable-  trsum0 @_ @r t = case knownSTK @r of-    STKScalar ->  -- optimized-      Concrete . Nested.rscalar . Nested.rsumAllPrim . unConcrete $ t-    _ -> trsum . rflatten $ t-  {-# INLINE trdot0 #-}-  trdot0 u v =-    Concrete $ Nested.rscalar $ Nested.rdot (unConcrete u) (unConcrete v)-  trdot1In u v = Concrete $ Nested.rdot1Inner (unConcrete u) (unConcrete v)-  {-# INLINE trmatvecmul #-}-  trmatvecmul m v = trdot1In m (trreplicate (rwidth m) v)-  trmatmul2 m1 m2 = case rshape m2 of-    _ :$: width2 :$: ZSR ->-      trdot1In (trtranspose [1, 0] (trreplicate width2 m1))-               (trtranspose [0, 2, 1] (trreplicate (rwidth m1) m2))-  trreplicate @_ @r k | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.rreplicate (k :$: ZSR) . unConcrete-  trreplicate0N @_ @r sh | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.rreplicate sh . unConcrete-  trindex = tindexZR-  trindex0 = tindex0R-  trscatter = tscatterZR-  trscatter1 = tscatterZ1R-  trgather = tgatherZR-  trgather1 = tgatherZ1R-  trconcrete = Concrete-  trfloor = Concrete . liftVR (V.map floor) . unConcrete-  trfromIntegral = Concrete . liftVR (V.map fromIntegral) . unConcrete-  {-# INLINE trcast #-}-  trcast = Concrete . liftVR (V.map realToFrac) . unConcrete-  trminIndex = Concrete . tminIndexR . unConcrete-  trmaxIndex = Concrete . tmaxIndexR . unConcrete-  triota n = Concrete $ Nested.rfromList1 $ NonEmpty.map fromInteger-             $ NonEmpty.fromList [0 .. fromIntegral n - 1]-  trappend @_ @r u v | Dict <- eltDictRep (knownSTK @r) =-    Concrete $ Nested.rappend (unConcrete u) (unConcrete v)-  trslice @_ @r i n | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.rslice i n . unConcrete-  trreverse @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.rrev1 . unConcrete-  trtranspose @_ @r perm | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.rtranspose perm . unConcrete-  trreshape @_ @_ @r sh | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.rreshape sh . unConcrete-  trbuild1 @_ @r k f | Dict <- eltDictRep (knownSTK @r) =-    Concrete $ tbuild1R k (unConcrete . f . Concrete)-  trmap0N @_ @r @r1 f t = case (knownSTK @r1, knownSTK @r) of-    (STKScalar, STKScalar) ->-      Concrete $ tmap0NR (unConcrete . f . Concrete) (unConcrete t)-    _ ->  -- this is the default implementation from the class-      rbuild (rshape t) (f . trindex t)-  trzipWith0N @_ @r1 @r2 @r f t u =-    case (knownSTK @r1, knownSTK @r2, knownSTK @r) of-      (STKScalar, STKScalar, STKScalar) ->-        Concrete-        $ tzipWith0NR (\v w -> unConcrete $ f (Concrete v) (Concrete w))-                      (unConcrete t) (unConcrete u)-      _ ->  -- this is the default implementation from the class-        rbuild (rshape u) (\ix -> f (trindex t ix) (trindex u ix))--  -- Shaped ops-  sshape @_ @r | Dict <- eltDictRep (knownSTK @r) = Nested.sshape . unConcrete-  tsfromVector @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.sfromListOuter SNat . NonEmpty.fromList . V.toList-    . fmapUnConcrete-  tsfromVector0N @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . tfromVector0NS . fmapUnConcrete-  tsunravelToList @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =-    fmapConcrete . Nested.stoListOuter . unConcrete-  tssum t = case tftk knownSTK t of-    FTKS _ FTKScalar ->  -- optimized-      Concrete . Nested.ssumOuter1 . unConcrete $ t-    FTKS _ x ->-      let l = tsunravelToList t-          sh = shsTail $ sshape t-      in foldr (taddTarget knownSTK) (tdefTarget (FTKS sh x)) l-  tssum0 @sh @r t | SNat <- shsProduct (knownShS @sh) = case knownSTK @r of-    STKScalar ->  -- optimized-      Concrete . Nested.sscalar . Nested.ssumAllPrim . unConcrete $ t-    _ -> tssum . sflatten $ t-  {-# INLINE tsdot0 #-}  -- this doesn't want to specialize-  tsdot0 u v  =-    Concrete $ Nested.sscalar $ Nested.sdot (unConcrete u) (unConcrete v)-  tsdot1In @_ (SNat @n) u v =-    Concrete $ Nested.sdot1Inner (Proxy @n) (unConcrete u) (unConcrete v)-  {-# INLINE tsmatvecmul #-}  -- this doesn't want to specialize-  tsmatvecmul m v = tsdot1In SNat m (tsreplicate SNat knownShS v)-  tsmatmul2 m1 m2 =-    tsdot1In SNat-             (tstranspose (Permutation.makePerm @'[1, 0])-                          (tsreplicate SNat knownShS m1))-             (tstranspose (Permutation.makePerm @'[0, 2, 1])-                          (tsreplicate SNat knownShS m2))-  tsindex = tindexZS-  tsindex0 = tindex0S-  -- Performance depends a lot on the number and size of tensors.-  -- If tensors are not tiny, memory taken by underlying vectors matters most-  -- and this implementation is probbaly optimal in this respect-  -- (the only new vectors are created by V.concat, but this is done on demand).-  -- TODO: optimize updateNS and make it consume and forget arguments-  -- one by one to make the above true-  ---  -- Note how ix being in bounds is checked. The semantics of the operation-  -- permits index out of bounds and then no tensors is added at such an index.-  tsscatter @shm @shn @shp t f =-    let shpshn = knownShS @shp `shsAppend` knownShS @shn-    in withKnownShS (knownShS @shm `shsAppend` knownShS @shn) $-       case tftk knownSTK t of-         FTKS _ x@FTKScalar ->  -- optimized-           gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn)-                                          :~: shp) $-           gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn)-                                          :~: shn) $-           let zero = tdefTarget (FTKS shpshn x)-               shm = knownShS @shm-               s = shsSize shm-               g ix =-                 let ix2 = f $ fmapConcrete ix-                 in if ixInBounds (fmapUnConcrete $ toList ix2)-                                  (shsToList shpshn)-                    then M.insertWith (V.zipWith (+)) ix2-                           (Nested.stoVector-                            $ tindexNS @_ @shm @shn (unConcrete t) ix)-                    else id-               ivs = foldr g M.empty [ fromLinearIdxS fromIntegral shm-                                       $ fromIntegral i-                                     | i <- [0 .. s - 1] ]-           in withKnownShS shpshn $-              updateNS @(Rank shp) zero-              $ map (second $ Concrete . Nested.sfromVector (knownShS @shn))-              $ M.assocs ivs-         FTKS _ x | Dict <- eltDictRep (ftkToSTK x) ->-           gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn)-                                       :~: shp) $-           gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn)-                                       :~: shn) $-           let zero = tdefTarget (FTKS shpshn x)-               shm = knownShS @shm-               s = shsSize shm-               g ix =-                 let ix2 = f $ fmapConcrete ix-                 in if ixInBounds (fmapUnConcrete $ toList ix2)-                                  (shsToList shpshn)-                    then M.insertWith (taddTarget knownSTK) ix2-                           (Concrete-                            $ tindexNS @_ @shm @shn (unConcrete t) ix)-                    else id-               ivs = foldr g M.empty [ fromLinearIdxS fromIntegral shm-                                       $ fromIntegral i-                                     | i <- [0 .. s - 1] ]-           in withKnownShS shpshn $-              updateNS @(Rank shp) zero-              $ M.assocs ivs-  tsscatter1 = tscatterZ1S-  -- The semantics of the operation permits index out of bounds-  -- and the result of such indexing is def, which is 0.-  -- TODO: are bounds checked in the optimized case?-  -- The same question also elsewhere.-  tsgather @shm @shn @_ @r t f =-    gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $-    gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $-    case knownSTK @r of-      STKScalar ->  -- optimized-        let shm = knownShS @shm-            s = shsSize shm-            l = [ stoVector-                  $ tsindex @_ @_ @shn-                      t (f (fmapConcrete-                            $ fromLinearIdxS fromIntegral shm i))-                | i <- [0 .. fromIntegral s - 1] ]-        in Concrete-           $ Nested.sfromVector (knownShS @shm `shsAppend` knownShS @shn)-           $ V.concat l-      _ ->-        withKnownShS (knownShS @shm `shsAppend` knownShS @shn) $-        sbuild @(Rank shm) (\ix -> t `tsindex` f ix)-  tsgather1 = tgatherZ1S-  tsconcrete = Concrete-  tsfloor = Concrete . liftVS (V.map floor) . unConcrete-  tsfromIntegral = Concrete . tfromIntegralS . unConcrete-  {-# INLINE tscast #-}  -- this doesn't want to specialize-  tscast = Concrete . liftVS (V.map realToFrac) . unConcrete-  tsminIndex = Concrete . tminIndexS . unConcrete-  tsmaxIndex = Concrete . tmaxIndexS . unConcrete-  tsiota @n = case NonEmpty.nonEmpty [0 .. valueOf @n - 1] of-    Nothing -> case sameNat (Proxy @n) (Proxy @0) of-      Just Refl -> Concrete $ Nested.semptyArray ZSS-      Nothing -> error "siota: wrong rank"-    Just l -> Concrete $ Nested.sfromList1 SNat $ NonEmpty.map fromInteger l-  tsappend @_ @_ @_ @r u v | Dict <- eltDictRep (knownSTK @r) =-    Concrete $ Nested.sappend (unConcrete u) (unConcrete v)-  tsslice @_ @_ @_ @_ @r i n _ | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.sslice i n . unConcrete-  tsreverse @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.srev1 . unConcrete-  tsbuild1 @_ @_ @r f | Dict <- eltDictRep (knownSTK @r) =-    Concrete $ tbuild1S (unConcrete . f . Concrete)-  tsmap0N @sh @r @r1 f v = case (knownSTK @r1, knownSTK @r) of-    (STKScalar, STKScalar) ->-      Concrete $ tmap0NS (unConcrete . f . Concrete) (unConcrete v)-    _ | Refl <- lemAppNil @sh ->-      -- this is the default implementation from the class-      gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])-      $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)-      $ sbuild @(Rank sh) (f . tsindex v)-  tszipWith0N @sh @r1 @r2 @r f t u =-    case (knownSTK @r1, knownSTK @r2, knownSTK @r) of-      (STKScalar, STKScalar, STKScalar) ->-        Concrete-        $ tzipWith0NS (\v w -> unConcrete $ f (Concrete v) (Concrete w))-                      (unConcrete t) (unConcrete u)-      _ | Refl <- lemAppNil @sh ->-        -- this is the default implementation from the class-        gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])-        $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)-        $ sbuild @(Rank sh) (\ix -> f (tsindex t ix) (tsindex u ix))--  -- Mixed ops-  xshape @_ @r | Dict <- eltDictRep (knownSTK @r) = Nested.mshape . unConcrete-  txfromVector @n @sh @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mcast (Nested.SKnown (SNat @n) :!% knownShX @sh)-    . Nested.mfromListOuter . NonEmpty.fromList . V.toList-    . fmapUnConcrete-  txfromVector0N @_ @r sh | Dict <- eltDictRep (knownSTK @r) =-    Concrete . tfromVector0NX sh . fmapUnConcrete-  txunravelToList @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =-    fmapConcrete . Nested.mtoListOuter . unConcrete-  txsum t = case tftk knownSTK t of-    FTKX _ FTKScalar ->  -- optimized-      Concrete . Nested.msumOuter1 . unConcrete $ t-    FTKX _ x ->-      let l = txunravelToList t-          sh = shxTail $ xshape t-      in foldr (taddTarget knownSTK) (tdefTarget (FTKX sh x)) l-  txsum0 @_ @r t =-   case knownSTK @r of-    STKScalar ->  -- optimized-      Concrete . Nested.mscalar . Nested.msumAllPrim . unConcrete $ t-    _ -> withSNat (shxSize $ xshape t) $ \snat ->-      txsum (xmcast (Nested.SKnown snat :!% ZKX) $ xflatten t)-  {-# INLINE txdot0 #-}-  txdot0 u v =-    Concrete $ Nested.mscalar $ Nested.mdot (unConcrete u) (unConcrete v)-  txdot1In @_ (SNat @n) u v =-    Concrete $ Nested.mdot1Inner (Proxy @(Just n)) (unConcrete u) (unConcrete v)-  txmatvecmul mm mn m v =-    withKnownShX (ssxFromShX $ mn :$% ZSX) $-    withKnownShX (ssxFromShX $ mm :$% mn :$% ZSX) $-    withSNat (fromSMayNat' mm) $ \(SNat @m) ->-    withSNat (fromSMayNat' mn) $ \(SNat @n) ->-      xmcast (ssxFromShX (mm :$% ZSX))-      $ txsum (xtr (txreplicate (SNat @m) knownShX-                      (xmcast (ssxFromShX (Nested.SKnown (SNat @n)-                                             :$% ZSX)) v)-                    * xmcast (ssxFromShX (Nested.SKnown (SNat @m)-                                            :$% Nested.SKnown (SNat @n)-                                            :$% ZSX)) m))-  {-# INLINE txmatvecmul #-}-  txmatmul2 m1 m2 =-    txdot1In SNat-             (txtranspose (Permutation.makePerm @'[1, 0])-                          (txreplicate SNat knownShX m1))-             (txtranspose (Permutation.makePerm @'[0, 2, 1])-                          (txreplicate SNat knownShX m2))-  txreplicate @_ @_ @r snat _sh | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mreplicate (Nested.SKnown snat :$% ZSX) . unConcrete-  txreplicate0N @sh @r sh | Refl <- lemAppNil @sh-                          , Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mreplicate sh . unConcrete-  txindex = tindexZX-  txindex0 = tindex0X-  txscatter @shm @shn @shp sh t f =-    withKnownShX (knownShX @shm `ssxAppend` knownShX @shn) $-    gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) $-    gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) $-    case tftk knownSTK t of-      FTKX _ x@FTKScalar ->  -- optimized-        let zero = tdefTarget (FTKX sh x)-            shm = shxTakeSSX (Proxy @shn) (knownShX @shm) (xshape t)-            shDropP = shxDropSSX (knownShX @shm) (xshape t)-            s = shxSize shm-            g ix =-              let ix2 = f $ fmapConcrete ix-              in if ixInBounds (fmapUnConcrete $ toList ix2) (shxToList sh)-                 then M.insertWith (V.zipWith (+)) ix2-                        (Nested.mtoVector-                         $ tindexNX @_ @shm @shn (unConcrete t) ix)-                 else id-            ivs = foldr g M.empty [ fromLinearIdxX fromIntegral shm-                                    $ fromIntegral i-                                  | i <- [0 .. s - 1] ]-        in updateNX @(Rank shp) zero-           $ map (second $ Concrete . Nested.mfromVector shDropP)-           $ M.assocs ivs-      FTKX _ x | Dict <- eltDictRep (ftkToSTK x) ->-        let zero = tdefTarget (FTKX sh x)-            shm = shxTakeSSX (Proxy @shn) (knownShX @shm) (xshape t)-            s = shxSize shm-            g ix =-              let ix2 = f $ fmapConcrete ix-              in if ixInBounds (fmapUnConcrete $ toList ix2) (shxToList sh)-                 then M.insertWith (taddTarget knownSTK) ix2-                        (Concrete-                         $ tindexNX @_ @shm @shn (unConcrete t) ix)-                 else id-            ivs = foldr g M.empty [ fromLinearIdxX fromIntegral shm-                                    $ fromIntegral i-                                  | i <- [0 .. s - 1] ]-        in updateNX @(Rank shp) zero-           $ M.assocs ivs-  txscatter1 = tscatterZ1X-  txgather @shm @shn @_ @r sh t f =-    gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $-    gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $-    case knownSTK @r of-      STKScalar ->  -- optimized-        let shm = shxTakeSSX (Proxy @shn) (knownShX @shm) sh-            s = shxSize shm-            l = [ xtoVector-                  $ txindex @_ @_ @shn-                      t (f (fmapConcrete-                            $ fromLinearIdxX fromIntegral shm i))-                | i <- [0 .. fromIntegral s - 1] ]-        in Concrete $ Nested.mfromVector sh $ V.concat l-      _ ->-        withKnownShX (ssxFromShX sh) $-        xbuild @(Rank shm) sh (\ix -> t `txindex` f ix)-  txgather1 = tgatherZ1X-  txconcrete = Concrete-  txfloor = Concrete . liftVX (V.map floor) . unConcrete-  txfromIntegral = Concrete . liftVX (V.map fromIntegral) . unConcrete-  {-# INLINE txcast #-}-  txcast = Concrete . liftVX (V.map realToFrac) . unConcrete-  txminIndex = Concrete . tminIndexX . unConcrete-  txmaxIndex = Concrete . tmaxIndexX . unConcrete-  txiota @n = let n = valueOf @n-                  t = Nested.mfromList1 $ NonEmpty.map fromInteger-                      $ NonEmpty.fromList [0 .. n - 1]-              in Concrete $ Nested.mcast (Nested.SKnown (SNat @n) :!% ZKX) t-  txappend @_ @_ @_ @r u v | Dict <- eltDictRep (knownSTK @r) =-    Concrete $ Nested.mappend (unConcrete u) (unConcrete v)-  txslice @_ @_ @_ @_ @r i n _ | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mslice i n . unConcrete-  txreverse @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mrev1 . unConcrete-  txtranspose @_ @_ @r perm | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mtranspose perm . unConcrete-  txreshape @_ @_ @r sh | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mreshape sh . unConcrete-  txbuild1 @_ @_ @r f | Dict <- eltDictRep (knownSTK @r) =-    Concrete $ tbuild1X (unConcrete . f . Concrete)--  -- Scalar ops-  tkconcrete = Concrete-  tkfloor = Concrete . floor . unConcrete-  tkfromIntegral = Concrete . fromIntegral . unConcrete-  tkcast = Concrete . realToFrac . unConcrete--  -- General operations that don't require LetTensor nor ShareTensor-  tftk stk (Concrete t) = tftkG stk t-  tconcrete _ = id-  tpair !u !v = Concrete (unConcrete u, unConcrete v)-  tproject1 = Concrete . fst . unConcrete-  tproject2 = Concrete . snd . unConcrete-  tsreplicate @_ @_ @x snat@SNat _sh | Dict <- eltDictRep (knownSTK @x) =-    Concrete . Nested.sreplicate (snat :$$ ZSS) . unConcrete-  tsreplicate0N @sh @r sh | Refl <- lemAppNil @sh-                          , Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.sreplicate sh . unConcrete-  tstranspose @_ @_ @r perm | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.stranspose perm . unConcrete-  tsreshape @_ @_ @x sh | Dict <- eltDictRep (knownSTK @x) =-    Concrete . Nested.sreshape sh . unConcrete-  -- The eta-expansion below is needed for typing.-  tmapAccumRDer _ k _ bftk eftk f _df _rf acc0 es =-    oRtmapAccumR k bftk eftk (\ (Concrete a) (Concrete b) ->-                                Concrete $ f (a, b)) acc0 es-  tmapAccumLDer _ k _ bftk eftk f _df _rf acc0 es =-    oRtmapAccumL k bftk eftk (\ (Concrete a) (Concrete b) ->-                                Concrete $ f (a, b)) acc0 es-  tApply f x = Concrete $ f $ unConcrete x-  tlambda _ f x = unConcrete $ unHFun f $ Concrete x-  tcond _ b u v = if b then u else v-  tprimalPart = id-  tdualPart stk t = DummyDualTarget (tftk stk t)-  tfromPrimal _ t = t-  tfromDual (DummyDualTarget ftk) = tdefTarget ftk-  tScale _ _ t = t-  -- The code for tvjp and tjvp in this instance is similar as for the-  -- ADVal ranked instance, because the type family instance is the same.-  tgrad @x xftk h =-    let rf :: RepConcrete x -> RepConcrete (ADTensorKind x)-        rf !a = unConcrete $ fst $ crevOnParams Nothing (unHFun h)-                                                 xftk (Concrete a)-    in rf-  tvjp @x @z xftk h =-    let rf :: RepConcrete (TKProduct (ADTensorKind z) x)-           -> RepConcrete (ADTensorKind x)-        rf !db_a = unConcrete $ fst-                   $ crevOnParams (Just $ Concrete $ fst db_a) (unHFun h)-                                   xftk (Concrete $ snd db_a)-    in rf-  tjvp @x @z xftk h =-    let df :: RepConcrete (TKProduct (ADTensorKind x) x)-           -> RepConcrete (ADTensorKind z)-        df !da_a = unConcrete $ fst-                   $ cfwdOnParams xftk (Concrete $ snd da_a)-                                   (unHFun h) (Concrete $ fst da_a)-    in df--  tfromVector snat@SNat stk v = assert (V.length v == sNatValue snat)-                                $ case stk of-    STKScalar -> tsfromVector $ V.map sfromK v-    STKR SNat x | Dict <- lemKnownSTK x -> trfromVector v-    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsfromVector v-    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txfromVector v-    STKProduct stk1 stk2 ->-      let (v1, v2) = V.unzip $ V.map tunpair v-      in tpair (tfromVector snat stk1 v1) (tfromVector snat stk2 v2)--  treplTarget = replTarget-  tdefTarget = defTarget-  taddTarget = addTarget-  tmultTarget = multTarget-  tsum0Target = sum0Target-  tdot0Target = dot0Target--instance ConvertTensor Concrete where-  tconvert c astk a | Dict <- eltDictRep astk-                    , Dict <- eltDictRep (convertSTK c astk) =-    Concrete $ Nested.convert (interpretTKConversion c) (unConcrete a)--  kfromR = Concrete . Nested.runScalar . unConcrete-  kfromS = Concrete . Nested.sunScalar . unConcrete-  kfromX = Concrete . Nested.munScalar . unConcrete-  rfromK = Concrete . Nested.rscalar . unConcrete-  rfromS @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.stoRanked . unConcrete-  {-# SPECIALIZE rfromS :: KnownShS sh => Concrete (TKS sh Double) -> Concrete (TKR (Rank sh) Double) #-}-  {-# SPECIALIZE rfromS :: KnownShS sh => Concrete (TKS sh Float) -> Concrete (TKR (Rank sh) Float) #-}-  rfromX @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mtoRanked . unConcrete-  sfromK = Concrete . Nested.sscalar . unConcrete-  sfromR @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . flip Nested.rcastToShaped knownShS . unConcrete-  sfromX @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.mcastToShaped knownShS . unConcrete-  xfromK = Concrete . Nested.mscalar . unConcrete-  xfromR @sh @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.rcastToMixed (knownShX @sh) . unConcrete-  xfromS @_ @sh' @r | Dict <- eltDictRep (knownSTK @r) =-    Concrete . Nested.scastToMixed (knownShX @sh') . unConcrete--  rzip @y @z (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)-                               , Dict <- eltDictRep (knownSTK @z) =-    Concrete $ Nested.rzip a b-  runzip a = let (!a1, !a2) = Nested.runzip $ unConcrete a-             in Concrete (a1, a2)-  szip @y @z  (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)-                                , Dict <- eltDictRep (knownSTK @z) =-    Concrete $ Nested.szip a b-  sunzip a = let (!a1, !a2) = Nested.sunzip $ unConcrete a-             in Concrete (a1, a2)-  xzip @y @z  (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)-                                , Dict <- eltDictRep (knownSTK @z) =-    Concrete $ Nested.mzip a b-  xunzip a = let (!a1, !a2) = Nested.munzip $ unConcrete a-             in Concrete (a1, a2)--  xnestR @sh1 @m @x sh | Dict <- eltDictRep (knownSTK @x)-                       , Refl <- lemRankReplicate (SNat @m) =-    Concrete-    . Nested.convert-        @(Nested.Mixed sh1 (Nested.Mixed (Replicate m Nothing) (RepConcrete x)))-        (Nested.ConvXX Nested.ConvXR)-    . Nested.mnest sh-    . unConcrete-  xnestS @sh1 @sh2 @x sh | Dict <- eltDictRep (knownSTK @x) =-    Concrete-    . Nested.convert-        @(Nested.Mixed sh1 (Nested.Mixed (MapJust sh2) (RepConcrete x)))-        (Nested.ConvXX Nested.ConvXS)-    . Nested.mnest sh-    . unConcrete-  xnest @_ @_ @x sh | Dict <- eltDictRep (knownSTK @x) =-    Concrete . Nested.mnest sh . unConcrete-  xunNestR @sh1 @m @x | Dict <- eltDictRep (knownSTK @x) =-    Concrete-    . Nested.munNest-    . Nested.convert-        @(Nested.Mixed sh1 (Nested.Ranked m (RepConcrete x)))-        (Nested.ConvXX Nested.ConvRX)-    . unConcrete-  xunNestS @sh1 @sh2 @x | Dict <- eltDictRep (knownSTK @x) =-    Concrete-    . Nested.munNest-    . Nested.convert-        @(Nested.Mixed sh1 (Nested.Shaped sh2 (RepConcrete x)))-        (Nested.ConvXX Nested.ConvSX)-    . unConcrete-  xunNest = Concrete . Nested.munNest . unConcrete--  tpairConv = tpair-  tunpairConv = tunpair--interpretTKConversion :: TKConversion a b-                    -> Nested.Conversion (RepConcrete a) (RepConcrete b)-interpretTKConversion c0 = case c0 of-  ConvId -> Nested.ConvId-  ConvCmp c1 c2 -> Nested.ConvCmp (interpretTKConversion c1)-                                  (interpretTKConversion c2)-  ConvRX -> Nested.ConvRX-  ConvSX -> Nested.ConvSX-  ConvXR stk | Dict <- eltDictRep stk -> Nested.ConvXR-  ConvXS -> Nested.ConvXS-  ConvXS' (FTKS sh' ftk) | Dict <- eltDictRep (ftkToSTK ftk) ->-    Nested.ConvXS' sh'-  ConvXX' (FTKX shx ftk) | Dict <- eltDictRep (ftkToSTK ftk) ->-    Nested.ConvXX' (ssxFromShX shx)-  ConvRR c -> Nested.ConvRR (interpretTKConversion c)-  ConvSS c -> Nested.ConvSS (interpretTKConversion c)-  ConvXX c -> Nested.ConvXX (interpretTKConversion c)-  ConvT2 c1 c2 ->-    Nested.ConvT2 (interpretTKConversion c1) (interpretTKConversion c2)-  Conv0X stk | Dict <- eltDictRep stk -> Nested.Conv0X-  ConvX0 -> Nested.ConvX0-  ConvNest (STKX sh x) | Dict <- eltDictRep x -> Nested.ConvNest sh-  ConvUnnest -> Nested.ConvUnnest-  ConvZip stk1 stk2 | Dict <- eltDictRep stk1-                    , Dict <- eltDictRep stk2 -> Nested.ConvZip-  ConvUnzip stk1 stk2 | Dict <- eltDictRep stk1-                      , Dict <- eltDictRep stk2 -> Nested.ConvUnzip----- * MapAccum internal definitions--ravel :: forall k y.-         SNat k -> SingletonTK y -> [Concrete y]-      -> Concrete (BuildTensorKind k y)-ravel k stk l = tfromVector k stk (V.fromList l)--unravel :: forall k y.-           SNat k -> SingletonTK y -> Concrete (BuildTensorKind k y)-        -> [Concrete y]-unravel = tunravelToListShare--oRtmapAccumR-  :: forall k accy by ey.-     SNat k-  -> FullShapeTK by-  -> FullShapeTK ey-  -> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))-  -> Concrete accy-  -> Concrete (BuildTensorKind k ey)-  -> Concrete (TKProduct accy (BuildTensorKind k by))-oRtmapAccumR k bftk eftk f acc0 es = case sNatValue k of-  0 -> tpair acc0 (treplicate k (ftkToSTK bftk) (tdefTarget bftk))-  _ ->-    let g a b = let res = f a b-                in (tproject1 res, tproject2 res)-        (xout, lout) = mapAccumR g acc0 (unravel k (ftkToSTK eftk) es)-    in tpair xout (ravel k (ftkToSTK bftk) lout)-      -- TODO: reimplement not with Haskell's mapAccumR to avoid the ravels--oRtmapAccumL-  :: forall k accy by ey.-     SNat k-  -> FullShapeTK by-  -> FullShapeTK ey-  -> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))-  -> Concrete accy-  -> Concrete (BuildTensorKind k ey)-  -> Concrete (TKProduct accy (BuildTensorKind k by))-oRtmapAccumL k bftk eftk f acc0 es = case sNatValue k of-  0 -> tpair acc0 (treplicate k (ftkToSTK bftk) (tdefTarget bftk))-  _ ->-    let g a b = let res = f a b-                in (tproject1 res, tproject2 res)-        (xout, lout) = mapAccumL g acc0 (unravel k (ftkToSTK eftk) es)-    in tpair xout (ravel k (ftkToSTK bftk) lout)----- * Ranked internal definitions---- TODO: check what the following did in tsum0R and if worth emulating--- (also in sum1Inner and extremum and maybe tdot0R):--- LA.sumElements $ OI.toUnorderedVectorT sh t-{--tdot0R t u = OR.toVector t LA.<.> OR.toVector u-  -- TODO: if offset 0 and same strides, use toUnorderedVectorT-  -- TODO: if either has length 1 values, it may or may not be faster to do-  -- tsum0R (t * u) -}---- TODO: check what the following did in tsum0R and if worth emulating--- (also in sum1Inner and extremum and maybe tdot0R):--- LA.sumElements $ OI.toUnorderedVectorT sh t---- TODO: try to weave a similar magic as in tindex0R--- TODO: for the non-singleton case see--- https://github.com/Mikolaj/horde-ad/pull/81#discussion_r1096532164-updateNR :: forall n m x. (KnownNat n, KnownNat m, KnownSTK x)-         => Concrete (TKR2 (n + m) x)-         -> [(IxROf Concrete n, Concrete (TKR2 m x))]-         -> Concrete (TKR2 (n + m) x)-updateNR arr upd = case knownSTK @x of-  STKScalar ->  -- optimized-    let values = rtoVector arr-        sh = rshape arr-        f !t (ix, u) =-          let v = rtoVector u-              i = fromIntegral $ unConcrete-                  $ toLinearIdxR @n @m fromIntegral sh ix-          in V.concat [V.take i t, v, V.drop (i + V.length v) t]-    in Concrete $ Nested.rfromVector sh (foldl' f values upd)-  _ ->-    let arrNested = rnest (SNat @n) arr-        shNested = rshape arrNested-        f i v = case lookup (fromLinearIdxR-                               @n (Concrete . fromIntegral)-                               shNested ((Concrete . fromIntegral) i)) upd of-          Just u -> rnest (SNat @0) u-          Nothing -> v-    in runNest $ trfromVector0N shNested $ V.fromList-       $ imap f $ trunravelToList $ rflatten arrNested--tminIndexR-  :: forall r r2 n.-     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)-  => Nested.Ranked (1 + n) r -> Nested.Ranked n r2-tminIndexR v | SNat <- Nested.rrank v =-  let f :: Nested.Ranked 1 r -> Nested.Ranked 0 r2-      f = Nested.rscalar . fromIntegral . ixrHead . Nested.rminIndexPrim-  in Nested.rrerank SNat ZSR f v--tmaxIndexR-  :: forall r r2 n.-     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)-  => Nested.Ranked (1 + n) r -> Nested.Ranked n r2-tmaxIndexR v | SNat <- Nested.rrank v =-  let f :: Nested.Ranked 1 r -> Nested.Ranked 0 r2-      f = Nested.rscalar . fromIntegral . ixrHead . Nested.rmaxIndexPrim-  in Nested.rrerank SNat ZSR f v---- We could generalize by unwinding and only then doing the PrimElt things,--- but we'd need a type family that says "replace this underlying scalars--- by this one", which makes things too complicated.------ We could also expose `liftVR` in the user API, but in addition--- to the main function argument such as floor or cast, it'd need the function's--- derivative, just as with mapAccums. Maybe it's better to generalize even more--- and permit arbitrary extra ops if given their derivatives.-liftVR-  :: (Nested.PrimElt r1, Nested.PrimElt r2)-  => (VS.Vector r1 -> VS.Vector r2)-  -> Nested.Ranked n r1 -> Nested.Ranked n r2-liftVR f = Ranked.liftRanked1 (Mixed.mliftNumElt1 (`liftVEltwise1` f))-{-# SPECIALIZE liftVR :: (VS.Vector Double -> VS.Vector Double) -> Nested.Ranked n Double -> Nested.Ranked n Double #-}-{-# SPECIALIZE liftVR :: (VS.Vector Float -> VS.Vector Float) -> Nested.Ranked n Float -> Nested.Ranked n Float #-}-{-# SPECIALIZE liftVR :: (VS.Vector Double -> VS.Vector Float) -> Nested.Ranked n Double -> Nested.Ranked n Float #-}-{-# SPECIALIZE liftVR :: (VS.Vector Float -> VS.Vector Double) -> Nested.Ranked n Float -> Nested.Ranked n Double #-}--ixInBounds :: [Int64] -> [Int] -> Bool-ixInBounds ix sh =-  and $ zipWith (\i dim -> 0 <= i && i < fromIntegral dim) ix sh--tindexNR-  :: (Nested.Elt r, Show r, KnownNat m, KnownNat n)-  => Nested.Ranked (m + n) r -> IxR m Int64 -> Nested.Ranked n r-tindexNR v ix = let sh = Nested.rshape v-                    !_A = assert (ixInBounds (toList ix) (toList sh)-                                  `blame` (v, ix)) ()-                in Nested.rindexPartial v (fmap fromIntegral ix)-{- TODO: benchmark if this is faster enough for its complexity;-         probably not, becasue orthotope's index does no canonicalization either-tindexNR v@(RS.A (RG.A sh OI.T{strides, offset, values})) ix =-  let l = indexToList ix-      linear = offset + sum (zipWith (*) (map fromIntegral l) strides)-      plen = valueOf @m  -- length of prefix being indexed out of-      !_A = assert (ixInBounds l sh `blame` (ix, sh, v)) ()-  in-    RS.A (RG.A (drop plen sh) OI.T{ strides = drop plen strides-                                  , offset = linear-                                  , values })--}--tindexZR-  :: forall r m n. (KnownSTK r, KnownNat m, KnownNat n)-  => Concrete (TKR2 (m + n) r) -> IxROf Concrete m -> Concrete (TKR2 n r)-tindexZR v ixConcrete | Dict <- showDictRep (knownSTK @r)-                  , Dict <- eltDictRep (knownSTK @r) =-  let ix = fmapUnConcrete ixConcrete-  in case tftk knownSTK v of-    FTKR sh x ->-     if ixInBounds (Foldable.toList ix) (Foldable.toList sh)-     then Concrete $ tindexNR (unConcrete v) ix-     else tdefTarget (FTKR (shrDrop @m sh) x)--tindex0R-  :: forall r m. (KnownSTK r, KnownNat m)-  => Concrete (TKR2 m r) -> IxROf Concrete m -> Concrete (TKR2 0 r)-tindex0R v ixConcrete | Dict <- eltDictRep (knownSTK @r) =-  let ix = fmapUnConcrete ixConcrete-  in case tftk knownSTK v of-    FTKR sh x ->-      if ixInBounds (toList ix) (toList sh)-      then let arr = Nested.rscalar-                     $ Nested.rindex (unConcrete v) (fmap fromIntegral ix)-           in Concrete arr-      else tdefTarget (FTKR ZSR x)-{- TODO: see above-tindex0R (RS.A (RG.A _ OI.T{..})) ix =-  values V.! (offset + sum (zipWith (*) (map fromIntegral $ indexToList ix)-                                        strides))--}---- Performance depends a lot on the number and size of tensors.--- If tensors are not tiny, memory taken by underlying vectors matters most--- and this implementation is probbaly optimal in this respect--- (the only new vectors are created by V.concat, but this is done on demand).--- TODO: optimize updateNR and make it consume and forget arguments--- one by one to make the above true------ Note how ix being in bounds is checked. The semantics of the operation--- permits index out of bounds and then no tensors is added at such an index.-tscatterZR :: forall m p n r.-              (KnownNat p, KnownNat m, KnownNat n, KnownSTK r)-           => IShR (p + n) -> Concrete (TKR2 (m + n) r)-           -> (IxROf Concrete m -> IxROf Concrete p)-           -> Concrete (TKR2 (p + n) r)-tscatterZR sh t f- | Dict <- eltDictRep (knownSTK @r) = case tftk knownSTK t of-  FTKR _ x@FTKScalar ->  -- optimized-    let zero = tdefTarget (FTKR sh x)-        (shm, shDropP) = shrSplitAt @m $ rshape t-        s = shrSize shm-        g ix =-          let ix2 = f $ fmapConcrete ix-          in if ixInBounds (fmapUnConcrete $ toList ix2) (toList sh)-             then M.insertWith (V.zipWith (+)) ix2-                               (Nested.rtoVector $ unConcrete t `tindexNR` ix)-             else id-        ivs = foldr g M.empty [ fromLinearIdxR fromIntegral shm i-                              | i <- [0 .. fromIntegral s - 1] ]-    in updateNR zero-       $ map (second $ Concrete . Nested.rfromVector shDropP)-       $ M.assocs ivs-  FTKR _ x | Dict <- showDictRep (ftkToSTK x) ->-    let zero = tdefTarget (FTKR sh x)-        (shm, _) = shrSplitAt @m $ rshape t-        s = shrSize shm-        g ix =-          let ix2 = f $ fmapConcrete ix-          in if ixInBounds (fmapUnConcrete $ toList ix2) (toList sh)-             then M.insertWith (taddTarget knownSTK) ix2-                               (Concrete $ unConcrete t `tindexNR` ix)-             else id-        ivs = foldr g M.empty [ fromLinearIdxR fromIntegral shm i-                              | i <- [0 .. fromIntegral s - 1] ]-    in updateNR zero-       $ M.assocs ivs---- TODO: update in place in ST or with a vector builder, but that requires--- building the underlying value vector with crafty index computations--- and then freezing it and calling Nested.rfromVector--- or optimize tscatterNR and instantiate it instead-tscatterZ1R :: (KnownSTK r, KnownNat p, KnownNat n)-            => IShR (p + n) -> Concrete (TKR2 (1 + n) r)-            -> (IntOf Concrete -> IxROf Concrete p)-            -> Concrete (TKR2 (p + n) r)-tscatterZ1R sh t f = case tftk knownSTK t of-  FTKR _ x ->-    let zero = tdefTarget (FTKR sh x)-        lt = trunravelToList t-        g i ti = let ix2 = f $ Concrete $ fromIntegral i-                 in if ixInBounds (fmapUnConcrete $ toList ix2) (toList sh)-                    then updateNR zero [(ix2, ti)]-                    else zero-        lu = imap g lt-    in foldr (taddTarget knownSTK) zero lu--tfromVector0NR-  :: Nested.KnownElt r-  => IShR n -> Data.Vector.Vector (Nested.Ranked 0 r) -> Nested.Ranked n r-tfromVector0NR sh l = case NonEmpty.nonEmpty $ V.toList l of-  Nothing -> Nested.rreshape sh Nested.remptyArray-  Just nl -> Nested.rfromListLinear sh $ NonEmpty.map Nested.runScalar nl--tbuild1R-  :: forall r n. (Nested.KnownElt r, KnownNat n)-  => Int -> (Int64 -> Nested.Ranked n r) -> Nested.Ranked (1 + n) r-tbuild1R k f = case NonEmpty.nonEmpty [0 .. fromIntegral k - 1] of-  Nothing -> case sameNat (Proxy @n) (Proxy @0) of-    Just Refl -> Nested.remptyArray-    Nothing -> error "rbuild1: shape ambiguity"-  Just l -> Nested.rfromListOuter $ NonEmpty.map f l  -- hope this fuses--tmap0NR-  :: (Nested.PrimElt r1, Nested.PrimElt r)-  => (Nested.Ranked 0 r1 -> Nested.Ranked 0 r) -> Nested.Ranked n r1-  -> Nested.Ranked n r-tmap0NR f = Ranked.liftRanked1-              (Mixed.mliftPrim (Nested.runScalar . f . Nested.rscalar ))-  -- too slow: tbuildNR (Nested.rshape v) (\ix -> f $ v `tindexNR` ix)--tzipWith0NR-  :: (Nested.PrimElt r, Nested.PrimElt r1, Nested.PrimElt r2)-  => (Nested.Ranked 0 r1 -> Nested.Ranked 0 r2 -> Nested.Ranked 0 r)-  -> Nested.Ranked n r1 -> Nested.Ranked n r2 -> Nested.Ranked n r-tzipWith0NR f =-  Ranked.liftRanked2-    (Mixed.mliftPrim2-       (\x y -> Nested.runScalar $ f (Nested.rscalar x) (Nested.rscalar y)))---- The semantics of the operation permits index out of bounds--- and the result of such indexing is def, which is 0.-tgatherZR :: forall m p n r.-             (KnownNat m, KnownNat p, KnownNat n, KnownSTK r)-          => IShR (m + n) -> Concrete (TKR2 (p + n) r)-          -> (IxROf Concrete m -> IxROf Concrete p)-          -> Concrete (TKR2 (m + n) r)-tgatherZR sh t f = case knownSTK @r of-  STKScalar ->  -- optimized-    let shm = shrTake @m sh-        s = shrSize shm-        l = [ rtoVector-              $ t `trindex` f (fmapConcrete $ fromLinearIdxR fromIntegral shm i)-            | i <- [0 .. fromIntegral s - 1] ]-    in Concrete $ Nested.rfromVector sh $ V.concat l-  _ -> rbuild sh (\ix -> t `trindex` f ix)--tgatherZ1R :: forall p n r.-              (KnownNat p, KnownNat n, KnownSTK r)-           => Int -> Concrete (TKR2 (p + n) r)-           -> (IntOf Concrete -> IxROf Concrete p)-           -> Concrete (TKR2 (1 + n) r)-tgatherZ1R k t f = case knownSTK @r of-  STKScalar ->  -- optimized-    trfromVector $ V.fromList $ map (\i -> t `trindex` f (Concrete i))-                                    [0 .. fromIntegral k - 1]-  _ -> trbuild1 k (\ix -> t `trindex` f ix)----- * Shaped internal definitions---- TODO: try to weave a similar magic as in tindex0R--- TODO: for the non-singleton case see--- https://github.com/Mikolaj/horde-ad/pull/81#discussion_r1096532164-updateNS :: forall n sh r.-            ( KnownSTK r, KnownShS sh, KnownShS (Drop n sh)-            , KnownShS (Take n sh) )-         => Concrete (TKS2 sh r)-         -> [(IxSOf Concrete (Take n sh), Concrete (TKS2 (Drop n sh) r))]-         -> Concrete (TKS2 sh r)-updateNS arr upd = case knownSTK @r of-  STKScalar ->-    let values = stoVector arr-        sh = knownShS @sh-        f !t (ix, u) =-          let v = stoVector u-              i = gcastWith (unsafeCoerceRefl-                             :: sh :~: Take n sh ++ Drop n sh)-                  $ fromIntegral $ unConcrete-                  $ toLinearIdxS @(Take n sh) @(Drop n sh)-                                 fromIntegral sh ix-          in V.concat [V.take i t, v, V.drop (i + V.length v) t]-    in Concrete $ Nested.sfromVector knownShS (foldl' f values upd)-  _ -> case shsProduct (knownShS @(Take n sh)) of-    SNat ->-      gcastWith (unsafeCoerceRefl :: sh :~: Take n sh ++ Drop n sh) $-      let arrNested = snest (knownShS @(Take n sh)) arr-          shNested = sshape arrNested-          f i v = case lookup (fromLinearIdxS-                                 @(Take n sh) (Concrete . fromIntegral)-                                 shNested ((Concrete . fromIntegral) i)) upd of-            Just u -> snest (knownShS @'[]) u-            Nothing -> v-      in sunNest @_ @(Take n sh) $ tsfromVector0N $ V.fromList-         $ imap f $ tsunravelToList $ sflatten arrNested--tfromIntegralS :: (GoodScalar r1, Integral r1, GoodScalar r2)-               => Nested.Shaped sh r1 -> Nested.Shaped sh r2-tfromIntegralS = liftVS (V.map fromIntegral)--tminIndexS-  :: forall n sh r r2.-     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)-  => Nested.Shaped (n ': sh) r -> Nested.Shaped (Init (n ': sh)) r2-tminIndexS v | sh1@(_ :$$ sh) <- Nested.sshape v =-  let f :: Nested.Shaped '[m] r -> Nested.Shaped '[] r2-      f = Nested.sscalar . fromIntegral . ixsHead . Nested.sminIndexPrim-  in case sh of-    ZSS -> f @n v-    _ | SNat @m <- shsLast sh1 ->-      gcastWith (unsafeCoerceRefl-                 :: Init (n ': sh) ++ '[m] :~: n ': sh) $-      gcastWith (unsafeCoerceRefl-                 :: Init (n ': sh) :~: Init (n ': sh) ++ '[]) $-      Nested.srerank @'[m] @'[] @(Init (n ': sh))-                     (shsInit sh1) ZSS (f @m) v--tmaxIndexS-  :: forall n sh r r2.-     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)-  => Nested.Shaped (n ': sh) r -> Nested.Shaped (Init (n ': sh)) r2-tmaxIndexS v | sh1@(_ :$$ sh) <- Nested.sshape v =-  let f :: Nested.Shaped '[m] r -> Nested.Shaped '[] r2-      f = Nested.sscalar . fromIntegral . ixsHead . Nested.smaxIndexPrim-  in case sh of-    ZSS -> f @n v-    _ | SNat @m <- shsLast sh1 ->-      gcastWith (unsafeCoerceRefl-                 :: Init (n ': sh) ++ '[m] :~: n ': sh) $-      gcastWith (unsafeCoerceRefl-                 :: Init (n ': sh) :~: Init (n ': sh) ++ '[]) $-      Nested.srerank @'[m] @'[] @(Init (n ': sh))-                     (shsInit sh1) ZSS (f @m) v--liftVS-  :: (Nested.PrimElt r1, Nested.PrimElt r)-  => (VS.Vector r1 -> VS.Vector r)-  -> Nested.Shaped sh r1 -> Nested.Shaped sh r-liftVS f = Shaped.liftShaped1 (Mixed.mliftNumElt1 (`liftVEltwise1` f))-{-# SPECIALIZE liftVS :: (VS.Vector Double -> VS.Vector Double) -> Nested.Shaped sh Double -> Nested.Shaped sh Double #-}-{-# SPECIALIZE liftVS :: (VS.Vector Float -> VS.Vector Float) -> Nested.Shaped sh Float -> Nested.Shaped sh Float #-}-{-# SPECIALIZE liftVS :: (VS.Vector Double -> VS.Vector Float) -> Nested.Shaped sh Double -> Nested.Shaped sh Float #-}-{-# SPECIALIZE liftVS :: (VS.Vector Float -> VS.Vector Double) -> Nested.Shaped sh Float -> Nested.Shaped sh Double #-}--tindexNS-  :: Nested.Elt r-  => Nested.Shaped (sh1 ++ sh2) r -> IxS sh1 Int64 -> Nested.Shaped sh2 r-tindexNS v ix = Nested.sindexPartial v (fmap fromIntegral ix)-{- TODO-tindexNS (SS.A (SG.A OI.T{strides, offset, values})) ix =-  let l = ShapedList.indexToList ix-      linear = offset + sum (zipWith (*) (map fromIntegral l) strides)-      plen = length l  -- length of prefix being indexed out of-  in-    SS.A (SG.A OI.T{ strides = drop plen strides-                   , offset = linear-                   , values })--}---- Note that after vectorization, the index may not fit within--- the type-level shape, which we catch in the @ixInBounds@--- and return def, so it's fine. Similarly in gather and scatter.-tindexZS-  :: forall r sh1 sh2. (KnownSTK r, KnownShS sh1, KnownShS sh2)-  => Concrete (TKS2 (sh1 ++ sh2) r) -> IxSOf Concrete sh1-  -> Concrete (TKS2 sh2 r)-tindexZS v ixConcrete | Dict <- eltDictRep (knownSTK @r) =-  let ix = fmapUnConcrete ixConcrete-  in withKnownShS (knownShS @sh1 `shsAppend` knownShS @sh2) $-     case tftk knownSTK v of-       FTKS sh x ->-         if ixInBounds (Foldable.toList ix) (shsToList sh)-         then Concrete $ tindexNS (unConcrete v) ix-         else tdefTarget (FTKS knownShS x)--tindex0S-  :: forall r sh. (KnownSTK r, KnownShS sh)-  => Concrete (TKS2 sh r) -> IxSOf Concrete sh -> Concrete (TKS2 '[] r)-tindex0S v ixConcrete | Dict <- eltDictRep (knownSTK @r) =-  let ix = fmapUnConcrete ixConcrete-  in case tftk knownSTK v of-    FTKS sh x ->-      if ixInBounds (toList ix) (toList sh)-      then let arr = Nested.sscalar-                     $ Nested.sindex (unConcrete v) (fmap fromIntegral ix)-           in Concrete arr-      else tdefTarget (FTKS ZSS x)-{- TODO: benchmark if this is faster enough for its complexity;-         probably not, becasue orthotope's index does no canonicalization either-tindex0S (SS.A (SG.A OI.T{..})) ix =-  values V.! (offset + sum (zipWith (*) (map fromIntegral-                                         $ ShapedList.indexToList ix)-                                        strides))-    -- to avoid linearizing @values@, we do everything in unsized way--}---- TODO: update in place in ST or with a vector builder, but that requires--- building the underlying value vector with crafty index computations--- and then freezing it and calling OS.fromVector--- or optimize tscatterNS and instantiate it instead-tscatterZ1S-  :: forall r n2 shn shp.-     (KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp)-  => Concrete (TKS2 (n2 ': shn) r)-  -> (IntOf Concrete -> IxSOf Concrete shp)-  -> Concrete (TKS2 (shp ++ shn) r)-tscatterZ1S t f = case tftk knownSTK t of-  FTKS _ x ->-    gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) $-    gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) $-    let shpshn = knownShS @shp `shsAppend` knownShS @shn-        zero = tdefTarget (FTKS shpshn x)-        lt = tsunravelToList t-        g i ti = let ix2 = f $ Concrete $ fromIntegral i-                 in if ixInBounds (fmapUnConcrete $ Foldable.toList ix2)-                                  (shsToList shpshn)-                    then withKnownShS shpshn $-                         updateNS @(Rank shp) zero [(ix2, ti)]-                    else zero-        lu = imap g lt-    in foldr (taddTarget (STKS shpshn (knownSTK @r))) zero lu--tfromVector0NS-  :: forall r sh. (Nested.KnownElt r, KnownShS sh)-  => Data.Vector.Vector (Nested.Shaped '[] r) -> Nested.Shaped sh r-tfromVector0NS l = case NonEmpty.nonEmpty $ V.toList l of-  Nothing -> case testEquality (shsProduct (knownShS @sh)) (SNat @0) of-    Just Refl -> Nested.sreshape (knownShS @sh)-                 $ Nested.semptyArray (knownShS @sh)-    Nothing -> error "tfromVector0N: empty list, but not shape"-  Just nl -> Nested.sfromListLinear knownShS $ NonEmpty.map Nested.sunScalar nl--tbuild1S-  :: forall k sh r. (KnownNat k, KnownShS sh, Nested.KnownElt r)-  => (Int64 -> Nested.Shaped sh r) -> Nested.Shaped (k ': sh) r-tbuild1S f = case NonEmpty.nonEmpty [0 .. valueOf @k - 1] of-  Nothing -> gcastWith (unsafeCoerceRefl :: k :~: 0) $-             Nested.semptyArray knownShS-  Just l -> Nested.sfromListOuter SNat $ NonEmpty.map f l  -- hope this fuses--tmap0NS-  :: forall r1 r sh. (Nested.PrimElt r1, Nested.PrimElt r)-  => (Nested.Shaped '[] r1 -> Nested.Shaped '[] r) -> Nested.Shaped sh r1-  -> Nested.Shaped sh r-tmap0NS f =-  Shaped.liftShaped1-    (Mixed.mliftPrim (Nested.sunScalar . f . Nested.sscalar))-      -- too slow: tbuildNS (tshapeS v) (\ix -> f $ v `tindexNS` ix)--tzipWith0NS-  :: forall r1 r2 r sh. (Nested.PrimElt r, Nested.PrimElt r1, Nested.PrimElt r2)-  => (Nested.Shaped '[] r1 -> Nested.Shaped '[] r2 -> Nested.Shaped '[] r)-  -> Nested.Shaped sh r1 -> Nested.Shaped sh r2 -> Nested.Shaped sh r-tzipWith0NS f =-  Shaped.liftShaped2-    (Mixed.mliftPrim2-       (\x y -> Nested.sunScalar $ f (Nested.sscalar x) (Nested.sscalar y)))--tgatherZ1S-  :: forall r n2 shn shp.-     (KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp)-  => Concrete (TKS2 (shp ++ shn) r)-  -> (IntOf Concrete -> IxSOf Concrete shp)-  -> Concrete (TKS2 (n2 ': shn) r)-tgatherZ1S t f =-  case knownSTK @r of-    STKScalar ->  -- optimized-      tsfromVector $ V.fromList $ map (\i -> t `tsindex` f (Concrete i))-                                      [0 .. valueOf @n2 - 1]-    _ -> tsbuild1 (\ix -> t `tsindex` f ix)----- * Mixed internal definitions--updateNX :: forall n sh r.-            (KnownSTK r, KnownShX (Drop n sh), KnownShX (Take n sh))-         => Concrete (TKX2 sh r)-         -> [(IxXOf Concrete (Take n sh), Concrete (TKX2 (Drop n sh) r))]-         -> Concrete (TKX2 sh r)-updateNX arr upd = case knownSTK @r of-  STKScalar ->-    let values = xtoVector arr-        sh = xshape arr-        f !t (ix, u) =-          let v = xtoVector u-              i = gcastWith (unsafeCoerceRefl-                             :: sh :~: Take n sh ++ Drop n sh)-                  $ fromIntegral $ unConcrete-                  $ toLinearIdxX @(Take n sh) @(Drop n sh)-                                 fromIntegral sh ix-          in V.concat [V.take i t, v, V.drop (i + V.length v) t]-    in Concrete $ Nested.mfromVector (xshape arr) (foldl' f values upd)-  _ | Dict <- eltDictRep (knownSTK @r) ->-      gcastWith (unsafeCoerceRefl :: sh :~: Take n sh ++ Drop n sh) $-      let arrNested = xnest (knownShX @(Take n sh)) arr-          shNested = xshape arrNested-          f i v = case lookup (fromLinearIdxX-                                 @(Take n sh) (Concrete . fromIntegral)-                                 shNested ((Concrete . fromIntegral) i)) upd of-            Just u -> xnest ZKX u-            Nothing -> v-      in withSNat (shxSize shNested) $ \snat ->-           xunNest @_ @(Take n sh) $ txfromVector0N shNested $ V.fromList-           $ imap f $ txunravelToList-           $ Concrete $ Nested.mcast (Nested.SKnown snat :!% ZKX)-           $ unConcrete $ xflatten arrNested--tminIndexX-  :: forall mn sh r r2.-     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)-  => Nested.Mixed (mn ': sh) r -> Nested.Mixed (Init (mn ': sh)) r2-tminIndexX v | sh1@(_ :$% sh) <- Nested.mshape v =-  let f :: Nested.Mixed '[mm] r -> Nested.Mixed '[] r2-      f = Nested.mscalar . fromIntegral . ixxHead-          . Nested.mminIndexPrim-  in case sh of-    ZSX -> f @mn v-    _ -> withSNat (fromSMayNat' (shxLast sh1)) $ \(_ :: SNat m) ->-      gcastWith (unsafeCoerceRefl-                 :: Init (mn ': sh) ++ '[Just m] :~: mn ': sh) $-      gcastWith (unsafeCoerceRefl-                 :: Init (mn ': sh) :~: Init (mn ': sh) ++ '[]) $-      Nested.mrerank @'[Just m] @'[] @(Init (mn ': sh))-                     (ssxFromShX $ shxInit sh1) ZSX (f @(Just m)) v--tmaxIndexX-  :: forall mn sh r r2.-     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)-  => Nested.Mixed (mn ': sh) r -> Nested.Mixed (Init (mn ': sh)) r2-tmaxIndexX v | sh1@(_ :$% sh) <- Nested.mshape v =-  let f :: Nested.Mixed '[mm] r -> Nested.Mixed '[] r2-      f = Nested.mscalar . fromIntegral . ixxHead-          . Nested.mmaxIndexPrim-  in case sh of-    ZSX -> f @mn v-    _ -> withSNat (fromSMayNat' (shxLast sh1)) $ \(_ :: SNat m) ->-      gcastWith (unsafeCoerceRefl-                 :: Init (mn ': sh) ++ '[Just m] :~: mn ': sh) $-      gcastWith (unsafeCoerceRefl-                 :: Init (mn ': sh) :~: Init (mn ': sh) ++ '[]) $-      Nested.mrerank @'[Just m] @'[] @(Init (mn ': sh))-                     (ssxFromShX $ shxInit sh1) ZSX (f @(Just m)) v--liftVX-  :: (Nested.PrimElt r1, Nested.PrimElt r)-  => (VS.Vector r1 -> VS.Vector r)-  -> Nested.Mixed sh r1 -> Nested.Mixed sh r-liftVX f = Mixed.mliftNumElt1 (`liftVEltwise1` f)-{-# SPECIALIZE liftVX :: (VS.Vector Double -> VS.Vector Double) -> Nested.Mixed sh Double -> Nested.Mixed sh Double #-}-{-# SPECIALIZE liftVX :: (VS.Vector Float -> VS.Vector Float) -> Nested.Mixed sh Float -> Nested.Mixed sh Float #-}-{-# SPECIALIZE liftVX :: (VS.Vector Double -> VS.Vector Float) -> Nested.Mixed sh Double -> Nested.Mixed sh Float #-}-{-# SPECIALIZE liftVX :: (VS.Vector Float -> VS.Vector Double) -> Nested.Mixed sh Float -> Nested.Mixed sh Double #-}--tindexNX-  :: Nested.Elt r-  => Nested.Mixed (sh1 ++ sh2) r -> IxX sh1 Int64 -> Nested.Mixed sh2 r-tindexNX v ix = Nested.mindexPartial v (fmap fromIntegral ix)--tindexZX-  :: forall r sh1 sh2. (KnownSTK r, KnownShX sh1, KnownShX sh2)-  => Concrete (TKX2 (sh1 ++ sh2) r) -> IxXOf Concrete sh1-  -> Concrete (TKX2 sh2 r)-tindexZX v ixConcrete | Dict <- eltDictRep (knownSTK @r) =-  let ix = fmapUnConcrete ixConcrete-  in withKnownShX (knownShX @sh1 `ssxAppend` knownShX @sh2) $-     case tftk knownSTK v of-       FTKX sh x ->-         if ixInBounds (Foldable.toList ix) (shxToList sh)-         then Concrete $ tindexNX (unConcrete v) ix-         else tdefTarget (FTKX (shxDropSSX (knownShX @sh1) sh) x)--tindex0X-  :: forall r sh. (KnownSTK r, KnownShX sh)-  => Concrete (TKX2 sh r) -> IxXOf Concrete sh -> Concrete (TKX2 '[] r)-tindex0X v ixConcrete | Dict <- eltDictRep (knownSTK @r) =-  let ix = fmapUnConcrete ixConcrete-  in case tftk knownSTK v of-    FTKX sh x ->-      if ixInBounds (toList ix) (toList sh)-      then let arr = Nested.mscalar-                     $ Nested.mindex (unConcrete v) (fmap fromIntegral ix)-           in Concrete arr-      else tdefTarget (FTKX ZSX x)--tscatterZ1X-  :: forall r n2 shn shp.-     (KnownSTK r, KnownNat n2, KnownShX shn, KnownShX shp)-  => IShX (shp ++ shn) -> Concrete (TKX2 (Just n2 ': shn) r)-  -> (IntOf Concrete -> IxXOf Concrete shp)-  -> Concrete (TKX2 (shp ++ shn) r)-tscatterZ1X sh t f =-  case tftk knownSTK t of-    FTKX _ x ->-      withKnownShX (ssxFromShX sh) $-      gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) $-      gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) $-      let zero = tdefTarget (FTKX sh x)-          lt = txunravelToList t-          g i ti = let ix2 = f $ Concrete $ fromIntegral i-                   in if ixInBounds (fmapUnConcrete $ Foldable.toList ix2)-                                    (shxToList sh)-                      then updateNX @(Rank shp) zero [(ix2, ti)]-                      else zero-          lu = imap g lt-      in foldr (taddTarget knownSTK) zero lu--tfromVector0NX-  :: forall r sh. Nested.KnownElt r-  => IShX sh -> Data.Vector.Vector (Nested.Mixed '[] r) -> Nested.Mixed sh r-tfromVector0NX sh l = case NonEmpty.nonEmpty $ V.toList l of-  Nothing -> if shxSize sh == 0-             then Nested.mreshape sh $ Nested.memptyArray sh-             else error "tfromVector0N: empty list, but not shape"-  Just nl -> Nested.mfromListLinear sh $ NonEmpty.map Nested.munScalar nl--tbuild1X-  :: forall k sh r. (KnownNat k, KnownShX sh, Nested.KnownElt r)-  => (Int64 -> Nested.Mixed sh r)-  -> Nested.Mixed (Just k ': sh) r-tbuild1X f = case NonEmpty.nonEmpty [0 .. valueOf @k - 1] of-  Nothing -> case testEquality (knownShX @sh) ZKX of-    Just Refl -> gcastWith (unsafeCoerceRefl :: k :~: 0) $-                 Nested.memptyArray ZSX-    Nothing -> error "xbuild1: shape ambiguity"-  Just l -> Nested.mcast (Nested.SKnown (SNat @k) :!% knownShX)-            $ Nested.mfromListOuter $ NonEmpty.map f l  -- hope this fuses--tgatherZ1X-  :: forall r n2 shn shp.-     (KnownSTK r, KnownShX shn, KnownShX shp)-  => SNat n2 -> Concrete (TKX2 (shp ++ shn) r)-  -> (IntOf Concrete -> IxXOf Concrete shp)-  -> Concrete (TKX2 (Just n2 ': shn) r)-tgatherZ1X SNat t f =-  case knownSTK @r of-    STKScalar ->  -- optimized-      txfromVector $ V.fromList $ map (\i -> t `txindex` f (Concrete i))-                                      [0 .. valueOf @n2 - 1]-    _ -> txbuild1 @_ @n2 (\ix -> t `txindex` f ix)--fmapConcrete :: Coercible (f (RepConcrete y)) (f (Concrete y))-               => f (RepConcrete y) -> f (Concrete y)-fmapConcrete = coerce--fmapUnConcrete :: Coercible (f (Concrete y)) (f (RepConcrete y))-               => f (Concrete y) -> f (RepConcrete y)-fmapUnConcrete = coerce+{-# LANGUAGE AllowAmbiguousTypes, ViewPatterns #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Tensor class instances for concrete arrays backed+-- by 'Data.Vector.Storable.Vector'+-- and defined in @ox-arrays@ on the basis of @orthotope@.+module HordeAd.Core.OpsConcrete+  () where++import Prelude++import Control.Monad (forM_)+import Control.Monad.ST+import Data.Coerce (Coercible, coerce)+import Data.Default+import Data.Function ((&))+import Data.Functor qualified as Functor+import Data.Int (Int16, Int32, Int64, Int8)+import Data.IntMap.Strict qualified as IM+import Data.List (scanl')+import Data.List.NonEmpty (NonEmpty)+import Data.List.NonEmpty qualified as NonEmpty+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Data.Vector.Storable qualified as VS+import Data.Vector.Storable.Mutable qualified as VSM+import Foreign.C (CInt)+import GHC.TypeLits (KnownNat, Nat, type (*), type (+))+import Type.Reflection (Typeable, typeRep)++import Data.Array.Nested (MapJust, Replicate, type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (mcast, shxFromShR, shxFromShS)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed qualified as Mixed+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked qualified as Ranked+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped qualified as Shaped+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types+  (Init, fromSNat', pattern SS, pattern SZ, unsafeCoerceRefl)+import Data.Array.Strided.Orthotope (liftVEltwise1)++import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Conversion+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.OpsADVal+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.UnwindNum++-- * Tensor classes instance++instance LetTensor Concrete where+  ttlet = (&)  -- doesn't have to be strict, just as kcond and scond+  ttletPrimal = (&)+  ttletPlain = (&)+  toShare = id+  tunshare = id+  tD _stk t DummyDualTarget{} = t+  {-# INLINE tfold #-}+  tfold k _ stk f x0 es = foldl' f x0 (tunravelToListShare k stk es)+   {- This is worse than the above when the vector needs to be allocated+      due to complex strides. Apparently this happens often enough+      and checking strides is costly for small folds (but do we care?+      but then, how often do big folds work on rank 1 arrays anyway?).+   case stk of+    STKScalar ->+      let g !yn !ym = f yn (Concrete ym)+      in VS.foldl' g x0 (stoVector es)+    STKR SZ STKScalar ->+      let g !yn !ym = f yn (rfromK $ Concrete ym)+      in VS.foldl' g x0 (rtoVector es)+    STKS ZSS STKScalar ->+      let g !yn !ym = f yn (sfromK $ Concrete ym)+      in VS.foldl' g x0 (stoVector es)+    STKX ZKX STKScalar ->+      let g !yn !ym = f yn (xfromK $ Concrete ym)+      in VS.foldl' g x0 (xtoVector es)+    _ -> foldl' f x0 (tunravelToListShare k stk es) -}+  {-# INLINE tscan #-}+  tscan k nstk stk f x0 as =+    case NonEmpty.nonEmpty $ scanl' f x0 $ tunravelToListShare k stk as of+      Just nl -> tfromList (snatSucc k) nstk nl+      Nothing -> error "tscan: impossible"++instance ShareTensor Concrete where+  tshare = id+  {-# INLINE tunpair #-}+  tunpair (Concrete (t1, t2)) = (Concrete t1, Concrete t2)++instance BaseTensor Concrete where+  isConcreteInstance = True+  {-# INLINE rshape #-}+  rshape @_ @x | Dict <- eltDictRep (knownSTK @x) = Nested.rshape . unConcrete+  {-# INLINE sshape #-}+  sshape @_ @x | Dict <- eltDictRep (knownSTK @x) = Nested.sshape . unConcrete+  {-# INLINE xshape #-}+  xshape @_ @x | Dict <- eltDictRep (knownSTK @x) = Nested.mshape . unConcrete+  {-# INLINE tftk #-}+  tftk stk (Concrete t) = tftkG stk t+  {-# INLINE tpair #-}+  tpair u v = u `seq` v `seq` Concrete (unConcrete u, unConcrete v)+  {-# INLINE tproject1 #-}+  tproject1 = Concrete . fst . unConcrete+  {-# INLINE tproject2 #-}+  tproject2 = Concrete . snd . unConcrete+  {-# INLINE kcond #-}+  kcond b u v = if unConcrete b then u else v+    -- doesn't have to be strict, just as tlet+  {-# INLINE scond #-}+  scond b u v = if unConcrete b then u else v+  {-# INLINE tcond #-}+  tcond _ b u v = if unConcrete b then u else v+  tkconcrete = Concrete+  trconcrete = Concrete+  tsconcrete = Concrete+  txconcrete = Concrete+  tconcrete _ = id+  {-# INLINE trfromVector #-}+  trfromVector @_ @x v | Dict <- eltDictRep (knownSTK @x) =+    case NonEmpty.nonEmpty $ V.toList $ fmapUnConcrete v of+      Just l -> Concrete $ Nested.rfromListOuterN (V.length v) l+      Nothing -> error "rfromVector: empty vector"+  {-# INLINE trfromVectorN #-}+  trfromVectorN shm t = case V.uncons t of+    Just (v, _) -> trreshape (shm `shrAppend` rshape v)+                   $ trfromVector t+    Nothing -> error "trfromVectorN: empty vector"+  {-# INLINE trfromVectorLinear #-}+  trfromVectorLinear shm v =+    let l = V.toList $ fmapUnConcrete v+    in Concrete $ Nested.rfromListPrimLinear shm l+  {-# INLINE trunravelToList #-}+  trunravelToList @_ @x | Dict <- eltDictRep (knownSTK @x) =+    fmapConcrete . Nested.rtoListOuter . unConcrete+  {-# INLINE trtoListLinear #-}+  trtoListLinear = fmapConcrete . Nested.rtoListPrimLinear . unConcrete+  {-# INLINE tsfromVector #-}+  tsfromVector @_ @_ @x v | Dict <- eltDictRep (knownSTK @x) =+    case NonEmpty.nonEmpty $ V.toList $ fmapUnConcrete v of+      Just l -> Concrete $ Nested.sfromListOuter SNat l+      Nothing -> error "sfromVector: empty vector"+  {-# INLINE tsfromVectorN #-}+  tsfromVectorN @shm @shn shm t | SNat <- shsProduct shm = case V.uncons t of+    Just (v, _) ->+      gcastWith (unsafeCoerceRefl+                 :: Product (shm ++ shn) :~: Product shm * Product shn) $+      tsreshape (shm `shsAppend` sshape v) $ tsfromVector t+    Nothing -> error "tsfromVectorN: empty vector"+  {-# INLINE tsfromVectorLinear #-}+  tsfromVectorLinear shm v =+    let l = V.toList $ fmapUnConcrete v+    in Concrete $ Nested.sfromListPrimLinear shm l+  {-# INLINE tsunravelToList #-}+  tsunravelToList @_ @_ @x | Dict <- eltDictRep (knownSTK @x) =+    fmapConcrete . Nested.stoListOuter . unConcrete+  {-# INLINE tstoListLinear #-}+  tstoListLinear = fmapConcrete . Nested.stoListPrimLinear . unConcrete+  {-# INLINE txfromVector #-}+  txfromVector @n @_ @x v | Dict <- eltDictRep (knownSTK @x) =+    case NonEmpty.nonEmpty $ V.toList $ fmapUnConcrete v of+      Just l -> Concrete $ Nested.mfromListOuterSN (SNat @n) l+      Nothing -> error "xfromVector: empty vector"+  {-# INLINE txfromVectorN #-}+  txfromVectorN shm t = case V.uncons t of+    Just (v, _) ->+      withSNat (shxSize shm) $ \(SNat @n) ->+      txreshape (shm `shxAppend` xshape v) $ txfromVector @_ @n t+    Nothing -> error "trfromVectorN: empty vector"+  {-# INLINE txfromVectorLinear #-}+  txfromVectorLinear shm v =+    let l = V.toList $ fmapUnConcrete v+    in Concrete $ Nested.mfromListPrimLinear shm l+  {-# INLINE txunravelToList #-}+  txunravelToList @_ @_ @x | Dict <- eltDictRep (knownSTK @x) =+    fmapConcrete . Nested.mtoListOuter . unConcrete+  {-# INLINE txtoListLinear #-}+  txtoListLinear = fmapConcrete . Nested.mtoListPrimLinear . unConcrete+  trsum @_ @x t = case knownSTK @x of+    STKScalar @r | Dict0 <- numFromTKAllNum (Proxy @r) ->+      Concrete . Nested.rsumOuter1Prim . unConcrete $ t  -- optimized+    _ -> case tftk knownSTK t of+      FTKR sh x ->  -- GHC 9.14 says (_ :$: rest) not exhaustive+        let l = trunravelToList t+        in foldl' (taddTarget knownSTK) (tdefTarget (FTKR (shrTail sh) x)) l+          -- Concrete has a ShareTensor instance, so taddTarget arguments+          -- don't need to be duplicable+  trsumN @_ @n @x =+    let go :: SNat m2 -> Concrete (TKR2 (m2 + n) x) -> Concrete (TKR2 n x)+        go SZ v = v+        go (SS k@SNat) v = go k (trsum v)+    in go SNat+  {-# INLINE trsum0 #-}+  trsum0 = Concrete . Nested.rsumAllPrim . unConcrete+  {-# INLINE trdot0 #-}+  trdot0 u v = Concrete $ Nested.rdot (unConcrete u) (unConcrete v)+  {-# INLINE trdot1In #-}+  trdot1In u v = Concrete $ Nested.rdot1Inner (unConcrete u) (unConcrete v)+  trmatvecmul m v = trdot1In m (trreplicate (rwidth m) v)+  trmatmul2 m1 m2 = case rshape m2 of+    _ :$: width2 :$: ZSR ->+      trdot1In (trtranspose [1, 0] (trreplicate width2 m1))+               (trtranspose [0, 2, 1] (trreplicate (rwidth m1) m2))+  {-# INLINE trreplicateN #-}+  trreplicateN @_ @_ @x shm | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.rreplicate shm . unConcrete+  {-# INLINE trreplicate0N #-}+  trreplicate0N sh = Concrete . Nested.rreplicatePrim sh . unConcrete+  tssum @_ @_ @x t = case knownSTK @x of+    STKScalar @r | Dict0 <- numFromTKAllNum (Proxy @r) ->+      Concrete . Nested.ssumOuter1Prim . unConcrete $ t  -- optimized+    _ -> case tftk knownSTK t of+      FTKS (_ :$$ rest) x ->+        let l = tsunravelToList t+        in foldl' (taddTarget knownSTK) (tdefTarget (FTKS rest x)) l+  tssumN @shm @shn @x =+    let go :: ShS shm2 -> Concrete (TKS2 (shm2 ++ shn) x)+           -> Concrete (TKS2 shn x)+        go ZSS v = v+        go (SNat :$$ rest) v =+          go rest (withKnownShS (rest `shsAppend` knownShS @shn) $ tssum v)+    in go (knownShS @shm)+  {-# INLINE tssum0 #-}+  tssum0 = Concrete . Nested.ssumAllPrim . unConcrete+  {-# INLINE tsdot0 #-}+  tsdot0 u v = Concrete $ Nested.sdot (unConcrete u) (unConcrete v)+  {-# INLINE tsdot1In #-}+  tsdot1In @_ (SNat @n) u v =+    Concrete $ Nested.sdot1Inner (Proxy @n) (unConcrete u) (unConcrete v)+  tsmatvecmul m v = tsdot1In SNat m (tsreplicate SNat v)+  tsmatmul2 m1 m2 =+    tsdot1In SNat+             (tstranspose (Permutation.makePerm @'[1, 0])+                          (tsreplicate SNat m1))+             (tstranspose (Permutation.makePerm @'[0, 2, 1])+                          (tsreplicate SNat m2))+  tsreplicateN @_ @_ @x shm | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.sreplicate shm . unConcrete+  {-# INLINE tsreplicate0N #-}+  tsreplicate0N sh = Concrete . Nested.sreplicatePrim sh . unConcrete+  txsum @_ @_ @x t = case knownSTK @x of+    STKScalar @r | Dict0 <- numFromTKAllNum (Proxy @r) ->+      Concrete . Nested.msumOuter1Prim . unConcrete $ t  -- optimized+    _ -> case tftk knownSTK t of+      FTKX (_ :$% rest) x ->+        let l = txunravelToList t+        in foldl' (taddTarget knownSTK) (tdefTarget (FTKX rest x)) l+  {-# INLINE txsum0 #-}+  txsum0 = Concrete . Nested.msumAllPrim . unConcrete+  txsumN @shm @shn @x t | SNat <- ssxRank (knownShX @shm) =+    let shmshn = xshape t+        go :: IShX shm2 -> Concrete (TKX2 (shm2 ++ shn) x)+           -> Concrete (TKX2 shn x)+        go ZSX v = v+        go (SKnown SNat :$% rest) v =+          go rest (withKnownShX (ssxFromShX rest `ssxAppend` knownShX @shn)+                   $ txsum v)+        go (SUnknown i :$% rest) v =+          withSNat i $ \snat ->+            go rest (withKnownShX (ssxFromShX rest `ssxAppend` knownShX @shn)+                     $ txsum $ xmcast (SKnown snat :!% ssxFromShX rest+                                       `ssxAppend` knownShX @shn) v)+    in gcastWith (unsafeCoerceRefl+                  :: Take (Rank shm) (shm ++ shn) :~: shm) $+       go (shxTake @(Rank shm) shmshn) t+  {-# INLINE txdot0 #-}+  txdot0 u v = Concrete $ Nested.mdot (unConcrete u) (unConcrete v)+  {-# INLINE txdot1In #-}+  txdot1In @_ (SNat @n) u v =+    Concrete $ Nested.mdot1Inner (Proxy @(Just n)) (unConcrete u) (unConcrete v)+  txmatvecmul mm mn m v =+    withKnownShX (ssxFromShX $ mn :$% ZSX) $+    withKnownShX (ssxFromShX $ mm :$% mn :$% ZSX) $+    withSNat (fromSMayNat' mm) $ \(SNat @m) ->+    withSNat (fromSMayNat' mn) $ \(SNat @n) ->+      xmcast (ssxFromShX (mm :$% ZSX))+      $ txsum (xtr (txreplicate (SNat @m)+                      (xmcast (ssxFromShX (SKnown (SNat @n) :$% ZSX)) v)+                    * xmcast (ssxFromShX (SKnown (SNat @m)+                                          :$% SKnown (SNat @n)+                                          :$% ZSX)) m))+  txmatmul2 m1 m2 =+    txdot1In SNat+             (txtranspose (Permutation.makePerm @'[1, 0])+                          (txreplicate SNat m1))+             (txtranspose (Permutation.makePerm @'[0, 2, 1])+                          (txreplicate SNat m2))+  {-# INLINE txreplicateN #-}+  txreplicateN @_ @_ @x shm | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mreplicate shm . unConcrete+  {-# INLINE txreplicate0N #-}+  txreplicate0N sh = Concrete . Nested.mreplicatePrim sh . unConcrete+  {-# INLINE trindex #-}+  trindex = tindexZR+  {-# INLINE trindex0 #-}+  trindex0 = tindex0R+  {-# INLINE troneHot #-}+  troneHot = toneHotR+  {-# INLINE trscatter #-}+  trscatter = tscatterZR+  -- no meaningful optimization here so far: trscatter1 = tscatterZ1R+  {-# INLINE trgather #-}+  trgather = tgatherZR+  {-# INLINE trgather1 #-}+  trgather1 = tgatherZ1R+  {-# INLINE tsindex #-}+  tsindex @_ @shn = tindexZS (knownShS @shn)+  {-# INLINE tsindex0 #-}+  tsindex0 = tindex0S+  {-# INLINE tsoneHot #-}+  tsoneHot @shp @shn = toneHotS (knownShS @shn) (knownShS @shp)+  {-# INLINE tsscatter #-}+  tsscatter @shm @shn @shp =+    tscatterZS (knownShS @shm) (knownShS @shn) (knownShS @shp)+  {-# INLINE tsgather #-}+  tsgather @shm @shn =+    tgatherZS (knownShS @shm) (knownShS @shn)+  {-# INLINE tsgather1 #-}+  tsgather1 @_ @shn = tgatherZ1S (knownShS @shn)+  {-# INLINE txindex #-}+  txindex = tindexZX+  {-# INLINE txindex0 #-}+  txindex0 = tindex0X+  {-# INLINE txoneHot #-}+  txoneHot = toneHotX+  {-# INLINE txscatter #-}+  txscatter @shm @shn = tscatterZX @shm @shn+  {-# INLINE txgather #-}+  txgather @shm @shn = tgatherZX @shm @shn+  {-# INLINE txgather1 #-}+  txgather1 = tgatherZ1X+  {-# INLINE tkfloor #-}+  tkfloor = Concrete . floor . unConcrete+  {-# INLINE tkfromIntegral #-}+  tkfromIntegral = fromIntegral . unConcrete+    -- Here runtime specialization would be questionable. It can take several+    -- comparisons to make a single scalar operation (much) faster.+    -- The only big gain is when/if this gets inlined into the interpreter+    -- and so the interpretation resulting in @a@ is performed at a concrete+    -- type, e.g., it may be a big arithmetic expression, which we never+    -- runtime-specialize. Or when/if this op is repeated in a loop+    -- and the comparison is floated up out of the loop.+    --+    -- Benchmarks indicate this lowers allocation considerably, but increases+    -- runtime just as considerably, so it's disabled for now.+  {-# INLINE tkcast #-}+  tkcast @r1 @r2 a =+    let cast :: (Differentiable r1', Differentiable r2')+             => Concrete (TKScalar r1') -> Concrete (TKScalar r2')+        {-# INLINE cast #-}+        cast = Concrete . realToFrac . unConcrete+    -- Specializing just for the cases covered by realToFrac rules+    -- in GHC.Internal.Float, except for the Int cases that the RealFrac+    -- constraint required by reverse differenciation precludes.+    in case typeRep @r1 of+      Is @Double -> case typeRep @r2 of+        Is @Float -> cast @Double @Float a+        _ -> cast a+      Is @Float -> case typeRep @r2 of+        Is @Double -> cast @Float @Double a+        _ -> cast a+      _ -> cast a+  tkargMin = Concrete . targMinK . unConcrete+  tkargMax = Concrete . targMaxK . unConcrete+  {-# INLINE trfloor #-}+  trfloor = Concrete . liftVR (V.map floor) . unConcrete+  {-# INLINE trfromIntegral #-}+  trfromIntegral @r1 @r2 a =+    let cast :: (GoodScalar r1', Integral r1', NumScalar r2')+             => Concrete (TKR n r1') -> Concrete (TKR n r2')+        {-# INLINE cast #-}+        cast = Concrete . liftVR (V.map fromIntegral) . unConcrete+    in case typeRep @r1 of+        Is @Int -> case typeRep @r2 of+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int8 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int16 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int32 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int64 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @CInt -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          _ -> cast a+        _ -> cast a+  {-# INLINE trcast #-}+  trcast @r1 @r2 a =+    let cast :: ( Differentiable r1', NumScalar r1'+                , Differentiable r2', NumScalar r2' )+             => Concrete (TKR n r1') -> Concrete (TKR n r2')+        {-# INLINE cast #-}+        cast = Concrete . liftVR (V.map realToFrac) . unConcrete+    in case typeRep @r1 of+      Is @Double -> case typeRep @r2 of+        Is @Float -> cast @Double @Float a+        _ -> cast a+      Is @Float -> case typeRep @r2 of+        Is @Double -> cast @Float @Double a+        _ -> cast a+      _ -> cast a+  trargMin = Concrete . targMinR . unConcrete+  trargMax = Concrete . targMaxR . unConcrete+  {-# INLINE triota #-}+  triota n = trfromIntegral $ Concrete $ Nested.riota @Int n+  {-# INLINE tsfloor #-}+  tsfloor = Concrete . liftVS (V.map floor) . unConcrete+  {-# INLINE tsfromIntegral #-}+  tsfromIntegral @r1 @r2 a =+    let cast :: (GoodScalar r1', Integral r1', NumScalar r2')+             => Concrete (TKS sh r1') -> Concrete (TKS sh r2')+        {-# INLINE cast #-}+        cast = Concrete . liftVS (V.map fromIntegral) . unConcrete+    in case typeRep @r1 of+        Is @Int -> case typeRep @r2 of+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int8 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int16 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int32 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int64 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @CInt -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          _ -> cast a+        _ -> cast a+  {-# INLINE tscast #-}+  tscast @r1 @r2 a =+    let cast :: ( Differentiable r1', NumScalar r1'+                , Differentiable r2', NumScalar r2' )+             => Concrete (TKS sh r1') -> Concrete (TKS sh r2')+        {-# INLINE cast #-}+        cast = Concrete . liftVS (V.map realToFrac) . unConcrete+    in case typeRep @r1 of+      Is @Double -> case typeRep @r2 of+        Is @Float -> cast @Double @Float a+        _ -> cast a+      Is @Float -> case typeRep @r2 of+        Is @Double -> cast @Float @Double a+        _ -> cast a+      _ -> cast a+  tsargMin = Concrete . targMinS . unConcrete+  tsargMax = Concrete . targMaxS . unConcrete+  {-# INLINE tsiota #-}+  tsiota @n = tsfromIntegral $ Concrete $ Nested.siota @Int (SNat @n)+  {-# INLINE txfloor #-}+  txfloor = Concrete . liftVX (V.map floor) . unConcrete+  {-# INLINE txfromIntegral #-}+  txfromIntegral @r1 @r2 a =+    let cast :: (GoodScalar r1', Integral r1', NumScalar r2')+             => Concrete (TKX sh r1') -> Concrete (TKX sh r2')+        {-# INLINE cast #-}+        cast = Concrete . liftVX (V.map fromIntegral) . unConcrete+    in case typeRep @r1 of+        Is @Int -> case typeRep @r2 of+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int8 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int16 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int32 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int64 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @Int64 -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @CInt -> cast a+          _ -> cast a+        Is @CInt -> case typeRep @r2 of+          Is @Int -> cast a+          Is @Double -> cast a+          Is @Float -> cast a+          Is @Int8 -> cast a+          Is @Int16 -> cast a+          Is @Int32 -> cast a+          Is @Int64 -> cast a+          _ -> cast a+        _ -> cast a+  {-# INLINE txcast #-}+  txcast @r1 @r2 a =+    let cast :: ( Differentiable r1', NumScalar r1'+                , Differentiable r2', NumScalar r2' )+             => Concrete (TKX sh r1') -> Concrete (TKX sh r2')+        {-# INLINE cast #-}+        cast = Concrete . liftVX (V.map realToFrac) . unConcrete+    in case typeRep @r1 of+      Is @Double -> case typeRep @r2 of+        Is @Float -> cast @Double @Float a+        _ -> cast a+      Is @Float -> case typeRep @r2 of+        Is @Double -> cast @Float @Double a+        _ -> cast a+      _ -> cast a+  txargMin = Concrete . targMinX . unConcrete+  txargMax = Concrete . targMaxX . unConcrete+  {-# INLINE txiota #-}+  txiota @n = txfromIntegral $ Concrete $ Nested.miota @Int (SNat @n)+  {-# INLINE trappend #-}+  trappend @_ @x u v | Dict <- eltDictRep (knownSTK @x) =+    Concrete $ Nested.rappend (unConcrete u) (unConcrete v)+  {-# INLINE trconcat #-}+  trconcat @_ @x l | Dict <- eltDictRep (knownSTK @x) =+    Concrete $ Nested.rconcat (fmapUnConcrete l)+  {-# INLINE trslice #-}+  trslice @_ @x i n | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.rslice i n . unConcrete+  {-# INLINE trreverse #-}+  trreverse @_ @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.rrev1 . unConcrete+  {-# INLINE trtranspose #-}+  trtranspose @_ @x perm | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.rtranspose perm . unConcrete+  {-# INLINE trreshape #-}+  trreshape @_ @_ @x sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.rreshape sh . unConcrete+  {-# INLINE tsappend #-}+  tsappend @_ @_ @_ @x u v | Dict <- eltDictRep (knownSTK @x) =+    Concrete $ Nested.sappend (unConcrete u) (unConcrete v)+  {-# INLINE tsslice #-}+  tsslice @_ @_ @_ @_ @x i n _ | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.sslice i n . unConcrete+  {-# INLINE tsreverse #-}+  tsreverse @_ @_ @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.srev1 . unConcrete+  {-# INLINE tstranspose #-}+  tstranspose @_ @_ @x perm | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.stranspose perm . unConcrete+  {-# INLINE tsreshape #-}+  tsreshape @_ @_ @x sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.sreshape sh . unConcrete+  {-# INLINE txappend #-}+  txappend @_ @_ @_ @x u v | Dict <- eltDictRep (knownSTK @x) =+    Concrete $ Nested.mappend (unConcrete u) (unConcrete v)+  {-# INLINE txconcat #-}+  txconcat @_ @x l | Dict <- eltDictRep (knownSTK @x) =+    Concrete $ Nested.mconcat (fmapUnConcrete l)+  {-# INLINE txslice #-}+  txslice @_ @_ @_ @_ @x i n k | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mslice i n k . unConcrete+  {-# INLINE txreverse #-}+  txreverse @_ @_ @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mrev1 . unConcrete+  {-# INLINE txtranspose #-}+  txtranspose @_ @_ @x perm | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mtranspose perm . unConcrete+  {-# INLINE txreshape #-}+  txreshape @_ @_ @x sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mreshape sh . unConcrete+  {-# INLINE tkbuild1 #-}+  tkbuild1 = tbuild1K+  {-# INLINE tkbuild #-}+  tkbuild @sh = tbuildK (knownShS @sh)+  {-# INLINE trbuild1 #-}+  trbuild1 @n @x k f =+    let g :: Int -> RepConcrete (TKR2 n x)+        g i = unConcrete $ f (Concrete i)+    in case knownSTK @x of+      STKScalar | SZ <- SNat @n ->+        Concrete $ Nested.rfromVector (k :$: ZSR)+        $ VS.generate k (Nested.runScalar . g)+      _ | Dict <- eltDictRep (knownSTK @x) -> case k of+        0 ->+          Concrete $ Nested.runNest Nested.remptyArray+        _ ->+          Concrete $ Nested.rfromListOuterN k $ NonEmpty.fromList+          $ map g [0 .. k - 1]+  {-# INLINE trbuild #-}+  trbuild @_ @n @x shm f =+    let g ix = unConcrete $ f (fmapConcrete ix)+    in case knownSTK @x of+      STKScalar | SZ <- SNat @n ->+        Concrete $ Nested.rgeneratePrim shm (Nested.runScalar . g)+      _ | Dict <- eltDictRep (knownSTK @x) ->+        Concrete $ Nested.runNest $ Nested.rgenerate shm g+  {-# INLINE trmap0N #-}+  trmap0N f t = Concrete $ tmap0NR (unConcrete . f . Concrete) (unConcrete t)+  {-# INLINE trzipWith0N #-}+  trzipWith0N f t u =+    Concrete+    $ tzipWith0NR (\v w -> unConcrete $ f (Concrete v) (Concrete w))+                  (unConcrete t) (unConcrete u)+  {-# INLINE tsbuild1 #-}+  tsbuild1 @_ @sh  = tbuild1S (knownShS @sh)+  {-# INLINE tsbuild #-}+  tsbuild @shm @shn  = tbuildS (knownShS @shm) (knownShS @shn)+  {-# INLINE tsmap0N #-}+  tsmap0N f v = Concrete $ tmap0NS (unConcrete . f . Concrete) (unConcrete v)+  {-# INLINE tszipWith0N #-}+  tszipWith0N f t u =+    Concrete+    $ tzipWith0NS (\v w -> unConcrete $ f (Concrete v) (Concrete w))+                  (unConcrete t) (unConcrete u)+  {-# INLINE txbuild1 #-}+  txbuild1 @k @sh @x f =+    let g :: Int -> RepConcrete (TKX2 sh x)+        g i = unConcrete $ f (Concrete i)+    in case knownSTK @x of+      STKScalar | ZKX <- knownShX @sh ->+        Concrete $ Nested.mfromVector (SKnown SNat :$% ZSX)+        $ VS.generate (valueOf @k) (Nested.munScalar . g)+          -- this is somewhat faster and not much more complex than if+          -- done with mgeneratePrim+      _ | Dict <- eltDictRep (knownSTK @x) -> case SNat @k of+        SZ ->+          Concrete $ Nested.munNest $ Nested.memptyArray ZSX+        _ ->+          Concrete $ Nested.mfromListOuterSN SNat $ NonEmpty.fromList+          $ map g [0 .. valueOf @k - 1]+  {-# INLINE txbuild #-}+  txbuild @shm @shn @x shm f =+    let g ix = unConcrete $ f (fmapConcrete ix)+    in case knownSTK @x of+      STKScalar | ZKX <- knownShX @shn+                , Refl <- lemAppNil @shm ->+        Concrete $ Nested.mgeneratePrim shm (Nested.munScalar . g)+      _ | Dict <- eltDictRep (knownSTK @x) ->+        Concrete $ Nested.munNest $ Nested.mgenerate shm g+  {-# INLINE tmapAccumLDer #-}+  tmapAccumLDer _ k _ bftk eftk (ConcreteFun f) _df _rf =+    tmapAccumLC k bftk eftk f+  {-# INLINE tapply #-}+  tapply (ConcreteFun f) = Concrete . f . unConcrete+  {-# INLINE tlambda #-}+  tlambda _ f = ConcreteFun $ unConcrete . unHFun f . Concrete+  -- The code for tvjp and tjvp in this instance is similar as for the+  -- ADVal ranked instance, because the type family instance is the same.+  {-# INLINE tgrad #-}+  tgrad @_ @r xftk h | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =+    ConcreteFun+    $ unConcrete . snd . crevOnParams Nothing (unHFun h) xftk . Concrete+  {-# INLINE tvjp #-}+  tvjp xftk h = ConcreteFun $ \db_a ->+    unConcrete $ snd+    $ crevOnParamsDt (Concrete $ fst db_a) (unHFun h) xftk (Concrete $ snd db_a)+  {-# INLINE tjvp #-}+  tjvp xftk h = ConcreteFun $ \da_a ->+    unConcrete $ snd+    $ cfwdOnParams xftk (Concrete $ snd da_a) (unHFun h) (Concrete $ fst da_a)+  tprimalPart = id+  {-# INLINE tdualPart #-}+  tdualPart stk t = DummyDualTarget (tftk stk t)+  tplainPart = id+  tfromPrimal _ t = t+  {-# INLINE tfromDual #-}+  tfromDual (DummyDualTarget ftk) = tdefTarget ftk+  tfromPlain _ t = t+  tScale _ _ t = t+  taddTarget = addTarget+  tmultTarget = multTarget+  tsum0Target = sum0Target+  tdot0Target = dot0Target+  xmcast @x sh2 | Dict <- eltDictRep (knownSTK @x) =+    Concrete . mcast sh2 . unConcrete++instance ConvertTensor Concrete where+  {-# INLINE tconvert #-}+  tconvert c astk a | Dict <- eltDictRep astk+                    , Dict <- eltDictRep (convertSTK c astk) =+    Concrete $ Nested.convert (interpretTKConversion c) (unConcrete a)++  {-# INLINE kfromR #-}+  kfromR = Concrete . Nested.runScalar . unConcrete+  {-# INLINE kfromS #-}+  kfromS = Concrete . Nested.sunScalar . unConcrete+  {-# INLINE kfromX #-}+  kfromX = Concrete . Nested.munScalar . unConcrete+  {-# INLINE rfromK #-}+  rfromK = Concrete . Nested.rscalar . unConcrete+  {-# INLINE rfromS #-}+  rfromS @_ @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.stoRanked . unConcrete+  {-# INLINE rfromX #-}+  rfromX @_ @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mtoRanked . unConcrete+  {-# INLINE sfromK #-}+  sfromK = Concrete . Nested.sscalar . unConcrete+  {-# INLINE sfromR #-}+  sfromR @_ @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . flip Nested.rcastToShaped knownShS . unConcrete+  {-# INLINE sfromX #-}+  sfromX @_ @_ @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mcastToShaped knownShS . unConcrete+  {-# INLINE xfromK #-}+  xfromK = Concrete . Nested.mscalar . unConcrete+  {-# INLINE xfromR #-}+  xfromR @sh @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.rcastToMixed (knownShX @sh) . unConcrete+  {-# INLINE xfromS #-}+  xfromS @_ @sh' @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.scastToMixed (knownShX @sh') . unConcrete++  {-# INLINE rzip #-}+  rzip @y @z (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)+                               , Dict <- eltDictRep (knownSTK @z) =+    Concrete $ Nested.rzip a b+  {-# INLINE runzip #-}+  runzip a = let (!a1, !a2) = Nested.runzip $ unConcrete a+             in Concrete (a1, a2)+  {-# INLINE szip #-}+  szip @y @z (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)+                               , Dict <- eltDictRep (knownSTK @z) =+    Concrete $ Nested.szip a b+  {-# INLINE sunzip #-}+  sunzip a = let (!a1, !a2) = Nested.sunzip $ unConcrete a+             in Concrete (a1, a2)+  {-# INLINE xzip #-}+  xzip @y @z (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)+                               , Dict <- eltDictRep (knownSTK @z) =+    Concrete $ Nested.mzip a b+  {-# INLINE xunzip #-}+  xunzip a = let (!a1, !a2) = Nested.munzip $ unConcrete a+             in Concrete (a1, a2)++  {-# INLINE xnestR #-}+  xnestR @sh1 @m @x sh | Dict <- eltDictRep (knownSTK @x)+                       , Refl <- lemRankReplicate (SNat @m) =+    Concrete+    . Nested.convert+        @(Nested.Mixed sh1 (Nested.Mixed (Replicate m Nothing) (RepConcrete x)))+        (Nested.ConvXX Nested.ConvXR)+    . Nested.mnest sh+    . unConcrete+  {-# INLINE xnestS #-}+  xnestS @sh1 @sh2 @x sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete+    . Nested.convert+        @(Nested.Mixed sh1 (Nested.Mixed (MapJust sh2) (RepConcrete x)))+        (Nested.ConvXX Nested.ConvXS)+    . Nested.mnest sh+    . unConcrete+  {-# INLINE xnest #-}+  xnest @_ @_ @x sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mnest sh . unConcrete+  {-# INLINE xunNestR #-}+  xunNestR @sh1 @m @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete+    . Nested.munNest+    . Nested.convert+        @(Nested.Mixed sh1 (Nested.Ranked m (RepConcrete x)))+        (Nested.ConvXX Nested.ConvRX)+    . unConcrete+  {-# INLINE xunNestS #-}+  xunNestS @sh1 @sh2 @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete+    . Nested.munNest+    . Nested.convert+        @(Nested.Mixed sh1 (Nested.Shaped sh2 (RepConcrete x)))+        (Nested.ConvXX Nested.ConvSX)+    . unConcrete+  {-# INLINE xunNest #-}+  xunNest = Concrete . Nested.munNest . unConcrete++  tpairConv = tpair+  tunpairConv = tunpair++interpretTKConversion :: TKConversion a b+                      -> Nested.Conversion (RepConcrete a) (RepConcrete b)+interpretTKConversion c0 = case c0 of+  ConvId -> Nested.ConvId+  ConvCmp c1 c2 -> Nested.ConvCmp (interpretTKConversion c1)+                                  (interpretTKConversion c2)+  ConvRX -> Nested.ConvRX+  ConvSX -> Nested.ConvSX+  ConvXR stk | Dict <- eltDictRep stk -> Nested.ConvXR+  ConvXS -> Nested.ConvXS+  ConvXS' (FTKS sh' ftk) | Dict <- eltDictRep (ftkToSTK ftk) ->+    Nested.ConvXS' sh'+  ConvXX' (FTKX shx ftk) | Dict <- eltDictRep (ftkToSTK ftk) ->+    Nested.ConvXX' (ssxFromShX shx)+  ConvRR c -> Nested.ConvRR (interpretTKConversion c)+  ConvSS c -> Nested.ConvSS (interpretTKConversion c)+  ConvXX c -> Nested.ConvXX (interpretTKConversion c)+  ConvT2 c1 c2 ->+    Nested.ConvT2 (interpretTKConversion c1) (interpretTKConversion c2)+  Conv0X stk | Dict <- eltDictRep stk -> Nested.Conv0X+  ConvX0 -> Nested.ConvX0+  ConvNest (STKX sh x) | Dict <- eltDictRep x -> Nested.ConvNest sh+  ConvUnnest -> Nested.ConvUnnest+  ConvZip stk1 stk2 | Dict <- eltDictRep stk1+                    , Dict <- eltDictRep stk2 -> Nested.ConvZip+  ConvUnzip stk1 stk2 | Dict <- eltDictRep stk1+                      , Dict <- eltDictRep stk2 -> Nested.ConvUnzip+++-- * Misc++fmapConcrete :: Coercible (f (RepConcrete y)) (f (Concrete y))+               => f (RepConcrete y) -> f (Concrete y)+fmapConcrete = coerce++fmapUnConcrete :: Coercible (f (Concrete y)) (f (RepConcrete y))+               => f (Concrete y) -> f (RepConcrete y)+fmapUnConcrete = coerce++-- Depite the warning, the pattern match is exhaustive.+ixInBoundsR :: IShR n -> IxROf Concrete n -> Bool+ixInBoundsR ZSR ZIR = True+ixInBoundsR (n :$: sh) (Concrete i :.: ix) =+  0 <= i && i < n && ixInBoundsR sh ix++ixInBoundsS :: ShS sh -> IxSOf Concrete sh -> Bool+ixInBoundsS ZSS ZIS = True+ixInBoundsS ((fromSNat' -> n) :$$ sh) (Concrete i :.$ ix) =+  0 <= i && i < n && ixInBoundsS sh ix++ixInBoundsX :: IShX sh -> IxXOf Concrete sh -> Bool+ixInBoundsX ZSX ZIX = True+ixInBoundsX ((fromSMayNat' -> n) :$% sh) (Concrete i :.% ix) =+  0 <= i && i < n && ixInBoundsX sh ix++-- Depite the warning, the pattern match is exhaustive.+ixrToLinearMaybe :: IShR n -> IxROf Concrete n -> Maybe Int+ixrToLinearMaybe = \sh ix -> goR sh ix 0+  where+    goR :: IShR n -> IxROf Concrete n -> Int -> Maybe Int+    goR ZSR ZIR !a = Just a+    goR (n :$: sh) (Concrete i :.: ix) a =+      if 0 <= i && i < n then goR sh ix (n * a + i) else Nothing++-- This would be shorter, but a bit more expensive:+--   ixxToLinearMaybe (ixxFromIxS ix)+ixsToLinearMaybe :: ShS sh -> IxSOf Concrete sh -> Maybe Int+ixsToLinearMaybe = \sh ix -> goS sh ix 0+  where+    goS :: ShS sh -> IxSOf Concrete sh -> Int -> Maybe Int+    goS ZSS ZIS !a = Just a+    goS ((fromSNat' -> n) :$$ sh) (Concrete i :.$ ix) a =+      if 0 <= i && i < n then goS sh ix (n * a + i) else Nothing++ixxToLinearMaybe :: IShX sh -> IxXOf Concrete sh -> Maybe Int+ixxToLinearMaybe = \sh ix -> goX sh ix 0+  where+    goX :: IShX sh -> IxXOf Concrete sh -> Int -> Maybe Int+    goX ZSX ZIX !a = Just a+    goX ((fromSMayNat' -> n) :$% sh) (Concrete i :.% ix) a =+      if 0 <= i && i < n then goX sh ix (n * a + i) else Nothing++tfromList :: forall y k.+             SNat k -> SingletonTK y -> NonEmpty (Concrete y)+          -> Concrete (BuildTensorKind k y)+{-# INLINE tfromList #-}+tfromList snat@SNat stk l = case stk of+  STKScalar ->+    Concrete $ Nested.sfromList1Prim snat $ fmapUnConcrete $ NonEmpty.toList l+  STKR SNat x | Dict <- eltDictRep x ->+    Concrete $ Nested.rfromListOuterN (fromSNat' snat) $ fmapUnConcrete l+  STKS _sh x | Dict <- eltDictRep x ->+    Concrete $ Nested.sfromListOuter snat $ fmapUnConcrete l+  STKX _sh x | Dict <- eltDictRep x ->+    Concrete $ Nested.mfromListOuterSN snat $ fmapUnConcrete l+  STKProduct stk1 stk2 ->+    let (l1, l2) = Functor.unzip $ coerce l  -- NonEmpty.map tunpair l+        a1 = tfromList snat stk1 l1+        a2 = tfromList snat stk2 l2+    in -- This processes a list of trivial primitive elements first,+       -- which does not force the list, which prevents forcing+       -- the other (tuple of) list prematurely, which makes streaming+       -- possible (sometimes).+       case stk2 of+         STKScalar @r2 | Just Refl <- testEquality (typeRep @r2)+                                                   (typeRep @Z1) ->+           a2 `seq` tpair a1 a2+         _ -> tpair a1 a2+           -- TODO: instead construct both (tuples of) tensors at once,+           -- element by element to stream even when more than one+           -- list of nontrivial elements is present; use mvecsWrite?++tmapAccumLC+  :: forall k accy by ey.+     SNat k+  -> FullShapeTK by+  -> FullShapeTK ey+  -> (RepConcrete (TKProduct accy ey) -> RepConcrete (TKProduct accy by))+  -> Concrete accy+  -> Concrete (BuildTensorKind k ey)+  -> Concrete (TKProduct accy (BuildTensorKind k by))+{-# INLINE tmapAccumLC #-}+tmapAccumLC k (FTKScalar @z1) eftk f acc0 es+  | Just Refl <- testEquality (typeRep @z1) (typeRep @Z1) =+    let h :: Concrete accy -> Concrete ey -> Concrete accy+        h !acc !e = Concrete $ fst $ f (unConcrete acc, unConcrete e)+        xout = foldl' h acc0 (tunravelToListShare k (ftkToSTK eftk) es)+        lout2 = tsreplicate0N (k :$$ ZSS) (Concrete Z1)+    in tpair xout lout2+tmapAccumLC k bftk eftk f !acc0 !es =+  let (xout, lout) =+        mapAccumL' (curry $ coerce f) acc0+                   (tunravelToListShare k (ftkToSTK eftk) es)+  in case NonEmpty.nonEmpty lout of+    Just nl ->+      -- The bang is needed to stream the list, which may still partially+      -- fail if the output is a tuple of non-Z1 lists and then only+      -- the first of them is streamed. Such tuples are common+      -- in gradients of non-fold mapAccums.+      let !lout2 = tfromList k (ftkToSTK bftk) nl+      in tpair xout lout2+    Nothing -> tpair xout (tdefTarget (buildFTK k bftk))++-- The explicit dictionary is needed to trick GHC into specializing f at types+-- Int, Double, etc. insteasd of at type r, to simpify away the dictionaries+-- emerging from the constraints in the signature of f.+--+-- Despite what GHC says, TKAllNum (TKScalar r) is not redundant,+-- because it ensures the error case can't appear.+contFromTKAllNum :: forall r a. (Typeable r, TKAllNum (TKScalar r))+                 => (Dict0 (Num r, Nested.NumElt r, GoodScalar r) -> a) -> a+{-# INLINE contFromTKAllNum #-}  -- needed for the specialization hack+contFromTKAllNum f = case typeRep @r of+  Is @Int -> f Dict0+  Is @Double -> f Dict0+  Is @Float -> f Dict0+  Is @Z1 -> f Dict0+  Is @Int64 -> f Dict0+  Is @Int32-> f Dict0+  Is @Int16 -> f Dict0+  Is @Int8 -> f Dict0+  Is @CInt -> f Dict0+  _ -> error "contFromTKAllNum: impossible type"++-- See above. The list comes from ox-arrays at [PRIMITIVE ELEMENT TYPES LIST].+contFromTypeable :: forall r a. Typeable r+                 => (Dict GoodScalar r -> a) -> a+{-# INLINE contFromTypeable #-}  -- needed for the specialization hack+contFromTypeable f = case typeRep @r of+  Is @Int -> f Dict+  Is @Double -> f Dict+  Is @Float -> f Dict+  Is @Z1 -> f Dict+  Is @Int64 -> f Dict+  Is @Int32 -> f Dict+  Is @Int16-> f Dict+  Is @Int8 -> f Dict+  Is @CInt -> f Dict+  Is @Bool -> f Dict+  Is @() -> f Dict+  _ -> error "contFromTypeable: unexpected type"++targMinK+  :: forall n r. (Nested.PrimElt r, Nested.NumElt r)+  => Nested.Shaped '[n] r -> Int+{-# INLINE targMinK #-}+targMinK = ixsHead . Nested.sminIndexPrim++targMaxK+  :: forall n r. (Nested.PrimElt r, Nested.NumElt r)+  => Nested.Shaped '[n] r -> Int+{-# INLINE targMaxK #-}+targMaxK = ixsHead . Nested.smaxIndexPrim++tbuild1K :: (KnownNat k, GoodScalar r)+         => (IntOf Concrete -> Concrete (TKScalar r))+         -> Concrete (TKS '[k] r)+{-# INLINE tbuild1K #-}+tbuild1K @k f =+  let g i = unConcrete $ f (Concrete i)+  in Concrete $ Nested.sfromVector (SNat :$$ ZSS)+     $ VS.generate (valueOf @k) g++tbuildK :: GoodScalar r+        => ShS sh -> (IxSOf Concrete sh -> Concrete (TKScalar r))+        -> Concrete (TKS sh r)+{-# INLINE tbuildK #-}+tbuildK sh f =+  let g ix = unConcrete $ f (fmapConcrete ix)+  in Concrete $ Nested.sgeneratePrim sh g+++-- * Ranked internal definitions++{-+tdot0R t u = OR.toVector t LA.<.> OR.toVector u+  -- TODO: if offset 0 and same strides, use toUnorderedVectorT+  -- TODO: if either has length 1 values, it may or may not be faster to do+  -- tsum0R (t * u) -}++-- We could generalize by unwinding and only then doing the PrimElt things,+-- but we'd need a type family that says "replace this underlying scalars+-- by this one", which makes things too complicated.+--+-- We could also expose `liftVR` in the user API, but in addition+-- to the main function argument such as floor or cast, it'd need the function's+-- derivative, just as with mapAccums. Maybe it's better to generalize even more+-- and permit arbitrary extra ops if given their derivatives.+liftVR+  :: (Nested.PrimElt r1, Nested.PrimElt r2)+  => (VS.Vector r1 -> VS.Vector r2)+  -> Nested.Ranked n r1 -> Nested.Ranked n r2+{-# INLINE liftVR #-}+liftVR f = Ranked.liftRanked1 (Mixed.mliftNumElt1 (`liftVEltwise1` f))++manyHotNR :: forall m n x. (KnownNat m, KnownSTK x)+          => FullShapeTK (TKR2 (m + n) x)+          -> [(Int, Concrete (TKR2 n x))]+          -> Concrete (TKR2 (m + n) x)+{-# INLINE manyHotNR #-}+manyHotNR (FTKR shRanked x) upd | Dict <- eltDictRep (knownSTK @x)+                                , Refl <- lemRankReplicate (Proxy @(m + n))+                                , Refl <- lemReplicatePlusApp+                                            (SNat @m)+                                            (Proxy @n)+                                            (Proxy @(Nothing @Nat)) = runST $ do+  let zero = unConcrete $ tdefTarget x+      sh = shxFromShR shRanked+  vecs <- Mixed.mvecsReplicate sh zero+    -- this avoids the slow case of mvecsReplicate if x is TKScalar+  forM_ upd $ \(ix, Concrete v) ->+    Mixed.mvecsWritePartialLinear+      (Proxy @(Replicate m Nothing)) ix (Nested.rtoMixed v) vecs+  Concrete . Nested.mtoRanked <$> Mixed.mvecsUnsafeFreeze sh vecs++tindexZR :: forall m n x. (KnownNat n, KnownSTK x)+         => Concrete (TKR2 (m + n) x) -> IxROf Concrete m+         -> Concrete (TKR2 n x)+{-# INLINE tindexZR #-}+tindexZR = case knownSTK @x of+  STKScalar @r -> contFromTypeable @r tindexZRDict+  _ -> tindexZRSlow++tindexZRSlow :: forall m n x. KnownSTK x+             => Concrete (TKR2 (m + n) x) -> IxROf Concrete m+             -> Concrete (TKR2 n x)+{-# NOINLINE tindexZRSlow #-}  -- the rare slow case+tindexZRSlow (Concrete v) ix | Dict <- eltDictRep (knownSTK @x) =+  if ixInBoundsR (shrTakeIx @m @n Proxy ix (Nested.rshape v))  ix+  then Concrete $ Nested.rindexPartial v (fmapUnConcrete ix)+  else case tftkG (STKR (Nested.rrank v) (knownSTK @x)) v of+         FTKR _ x -> tdefTarget (FTKR (shrDropIx @m @n ix (Nested.rshape v)) x)++-- See the comment about tscatterZSDict.+tindexZRDict :: forall m n r. KnownNat n+             => Dict GoodScalar r+             -> Concrete (TKR (m + n) r)+             -> IxROf Concrete m+             -> Concrete (TKR n r)+{-# INLINE tindexZRDict #-}+tindexZRDict Dict = tindexZRScalar++tindexZRScalar :: forall m n r. (KnownNat n, GoodScalar r)+               => Concrete (TKR (m + n) r) -> IxROf Concrete m+               -> Concrete (TKR n r)+{-# INLINE tindexZRScalar #-}+tindexZRScalar (Concrete v) ix = case SNat @n of+  SZ ->  -- an optimized common case+    rfromK $ Concrete v `tindex0RImpl` ix+  _ ->+    Concrete+    $ if ixInBoundsR (shrTakeIx @m @n Proxy ix (Nested.rshape v)) ix+      then Nested.rindexPartial v (fmapUnConcrete ix)+      else Nested.rreplicatePrim (shrDropIx @m @n ix (Nested.rshape v)) def++tindex0R :: forall m r. GoodScalar r+         => Concrete (TKR m r) -> IxROf Concrete m+         -> Concrete (TKScalar r)+{-# INLINE tindex0R #-}+tindex0R = contFromTypeable @r tindex0RDict++tindex0RImpl :: forall m r. GoodScalar r+             => Concrete (TKR m r) -> IxROf Concrete m+             -> Concrete (TKScalar r)+{-# INLINE tindex0RImpl #-}+tindex0RImpl (Concrete v) ix =+  Concrete+  $ if ixInBoundsR (Nested.rshape v) ix+    then Nested.rindex v (fmapUnConcrete ix)+    else def++tindex0RDict :: forall m r.+                Dict GoodScalar r+             -> Concrete (TKR m r)+             -> IxROf Concrete m+             -> Concrete (TKScalar r)+{-# INLINE tindex0RDict #-}+tindex0RDict Dict = tindex0RImpl++toneHotR :: forall m n x. (KnownNat m, KnownNat n, KnownSTK x)+         => IShR m -> Concrete (TKR2 n x) -> IxROf Concrete m+         -> Concrete (TKR2 (m + n) x)+{-# INLINE toneHotR #-}+toneHotR sh1 !v ix = case tftk knownSTK v of+  FTKR sh2 x ->+    let ftk = FTKR (sh1 `shrAppend` sh2) x+    in case ixrToLinearMaybe sh1 ix of+         Nothing -> tdefTarget ftk+         Just i2 -> manyHotNR ftk [(i2, v)]++-- Note how ix being in bounds is checked. The semantics of the operation+-- permits index out of bounds and then no tensor is added at such an index.+tscatterZR+  :: forall m n p x.+     (KnownNat m, KnownNat n, KnownNat p, TKAllNum x, KnownSTK x)+  => IShR p -> Concrete (TKR2 (m + n) x)+  -> (IxROf Concrete m -> IxROf Concrete p)+  -> Concrete (TKR2 (p + n) x)+{-# INLINE tscatterZR #-}+tscatterZR = case knownSTK @x of+  STKScalar @r ->  -- we don't use full dictionary from FTKScalar+    contFromTKAllNum @r tscatterZRDict+      -- Optimized: using (+) instead of taddTarget+  _ -> tscatterZRSlow++tscatterZRSlow+  :: forall m n p x.+     (KnownNat m, KnownNat n, KnownNat p, TKAllNum x, KnownSTK x)+  => IShR p -> Concrete (TKR2 (m + n) x)+  -> (IxROf Concrete m -> IxROf Concrete p)+  -> Concrete (TKR2 (p + n) x)+{-# NOINLINE tscatterZRSlow #-}  -- the rare slow case+tscatterZRSlow shp (Concrete v) f | Dict <- eltDictRep (knownSTK @x) =+  case tftkG (STKR SNat (knownSTK @x)) v of+    FTKR sht x ->+      let shm = shrTake @m sht+          shn = shrDrop @m sht+          ftk = FTKR (shp `shrAppend` shn) x+          g :: IIxR m+            -> IM.IntMap (Concrete (TKR2 n x))+            -> IM.IntMap (Concrete (TKR2 n x))+          g ix =+            let ix2 = f $ fmapConcrete ix+            in case ixrToLinearMaybe shp ix2 of+              Nothing -> id+              Just i2 ->+                IM.insertWith (taddTarget knownSTK) i2+                  (Concrete $ Nested.rindexPartial v ix)+          !ivs = foldl' (flip g) IM.empty (shrEnum shm)+      in manyHotNR ftk $ IM.assocs ivs++-- See the comment about tscatterZSDict.+tscatterZRDict+  :: forall m n p r. (KnownNat m, KnownNat n)+  => Dict0 (Num r, Nested.NumElt r, GoodScalar r)+  -> IShR p -> Concrete (TKR (m + n) r)+  -> (IxROf Concrete m -> IxROf Concrete p)+  -> Concrete (TKR (p + n) r)+{-# INLINE [1] tscatterZRDict #-}+tscatterZRDict Dict0 = tscatterZRScalar++tscatterZRScalar+  :: forall m n p r.+     (KnownNat m, KnownNat n, GoodScalar r, Num r, Nested.NumElt r)+  => IShR p -> Concrete (TKR (m + n) r)+  -> (IxROf Concrete m -> IxROf Concrete p)+  -> Concrete (TKR (p + n) r)+{-# INLINE tscatterZRScalar #-}+tscatterZRScalar shp (Concrete v) f =+  let sht = Nested.rshape v+      shm = shrTake @m sht+  in case SNat @n of+       SZ -> runST $ do+         -- Optimized: using Nested.rindex instesad of Nested.rindexPartial.+         vec <- VSM.replicate (shrSize shp) 0+         forM_ (shrEnum shm) $ \ix -> do+           let ix2 = f $ fmapConcrete ix+           case ixrToLinearMaybe shp ix2 of+             Nothing -> return ()+             Just i2 -> do+               let u = Nested.rindex v ix+               u2 <- VSM.read vec i2+               VSM.write vec i2 (u + u2)+         Concrete . Nested.rfromVector shp <$> VS.unsafeFreeze vec+       _ -> runST $ do+         let g :: IIxR m+               -> IM.IntMap (Nested.Ranked n r)+               -> IM.IntMap (Nested.Ranked n r)+             g ix =+               let ix2 = f $ fmapConcrete ix+               in case ixrToLinearMaybe shp ix2 of+                 Nothing -> id+                 Just i2 ->+                   IM.insertWith (+) i2 (Nested.rindexPartial v ix)+             ivs = foldl' (flip g) IM.empty (shrEnum shm)+             shn = shrDrop @m sht+             shnSize = shrSize shn+         vec <- VSM.replicate (shrSize shp * shnSize) 0+         forM_ (IM.assocs ivs) $ \(i, u) ->+           VS.copy (VSM.slice (i * shnSize) shnSize vec) (Nested.rtoVector u)+             -- TODO: use toVectorList instead+         Concrete . Nested.rfromVector (shp `shrAppend` shn)+           <$> VS.unsafeFreeze vec++-- The semantics of the operation permits index out of bounds+-- and the result of such indexing is def, which is 0.+tgatherZR+  :: forall m n p x. (KnownNat m, KnownNat n, KnownSTK x)+  => IShR m -> Concrete (TKR2 (p + n) x)+  -> (IxROf Concrete m -> IxROf Concrete p)+  -> Concrete (TKR2 (m + n) x)+{-# INLINE tgatherZR #-}+tgatherZR = case knownSTK @x of+  STKScalar @r ->  -- we don't use full dictionary from FTKScalar+    contFromTypeable @r tgatherZRDict+      -- Code gets specialized to a particular underlying scalar.+  _ -> tgatherZRSlow++tgatherZRSlow+  :: forall m n p x. (KnownNat m, KnownNat n, KnownSTK x)+  => IShR m -> Concrete (TKR2 (p + n) x)+  -> (IxROf Concrete m -> IxROf Concrete p)+  -> Concrete (TKR2 (m + n) x)+{-# NOINLINE tgatherZRSlow #-}+tgatherZRSlow shm !t f =+  trbuild shm (\ix -> t `tindexZRSlow` f ix)++tgatherZRDict+  :: forall m n p r. (KnownNat m, KnownNat n)+  => Dict GoodScalar r+  -> IShR m -> Concrete (TKR (p + n) r)+  -> (IxROf Concrete m -> IxROf Concrete p)+  -> Concrete (TKR (m + n) r)+{-# INLINE tgatherZRDict #-}+tgatherZRDict Dict = tgatherZRScalar++tgatherZRScalar+  :: forall m n p r. (KnownNat m, KnownNat n, GoodScalar r)+  => IShR m -> Concrete (TKR (p + n) r)+  -> (IxROf Concrete m -> IxROf Concrete p)+  -> Concrete (TKR (m + n) r)+{-# INLINE tgatherZRScalar #-}+tgatherZRScalar shm !t f = case SNat @n of+  SZ ->  -- an optimized common case+    let g ix = unConcrete $ t `tindex0RImpl` f (fmapConcrete ix)+    in Concrete $ Nested.rgeneratePrim shm g+  _ -> trbuild shm (\ix -> t `tindexZRScalar` f ix)++tgatherZ1R+  :: forall n p x. (KnownNat n, KnownNat p, KnownSTK x)+  => Int -> Concrete (TKR2 (p + n) x)+  -> (IntOf Concrete -> IxROf Concrete p)+  -> Concrete (TKR2 (1 + n) x)+{-# INLINE tgatherZ1R #-}+tgatherZ1R = case knownSTK @x of+  STKScalar @r ->  -- we don't use full dictionary from FTKScalar+    contFromTypeable @r tgatherZ1RDict+      -- Code gets specialized to a particular underlying scalar.+  _ -> tgatherZ1RSlow++tgatherZ1RSlow+  :: forall n p x. (KnownNat n, KnownSTK x)+  => Int -> Concrete (TKR2 (p + n) x)+  -> (IntOf Concrete -> IxROf Concrete p)+  -> Concrete (TKR2 (1 + n) x)+{-# NOINLINE tgatherZ1RSlow #-}+tgatherZ1RSlow k !t f =+  trbuild1 k (\ix -> t `tindexZRSlow` f ix)++tgatherZ1RDict+  :: forall n p r. (KnownNat n, KnownNat p)+  => Dict GoodScalar r+  -> Int -> Concrete (TKR (p + n) r)+  -> (IntOf Concrete -> IxROf Concrete p)+  -> Concrete (TKR (1 + n) r)+{-# INLINE tgatherZ1RDict #-}+tgatherZ1RDict Dict = tgatherZ1RScalar++tgatherZ1RScalar+  :: forall n p r. (KnownNat n, KnownNat p, GoodScalar r)+  => Int -> Concrete (TKR (p + n) r)+  -> (IntOf Concrete -> IxROf Concrete p)+  -> Concrete (TKR (1 + n) r)+{-# INLINE tgatherZ1RScalar #-}+tgatherZ1RScalar k !t f = case SNat @n of+  SZ ->  -- an optimized common case+    let shm = k :$: shrDrop (rshape t)+        g i = unConcrete $ t `tindex0RImpl` f (Concrete i)+    in Concrete $ Nested.rfromVector shm $ VS.generate k g+  _ -> trbuild1 k (\ix -> t `tindexZRScalar` f ix)++targMinR+  :: forall r n. (Nested.PrimElt r, Nested.NumElt r)+  => Nested.Ranked (1 + n) r -> Nested.Ranked n Int+{-# INLINE targMinR #-}+targMinR v | SNat <- Nested.rrank v =+  let f :: Nested.Ranked 1 r -> Nested.Ranked 0 Int+      f = Nested.rscalar . ixrHead . Nested.rminIndexPrim+  in Nested.runNest $ Nested.rrerankPrim ZSR f (Nested.rnest SNat v)++targMaxR+  :: forall r n. (Nested.PrimElt r, Nested.NumElt r)+  => Nested.Ranked (1 + n) r -> Nested.Ranked n Int+{-# INLINE targMaxR #-}+targMaxR v | SNat <- Nested.rrank v =+  let f :: Nested.Ranked 1 r -> Nested.Ranked 0 Int+      f = Nested.rscalar . ixrHead . Nested.rmaxIndexPrim+  in Nested.runNest $ Nested.rrerankPrim ZSR f (Nested.rnest SNat v)++tmap0NR+  :: (Nested.PrimElt r1, Nested.PrimElt r)+  => (r1 -> r) -> Nested.Ranked n r1 -> Nested.Ranked n r+{-# INLINE tmap0NR #-}+tmap0NR f = Ranked.liftRanked1 (Mixed.mliftPrim f)++tzipWith0NR+  :: (Nested.PrimElt r, Nested.PrimElt r1, Nested.PrimElt r2)+  => (r1 -> r2 -> r) -> Nested.Ranked n r1 -> Nested.Ranked n r2+  -> Nested.Ranked n r+{-# INLINE tzipWith0NR #-}+tzipWith0NR f = Ranked.liftRanked2 (Mixed.mliftPrim2 f)+++-- * Shaped internal definitions++liftVS+  :: (Nested.PrimElt r1, Nested.PrimElt r)+  => (VS.Vector r1 -> VS.Vector r)+  -> Nested.Shaped sh r1 -> Nested.Shaped sh r+{-# INLINE liftVS #-}+liftVS f = Shaped.liftShaped1 (Mixed.mliftNumElt1 (`liftVEltwise1` f))++manyHotNS :: forall shn shp x.+             ShS shn -> ShS shp -> FullShapeTK x+          -> [(Int, Concrete (TKS2 shn x))]+          -> Concrete (TKS2 (shp ++ shn) x)+{-# INLINE manyHotNS #-}+manyHotNS shn shp x upd | Dict <- eltDictRep (ftkToSTK x)+                        , let shShaped = shp `shsAppend` shn+                        , Refl <- lemRankMapJust shShaped+                        , Refl <- lemMapJustApp shp (Proxy @shn) = runST $ do+  let zero = unConcrete $ tdefTarget x+      sh = shxFromShS shShaped+  vecs <- Mixed.mvecsReplicate sh zero+    -- this avoids the slow case of mvecsReplicate if x is TKScalar+  forM_ upd $ \(ix, Concrete v) ->+    Mixed.mvecsWritePartialLinear+      (Proxy @(MapJust shp)) ix (Nested.stoMixed v) vecs+  Concrete . Nested.mcastToShaped shShaped <$> Mixed.mvecsUnsafeFreeze sh vecs++-- Note that after vectorization, the index may not fit within+-- the type-level shape, which we catch in the @ixInBounds@+-- and return def, so it's fine. Similarly in gather and scatter.+tindexZS :: forall shm shn x. KnownSTK x+         => ShS shn -> Concrete (TKS2 (shm ++ shn) x) -> IxSOf Concrete shm+         -> Concrete (TKS2 shn x)+{-# INLINE tindexZS #-}+tindexZS = case knownSTK @x of+  STKScalar @r -> contFromTypeable @r tindexZSDict+  _ -> tindexZSSlow++tindexZSSlow :: forall shm shn x. KnownSTK x+             => ShS shn -> Concrete (TKS2 (shm ++ shn) x) -> IxSOf Concrete shm+             -> Concrete (TKS2 shn x)+{-# NOINLINE tindexZSSlow #-}  -- the rare slow case+tindexZSSlow shn (Concrete v) ix | Dict <- eltDictRep (knownSTK @x) =+  if ixInBoundsS (shsTakeIx @shm @shn Proxy ix (Nested.sshape v)) ix+  then Concrete $ Nested.sindexPartial v (fmapUnConcrete ix)+  else case tftkG (STKS (Nested.sshape v) (knownSTK @x)) v of+         FTKS _sh x -> tdefTarget (FTKS shn x)++tindexZSDict :: forall shm shn r.+                Dict GoodScalar r+             -> ShS shn -> Concrete (TKS (shm ++ shn) r)+             -> IxSOf Concrete shm+             -> Concrete (TKS shn r)+{-# INLINE tindexZSDict #-}+tindexZSDict Dict = tindexZSScalar++tindexZSScalar :: forall shm shn r. GoodScalar r+               => ShS shn -> Concrete (TKS (shm ++ shn) r) -> IxSOf Concrete shm+               -> Concrete (TKS shn r)+{-# INLINE tindexZSScalar #-}+tindexZSScalar shn (Concrete v) ix = case shn of+  ZSS | Refl <- lemAppNil @shm ->  -- an optimized common case+    sfromK $ tindex0SImpl (Concrete v) ix+      -- TODO: benchmark and possibly duplicate the Concrete instance+      -- to avoid this overhead in symbolic pipeline; also in tindexZRScalar+  _ ->+    Concrete+    $ if ixInBoundsS (shsTakeIx @shm @shn Proxy ix (Nested.sshape v)) ix+      then Nested.sindexPartial v (fmapUnConcrete ix)+      else Nested.sreplicatePrim shn def++tindex0S :: forall sh1 r. GoodScalar r+         => Concrete (TKS sh1 r) -> IxSOf Concrete sh1+         -> Concrete (TKScalar r)+{-# INLINE tindex0S #-}+tindex0S = contFromTypeable @r tindex0SDict++tindex0SImpl :: forall sh1 r. GoodScalar r+             => Concrete (TKS sh1 r) -> IxSOf Concrete sh1+             -> Concrete (TKScalar r)+{-# INLINE tindex0SImpl #-}+tindex0SImpl (Concrete v) ix =+  Concrete+  $ if ixInBoundsS (Nested.sshape v) ix+    then Nested.sindex v (fmapUnConcrete ix)+    else def++tindex0SDict :: forall sh1 r.+                Dict GoodScalar r+             -> Concrete (TKS sh1 r)+             -> IxSOf Concrete sh1+             -> Concrete (TKScalar r)+{-# INLINE tindex0SDict #-}+tindex0SDict Dict = tindex0SImpl++toneHotS :: forall shp shn x. KnownSTK x+         => ShS shn -> ShS shp -> Concrete (TKS2 shn x) -> IxSOf Concrete shp+         -> Concrete (TKS2 (shp ++ shn) x)+{-# INLINE toneHotS #-}+toneHotS shn shp !v ix = case tftk (STKS shn knownSTK) v of+  FTKS _ x ->+    let ftk = FTKS (shp `shsAppend` shn) x+    in case ixsToLinearMaybe shp ix of+         Nothing -> tdefTarget ftk+         Just i2 -> manyHotNS shn shp x [(i2, v)]++-- Note how ix being in bounds is checked. The semantics of the operation+-- permits index out of bounds and then no tensor is added at such an index.+tscatterZS+  :: forall shm shn shp x. (TKAllNum x, KnownSTK x)+  => ShS shm -> ShS shn -> ShS shp+  -> Concrete (TKS2 (shm ++ shn) x)+  -> (IxSOf Concrete shm -> IxSOf Concrete shp)+  -> Concrete (TKS2 (shp ++ shn) x)+{-# INLINE tscatterZS #-}+tscatterZS = case knownSTK @x of+  STKScalar @r ->  -- we don't use full dictionary from FTKScalar+    contFromTKAllNum @r (tscatterZSDict @shm @shn)+      -- Optimized: using (+) instead of taddTarget+  _ -> tscatterZSSlow @shm @shn++tscatterZSSlow+  :: forall shm shn shp x. (TKAllNum x, KnownSTK x)+  => ShS shm -> ShS shn -> ShS shp+  -> Concrete (TKS2 (shm ++ shn) x)+  -> (IxSOf Concrete shm -> IxSOf Concrete shp)+  -> Concrete (TKS2 (shp ++ shn) x)+{-# NOINLINE tscatterZSSlow #-}  -- the rare slow case+tscatterZSSlow shm shn shp (Concrete v) f | Dict <- eltDictRep (knownSTK @x) =+  case tftkG (STKS (shm `shsAppend` shn) (knownSTK @x)) v of+    FTKS _sht x ->+      let g :: IIxS shm+            -> IM.IntMap (Concrete (TKS2 shn x))+            -> IM.IntMap (Concrete (TKS2 shn x))+          g ix =+            let ix2 = f $ fmapConcrete ix+            in case ixsToLinearMaybe shp ix2 of+              Nothing -> id+              Just i2 ->+                IM.insertWith (taddTarget $ STKS shn (ftkToSTK x)) i2+                  (Concrete $ Nested.sindexPartial @shm @shn v ix)+          !ivs = foldl' (flip g) IM.empty (shsEnum shm)+      in manyHotNS shn shp x $ IM.assocs ivs++-- The inlining phase control and the explicit dictionaryy argument+-- are required for GHC to consistently specialize the inlined+-- tscatterZSDict code to Int, Double, etc.+-- Phase 99 suffices in GHC 9.14.1, but a later phase is set for a good measure.+-- Maybe what needs to happen is that tscatterZSDict gets specialized+-- before it's inlined.+tscatterZSDict+  :: forall shm shn shp r.+     Dict0 (Num r, Nested.NumElt r, GoodScalar r)+  -> ShS shm -> ShS shn -> ShS shp+  -> Concrete (TKS (shm ++ shn) r)+  -> (IxSOf Concrete shm -> IxSOf Concrete shp)+  -> Concrete (TKS (shp ++ shn) r)+{-# INLINE [1] tscatterZSDict #-}+tscatterZSDict Dict0 = tscatterZSScalar @shm @shn++tscatterZSScalar+  :: forall shm shn shp r. (GoodScalar r, Num r, Nested.NumElt r)+  => ShS shm -> ShS shn -> ShS shp+  -> Concrete (TKS (shm ++ shn) r)+  -> (IxSOf Concrete shm -> IxSOf Concrete shp)+  -> Concrete (TKS (shp ++ shn) r)+{-# INLINE tscatterZSScalar #-}+tscatterZSScalar shm shn shp (Concrete v) f =+  case shn of+    ZSS | Refl <- lemAppNil @shp+        , Refl <- lemAppNil @shm -> runST $ do+      -- Optimized: using Nested.sindex instead of Nested.sindexPartial.+      vec <- VSM.replicate (shsSize shp) 0+      forM_ (shsEnum shm) $ \ix -> do+        let ix2 = f $ fmapConcrete ix+        case ixsToLinearMaybe shp ix2 of+          Nothing -> return ()+          Just i2 -> do+            let u = Nested.sindex v ix+            u2 <- VSM.read vec i2+            VSM.write vec i2 (u + u2)+      Concrete . Nested.sfromVector shp <$> VS.unsafeFreeze vec+    _ -> runST $ do+      let g :: IIxS shm+            -> IM.IntMap (Nested.Shaped shn r)+            -> IM.IntMap (Nested.Shaped shn r)+          g ix =+            let ix2 = f $ fmapConcrete ix+            in case ixsToLinearMaybe shp ix2 of+              Nothing -> id+              Just i2 ->+                IM.insertWith (+) i2 (Nested.sindexPartial @shm @shn v ix)+          ivs = foldl' (flip g) IM.empty (shsEnum shm)+          shnSize = shsSize shn+      vec <- VSM.replicate (shsSize shp * shnSize) 0+      forM_ (IM.assocs ivs) $ \(i, u) ->+        VS.copy (VSM.slice (i * shnSize) shnSize vec) (Nested.stoVector u)+      Concrete . Nested.sfromVector (shp `shsAppend` shn)+        <$> VS.unsafeFreeze vec++tgatherZS+  :: forall shm shn shp x. KnownSTK x+  => ShS shm -> ShS shn+  -> Concrete (TKS2 (shp ++ shn) x)+  -> (IxSOf Concrete shm -> IxSOf Concrete shp)+  -> Concrete (TKS2 (shm ++ shn) x)+{-# INLINE tgatherZS #-}+tgatherZS = case knownSTK @x of+  STKScalar @r ->  -- we don't use full dictionary from FTKScalar+    contFromTypeable @r (tgatherZSDict @shm @shn)+      -- Code gets specialized to a particular underlying scalar.+  _ -> tgatherZSSlow @shm @shn++tgatherZSSlow+  :: forall shm shn shp x. KnownSTK x+  => ShS shm -> ShS shn+  -> Concrete (TKS2 (shp ++ shn) x)+  -> (IxSOf Concrete shm -> IxSOf Concrete shp)+  -> Concrete (TKS2 (shm ++ shn) x)+{-# NOINLINE tgatherZSSlow #-}+tgatherZSSlow shm shn !t f =+  tbuildS shm shn (tindexZSSlow shn t . f)++tgatherZSDict+  :: forall shm shn shp r.+     Dict GoodScalar r+  -> ShS shm -> ShS shn+  -> Concrete (TKS (shp ++ shn) r)+  -> (IxSOf Concrete shm -> IxSOf Concrete shp)+  -> Concrete (TKS (shm ++ shn) r)+{-# INLINE tgatherZSDict #-}+tgatherZSDict Dict = tgatherZSScalar @shm @shn++tgatherZSScalar+  :: forall shm shn shp r. GoodScalar r+  => ShS shm -> ShS shn+  -> Concrete (TKS (shp ++ shn) r)+  -> (IxSOf Concrete shm -> IxSOf Concrete shp)+  -> Concrete (TKS (shm ++ shn) r)+{-# INLINE tgatherZSScalar #-}+tgatherZSScalar shm shn !t f = case shn of+  ZSS | Refl <- lemAppNil @shm+      , Refl <- lemAppNil @shp ->  -- an optimized common case+    tbuildK shm (\ix -> t `tindex0SImpl` f ix)+  _ -> tbuildS shm shn (tindexZSScalar shn t . f)++tgatherZ1S+  :: forall k shn shp x. (KnownNat k, KnownSTK x)+  => ShS shn+  -> Concrete (TKS2 (shp ++ shn) x)+  -> (IntOf Concrete -> IxSOf Concrete shp)+  -> Concrete (TKS2 (k ': shn) x)+{-# INLINE tgatherZ1S #-}+tgatherZ1S = case knownSTK @x of+  STKScalar @r ->  -- we don't use full dictionary from FTKScalar+    contFromTypeable @r (tgatherZ1SDict @k @shn)+      -- Code gets specialized to a particular underlying scalar.+  _ -> tgatherZ1SSlow @k @shn++tgatherZ1SSlow+  :: forall k shn shp x. (KnownNat k, KnownSTK x)+  => ShS shn+  -> Concrete (TKS2 (shp ++ shn) x)+  -> (IntOf Concrete -> IxSOf Concrete shp)+  -> Concrete (TKS2 (k ': shn) x)+{-# NOINLINE tgatherZ1SSlow #-}+tgatherZ1SSlow shn !t f =+  tbuild1S shn (tindexZSSlow shn t . f)++tgatherZ1SDict+  :: forall k shn shp r. KnownNat k+  => Dict GoodScalar r+  -> ShS shn+  -> Concrete (TKS (shp ++ shn) r)+  -> (IntOf Concrete -> IxSOf Concrete shp)+  -> Concrete (TKS (k ': shn) r)+{-# INLINE tgatherZ1SDict #-}+tgatherZ1SDict Dict = tgatherZ1SScalar @k @shn++tgatherZ1SScalar+  :: forall k shn shp r. (KnownNat k, GoodScalar r)+  => ShS shn+  -> Concrete (TKS (shp ++ shn) r)+  -> (IntOf Concrete -> IxSOf Concrete shp)+  -> Concrete (TKS (k ': shn) r)+{-# INLINE tgatherZ1SScalar #-}+tgatherZ1SScalar shn !t f = case shn of+  ZSS | Refl <- lemAppNil @shp ->  -- an optimized common case+    tkbuild1 @_ @k (\ix -> t `tindex0SImpl` f ix)+  _ -> tbuild1S shn (tindexZSScalar shn t . f)++targMinS+  :: forall n sh r. (Nested.PrimElt r, Nested.NumElt r)+  => Nested.Shaped (n ': sh) r -> Nested.Shaped (Init (n ': sh)) Int+{-# INLINE targMinS #-}+targMinS v | sh1@(_ :$$ sh) <- Nested.sshape v =+  let f :: Nested.Shaped '[m] r -> Nested.Shaped '[] Int+      f = Nested.sscalar . ixsHead . Nested.sminIndexPrim+  in case sh of+    ZSS -> f @n v+    _ | SNat @m <- shsLast sh1 ->+      gcastWith (unsafeCoerceRefl+                 :: Init (n ': sh) ++ '[m] :~: n ': sh) $+      gcastWith (unsafeCoerceRefl+                 :: Init (n ': sh) :~: Init (n ': sh) ++ '[]) $+      Nested.sunNest $+        Nested.srerankPrim @'[m] @'[] @(Init (n ': sh)) ZSS (f @m) $+          Nested.snest (shsInit sh1) v++targMaxS+  :: forall n sh r. (Nested.PrimElt r, Nested.NumElt r)+  => Nested.Shaped (n ': sh) r -> Nested.Shaped (Init (n ': sh)) Int+{-# INLINE targMaxS #-}+targMaxS v | sh1@(_ :$$ sh) <- Nested.sshape v =+  let f :: Nested.Shaped '[m] r -> Nested.Shaped '[] Int+      f = Nested.sscalar . ixsHead . Nested.smaxIndexPrim+  in case sh of+    ZSS -> f @n v+    _ | SNat @m <- shsLast sh1 ->+      gcastWith (unsafeCoerceRefl+                 :: Init (n ': sh) ++ '[m] :~: n ': sh) $+      gcastWith (unsafeCoerceRefl+                 :: Init (n ': sh) :~: Init (n ': sh) ++ '[]) $+      Nested.sunNest $+        Nested.srerankPrim @'[m] @'[] @(Init (n ': sh)) ZSS (f @m) $+          Nested.snest (shsInit sh1) v++tmap0NS+  :: forall r1 r sh. (Nested.PrimElt r1, Nested.PrimElt r)+  => (r1 -> r) -> Nested.Shaped sh r1 -> Nested.Shaped sh r+{-# INLINE tmap0NS #-}+tmap0NS f = Shaped.liftShaped1 (Mixed.mliftPrim f)++tzipWith0NS+  :: forall r1 r2 r sh. (Nested.PrimElt r, Nested.PrimElt r1, Nested.PrimElt r2)+  => (r1 -> r2 -> r) -> Nested.Shaped sh r1 -> Nested.Shaped sh r2+  -> Nested.Shaped sh r+{-# INLINE tzipWith0NS #-}+tzipWith0NS f = Shaped.liftShaped2 (Mixed.mliftPrim2 f)++tbuild1S :: forall k sh x. (KnownNat k, KnownSTK x)+         => ShS sh -> (IntOf Concrete -> Concrete (TKS2 sh x))+         -> Concrete (TKS2 (k ': sh) x)+{-# INLINE tbuild1S #-}+tbuild1S sh f = case knownSTK @x of+  STKScalar | ZSS <- sh ->+    tbuild1K (Concrete . Nested.sunScalar . unConcrete . f)+  _ | Dict <- eltDictRep (knownSTK @x) -> case SNat @k of+    SZ ->+      Concrete $ Nested.semptyArray sh+    _ ->+      let g i = unConcrete $ f (Concrete i)+      in Concrete $ Nested.sfromListOuter SNat $ NonEmpty.fromList+         $ map g [0 .. valueOf @k - 1]++tbuildS :: forall shm shn x. KnownSTK x+        => ShS shm -> ShS shn -> (IxSOf Concrete shm -> Concrete (TKS2 shn x))+        -> Concrete (TKS2 (shm ++ shn) x)+{-# INLINE tbuildS #-}+tbuildS shm shn f = case knownSTK @x of+  STKScalar | ZSS <- shn+            , Refl <- lemAppNil @shm ->+    tbuildK shm (Concrete . Nested.sunScalar . unConcrete . f)+  _ | Dict <- eltDictRep (knownSTK @x) ->+    withKnownShS shn $+      -- this is needed due to+      -- instance (KnownShS sh, KnownElt a) => KnownElt (Shaped sh a) where+    let h ix = unConcrete $ f (fmapConcrete ix)+    in Concrete $ Nested.sunNest $ Nested.sgenerate shm h+++-- * Mixed internal definitions++liftVX+  :: (Nested.PrimElt r1, Nested.PrimElt r)+  => (VS.Vector r1 -> VS.Vector r)+  -> Nested.Mixed sh r1 -> Nested.Mixed sh r+{-# INLINE liftVX #-}+liftVX f = Mixed.mliftNumElt1 (`liftVEltwise1` f)++manyHotNX :: forall sh1 sh2 x. KnownSTK x+          => FullShapeTK (TKX2 (sh1 ++ sh2) x)+          -> [(Int, Concrete (TKX2 sh2 x))]+          -> Concrete (TKX2 (sh1 ++ sh2) x)+{-# INLINE manyHotNX #-}+manyHotNX (FTKX sh x) upd | Dict <- eltDictRep (knownSTK @x) = runST $ do+  let zero = unConcrete $ tdefTarget x+  vecs <- Mixed.mvecsReplicate sh zero+    -- this avoids the slow case of mvecsReplicate if x is TKScalar+  forM_ upd $ \(ix, Concrete v) ->+    Mixed.mvecsWritePartialLinear (Proxy @sh1) ix v vecs+  Concrete <$> Mixed.mvecsUnsafeFreeze sh vecs++tindexZX :: forall shm shn x. KnownSTK x+         => Concrete (TKX2 (shm ++ shn) x) -> IxXOf Concrete shm+         -> Concrete (TKX2 shn x)+{-# INLINE tindexZX #-}+tindexZX (Concrete v) ix | Dict <- eltDictRep (knownSTK @x) =+  if ixInBoundsX (shxTakeIx @shm @shn Proxy ix (Nested.mshape v)) ix+  then Concrete $ Nested.mindexPartial v (fmapUnConcrete ix)+  else case tftkG (STKX (ssxFromShX $ Nested.mshape v) (knownSTK @x)) v of+         FTKX sh x -> tdefTarget (FTKX (shxDropIx @shm @shn ix sh) x)++tindex0X :: GoodScalar r+         => Concrete (TKX sh1 r) -> IxXOf Concrete sh1 -> Concrete (TKScalar r)+{-# INLINE tindex0X #-}+tindex0X (Concrete v) ix =+  Concrete+  $ if ixInBoundsX (Nested.mshape v) ix+    then Nested.mindex v (fmapUnConcrete ix)+    else def++toneHotX :: forall sh1 sh2 x. (KnownShX sh2, KnownSTK x)+         => IShX sh1 -> Concrete (TKX2 sh2 x) -> IxXOf Concrete sh1+         -> Concrete (TKX2 (sh1 ++ sh2) x)+{-# INLINE toneHotX #-}+toneHotX sh1 !v ix = case tftk knownSTK v of+  FTKX sh2 x ->+    let ftk = FTKX (sh1 `shxAppend` sh2) x+    in case ixxToLinearMaybe sh1 ix of+         Nothing -> tdefTarget ftk+         Just i2 -> manyHotNX @sh1 ftk [(i2, v)]++tscatterZX :: forall shm shn shp x.+              (KnownShX shm, KnownShX shn, TKAllNum x, KnownSTK x)+           => IShX shp -> Concrete (TKX2 (shm ++ shn) x)+           -> (IxXOf Concrete shm -> IxXOf Concrete shp)+           -> Concrete (TKX2 (shp ++ shn) x)+{-# INLINE tscatterZX #-}  -- this function takes a function as an argument+tscatterZX shp v0@(Concrete v) f | Dict <- eltDictRep (knownSTK @x) =+  let sht = Nested.mshape v+      shm = shxTakeSSX (Proxy @shn) (knownShX @shm) sht+      shn = shxDropSSX @_ @shn (knownShX @shm) sht+  in withKnownShX (knownShX @shm `ssxAppend` knownShX @shn) $+     case (knownShX @shn, tftk knownSTK v0) of+       (ZKX, FTKX _ (FTKScalar @r)) | Dict0 <- numFromTKAllNum (Proxy @r)+                                    , Refl <- lemAppNil @shp+                                    , Refl <- lemAppNil @shm -> runST $ do+         -- Optimized: using (+) instead of taddTarget and using mindex.+         vec <- VSM.replicate (shxSize shp) 0+         forM_ (shxEnum shm) $ \ix -> do+           let ix2 = f $ fmapConcrete ix+           case ixxToLinearMaybe shp ix2 of+             Nothing -> return ()+             Just i2 -> do+               let u = Nested.mindex v ix+               u2 <- VSM.read vec i2+               VSM.write vec i2 (u + u2)+         Concrete . Nested.mfromVector shp <$> VS.unsafeFreeze vec+       (_, FTKX _ (FTKScalar @r)) | Dict0 <- numFromTKAllNum (Proxy @r) ->+         -- Optimized: using (+) instead of taddTarget.+         -- TODO: write to vecs and use a bitmap to record the written indexes+         -- and the intmap only for subsequent writes+         let ftk = FTKX (shp `shxAppend` shn) FTKScalar+             g :: IIxX shm+               -> IM.IntMap (Concrete (TKX2 shn x))+               -> IM.IntMap (Concrete (TKX2 shn x))+             g ix =+               let ix2 = f $ fmapConcrete ix+               in case ixxToLinearMaybe shp ix2 of+                 Nothing -> id+                 Just i2 ->+                   IM.insertWith (+) i2+                     (Concrete $ Nested.mindexPartial @_ @shm @shn v ix)+             !ivs = foldl' (flip g) IM.empty (shxEnum shm)+         in manyHotNX @shp ftk $ IM.assocs ivs+       (_, FTKX _ x) ->+         -- TODO: write to vecs and use a bitmap to record the written indexes+         -- and the intmap only for subsequent writes+         let ftk = FTKX (shp `shxAppend` shn) x+             g :: IIxX shm+               -> IM.IntMap (Concrete (TKX2 shn x))+               -> IM.IntMap (Concrete (TKX2 shn x))+             g ix =+               let ix2 = f $ fmapConcrete ix+               in case ixxToLinearMaybe shp ix2 of+                 Nothing -> id+                 Just i2 ->+                   IM.insertWith (taddTarget knownSTK) i2+                     (Concrete $ Nested.mindexPartial @_ @shm @shn v ix)+             !ivs = foldl' (flip g) IM.empty (shxEnum shm)+         in manyHotNX @shp ftk $ IM.assocs ivs++tgatherZX :: forall shm shn shp x. (KnownShX shm, KnownShX shn, KnownSTK x)+          => IShX shm+          -> Concrete (TKX2 (shp ++ shn) x)+          -> (IxXOf Concrete shm -> IxXOf Concrete shp)+          -> Concrete (TKX2 (shm ++ shn) x)+{-# INLINE tgatherZX #-}  -- this function takes a function as an argument+tgatherZX shm !t f =+  gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+  case (knownShX @shn, knownSTK @x) of+    (ZKX, STKScalar) | Refl <- lemAppNil @shm+                     , Refl <- lemAppNil @shp ->  -- an optimized common case+      let g ix = unConcrete $ t `txindex0` f (fmapConcrete ix)+      in Concrete $ Nested.mgeneratePrim shm g+    _ ->+      case ssxRank (knownShX @shm) of+        SNat -> -- needed only for GHC 9.10+          txbuild @_ @shm @shn shm (\ix -> t `txindex` f ix)++tgatherZ1X :: forall k shn shp x.+              (KnownNat k, KnownShX shn, KnownShX shp, KnownSTK x)+           => SNat k -> Concrete (TKX2 (shp ++ shn) x)+           -> (IntOf Concrete -> IxXOf Concrete shp)+           -> Concrete (TKX2 (Just k ': shn) x)+{-# INLINE tgatherZ1X #-}   -- this function takes a function as an argument+tgatherZ1X _ !t f = case (knownShX @shn, knownSTK @x) of+  (ZKX, STKScalar) | Refl <- lemAppNil @shp ->  -- an optimized common case+    let shm = SKnown SNat :$% shxDropSSX (knownShX @shp)+                                         (Nested.mshape $ unConcrete t)+        g i = unConcrete $ t `txindex0` f (Concrete i)+    in Concrete $ Nested.mfromVector shm $ VS.generate (valueOf @k) g+  _ -> txbuild1 @_ @k (\ix -> t `txindex` f ix)++targMinX+  :: forall mn sh r. (Nested.PrimElt r, Nested.NumElt r)+  => Nested.Mixed (mn ': sh) r -> Nested.Mixed (Init (mn ': sh)) Int+{-# INLINE targMinX #-}+targMinX v | sh1@(_ :$% sh) <- Nested.mshape v =+  let f :: Nested.Mixed '[mm] r -> Nested.Mixed '[] Int+      f = Nested.mscalar . ixxHead . Nested.mminIndexPrim+  in case sh of+    ZSX -> f @mn v+    _ -> withSNat (fromSMayNat' (shxLast sh1)) $ \(_ :: SNat m) ->+      gcastWith (unsafeCoerceRefl+                 :: Init (mn ': sh) ++ '[Just m] :~: mn ': sh) $+      gcastWith (unsafeCoerceRefl+                 :: Init (mn ': sh) :~: Init (mn ': sh) ++ '[]) $+      Nested.munNest $+        Nested.mrerankPrim @'[Just m] @'[] @(Init (mn ': sh))+                           ZSX (f @(Just m)) $+          Nested.mnest (ssxFromShX $ shxInit sh1) v++targMaxX+  :: forall mn sh r. (Nested.PrimElt r, Nested.NumElt r)+  => Nested.Mixed (mn ': sh) r -> Nested.Mixed (Init (mn ': sh)) Int+{-# INLINE targMaxX #-}+targMaxX v | sh1@(_ :$% sh) <- Nested.mshape v =+  let f :: Nested.Mixed '[mm] r -> Nested.Mixed '[] Int+      f = Nested.mscalar . ixxHead . Nested.mmaxIndexPrim+  in case sh of+    ZSX -> f @mn v+    _ -> withSNat (fromSMayNat' (shxLast sh1)) $ \(_ :: SNat m) ->+      gcastWith (unsafeCoerceRefl+                 :: Init (mn ': sh) ++ '[Just m] :~: mn ': sh) $+      gcastWith (unsafeCoerceRefl+                 :: Init (mn ': sh) :~: Init (mn ': sh) ++ '[]) $+      Nested.munNest $+        Nested.mrerankPrim @'[Just m] @'[] @(Init (mn ': sh))+                           ZSX (f @(Just m)) $+          Nested.mnest (ssxFromShX $ shxInit sh1) v
+ src/HordeAd/Core/PPEngine.hs view
@@ -0,0 +1,96 @@+-- | Predefined common functions for pretty-printing AST terms+-- in a few useful configurations.+module HordeAd.Core.PPEngine+  ( printAstVarName+  , printAstSimple, printAstPretty, printAstPrettyButNested+  , printArtifactSimple, printArtifactPretty+  , printArtifactPrimalSimple, printArtifactPrimalPretty+  ) where++import Prelude++import Data.IntMap.Strict (IntMap)+import Data.IntMap.Strict qualified as IM++import HordeAd.Core.Ast+import HordeAd.Core.PPTools++-- * Pretty-printing terms in a few useful configurations++{-# INLINE printAstVarRename #-}+printAstVarRename :: IntMap String -> AstVarName '(s, y) -> String+printAstVarRename renames var =+  printAstVar (defaulPrintConfig {varRenames = renames}) var ""++{-# INLINE printAstSimpleRename #-}+printAstSimpleRename :: KnownSpan s+                     => IntMap String -> AstTensor ms s y -> String+printAstSimpleRename renames t =+  printAst+    (defaulPrintConfig {loseRoudtrip = False, varRenames = renames}) 0 t ""++{-# INLINE printAstPrettyRename #-}+printAstPrettyRename :: KnownSpan s+                     => IntMap String -> AstTensor ms s y -> String+printAstPrettyRename renames t =+  printAst (defaulPrintConfig {varRenames = renames}) 0 t ""++{-# INLINE printAstVarName #-}+printAstVarName :: AstVarName '(s, y) -> String+printAstVarName var =+  printAstVar defaulPrintConfig var ""++-- | Print an AST term in a form close to being able to roundtrip,+-- including explicit sharing preservation.+{-# INLINE printAstSimple #-}+printAstSimple :: KnownSpan s+               => AstTensor ms s y -> String+printAstSimple t =+  printAst (defaulPrintConfig {loseRoudtrip = False}) 0 t ""++-- | Print an AST term in a readable form that does not roundtrip,+-- and where Haskell @let@ (sharing on Haskell heap) is used instead+-- of explicit sharing of subterms.+{-# INLINE printAstPretty #-}+printAstPretty :: KnownSpan s+               => AstTensor ms s y -> String+printAstPretty t =+  printAst defaulPrintConfig 0 t ""++{-# INLINE printAstPrettyButNested #-}+printAstPrettyButNested :: KnownSpan s+                        => AstTensor ms s y -> String+printAstPrettyButNested t =+  printAst (defaulPrintConfig {ignoreNestedLambdas = False}) 0 t ""++{-# INLINE printArtifactSimple #-}+printArtifactSimple :: AstArtifactRev x z -> String+printArtifactSimple AstArtifactRev{..} =+  let nDt = fromEnum (varNameToAstVarId artVarDtRev) - 100000000+      renames = IM.singleton nDt "dret"+      varsPP = [ printAstVarRename renames artVarDtRev+               , printAstVarRename renames artVarDomainRev ]+  in "\\" ++ unwords varsPP+          ++ " -> " ++ printAstSimpleRename renames artDerivativeRev++{-# INLINE printArtifactPretty #-}+printArtifactPretty :: AstArtifactRev x z -> String+printArtifactPretty AstArtifactRev{..} =+  let nDt = fromEnum (varNameToAstVarId artVarDtRev) - 100000000+      renames = IM.singleton nDt "dret"+      varsPP = [ printAstVarRename renames artVarDtRev+               , printAstVarRename renames artVarDomainRev ]+  in "\\" ++ unwords varsPP+          ++ " -> " ++ printAstPrettyRename renames artDerivativeRev++{-# INLINE printArtifactPrimalSimple #-}+printArtifactPrimalSimple :: AstArtifactRev x z -> String+printArtifactPrimalSimple AstArtifactRev{..} =+  "\\" ++ printAstVarName artVarDomainRev+       ++ " -> " ++ printAstSimple artPrimalRev++{-# INLINE printArtifactPrimalPretty #-}+printArtifactPrimalPretty :: AstArtifactRev x z -> String+printArtifactPrimalPretty AstArtifactRev{..} =+  "\\" ++ printAstVarName artVarDomainRev+       ++ " -> " ++ printAstPretty artPrimalRev
+ src/HordeAd/Core/PPTools.hs view
@@ -0,0 +1,669 @@+{-# LANGUAGE CPP #-}+#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)+{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}+#endif+-- | Pretty-printing of the AST. Some of the variants of pretty-printing+-- almost roundtrip, while others are more readable but less faithful.+module HordeAd.Core.PPTools+  ( PrintConfig(..), defaulPrintConfig+  , printAstVar, printAst+  ) where++import Prelude++import Data.Foldable qualified as Foldable+import Data.IntMap.Strict (IntMap)+import Data.IntMap.Strict qualified as IM+import Data.List (intersperse)+import Data.Type.Equality (testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Type.Reflection (typeRep)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Permutation (Perm (..), permToList)+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat')++import HordeAd.Core.Ast+import HordeAd.Core.AstTools+import HordeAd.Core.Conversion+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Pretty-printing config++-- Modeled after https://github.com/VMatthijs/CHAD/blob/755fc47e1f8d1c3d91455f123338f44a353fc265/src/TargetLanguage.hs#L335.++-- TODO: ensure that terms roundtrip if neither loseRoudtrip+-- nor ignoreNestedLambdas is set and that explicit sharing is then preserved+-- as opposed to displaying sharing as Haskell lets.+-- The operations in OpsTensor should suffice to roundtrip. Adjust both+-- OpsTensor and the pretty-printing until they do.+-- Note that disabling ignoreNestedLambdas causes derivatives to be computed,+-- so pretty-printing in this way can be very expensive.+data PrintConfig = PrintConfig+  { loseRoudtrip        :: Bool+  , ignoreNestedLambdas :: Bool+  , varRenames          :: IntMap String+  }++defaulPrintConfig :: PrintConfig+defaulPrintConfig = PrintConfig+  { loseRoudtrip        = True+  , ignoreNestedLambdas = True+  , varRenames          = IM.empty+  }+++-- * Pretty-printing of variables++printAstVarId :: String -> PrintConfig -> AstVarId -> ShowS+printAstVarId prefix cfg var =+  let n = fromEnum var - 100000000+  in showString $ case IM.lookup n (varRenames cfg) of+    Just name | name /= "" -> name+    _ | n < 0 -> prefix ++ "m" ++ show (- n)+    _ -> prefix ++ show n++printAstVar :: PrintConfig -> AstVarName '(s, y) -> ShowS+printAstVar cfg var =+  case varNameToFTK var of+    FTKScalar @r | SPlainSpan <- varNameToSpan var+                 , Just Refl <- testEquality (typeRep @r) (typeRep @Int) ->+      printAstIntVar cfg var+    _ -> let prefix = case lengthSTK (ftkToSTK $ varNameToFTK var) of+               0 -> "x"+               1 -> "v"+               2 -> "m"+               3 -> "t"+               4 -> "u"+               _ -> "w"+         in printAstVarId prefix cfg (varNameToAstVarId var)++printAstIntVar :: PrintConfig -> IntVarName -> ShowS+printAstIntVar cfg var = printAstVarId "i" cfg (varNameToAstVarId var)+++-- * Pretty-printing of AST terms++-- Precedences used are as in Haskell.+printAst :: forall s y ms. KnownSpan s+         => PrintConfig -> Int -> AstTensor ms s y -> ShowS+printAst cfg d = \case+  AstPair t1 t2 ->+    showParen (d > 10)+    $ showString "tpair "+      . printAst cfg 11 t1+      . showString " "+      . printAst cfg 11 t2+  AstProject1 t -> printPrefixOp printAst cfg d "tproject1" [t]+  AstProject2 t -> printPrefixOp printAst cfg d "tproject2" [t]+  AstMapAccumLDer k bftk eftk f df rf acc0 es ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "tmapAccumLDer "+      . showParen True (shows k)+      . showString " "+      . printAstHFun cfg 10 f+      . showString " "+      . printAstHFun cfg 10 df+      . showString " "+      . printAstHFun cfg 10 rf+      . showString " "+      . printAst cfg 11 acc0+      . showString " "+      . printAst cfg 11 es+   else+    showParen (d > 10)+    $ showString "tmapAccumLDer "+      . showParen True (shows k)+      . showString " "+      . showParen True (shows (ftkAst acc0))+      . showString " "+      . showParen True (shows bftk)+      . showString " "+      . showParen True (shows eftk)+      . showString " "+      . printAstHFun cfg 10 f+      . showString " "+      . printAstHFun cfg 10 df+      . showString " "+      . printAstHFun cfg 10 rf+      . showString " "+      . printAst cfg 11 acc0+      . showString " "+      . printAst cfg 11 es+  AstApply t ll -> showParen (d > 10)+                   $ showString "tapply "+                     . printAstHFunOneUnignore cfg 10 t+                         -- this is a lambda, but not nested, so always printed+                     . showString " "+                     . printAst cfg 11 ll+  AstVar var -> printAstVar cfg var+  AstBuild1 k stk (var, v) ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "tbuild1 ("+      . shows k+      . showString ") "+      . (showParen True+         $ showString "\\"+           . printAstIntVar cfg var+           . showString " -> "+           . printAst cfg 0 v)+   else+    showParen (d > 10)+    $ showString "tbuild1 ("+      . shows k+      . showString ") "+      . showParen True (shows stk)+      . showString " "+      . (showParen True+         $ showString "\\"+           . printAstIntVar cfg var+           . showString " -> "+           . printAst cfg 0 v)++  t@(AstLet @_ @_ @s0 var0 u0 v0) ->+    withKnownSpan (varNameToSpan var0) $+    if loseRoudtrip cfg+    then let collect :: AstTensor AstMethodLet s y -> ([(ShowS, ShowS)], ShowS)+             collect (AstLet var u v) =+               let name = printAstVar cfg var+                   uPP = withKnownSpan (varNameToSpan var) $ printAst cfg 0 u+                   (rest, corePP) = collect v+               in ((name, uPP) : rest, corePP)+             collect v = ([], printAst cfg 0 v)+             (pairs, core) = collect t+         in showParen (d > 0)+            $ showString "let "+              . foldr (.) id (intersperse (showString " ; ")+                  [name . showString " = " . uPP | (name, uPP) <- pairs])+              . showString " in "+              . core+    else let keyword = case (knownSpan @s0, knownSpan @s) of+               (SPrimalStepSpan SFullSpan, SFullSpan) -> "tletPrimal "+               (SPlainSpan, SFullSpan) -> "tletPlain "+               _ -> "tlet "+         in showParen (d > 10)+            $ showString keyword+              . printAst cfg 11 u0+              . showString " "+              . (showParen True+                 $ showString "\\"+                   . printAstVar cfg var0+                   . showString " -> "+                   . printAst cfg 0 v0)+  AstShare _var v -> printPrefixOp printAst cfg d "tshare" [v]+  AstToShare v -> printPrefixOp printAst cfg d "toShare" [v]++  AstConcreteK k -> showNumber k+  AstFromPlain (AstConcreteK k) | loseRoudtrip cfg -> showNumber k+  AstFromPlain t@(AstConvert _ AstConcreteK{}) | loseRoudtrip cfg ->+    printAst cfg d t+  AstFromPlain t@(AstPair AstConcreteK{} AstConcreteK{}) | loseRoudtrip cfg ->+    printAst cfg d t+  AstConcreteS a -> case Nested.sshape a of+    ZSS -> showParen (d > 10)+           $ showString "sscalar "+             . showNumber (Nested.sunScalar a)+    _ -> showParen (d > 10)+         $ showString "sconcrete "+           . (showParen True+              $ shows a)+  AstFromPlain (AstConcreteS a) | loseRoudtrip cfg -> case Nested.sshape a of+    ZSS -> showParen (d > 10)+           $ showString "sscalar "+             . showNumber (Nested.sunScalar a)+    _ -> showParen (d > 10)+         $ showString "sconcrete "+           . (showParen True+              $ shows a)+  AstFromPlain t@(AstConvert _ AstConcreteS{}) | loseRoudtrip cfg ->+    printAst cfg d t+  AstFromPlain t@(AstPair AstConcreteS{} AstConcreteS{}) | loseRoudtrip cfg ->+    printAst cfg d t++  AstPrimalPart a ->+    if loseRoudtrip cfg+    then case ftkAst a of+      FTKScalar{} -> printPrefixOp printAst cfg d "kprimalPart" [a]+      FTKR{} -> printPrefixOp printAst cfg d "rprimalPart" [a]+      FTKS{} -> printPrefixOp printAst cfg d "sprimalPart" [a]+      FTKX{} -> printPrefixOp printAst cfg d "xprimalPart" [a]+      _      -> printPrefixOp printAst cfg d "tprimalPart" [a]+    else printPrefixOp printAst cfg d "tprimalPart" [a]+  AstDualPart a ->+    if loseRoudtrip cfg+    then case ftkAst a of+      FTKScalar{} -> printPrefixOp printAst cfg d "kdualPart" [a]+      FTKR{} -> printPrefixOp printAst cfg d "rdualPart" [a]+      FTKS{} -> printPrefixOp printAst cfg d "sdualPart" [a]+      FTKX{} -> printPrefixOp printAst cfg d "xdualPart" [a]+      _      -> printPrefixOp printAst cfg d "tdualPart" [a]+    else printPrefixOp printAst cfg d+                       ("tdualPart (" ++ show (ftkToSTK (ftkAst a)) ++ ")") [a]+  AstPlainPart a ->+    if loseRoudtrip cfg+    then case ftkAst a of+      FTKScalar{} -> printPrefixOp printAst cfg d "kplainPart" [a]+      FTKR{} -> printPrefixOp printAst cfg d "rplainPart" [a]+      FTKS{} -> printPrefixOp printAst cfg d "splainPart" [a]+      FTKX{} -> printPrefixOp printAst cfg d "xplainPart" [a]+      _      -> printPrefixOp printAst cfg d "tplainPart" [a]+    else printPrefixOp printAst cfg d "tplainPart" [a]+  AstFromPrimal a ->+    if loseRoudtrip cfg+    then case ftkAst a of+      FTKScalar{} -> printPrefixOp printAst cfg d "kfromPrimal" [a]+      FTKR{} -> printPrefixOp printAst cfg d "rfromPrimal" [a]+      FTKS{} -> printPrefixOp printAst cfg d "sfromPrimal" [a]+      FTKX{} -> printPrefixOp printAst cfg d "xfromPrimal" [a]+      _      -> printPrefixOp printAst cfg d "tfromPrimal" [a]+    else printPrefixOp+           printAst cfg d+           ("tfromPrimal (" ++ show (ftkToSTK (ftkAst a)) ++ ")") [a]+  AstFromDual a ->+    if loseRoudtrip cfg+    then case ftkAst a of+      FTKScalar{} -> printPrefixOp printAst cfg d "kfromDual" [a]+      FTKR{} -> printPrefixOp printAst cfg d "rfromDual" [a]+      FTKS{} -> printPrefixOp printAst cfg d "sfromDual" [a]+      FTKX{} -> printPrefixOp printAst cfg d "xfromDual" [a]+      _      -> printPrefixOp printAst cfg d "tfromDual" [a]+    else printPrefixOp printAst cfg d "tfromDual" [a]+  AstFromPlain a ->+    if loseRoudtrip cfg+    then case ftkAst a of+      FTKScalar{} -> printPrefixOp printAst cfg d "kfromPlain" [a]+      FTKR{} -> printPrefixOp printAst cfg d "rfromPlain" [a]+      FTKS{} -> printPrefixOp printAst cfg d "sfromPlain" [a]+      FTKX{} -> printPrefixOp printAst cfg d "xfromPlain" [a]+      _      -> printPrefixOp printAst cfg d "tfromPlain" [a]+    else printPrefixOp+           printAst cfg d+           ("tfromPlain (" ++ show (ftkToSTK (ftkAst a)) ++ ")") [a]++  AstPlusK u v -> printBinaryOp printAst cfg d u (6, "+") v+  AstTimesK u v -> printBinaryOp printAst cfg d u (7, "*") v+  AstN1K opCode u -> printAstN1R printAst cfg d opCode u+  AstR1K opCode u -> printAstR1R printAst cfg d opCode u+  AstR2K opCode u v -> printAstR2R printAst cfg d opCode u v+  AstI2K opCode u v -> printAstI2R printAst cfg d opCode u v+  AstFloorK v -> printPrefixOp printAst cfg d "kfloor" [v]+  AstFromIntegralK v -> printPrefixOp printAst cfg d "kfromIntegral" [v]+  AstCastK v -> printPrefixOp printAst cfg d "kcast" [v]+  AstArgMinK v -> printPrefixOp printAst cfg d "kargMin" [v]+  AstArgMaxK v -> printPrefixOp printAst cfg d "kargMax" [v]+  AstIndexK v ix ->+    showParen (d > 9)+    $ printAst cfg 10 v+      . showString " `sindex0` "+      . showListWith (printAst cfg 0) (Foldable.toList ix)++  AstPlusS u v -> printBinaryOp printAst cfg d u (6, "+") v+  AstTimesS u v -> printBinaryOp printAst cfg d u (7, "*") v+  AstN1S opCode u -> printAstN1R printAst cfg d opCode u+  AstR1S opCode u -> printAstR1R printAst cfg d opCode u+  AstR2S opCode u v -> printAstR2R printAst cfg d opCode u v+  AstI2S opCode u v -> printAstI2R printAst cfg d opCode u v+  AstFloorS a -> printPrefixOp printAst cfg d "sfloor" [a]+  AstFromIntegralS a -> printPrefixOp printAst cfg d "sfromIntegral" [a]+  AstCastS a -> printPrefixOp printAst cfg d "scast" [a]+  AstArgMinS a -> printPrefixOp printAst cfg d "sargMin" [a]+  AstArgMaxS a -> printPrefixOp printAst cfg d "sargMax" [a]+  AstIndexS _ v ix ->+    showParen (d > 9)+    $ printAst cfg 10 v+      . showString " !$ "+      . showListWith (printAst cfg 0) (Foldable.toList ix)++  AstCondK b a1 a2 ->+    showParen (d > 10)+    $ showString "ifH "+      . printAst cfg 11 b+      . showString " "+      . printAst cfg 11 a1+      . showString " "+      . printAst cfg 11 a2+  AstCondS b a1 a2 ->+    showParen (d > 10)+    $ showString "ifH "+      . printAst cfg 11 b+      . showString " "+      . printAst cfg 11 a1+      . showString " "+      . printAst cfg 11 a2+  AstFromVectorK shm l ->+    let s = "sfromVector0N @" ++ showListWith shows (shsToList shm) " "+    in showParen (d > 10)+       $ showString s+         . (showParen True+            $ showString "fromList "+              . showListWith (printAst cfg 0) (V.toList l))+  AstFromVectorS (snat :$$ ZSS) l ->+    let s = "sfromVector @" ++ show (fromSNat' snat) ++ " "+    in showParen (d > 10)+       $ showString s+         . (showParen True+            $ showString "fromList "+              . showListWith (printAst cfg 0) (V.toList l))+  AstFromVectorS shm l ->+    let s = "sfromVectorN @" ++ showListWith shows (shsToList shm) " "+    in showParen (d > 10)+       $ showString s+         . (showParen True+            $ showString "fromList "+              . showListWith (printAst cfg 0) (V.toList l))+  AstSumK v | FTKS shm _ <- ftkAst v ->+    let s = "ssum0 @" ++ showListWith shows (shsToList shm) ""+    in printPrefixOp printAst cfg d s [v]+  AstSumS (snat :$$ ZSS) v ->+    let s = "ssum @" ++ show (fromSNat' snat)+    in printPrefixOp printAst cfg d s [v]+  AstSumS shm v ->+    let s = "ssumN @" ++ showListWith shows (shsToList shm) ""+    in printPrefixOp printAst cfg d s [v]+  AstScatterS _ _ _ v (ZS, ix) ->+    showParen (d > 9)+    $ showString "soneHot "+      . printAst cfg 11 v+      . showString " "+      . showListWith (printAst cfg 0) (Foldable.toList ix)+  AstScatterS (snat :$$ _) _ _ v (var ::$ ZS, ix) ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString ("sscatter1 @" ++ show (fromSNat' snat) ++ " ")+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . printAstIntVar cfg var+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+   else+    showParen (d > 10)+    $ showString ("sscatter1 @" ++ show (fromSNat' snat) ++ " ")+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . printAstIntVar cfg var+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+  AstScatterS shm _ _ v (vars, ix) ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "sscatter @"+      . showListWith shows (shsToList shm)+      . showString " "+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . showListWith (printAstIntVar cfg) (Foldable.toList vars)+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+   else+    showParen (d > 10)+    $ showString ("sscatter @" ++ show shm ++ " ")+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . showListWith (printAstIntVar cfg) (Foldable.toList vars)+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+  -- This is too common to be verbose even in no loseRoudtrip mode.+  AstReplicateK shm v ->+    let s = "sreplicate0N @" ++ showListWith shows (shsToList shm) ""+    in printPrefixOp printAst cfg d s [v]+  AstReplicateS (snat :$$ ZSS) v ->+    let s = "sreplicate @" ++ show (fromSNat' snat)+    in printPrefixOp printAst cfg d s [v]+  AstReplicateS shm v ->+    let s = "sreplicateN @" ++ showListWith shows (shsToList shm) ""+    in printPrefixOp printAst cfg d s [v]+  {- Let's re-enable this when/if we remove AstIndexS altogether+     or at least stop rewriting this to AstIndexS but instead optimize+     the instances for this case:+  AstGatherS _ v (ZS, ix) ->+    showParen (d > 9)+    $ printAst cfg 10 v+      . showString " !$ "+      . showListWith (printAst cfg 0) (Foldable.toList ix) -}+  AstGatherS (snat :$$ _) _ _ v (var ::$ ZS, ix) ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString ("sgather1 @" ++ show (fromSNat' snat) ++ " ")+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . printAstIntVar cfg var+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+   else+    showParen (d > 10)+    $ showString ("sgather1 @" ++ show (fromSNat' snat) ++ " ")+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . printAstIntVar cfg var+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+  AstGatherS shm _ _ v (vars, ix) ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "sgather @"+      . showListWith shows (shsToList shm)+      . showString " "+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . showListWith (printAstIntVar cfg)+                          (Foldable.toList vars)+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+   else+    showParen (d > 10)+    $ showString ("sgather @" ++ show shm ++ " ")+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . showListWith (printAstIntVar cfg)+                          (Foldable.toList vars)+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+  AstIotaS snat ->+    showParen (d > 10)+    $ showString ("siota (" ++ show snat ++ ")")+  AstAppendS t1 t2 ->+    showParen (d > 10)+    $ showString "sappend "+      . printAst cfg 11 t1+      . showString " "+      . printAst cfg 11 t2+  AstSliceS i n _k v ->+    printPrefixOp printAst cfg d+                  ("sslice (" ++ show i ++ ") (" ++ show n ++ ")") [v]+  AstReverseS v -> printPrefixOp printAst cfg d "sreverse" [v]+  AstTransposeS (SNat' @1 `PCons` SNat' @0 `PCons` PNil) v ->+    printPrefixOp printAst cfg d "str" [v]+  AstTransposeS perm v ->+   if loseRoudtrip cfg+   then printPrefixOp printAst cfg d+                      ("stranspose @"+                       ++ showListWith shows (permToList perm) "") [v]+   else printPrefixOp+          printAst cfg d+          ("ttranspose (makePerm @"+                        ++ showListWith shows (permToList perm) "" ++ ")") [v]+  AstReshapeS sh2 v ->+    printPrefixOp printAst cfg d ("sreshape @"+                                  ++ showListWith shows (shsToList sh2) "") [v]++  -- TODO: pretty-print correctly szip, sunzip, snestS, sunNestS+  -- or at least make sure they get printed as tconvert, not as the others+  AstConvert c t -> case (ftkToSTK (ftkAst t), convertFTK c (ftkAst t)) of+    (STKScalar, FTKR{}) -> printPrefixOp printAst cfg d "rfromK" [t]+    (STKScalar, FTKS{}) -> printPrefixOp printAst cfg d "sfromK" [t]+    (STKScalar, FTKX{}) -> printPrefixOp printAst cfg d "xfromK" [t]+    (STKR{}, FTKScalar) -> printPrefixOp printAst cfg d "kfromR" [t]+    (STKR{}, FTKS{}) -> printPrefixOp printAst cfg d "sfromR" [t]+    (STKR{}, FTKX{}) -> printPrefixOp printAst cfg d "xfromR" [t]+    (STKS{}, FTKScalar) -> printPrefixOp printAst cfg d "kfromS" [t]+    (STKS{}, FTKR{}) -> printPrefixOp printAst cfg d "rfromS" [t]+    (STKS{}, FTKX{}) -> printPrefixOp printAst cfg d "xfromS" [t]+    (STKX{}, FTKScalar) -> printPrefixOp printAst cfg d "kfromX" [t]+    (STKX{}, FTKR{}) -> printPrefixOp printAst cfg d "rfromX" [t]+    (STKX{}, FTKS{}) -> printPrefixOp printAst cfg d "sfromX" [t]+    (ystk, _) -> let s = "tconvert (" ++ show c ++ ") (" ++ show ystk ++ ")"+                 in printPrefixOp printAst cfg d s [t]++  AstDot0 u v ->+    printPrefixOp printAst cfg d "sdot0" [u, v]+  AstDot1InS _ _ u v ->+    printPrefixOp printAst cfg d "sdot1In" [u, v]+  AstMatmul2S SNat SNat SNat u v ->+    showParen (d > 10)+    $ showString "smatmul2 "+      . printAst cfg 11 u+      . showString " "+      . printAst cfg 11 v++  AstBoolNotK u -> printPrefixOp printAst cfg d "notB" [u]+  AstBoolNotS u -> printPrefixOp printAst cfg d "notS" [u]  -- TODO+  AstBoolAndK u v -> printBinaryOp printAst cfg d u (3, "&&*") v+  AstBoolAndS u v -> printBinaryOp printAst cfg d u (3, "`andS`") v  -- TODO+  AstLeqK u v -> printBinaryOp printAst cfg d u (4, "<=.") v+  AstLeq u v -> printBinaryOp printAst cfg d u (4, "<=.") v+  AstLeqS _ _ u v -> printBinaryOp printAst cfg d u (4, "`leqS`") v  -- TODO++showNumber :: Show a => a -> ShowS+{-# INLINE showNumber #-}+showNumber a = case show a of+  "-0.0" -> showString "0.0"  -- conserve sanity+  s@('-' : _) -> showParen True (showString s)+  s -> showString s++-- Differs from standard only in the space after comma.+showListWith :: (a -> ShowS) -> [a] -> ShowS+{-# INLINE showListWith #-}+showListWith = showCollectionWith "[" ", " "]"++showCollectionWith :: String -> String -> String -> (a -> ShowS) -> [a] -> ShowS+{-# INLINE showCollectionWith #-}+showCollectionWith start _   end _     []     s = start ++ end ++ s+showCollectionWith start sep end showx (x:xs) s = start ++ showx x (showl xs)+ where+  showl []     = end ++ s+  showl (y:ys) = sep ++ showx y (showl ys)++printAstHFun :: KnownSpan s+             => PrintConfig -> Int -> AstHFun s x y -> ShowS+printAstHFun cfg d = \case+  AstLambda !var !l ->+    if loseRoudtrip cfg+    then if ignoreNestedLambdas cfg+         then showString "<lambda>"+         else showParen (d > 0)+              $ showString "\\"+                . printAstVar cfg var+                . showString " -> "+                . printAst cfg 0 l+    else showParen (d > 0)+         $ showString "tlambda $ \\"+           . printAstVar cfg var+           . showString " -> "+           . printAst cfg 0 l++printAstHFunOneUnignore :: KnownSpan s+                        => PrintConfig -> Int -> AstHFun s x y -> ShowS+printAstHFunOneUnignore cfg d = \case+  AstLambda !var !l ->+    if loseRoudtrip cfg+    then showParen (d > 0)+         $ showString "\\"+           . printAstVar cfg var+           . showString " -> "+           . printAst cfg 0 l+    else showParen (d > 0)+         $ showString "tlambda $ \\"+           . printAstVar cfg var+           . showString " -> "+           . printAst cfg 0 l++printAstN1R :: (PrintConfig -> Int -> a -> ShowS)+            -> PrintConfig -> Int -> OpCodeNum1 -> a -> ShowS+{-# INLINE printAstN1R #-}+printAstN1R pr cfg d opCode u = case opCode of+  NegateOp -> printPrefixOp pr cfg d "negate" [u]+  AbsOp -> printPrefixOp pr cfg d "abs" [u]+  SignumOp -> printPrefixOp pr cfg d "signum" [u]++printAstR1R :: (PrintConfig -> Int -> a -> ShowS)+            -> PrintConfig -> Int -> OpCode1 -> a -> ShowS+{-# INLINE printAstR1R #-}+printAstR1R pr cfg d opCode u = case opCode of+  RecipOp -> printPrefixOp pr cfg d "recip" [u]+  ExpOp -> printPrefixOp pr cfg d "exp" [u]+  LogOp -> printPrefixOp pr cfg d "log" [u]+  SqrtOp -> printPrefixOp pr cfg d "sqrt" [u]+  SinOp -> printPrefixOp pr cfg d "sin" [u]+  CosOp -> printPrefixOp pr cfg d "cos" [u]+  TanOp -> printPrefixOp pr cfg d "tan" [u]+  AsinOp -> printPrefixOp pr cfg d "asin" [u]+  AcosOp -> printPrefixOp pr cfg d "acos" [u]+  AtanOp -> printPrefixOp pr cfg d "atan" [u]+  SinhOp -> printPrefixOp pr cfg d "sinh" [u]+  CoshOp -> printPrefixOp pr cfg d "cosh" [u]+  TanhOp -> printPrefixOp pr cfg d "tanh" [u]+  AsinhOp -> printPrefixOp pr cfg d "asinh" [u]+  AcoshOp -> printPrefixOp pr cfg d "acosh" [u]+  AtanhOp -> printPrefixOp pr cfg d "atanh" [u]++printAstR2R :: (PrintConfig -> Int -> a -> ShowS)+            -> PrintConfig -> Int -> OpCode2 -> a -> a -> ShowS+{-# INLINE printAstR2R #-}+printAstR2R pr cfg d opCode u v = case opCode of+  DivideOp -> printBinaryOp pr cfg d u (7, "/") v+  PowerOp -> printBinaryOp pr cfg d u (8, "**") v+  LogBaseOp -> printPrefixOp pr cfg d "logBase" [u, v]+  Atan2Op -> printPrefixOp pr cfg d "atan2H" [u, v]++printAstI2R :: (PrintConfig -> Int -> a -> ShowS)+            -> PrintConfig -> Int -> OpCodeIntegral2 -> a -> a -> ShowS+{-# INLINE printAstI2R #-}+printAstI2R pr cfg d opCode u v = case opCode of+  QuotOp -> printPrefixOp pr cfg d "quotH" [u, v]+  RemOp -> printPrefixOp pr cfg d "remH" [u, v]++printPrefixOp :: (PrintConfig -> Int -> a -> ShowS)+              -> PrintConfig -> Int -> String -> [a]+              -> ShowS+{-# INLINE printPrefixOp #-}+printPrefixOp pr cfg d funcname args =+  let rs = map (\arg -> showString " " . pr cfg 11 arg) args+  in showParen (d > 10)+     $ showString funcname+       . foldr (.) id rs++printBinaryOp :: (PrintConfig -> Int -> a -> ShowS)+              -> PrintConfig -> Int -> a -> (Int, String) -> a+              -> ShowS+{-# INLINE printBinaryOp #-}+printBinaryOp pr cfg d left (prec, opstr) right =+  showParen (d > prec)+  $ pr cfg (prec + 1) left+    . showString (" " ++ opstr ++ " ")+    . pr cfg (prec + 1) right
src/HordeAd/Core/TensorKind.hs view
@@ -1,3 +1,4 @@+{-# LANGUAGE LambdaCase #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} -- | Two kinds of singletons for tensor kindss and constraints@@ -5,32 +6,32 @@ module HordeAd.Core.TensorKind   ( -- * Tensor kind singletons     SingletonTK(..), KnownSTK(..)-  , TKConversion(..), convertSTK, convertFTK, buildTKConversion-  , withKnownSTK, lemKnownSTK, sameKnownSTK, sameSTK+  , withKnownSTK, lemKnownSTK, sameSTK+  , lemTKAllNumAD, lemTKScalarAllNumAD, numFromTKAllNum   , stkUnit, buildSTK, razeSTK, adSTK-  , lemKnownSTKOfBuild, lemKnownSTKOfAD, lemBuildOfAD, lengthSTK, widthSTK+  , lemKnownSTKOfBuild, lemBuildOfAD, lengthSTK, widthSTK     -- * Full shape tensor kind quasi-singletons   , FullShapeTK(..)   , matchingFTK, ftkToSTK, ftkUnit, buildFTK, razeFTK, adFTK, differentiableFTK-  , DummyDualTarget(..)   ) where  import Prelude hiding ((.))  import Control.Category-import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Proxy (Proxy)+import Data.Type.Equality (testEquality, (:~:) (Refl))+import Foreign.C (CInt) import GHC.Exts (withDict)-import GHC.TypeLits (KnownNat, OrderingI (..), cmpNat, fromSNat, type (+))-import Type.Reflection (typeRep)+import GHC.TypeLits (KnownNat)+import Type.Reflection (Typeable, typeRep) -import Data.Array.Nested (MapJust, Replicate, type (++))-import Data.Array.Nested.Convert (shrFromShX, shsFromShX, shxFromShS, shxFromShR, shsFromSSX)-import Data.Array.Nested.Lemmas+import Data.Array.Nested qualified as Nested import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (unsafeCoerceRefl)+import Data.Array.Nested.Types+  (fromSNat', pattern SS, pattern SZ, unsafeCoerceRefl)  import HordeAd.Core.Types @@ -72,11 +73,11 @@          => KnownSTK (TKProduct y z) where   knownSTK = STKProduct (knownSTK @y) (knownSTK @z) --- | Turning a singleton into a constraint via a continuation.+-- | Turn a singleton into a constraint via a continuation. withKnownSTK :: forall y r. SingletonTK y -> (KnownSTK y => r) -> r withKnownSTK = withDict @(KnownSTK y) --- | Turning a singleton into a dictionary containing constraint.+-- | Turn a singleton into a dictionary. lemKnownSTK :: SingletonTK y -> Dict KnownSTK y lemKnownSTK = \case   STKScalar -> Dict@@ -86,30 +87,62 @@   STKProduct stk1 stk2 | Dict <- lemKnownSTK stk1                        , Dict <- lemKnownSTK stk2 -> Dict -sameKnownSTK :: forall y1 y2. (KnownSTK y1, KnownSTK y2)-             => Maybe (y1 :~: y2)-sameKnownSTK = sameSTK (knownSTK @y1) (knownSTK @y2)- -- | A plausible implementation of `testEquality` on `SingletonTK`. sameSTK :: SingletonTK y1 -> SingletonTK y2 -> Maybe (y1 :~: y2)-sameSTK stk1 stk2 = case (stk1, stk2) of-  (STKScalar @r1, STKScalar @r2)+{-# INLINE sameSTK #-}+sameSTK = \cases+  (STKScalar @r1) (STKScalar @r2)     | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->       Just Refl-  (STKR snat1 x1, STKR snat2 x2)+  -- A common non-recursive case as a shorthand:+  (STKS sh1 (STKScalar @r1)) (STKS sh2 (STKScalar @r2))+    | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+    , Just Refl <- testEquality sh1 sh2 ->+      Just Refl+  (STKR snat1 x1) (STKR snat2 x2)     | Just Refl <- sameSTK x1 x2, Just Refl <- testEquality snat1 snat2 ->       Just Refl-  (STKS sh1 x1, STKS sh2 x2)+  (STKS sh1 x1) (STKS sh2 x2)     | Just Refl <- sameSTK x1 x2, Just Refl <- testEquality sh1 sh2 ->       Just Refl-  (STKX sh1 x1, STKX sh2 x2)+  (STKX sh1 x1) (STKX sh2 x2)     | Just Refl <- sameSTK x1 x2, Just Refl <- testEquality sh1 sh2 ->       Just Refl-  (STKProduct x1 y1, STKProduct x2 y2)+  (STKProduct x1 y1) (STKProduct x2 y2)     | Just Refl <- sameSTK x1 x2, Just Refl <- sameSTK y1 y2 ->       Just Refl-  _ -> Nothing+  _ _ -> Nothing +lemTKAllNumAD :: SingletonTK y -> Dict0 (TKAllNum (ADTensorKind y))+lemTKAllNumAD = \case+  STKScalar @r -> ifDifferentiable @r Dict0 Dict0+  STKR _ x | Dict0 <- lemTKAllNumAD x -> Dict0+  STKS _ x | Dict0 <- lemTKAllNumAD x -> Dict0+  STKX _ x | Dict0 <- lemTKAllNumAD x -> Dict0+  STKProduct stk1 stk2 | Dict0 <- lemTKAllNumAD stk1+                       , Dict0 <- lemTKAllNumAD stk2 -> Dict0++lemTKScalarAllNumAD :: forall r. Typeable r+                    => Proxy r+                    -> Dict0 (TKAllNum (TKScalar (ADTensorScalar r)))+lemTKScalarAllNumAD _ = ifDifferentiable @r Dict0 Dict0++-- Despite what GHC says, TKAllNum (TKScalar r) is not redundant,+-- because it ensures the error case can't appear.+numFromTKAllNum :: forall r. (Typeable r, TKAllNum (TKScalar r))+                => Proxy r -> Dict0 (Num r, Nested.NumElt r)+numFromTKAllNum _ = case typeRep @r of+  Is @Int -> Dict0+  Is @Double -> Dict0+  Is @Float -> Dict0+  Is @Int8 -> Dict0+  Is @Int16 -> Dict0+  Is @Int32-> Dict0+  Is @Int64 -> Dict0+  Is @CInt -> Dict0+  Is @Z1 -> Dict0+  _ -> error "numFromTKAllNum: impossible type"+ stkUnit :: SingletonTK TKUnit stkUnit = STKScalar @@ -121,16 +154,19 @@   STKX sh x -> STKX (SKnown snat :!% sh) x   STKProduct stk1 stk2 -> STKProduct (buildSTK snat stk1) (buildSTK snat stk2) +-- The argument is assumed to be the result of buildSTK. razeSTK :: SingletonTK z -> SingletonTK (RazeTensorKind z)-razeSTK = \case+razeSTK stk = case stk of   STKScalar -> error "razeSTK: impossible argument"-  STKR snat@SNat x ->-    case cmpNat (SNat @1) snat of-      LTI -> STKR SNat x-      EQI -> STKR SNat x-      _ -> error "razeSTK: impossible argument"+  STKR SZ _ -> error "razeSTK: impossible argument"+  STKR (SS snat) x -> STKR snat x   STKS ZSS _ -> error "razeSTK: impossible argument"-  STKS (_ :$$ sh) x -> STKS sh x+  STKS (_ :$$ ZSS) STKScalar -> STKScalar+  STKS (_ :$$ ZSS) x@STKR{} -> STKS ZSS x+  STKS (_ :$$ ZSS) x@STKS{}  -> STKS ZSS x+  STKS (_ :$$ ZSS) x@STKX{}  -> STKS ZSS x+  STKS (_ :$$ ZSS) x@STKProduct{}  -> STKS ZSS x+  STKS (_ :$$ m :$$ sh) x -> STKS (m :$$ sh) x   STKX ZKX _ -> error "razeSTK: impossible argument"   STKX (SUnknown _ :!% _) _ -> error "razeSTK: impossible argument"   STKX (SKnown _ :!% sh) x -> STKX sh x@@ -138,25 +174,16 @@  adSTK :: SingletonTK y -> SingletonTK (ADTensorKind y) adSTK = \case-  t@(STKScalar @r) -> case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> t-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> t-      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1)-           STKScalar+  t@(STKScalar @r) -> ifDifferentiable @r t STKScalar   STKR sh x -> STKR sh $ adSTK x   STKS sh x -> STKS sh $ adSTK x   STKX sh x -> STKX sh $ adSTK x   STKProduct stk1 stk2 -> STKProduct (adSTK stk1) (adSTK stk2)  lemKnownSTKOfBuild :: SNat k -> SingletonTK y-                     -> Dict KnownSTK (BuildTensorKind k y)+                   -> Dict KnownSTK (BuildTensorKind k y) lemKnownSTKOfBuild snat = lemKnownSTK . buildSTK snat -lemKnownSTKOfAD :: SingletonTK y-                  -> Dict KnownSTK (ADTensorKind y)-lemKnownSTKOfAD = lemKnownSTK . adSTK- lemBuildOfAD :: SNat k -> SingletonTK y              -> BuildTensorKind k (ADTensorKind y)                 :~: ADTensorKind (BuildTensorKind k y)@@ -170,7 +197,7 @@  lengthSTK :: SingletonTK x -> Int lengthSTK STKScalar = 0-lengthSTK (STKR snat _) = fromInteger $ fromSNat snat+lengthSTK (STKR snat _) = fromSNat' snat lengthSTK (STKS sh _) = shsLength sh lengthSTK (STKX sh _) = ssxLength sh lengthSTK (STKProduct sy sz) = lengthSTK sy `max` lengthSTK sz@@ -185,169 +212,7 @@   STKX{} -> 1   STKProduct stk1 stk2 -> widthSTK stk1 + widthSTK stk2 --- | This is copied, with modifications, from ox-arrays.------ This is a recipe for converting arrays, not always followed,--- and a proof a conversion is possible, with some proof obligations--- delayed to runtime (in ConvXS' and ConvXX', where not only the ranks--- of the shapes need to agree, but also the dimensions of the input--- array and of the output shape, which is not all captured in the type).--- As in ox-arrays, conversions only change the meta-data, not the underlying--- vector representation of the array.-type role TKConversion nominal nominal-data TKConversion (a :: TK) (b :: TK) where-  ConvId  :: TKConversion a a-  ConvCmp :: TKConversion b c -> TKConversion a b -> TKConversion a c -  ConvRX  :: TKConversion (TKR2 n a) (TKX2 (Replicate n Nothing) a)-  ConvSX  :: TKConversion (TKS2 sh a) (TKX2 (MapJust sh) a)--  ConvXR  :: SingletonTK a -> TKConversion (TKX2 sh a) (TKR2 (Rank sh) a)-  ConvXS  :: TKConversion (TKX2 (MapJust sh) a) (TKS2 sh a)-  ConvXS' :: Rank sh ~ Rank sh'-          => FullShapeTK (TKS2 sh' a)-          -> TKConversion (TKX2 sh a) (TKS2 sh' a)--  ConvXX' :: Rank sh ~ Rank sh'-          => FullShapeTK (TKX2 sh' a)-          -> TKConversion (TKX2 sh a) (TKX2 sh' a)--  ConvRR  :: TKConversion a b -> TKConversion (TKR2 n a) (TKR2 n b)-  ConvSS  :: TKConversion a b -> TKConversion (TKS2 sh a) (TKS2 sh b)-  ConvXX  :: TKConversion a b -> TKConversion (TKX2 sh a) (TKX2 sh b)-  ConvT2  :: TKConversion a a'-          -> TKConversion b b'-          -> TKConversion (TKProduct a b) (TKProduct a' b')--  Conv0X  :: SingletonTK a -> TKConversion a (TKX2 '[] a)-  ConvX0  :: TKConversion (TKX2 '[] a) a--  ConvNest :: SingletonTK (TKX2 sh a)-           -> TKConversion (TKX2 (sh ++ sh') a) (TKX2 sh (TKX2 sh' a))-  ConvUnnest :: TKConversion (TKX2 sh (TKX2 sh' a)) (TKX2 (sh ++ sh') a)--  ConvZip   :: SingletonTK a -> SingletonTK b-            -> TKConversion (TKProduct (TKX2 sh a) (TKX2 sh b))-                            (TKX2 sh (TKProduct a b))-  ConvUnzip :: SingletonTK a -> SingletonTK b-            -> TKConversion (TKX2 sh (TKProduct a b))-                            (TKProduct (TKX2 sh a) (TKX2 sh b))--deriving instance Show (TKConversion a b)--instance Category TKConversion where-  id = ConvId-  (.) = ConvCmp--convertSTK :: TKConversion a b -> SingletonTK a -> SingletonTK b-convertSTK = \cases-  ConvId astk -> astk-  (ConvCmp c1 c2) astk -> convertSTK c1 (convertSTK c2 astk)-  ConvRX (STKR n a) -> STKX (ssxReplicate n) a-  ConvSX (STKS sh a) -> STKX (ssxFromShX $ shxFromShS sh) a-  (ConvXR _stk) (STKX ssx a) -> STKR (ssxRank ssx) a-  ConvXS (STKX ssx a) -> STKS (shsFromSSX ssx) a-  (ConvXS' (FTKS sh _x)) (STKX _ssx2 a) -> STKS sh a-  (ConvXX' (FTKX shx _x)) (STKX _ssx2 a) -> STKX (ssxFromShX shx) a-  (ConvRR c) (STKR n a) -> STKR n (convertSTK c a)-  (ConvSS c) (STKS sh a) -> STKS sh (convertSTK c a)-  (ConvXX c) (STKX ssx a) -> STKX ssx (convertSTK c a)-  (ConvT2 c1 c2) (STKProduct stk1 stk2) ->-    STKProduct (convertSTK c1 stk1) (convertSTK c2 stk2)-  (Conv0X _stk) stk -> STKX ZKX stk-  ConvX0 (STKX ZKX stk) -> stk-  (ConvNest (STKX ssx x)) (STKX shsh' _x) ->-    STKX ssx (STKX (ssxDropSSX ssx shsh') x)-  ConvUnnest (STKX sh (STKX sh' x)) -> STKX (sh `ssxAppend` sh') x-  (ConvZip _ _) (STKProduct (STKX sh a1) (STKX _sh a2)) ->-    STKX sh (STKProduct a1 a2)-  (ConvUnzip _ _) (STKX sh (STKProduct a1 a2)) ->-    STKProduct (STKX sh a1) (STKX sh a2)--convertFTK :: TKConversion a b -> FullShapeTK a -> FullShapeTK b-convertFTK = \cases-  ConvId aftk -> aftk-  (ConvCmp c1 c2) aftk -> convertFTK c1 (convertFTK c2 aftk)-  ConvRX (FTKR shr a) -> FTKX (shxFromShR shr) a-  ConvSX (FTKS sh a) -> FTKX (shxFromShS sh) a-  (ConvXR _stk) (FTKX shx a) -> FTKR (shrFromShX shx) a-  ConvXS (FTKX shx a) -> FTKS (shsFromShX shx) a-  (ConvXS' ftk) _ -> ftk-  (ConvXX' ftk) _ -> ftk-  (ConvRR c) (FTKR shr a) -> FTKR shr (convertFTK c a)-  (ConvSS c) (FTKS sh a) -> FTKS sh (convertFTK c a)-  (ConvXX c) (FTKX shx a) -> FTKX shx (convertFTK c a)-  (ConvT2 c1 c2) (FTKProduct ftk1 ftk2) ->-    FTKProduct (convertFTK c1 ftk1) (convertFTK c2 ftk2)-  (Conv0X _stk) ftk -> FTKX ZSX ftk-  ConvX0 (FTKX ZSX ftk) -> ftk-  (ConvNest @_ @_ @sh' (STKX ssx _x)) (FTKX shsh' x) ->-    FTKX (shxTakeSSX (Proxy @sh') ssx shsh') (FTKX (shxDropSSX ssx shsh') x)-  ConvUnnest (FTKX sh (FTKX sh' x)) -> FTKX (sh `shxAppend` sh') x-  (ConvZip _ _) (FTKProduct (FTKX sh a1) (FTKX _sh a2)) ->-    FTKX sh (FTKProduct a1 a2)-  (ConvUnzip _ _) (FTKX sh (FTKProduct a1 a2)) ->-    FTKProduct (FTKX sh a1) (FTKX sh a2)--buildTKConversion :: SNat k -> FullShapeTK a-                  -> TKConversion a b-                  -> TKConversion (BuildTensorKind k a) (BuildTensorKind k b)-buildTKConversion k aftk c0 = case c0 of-  ConvId -> ConvId-  ConvCmp c1 c2 -> ConvCmp (buildTKConversion k (convertFTK c2 aftk) c1)-                           (buildTKConversion k aftk c2)-  ConvRX | FTKR @n shr xstk <- aftk-         , Refl <- lemRankReplicate (Proxy @n)-         , Refl <- lemRankReplicate (Proxy @(1 + n)) ->-    ConvCmp (ConvXX' (FTKX (SKnown k :$% shxFromShR shr) xstk)) ConvRX-  ConvSX -> ConvSX-  ConvXR stk -> ConvXR stk-  ConvXS -> ConvXS-  ConvXS' ftk -> ConvXS' (buildFTK k ftk)-  ConvXX' ftk -> ConvXX' (buildFTK k ftk)-  ConvRR c -> ConvRR c-  ConvSS c -> ConvSS c-  ConvXX c -> ConvXX c-  ConvT2 c1 c2 | FTKProduct ftk1 ftk2 <- aftk ->-    ConvT2 (buildTKConversion k ftk1 c1) (buildTKConversion k ftk2 c2)-  Conv0X _astk -> case aftk of-    FTKScalar -> ConvSX-    FTKR @n shr x | Refl <- lemRankReplicate (Proxy @n)-                  , Refl <- lemRankReplicate (Proxy @(1 + n)) ->-      ConvCmp (ConvXX (ConvXR (ftkToSTK x)))-              (ConvCmp (ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x)))-                       (ConvCmp-                          (ConvXX' (FTKX (SKnown k :$% shxFromShR shr) x))-                          ConvRX))-    FTKS _sh x ->-      ConvCmp (ConvXX ConvXS)-              (ConvCmp (ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x)))-                       ConvSX)-    FTKX _ssx x -> ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x))-    FTKProduct aftk1 aftk2 ->-      buildTKConversion-        k aftk (ConvCmp (ConvZip (ftkToSTK aftk1) (ftkToSTK aftk2))-                        (ConvT2 (Conv0X (ftkToSTK aftk1))-                                (Conv0X (ftkToSTK aftk2))))-  ConvX0 -> case aftk of-    FTKX ZSX FTKScalar -> ConvXS-    FTKX ZSX (FTKR @n _n x) | Refl <- lemRankReplicate (Proxy @n) ->-      ConvCmp (ConvXR (ftkToSTK x))-              (ConvCmp ConvUnnest (ConvXX ConvRX))-    FTKX ZSX FTKS{} ->-      ConvCmp ConvXS-              (ConvCmp ConvUnnest (ConvXX ConvSX))-    FTKX ZSX FTKX{} -> ConvUnnest-    FTKX ZSX (FTKProduct aftk1 aftk2) ->-      buildTKConversion-        k aftk (ConvCmp (ConvT2 ConvX0 ConvX0)-                        (ConvUnzip (ftkToSTK aftk1) (ftkToSTK aftk2)))-  ConvNest (STKX sh x) -> ConvNest (STKX (SKnown k :!% sh) x)-  ConvUnnest -> ConvUnnest-  ConvZip astk1 astk2 -> ConvZip astk1 astk2-  ConvUnzip astk1 astk2 -> ConvUnzip astk1 astk2-- -- * Full shape tensor kind quasi-singletons  -- | Full shape tensor kind singleton type.@@ -368,26 +233,32 @@ -- take into account shape difference in ranked and mixed tensors -- that `FullShapeTK`, but not `SingletonTK`, captures. matchingFTK :: FullShapeTK y1 -> FullShapeTK y2 -> Maybe (y1 :~: y2)-matchingFTK ftk1 ftk2 = case (ftk1, ftk2) of-  (FTKScalar @r1, FTKScalar @r2)+{-# INLINE matchingFTK #-}+matchingFTK = \cases+  (FTKScalar @r1) (FTKScalar @r2)     | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->       Just Refl-  (FTKR sh1 x1, FTKR sh2 x2)+  -- A common non-recursive case as a shorthand:+  (FTKS sh1 (FTKScalar @r1)) (FTKS sh2 (FTKScalar @r2))+    | Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+    , Just Refl <- testEquality sh1 sh2 ->+      Just Refl+  (FTKR sh1 x1) (FTKR sh2 x2)     | Just Refl <- matchingFTK x1 x2     , Just Refl <- testEquality (shrRank sh1) (shrRank sh2) ->  -- weaker!!!       Just Refl-  (FTKS sh1 x1, FTKS sh2 x2)+  (FTKS sh1 x1) (FTKS sh2 x2)     | Just Refl <- matchingFTK x1 x2     , Just Refl <- testEquality sh1 sh2 ->       Just Refl-  (FTKX sh1 x1, FTKX sh2 x2)+  (FTKX sh1 x1) (FTKX sh2 x2)     | Just Refl <- matchingFTK x1 x2     , Just Refl <- testEquality (ssxFromShX sh1) (ssxFromShX sh2) ->  -- !!!       Just Refl-  (FTKProduct x1 y1, FTKProduct x2 y2)+  (FTKProduct x1 y1) (FTKProduct x2 y2)     | Just Refl <- matchingFTK x1 x2, Just Refl <- matchingFTK y1 y2 ->       Just Refl-  _ -> Nothing+  _ _ -> Nothing  -- | A conversion that is fully determined by the property that it -- commutes with the `testEquality` implementations.@@ -405,19 +276,19 @@ buildFTK :: SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y) buildFTK snat@SNat = \case   FTKScalar -> FTKS (snat :$$ ZSS) FTKScalar-  FTKR sh x -> FTKR (sNatValue snat :$: sh) x+  FTKR sh x -> FTKR (fromSNat' snat :$: sh) x   FTKS sh x -> FTKS (snat :$$ sh) x   FTKX sh x -> FTKX (SKnown snat :$% sh) x   FTKProduct ftk1 ftk2 -> FTKProduct (buildFTK snat ftk1) (buildFTK snat ftk2) +-- Depite the warning, the pattern match is exhaustive and if a dummy+-- pattern is added, GHC 9.14.1 complains about that, in turn. razeFTK :: forall y k.-           SNat k -> SingletonTK y-        -> FullShapeTK (BuildTensorKind k y)+           SNat k -> SingletonTK y -> FullShapeTK (BuildTensorKind k y)         -> FullShapeTK y-razeFTK snat@SNat stk ftk = case (stk, ftk) of+razeFTK snat stk ftk = case (stk, ftk) of   (STKScalar, FTKS (_ :$$ ZSS) FTKScalar) -> FTKScalar   (STKR{}, FTKR (_ :$: sh) x) -> FTKR sh x-  (STKR{}, FTKR ZSR _) -> error "razeFTK: impossible built tensor kind"   (STKS{}, FTKS (_ :$$ sh) x) -> FTKS sh x   (STKX{}, FTKX (_ :$% sh) x) -> FTKX sh x   (STKProduct stk1 stk2, FTKProduct ftk1 ftk2) ->@@ -426,12 +297,7 @@ adFTK :: FullShapeTK y       -> FullShapeTK (ADTensorKind y) adFTK = \case-  t@(FTKScalar @r) -> case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> t-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> t-      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $-           FTKScalar @Z1+  t@(FTKScalar @r) -> ifDifferentiable @r t FTKScalar   FTKR sh x -> FTKR sh $ adFTK x   FTKS sh x -> FTKS sh $ adFTK x   FTKX sh x -> FTKX sh $ adFTK x@@ -441,16 +307,8 @@ -- from any non-differentiable types, such as integers. differentiableFTK :: FullShapeTK y -> Bool differentiableFTK = \case-  FTKScalar @r -> case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> True-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> True-      _ -> False+  FTKScalar @r -> ifDifferentiable @r True False   FTKR _ x -> differentiableFTK x   FTKS _ x -> differentiableFTK x   FTKX _ x -> differentiableFTK x   FTKProduct ftk1 ftk2 -> differentiableFTK ftk1 && differentiableFTK ftk2--type role DummyDualTarget nominal-type DummyDualTarget :: Target-newtype DummyDualTarget y = DummyDualTarget (FullShapeTK y)
src/HordeAd/Core/Types.hs view
@@ -1,5 +1,5 @@-{-# LANGUAGE AllowAmbiguousTypes, DerivingVia, ImpredicativeTypes,-             UndecidableInstances, ViewPatterns #-}+{-# LANGUAGE AllowAmbiguousTypes, DerivingVia, UndecidableInstances,+             UndecidableSuperClasses, ViewPatterns #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} {-# OPTIONS_GHC -Wno-orphans #-}@@ -7,81 +7,80 @@ module HordeAd.Core.Types   ( -- * Re-exports and definitions to help express and manipulate type-level natural numbers     SNat, pattern SNat, pattern SNat'-  , withSNat, sNatValue, proxyFromSNat, valueOf+  , withSNat, valueOf, snatSucc     -- * Kinds of the parameterized types that determine the structure of a tensor-  , Target, TK (..), TKR, TKS, TKX, TKUnit+  , Target, TK (..), TKR, TKS, TKX, TKUnit, TKAllNum     -- * Some fundamental constraints and types related to tensors-  , GoodScalar, GoodScalarConstraint, Differentiable, IfDifferentiable(..)+  , GoodScalar, NumScalar, GoodScalarConstraint+  , Differentiable, ifDifferentiable   , BuildTensorKind, RazeTensorKind, ADTensorKind, ADTensorScalar     -- * Type families that tensors belong to-  , IntOf, HFunOf, PrimalOf, DualOf, ShareOf, BoolOf+  , IntOf, HFunOf, PrimalOf, DualOf, PlainOf, ShareOf, BoolOf   , IxROf, IxSOf, IxXOf     -- * The Z1 Num unit type and its instances   , Z1(..)     -- * Misc-  , Dict(..), IntegralH(..), RealFloatH(..), Boolean (..), EqH(..), OrdH(..)+  , pattern Is, mapAccumL'+  , Dict(..), Dict0(..)+  , IntegralH(..), RealFloatH(..), Boolean (..), EqH(..), OrdH(..)   , backpermutePrefixList-  , toLinearIdxR, fromLinearIdxR, toLinearIdxS, fromLinearIdxS-  , toLinearIdxX, fromLinearIdxX     -- * Feature requests for ox-arrays   , Take, Drop, UnMapSucc-  , listsTake, listsDrop, listsSplitAt, ixrTake, ixrDrop, ixrSplitAt-  , shrTake, shrDrop, shrSplitAt-  , listrSplitAt, ixsDrop, shsTake, shsDrop-  , shxTake, shxDrop, ixxTake, ixxDrop'-  , listsTakeLen, listsDropLen, shsDropLen-  , permRInverse, ssxPermutePrefix, shxPermutePrefix-  , shCastSX-  , shsFromIxS, shsFromListS-  , normalizePermutationHack, backpermCycle, permCycle+  -- TODO: ixrTake, ixrDrop, ixrSplitAt+  , shrTake, shrDrop+  , ixsTake, ixsDrop, ixsSplitAt, shsTake, shsDrop+  -- TODO: ixxTake, ixxDrop'+  , shxTake, shxDrop+  , permRInverse, ssxPermutePrefix, shxPermutePrefix, shCastSX+  , normalizePermutationHack+  , backpermCycle, permCycle, backpermCycleN, permCycleN   , permUnShift1-  , ssxTakeIx   ) where  import Prelude  import Control.DeepSeq (NFData (..))-import Data.Array.Internal.RankedS qualified as RS import Data.Boolean (Boolean (..))-import Data.Coerce (coerce) import Data.Default-import Data.Functor.Const-import Data.Int (Int64)-import Data.Kind (Type)+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Kind (Constraint, Type) import Data.List (dropWhileEnd, sort) import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, (:~:) (Refl))+import Data.Type.Equality (gcastWith, type (~~), (:~:) (Refl), (:~~:) (HRefl)) import Data.Vector.Storable qualified as V import Foreign.C (CInt) import Foreign.Storable (Storable (..))-import GHC.Exts (IsList (..))+import GHC.Exts (oneShot) import GHC.Generics (Generic) import GHC.TypeLits-  ( KnownNat+  ( ErrorMessage (..)+  , KnownNat   , Nat   , SNat-  , fromSNat+  , TypeError   , pattern SNat   , sameNat   , type (+)   , type (-)-  , withSomeSNat   )+import GHC.TypeNats qualified as TN import System.Random-import Type.Reflection (Typeable)+import Type.Reflection (TypeRep, Typeable, eqTypeRep, pattern TypeRep, typeRep) import Unsafe.Coerce (unsafeCoerce) -import Data.Array.Nested (type (++))+import Data.Array.Nested (MapJust, Replicate) import Data.Array.Nested qualified as Nested import Data.Array.Nested.Convert (shxFromShS) import Data.Array.Nested.Mixed qualified as Mixed import Data.Array.Nested.Mixed.Shape-import Data.Array.Nested.Permutation (DropLen, PermR, TakeLen)+import Data.Array.Nested.Permutation (PermR) import Data.Array.Nested.Permutation qualified as Permutation import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape-import Data.Array.Nested.Types (Dict (..), Tail, fromSNat', unsafeCoerceRefl)-import Data.Array.Strided.Orthotope (NumElt (..), fromO, toO)+import Data.Array.Nested.Types+  (Dict (..), Tail, fromSNat', pattern SS, pattern SZ, unsafeCoerceRefl)+import Data.Array.Strided.Arith (NumElt (..))+import Data.Array.Strided.Array as SA  -- * Definitions to help express and manipulate type-level natural numbers @@ -89,22 +88,15 @@   rnf _ = ()  withSNat :: Int -> (forall n. KnownNat n => (SNat n -> r)) -> r-withSNat i f = withSomeSNat (fromIntegral i) $ \case-  Just snat@SNat -> f snat-  Nothing -> error $ "withSNat: negative argument: " ++ show i--sNatValue :: forall n. SNat n -> Int-{-# INLINE sNatValue #-}-sNatValue = fromInteger . fromSNat--proxyFromSNat :: SNat n -> Proxy n-proxyFromSNat SNat = Proxy+{-# INLINE withSNat #-}+withSNat i f = TN.withSomeSNat (toEnum i) $ \case+  snat@SNat -> f snat  {-# INLINE valueOf #-} valueOf :: forall n r. (KnownNat n, Num r) => r-valueOf = fromInteger $ fromSNat (SNat @n)+valueOf = fromIntegral $ fromSNat' (SNat @n) -pattern SNat' :: forall n m. KnownNat n => (KnownNat m, n ~ m) => SNat m+pattern SNat' :: forall n m. KnownNat n => n ~ m => SNat m pattern SNat' <- (matchSNat (Proxy @n) -> Just (Refl :: n :~: m))   where SNat' = SNat @@ -112,7 +104,10 @@           => proxy n -> SNat m -> Maybe (n :~: m) matchSNat p m@SNat = sameNat p m +snatSucc :: SNat n -> SNat (1 + n)+snatSucc SNat = SNat + -- * Kinds of the functors that determine the structure of a tensor type  -- | The parameterized carrier types of the algebras@@ -137,73 +132,97 @@  type TKUnit = TKScalar Z1 +type family TKAllNum (y :: TK) :: Constraint where+  TKAllNum (TKScalar Int64) = ()+  TKAllNum (TKScalar Int32) = ()+  TKAllNum (TKScalar Int16) = ()+  TKAllNum (TKScalar Int8) = ()+  TKAllNum (TKScalar Int) = ()+  TKAllNum (TKScalar CInt) = ()+  TKAllNum (TKScalar Double) = ()+  TKAllNum (TKScalar Float) = ()+  TKAllNum (TKScalar Z1) = ()+  TKAllNum (TKR2 _ y) = TKAllNum y+  TKAllNum (TKS2 _ y) = TKAllNum y+  TKAllNum (TKX2 _ y) = TKAllNum y+  TKAllNum (TKProduct x y) = (TKAllNum x, TKAllNum y)+  TKAllNum a =+    TypeError (Text "TKAllNum: No Num instance for type " :<>: ShowType a) + -- * Some fundamental constraints and types  type GoodScalarConstraint r =-  ( Show r, Ord r, Num r, Typeable r, IfDifferentiable r, Default r-  , NFData r, Nested.PrimElt r, Nested.KnownElt r, Nested.NumElt r-  , forall sh. Show (Nested.Mixed sh r), forall sh. Ord (Nested.Mixed sh r)-  , forall sh. NFData (Nested.Mixed sh r))+  ( Show r, Ord r, Typeable r, Default r, NFData r+  , Nested.PrimElt r, Nested.KnownElt r )  -- Attempted optimization via storing one pointer to a class dictionary--- in existential datatypes instead of six pointers. No effect, strangely.--- As a side effect, this avoids ImpredicativeTypes.--- Also, the constraint can be represented by a single Dict.+-- in existential datatypes instead of six pointers. No effect.+-- As a side effect, this avoids ImpredicativeTypes and UndecidableInstances+-- in many modules. -- -- | The constraint that signifies a scalar type, e.g., a float or an integers, -- is a well-behaved cell content of tensors supported by horde-ad.-class GoodScalarConstraint r => GoodScalar r+class    GoodScalarConstraint r => GoodScalar r instance GoodScalarConstraint r => GoodScalar r +class    (GoodScalarConstraint r, Num r, Nested.NumElt r, TKAllNum (TKScalar r))+         => NumScalar r+instance (GoodScalarConstraint r, Num r, Nested.NumElt r, TKAllNum (TKScalar r))+         => NumScalar r+ -- | The constraint for scalars that can be non-trivially differentiated, -- e.g., floating numbers, but not integers. type Differentiable r =-  (RealFloatH r, Nested.FloatElt r, RealFrac r, RealFloat r, Random r)+  ( RealFloatH r, Nested.FloatElt r, RealFrac r, RealFloat r, Random r+  , ADTensorScalar r ~ r )  -- We white-list all types on which we permit differentiation (e.g., SGD) -- to work. This is for technical typing purposes and imposes updates -- (and buggy omissions) when new scalar types are added, but it has -- the advantage of giving more control and visiblity. Sadly the list -- is repeated in several other places.-class IfDifferentiable r where-  ifDifferentiable :: (Differentiable r => a) -> a -> a--instance {-# OVERLAPPABLE #-} IfDifferentiable r where-  ifDifferentiable _ a = a---- The white-listed differentiable types.-instance IfDifferentiable Double where-  ifDifferentiable ra _ = ra-instance IfDifferentiable Float where-  ifDifferentiable ra _ = ra+--+-- BTW, a lot of occurrences of the non-differentiable case appears+-- when the parameters of the objective functions contain a ListR,+-- in which nil is represented as Z1, which is not differentiable.+-- In other cases it should be rare.+ifDifferentiable :: forall r a. Typeable r+                 => ((NumScalar r, Differentiable r) => a)+                 -> (ADTensorScalar r ~ Z1 => a)+                 -> a+{-# INLINE ifDifferentiable #-}+ifDifferentiable diff nonDiff = case typeRep @r of+  Is @Double -> diff+  Is @Float -> diff+  _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) nonDiff  type family BuildTensorKind k tk where   BuildTensorKind k (TKScalar r) = TKS '[k] r-  BuildTensorKind k (TKR2 n r) = TKR2 (1 + n) r-  BuildTensorKind k (TKS2 sh r) = TKS2 (k : sh) r-  BuildTensorKind k (TKX2 sh r) = TKX2 (Just k : sh) r+  BuildTensorKind k (TKR2 n x) = TKR2 (1 + n) x+  BuildTensorKind k (TKS2 sh x) = TKS2 (k : sh) x+  BuildTensorKind k (TKX2 sh x) = TKX2 (Just k : sh) x   BuildTensorKind k (TKProduct y z) =     TKProduct (BuildTensorKind k y) (BuildTensorKind k z)  -- This is an inverse of BuildTensorKind.--- This could be more efficient---   RazeTensorKind (TKS2 '[m] (TKScalar r)) = TKScalar r--- but then we'd lose the simplifying property that razing does not--- change the tensor kind variant, which is important, e.g.,--- when rewriting AstFromS t and trying to use AstFromS of the razed t. type family RazeTensorKind tk where-  RazeTensorKind (TKR2 n r) = TKR2 (n - 1) r-  RazeTensorKind (TKS2 sh r) = TKS2 (Tail sh) r-  RazeTensorKind (TKX2 sh r) = TKX2 (Tail sh) r-  RazeTensorKind (TKProduct y z) =-    TKProduct (RazeTensorKind y) (RazeTensorKind z)+  RazeTensorKind (TKR2 n x) = TKR2 (n - 1) x+  RazeTensorKind (TKS2 '[n] (TKScalar r)) = TKScalar r+  RazeTensorKind (TKS2 '[n] (TKR2 k y)) = TKS2 '[] (TKR2 k y)+  RazeTensorKind (TKS2 '[n] (TKS2 sh y)) = TKS2 '[] (TKS2 sh y)+  RazeTensorKind (TKS2 '[n] (TKX2 sh y)) = TKS2 '[] (TKX2 sh y)+  RazeTensorKind (TKS2 '[n] (TKProduct y z)) = TKS2 '[] (TKProduct y z)+  RazeTensorKind (TKS2 (n ': m ': sh) x) = TKS2 (m ': sh) x+  RazeTensorKind (TKX2 sh x) = TKX2 (Tail sh) x+  RazeTensorKind (TKProduct z1 z2) =+    TKProduct (RazeTensorKind z1) (RazeTensorKind z2)  type family ADTensorKind tk where   ADTensorKind (TKScalar r) = TKScalar (ADTensorScalar r)-  ADTensorKind (TKR2 n r) = TKR2 n (ADTensorKind r)-  ADTensorKind (TKS2 sh r) = TKS2 sh (ADTensorKind r)-  ADTensorKind (TKX2 sh r) = TKX2 sh (ADTensorKind r)+  ADTensorKind (TKR2 n x) = TKR2 n (ADTensorKind x)+  ADTensorKind (TKS2 sh x) = TKS2 sh (ADTensorKind x)+  ADTensorKind (TKX2 sh x) = TKX2 sh (ADTensorKind x)   ADTensorKind (TKProduct y z) =     TKProduct (ADTensorKind y) (ADTensorKind z) @@ -220,8 +239,10 @@ -- but would lead to irregular tensors, especially after vectorization, -- and would prevent statically known shapes. -- | The (arbitrariliy chosen) type of scalars to be used for indexing.-type IntOf (f :: Target) = PrimalOf f (TKScalar Int64)+type IntOf (f :: Target) = PlainOf f (TKScalar Int) +type BoolOf (f :: Target) = PlainOf f (TKScalar Bool)+ -- | The type family is defined in order to give a special instance -- for AST that preserves sharing and, even more importantly, keeps -- the computation of dervative functions lazy. See the definition@@ -230,16 +251,15 @@ -- -- The type family can't easily be made injective, because the @ADVal f@ -- instance is independent of @f@.-type family HFunOf (f :: Target)-                   (x :: TK) (z :: TK) :: Type+type family HFunOf (f :: Target) :: TK -> TK -> Type  type family PrimalOf (f :: Target) :: Target  type family DualOf (f :: Target) :: Target -type family ShareOf (f :: Target) :: Target+type family PlainOf (f :: Target) :: Target -type family BoolOf (t :: Target) :: Type+type family ShareOf (f :: Target) :: Target  -- | Ranked index, that is, a sized list of individual indices into -- ranked tensors.@@ -307,33 +327,96 @@   numEltNeg _ arr = arr   numEltAbs _ arr = arr   numEltSignum _ arr = arr-  numEltSum1Inner _ arr = fromO (RS.index (toO arr) 0)-  numEltProduct1Inner _ arr = fromO (RS.index (toO arr) 0)+  numEltSum1Inner _ = indexZ1+  numEltProduct1Inner _ = indexZ1   numEltSumFull _ _arr = Z1   numEltProductFull _ _arr = Z1-  numEltMinIndex snat _arr = replicate (sNatValue snat) 0-  numEltMaxIndex snat _arr = replicate (sNatValue snat) 0-  numEltDotprodInner _ arr1 _arr2 = fromO (RS.index (toO arr1) 0)+  numEltMinIndex snat _arr = replicate (fromSNat' snat) 0+  numEltMaxIndex snat _arr = replicate (fromSNat' snat) 0+  numEltDotprodInner _ arr1 _arr2 = indexZ1 arr1 +indexZ1 :: SA.Array (1 + n) a -> SA.Array n a+indexZ1 SA.Array{ SA.arrShape = _ : sh+                , SA.arrStrides = _ : strides+                , SA.arrOffset = _+                , SA.arrValues } = SA.Array { SA.arrShape = sh+                                            , SA.arrStrides = strides+                                            , SA.arrOffset = 0+                                            , SA.arrValues = arrValues }+indexZ1 _ = error "indexZ1: impossible" + -- * Misc  -- TODO: move all these somewhere +-- | Evidence for the constraint @c@. Works for type families,+-- such as @TKAllNum@.+type role Dict0 representational+data Dict0 c where+  Dict0 :: c => Dict0 c++pattern Is :: forall a b. Typeable a => a ~~ b => TypeRep b+pattern Is <- (eqTypeRep (TypeRep @a) -> Just (HRefl :: a :~~: b))++-- All below is copied from Data.OldList but made strict in state.+mapAccumL' :: (acc -> x -> (acc, y)) -> acc -> [x] -> (acc, [y])+{-# INLINE [1] mapAccumL' #-}+mapAccumL' _ s [] = (s, [])+mapAccumL' f !s (x : xs) = (s'', y : ys)+ where (s', y) = f s x+       (s'', ys) = mapAccumL' f s' xs++-- This apparently increases performance, so it must somehow work.+{-# RULES+"mapAccumL'" [~1] forall f s xs . mapAccumL' f s xs = foldr (mapAccumLF' f) pairWithNil xs s+"mapAccumL'List" [1] forall f s xs . foldr (mapAccumLF' f) pairWithNil xs s = mapAccumL' f s xs+ #-}++pairWithNil :: acc -> (acc, [y])+{-# INLINE [0] pairWithNil #-}+pairWithNil x = (x, [])++mapAccumLF' :: (acc -> x -> (acc, y)) -> x -> (acc -> (acc, [y])) -> acc+            -> (acc, [y])+{-# INLINE [0] mapAccumLF' #-}+mapAccumLF' f = \x r -> oneShot (\ !s ->+                          let (s', y) = f s x+                              (s'', ys) = r s'+                          in (s'', y : ys))+ -- | A variant of Integral for horde-ad without the problematic operations:--- 'mod' and 'rem' that are very slow, pair-producing operations that don't fit+-- 'mod' and 'rem' that are very slow (and simplifying them in AST requires+-- constructing conditionals, etc), pair-producing operations that don't fit -- the AST GADT format and 'toInteger' that doesn't make sense with AST -- without an environment to look up variables in.+infixl 7 `quotH`, `remH` class Num a => IntegralH a where   quotH, remH :: a -> a -> a  instance IntegralH Int64 where-  quotH = quot-  remH = rem+  quotH a b = if b == 0 then a else quot a b+  remH a b = if b == 0 then a else rem a b +instance IntegralH Int32 where+  quotH a b = if b == 0 then a else quot a b+  remH a b = if b == 0 then a else rem a b++instance IntegralH Int16 where+  quotH a b = if b == 0 then a else quot a b+  remH a b = if b == 0 then a else rem a b++instance IntegralH Int8 where+  quotH a b = if b == 0 then a else quot a b+  remH a b = if b == 0 then a else rem a b++instance IntegralH Int where+  quotH a b = if b == 0 then a else quot a b+  remH a b = if b == 0 then a else rem a b+ instance IntegralH CInt where-  quotH = quot-  remH = rem+  quotH a b = if b == 0 then a else quot a b+  remH a b = if b == 0 then a else rem a b  -- | The standard 'RealFloat' brings in a lot of operations we are not -- interested in and 'RealFrac' and 'Real' that we can't faithfully implement@@ -349,8 +432,8 @@  {- TODO: these would be better, but everything then overlaps with everything: instance {-# OVERLAPPABLE #-} Integral r => IntegralH r where-  quotH = quot-  remH = rem+  quotH a b = if b == 0 then a else quot a b+  remH a b = if b == 0 then a else rem a b  instance {-# OVERLAPPABLE #-} (Floating r, RealFloat r) => RealFloatH r where   atan2H = atan2@@ -371,148 +454,7 @@ backpermutePrefixList :: PermR -> [i] -> [i] backpermutePrefixList p l = map (l !!) p ++ drop (length p) l --- I can't switch to ixxFromLinear and ixxToLinear from ox-arrays--- even just because IntegralH is not available in ox-arrays.------ | Given a multidimensional index, get the corresponding linear--- index into the buffer. Note that the index doesn't need to be pointing--- at a scalar. It may point at the start of a larger tensor instead.------ If any of the dimensions is 0 or if rank is 0, the result will be 0,--- which is fine, that's pointing at the start of the empty buffer.--- Note that the resulting 0 may be a complex term.------ Warning: @fromInteger@ of type @j@ cannot be used.-toLinearIdxR :: forall m n j. Num j-             => (Int -> j) -> ShR (m + n) Int -> IxR m j -> j-toLinearIdxR fromInt = \sh idx -> go sh idx (fromInt 0)-  where-    -- Additional argument: index, in the @m - m1@ dimensional array so far,-    -- of the @m - m1 + n@ dimensional tensor pointed to by the current-    -- @m - m1@ dimensional index prefix.-    go :: ShR (m1 + n) Int -> IxR m1 j -> j -> j-    go sh ZIR tensidx = fromInt (shrSize sh) * tensidx-    go (n :$: sh) (i :.: idx) tensidx = go sh idx (fromInt n * tensidx + i) --- | Given a linear index into the buffer, get the corresponding--- multidimensional index. Here we require an index pointing at a scalar.------ If any of the dimensions is 0, the linear index has to be 0--- (which we can't assert, because j may be a term and so == lies)--- and a fake index with correct length but lots of zeroes is produced,--- because it doesn't matter, because it's going to point at the start--- of the empty buffer anyway.------ Warning: @fromInteger@ of type @j@ cannot be used.-fromLinearIdxR :: forall n j. IntegralH j-               => (Int -> j) -> ShR n Int -> j -> IxR n j-fromLinearIdxR fromInt = \sh lin -> snd (go sh lin)-  where-    -- Returns (linear index into array of sub-tensors,-    -- multi-index within sub-tensor).-    go :: ShR n1 Int -> j -> (j, IxR n1 j)-    go ZSR n = (n, ZIR)-    go (k :$: sh) _ | signum k == 0 =-      (fromInt 0, fromInt 0 :.: zeroOfR fromInt sh)-    go (n :$: sh) lin =-      let (tensLin, idxInTens) = go sh lin-          tensLin' = tensLin `quotH` fromInt n-          i = tensLin `remH` fromInt n-      in (tensLin', i :.: idxInTens)---- | The zero index in this shape (not dependent on the actual integers).-zeroOfR :: Num j => (Int -> j) -> ShR n i -> IxR n j-zeroOfR _ ZSR = ZIR-zeroOfR fromInt (_ :$: sh) = fromInt 0 :.: zeroOfR fromInt sh---- | Given a multidimensional index, get the corresponding linear--- index into the buffer. Note that the index doesn't need to be pointing--- at a scalar. It may point at the start of a larger tensor instead.------ If any of the dimensions is 0 or if rank is 0, the result will be 0,--- which is fine, that's pointing at the start of the empty buffer.--- Note that the resulting 0 may be a complex term.------ Warning: @fromInteger@ of type @j@ cannot be used.-toLinearIdxS :: forall sh1 sh2 j. Num j-             => (Int -> j) -> ShS (sh1 ++ sh2) -> IxS sh1 j -> j-toLinearIdxS fromInt = \sh idx -> go sh idx (fromInt 0)-  where-    -- Additional argument: index, in the @m - m1@ dimensional array so far,-    -- of the @m - m1 + n@ dimensional tensor pointed to by the current-    -- @m - m1@ dimensional index prefix.-    go :: forall sh3. ShS (sh3 ++ sh2) -> IxS sh3 j -> j -> j-    go sh ZIS tensidx = fromInt (shsSize sh) * tensidx-    go ((:$$) n sh) (i :.$ idx) tensidx =-      go sh idx (fromInt (sNatValue n) * tensidx + i)-    go _ _ _ = error "toLinearIdxS: impossible pattern needlessly required"---- | Given a linear index into the buffer, get the corresponding--- multidimensional index. Here we require an index pointing at a scalar.------ If any of the dimensions is 0, the linear index has to be 0--- (which we can't assert, because j may be a term and so == lies)--- and a fake index with correct length but lots of zeroes is produced,--- because it doesn't matter, because it's going to point at the start--- of the empty buffer anyway.------ Warning: @fromInteger@ of type @j@ cannot be used.-fromLinearIdxS :: forall sh j. IntegralH j-               => (Int -> j) -> ShS sh -> j -> IxS sh j-fromLinearIdxS fromInt = \sh lin -> snd (go sh lin)-  where-    -- Returns (linear index into array of sub-tensors,-    -- multi-index within sub-tensor).-    go :: ShS sh1 -> j -> (j, IxS sh1 j)-    go ZSS n = (n, ZIS)-    go ((:$$) k sh) _ | sNatValue k == 0 =-      (fromInt 0, fromInt 0 :.$ zeroOfS fromInt sh)-    go ((:$$) n sh) lin =-      let (tensLin, idxInTens) = go sh lin-          tensLin' = tensLin `quotH` fromInt (sNatValue n)-          i = tensLin `remH` fromInt (sNatValue n)-      in (tensLin', i :.$ idxInTens)---- | The zero index in this shape (not dependent on the actual integers).-zeroOfS :: Num j => (Int -> j) -> ShS sh -> IxS sh j-zeroOfS _ ZSS = ZIS-zeroOfS fromInt ((:$$) _ sh) = fromInt 0 :.$ zeroOfS fromInt sh--toLinearIdxX :: forall sh1 sh2 j. Num j-             => (Int -> j) -> IShX (sh1 ++ sh2) -> IxX sh1 j -> j-toLinearIdxX fromInt = \sh idx -> go sh idx (fromInt 0)-  where-    -- Additional argument: index, in the @m - m1@ dimensional array so far,-    -- of the @m - m1 + n@ dimensional tensor pointed to by the current-    -- @m - m1@ dimensional index prefix.-    go :: forall sh3. IShX (sh3 ++ sh2) -> IxX sh3 j -> j -> j-    go sh ZIX tensidx = fromInt (shxSize sh) * tensidx-    go ((:$%) n sh) (i :.% idx) tensidx =-      go sh idx (fromInt (fromSMayNat' n) * tensidx + i)-    go _ _ _ = error "toLinearIdxX: impossible pattern needlessly required"--fromLinearIdxX :: forall sh j. IntegralH j-               => (Int -> j) -> IShX sh -> j -> IxX sh j-fromLinearIdxX fromInt = \sh lin -> snd (go sh lin)-  where-    -- Returns (linear index into array of sub-tensors,-    -- multi-index within sub-tensor).-    go :: IShX sh1 -> j -> (j, IxX sh1 j)-    go ZSX n = (n, ZIX)-    go ((:$%) k sh) _ | fromSMayNat' k == 0 =-      (fromInt 0, fromInt 0 :.% zeroOfX fromInt sh)-    go ((:$%) n sh) lin =-      let (tensLin, idxInTens) = go sh lin-          tensLin' = tensLin `quotH` fromInt (fromSMayNat' n)-          i = tensLin `remH` fromInt (fromSMayNat' n)-      in (tensLin', i :.% idxInTens)---- | The zero index in this shape (not dependent on the actual integers).-zeroOfX :: Num j => (Int -> j) -> IShX sh -> IxX sh j-zeroOfX _ ZSX = ZIX-zeroOfX fromInt ((:$%) _ sh) = fromInt 0 :.% zeroOfX fromInt sh-- -- * Shopping list for ox-arrays  -- All of the below should have better types and/or implementations,@@ -525,6 +467,7 @@ shCastSX :: Rank sh ~ Rank sh' => StaticShX sh' -> ShS sh -> IShX sh' shCastSX ssx sh = shxCast' ssx (shxFromShS sh) + -- ** Permutation-related operations  -- - Permutation.permInverse for ShS would be helpful in addition to@@ -546,18 +489,27 @@  -- TODO: can this be defined for Permutation.Perm using @Mod@? -- A representation of a cycle backpermutation that moves elements--- to indexes one less (the the left, to the back).-backpermCycle :: Int -> Permutation.PermR-backpermCycle 0 = []-backpermCycle 1 = []-backpermCycle n = [k `mod` n | k <- [1 .. n]]+-- to indexes @n@ less (to the left, to the back), wrapping.+backpermCycleN :: Int -> Int -> Permutation.PermR+backpermCycleN _ 0 = []+backpermCycleN _ 1 = []+backpermCycleN n k = [(m + n) `mod` k | m <- [0 .. k - 1]]  -- TODO: can this be defined for Permutation.Perm using @Mod@?+-- A representation of a cycle permutation that is reverse to @backpermCycleN@.+permCycleN :: Int -> Int -> Permutation.PermR+permCycleN _ 0 = []+permCycleN _ 1 = []+permCycleN n k = [(m - n) `mod` k | m <- [0 .. k - 1]]++-- A representation of a cycle backpermutation that moves elements+-- to indexes one less (to the left, to the back).+backpermCycle :: Int -> Permutation.PermR+backpermCycle = backpermCycleN 1+ -- A representation of a cycle permutation that is reverse to @backpermCycle@. permCycle :: Int -> Permutation.PermR-permCycle 0 = []-permCycle 1 = []-permCycle n = [k `mod` n | k <- [-1, 0 .. n - 2]]+permCycle = permCycleN 1  type family UnMapSucc is where   UnMapSucc '[] = '[]@@ -569,7 +521,7 @@ permUnShift1 :: Permutation.Perm (0 : l)              -> Permutation.Perm (UnMapSucc l) permUnShift1 (Permutation.PCons _ permRest) =-  Permutation.permFromList+  Permutation.permFromListCont     (permUnMapSucc (Permutation.permToList' permRest)) unsafeCoerce  where   permUnMapSucc :: [Int] -> [Int]@@ -602,134 +554,124 @@   Drop 0 xs = xs   Drop n (x ': xs) = Drop (n - 1) xs -listsTake :: forall len sh i.-             (KnownShS sh, KnownNat len, KnownShS (Take len sh))-          => ListS sh (Const i) -> ListS (Take len sh) (Const i)-listsTake ix = fromList $ take (valueOf @len) $ toList ix--listsDrop :: forall len sh i.-             (KnownShS sh, KnownNat len, KnownShS (Drop len sh))-          => ListS sh (Const i) -> ListS (Drop len sh) (Const i)-listsDrop ix = fromList $ drop (valueOf @len) $ toList ix+{- TODO:+listrTake :: forall len n i. KnownNat len+          => ListR (len + n) i -> ListR len i+listrTake -listsSplitAt-  :: (KnownShS sh, KnownNat len, KnownShS (Drop len sh), KnownShS (Take len sh))-  => ListS sh (Const i)-  -> (ListS (Take len sh) (Const i), ListS (Drop len sh) (Const i))-listsSplitAt ix = (listsTake ix, listsDrop ix)+listrDrop :: forall len n i. KnownNat len+          => ListR (len + n) i -> ListR n i+listrDrop -ixrTake :: forall m n i. (KnownNat m, KnownNat n)+ixrTake :: forall m n i. KnownNat m         => IxR (m + n) i -> IxR m i ixrTake (IxR ix) = IxR $ listrTake ix -ixrDrop :: forall m n i. (KnownNat m, KnownNat n)+ixrDrop :: forall m n i. KnownNat m         => IxR (m + n) i -> IxR n i ixrDrop (IxR ix) = IxR $ listrDrop ix -ixrSplitAt :: (KnownNat m, KnownNat n)+ixrSplitAt :: KnownNat m            => IxR (m + n) i -> (IxR m i, IxR n i) ixrSplitAt ix = (ixrTake ix, ixrDrop ix)--shrTake :: forall m n i. (KnownNat n, KnownNat m)-        => ShR (m + n) i -> ShR m i-shrTake (ShR ix) = ShR $ listrTake ix--shrDrop :: forall m n i. (KnownNat m, KnownNat n)-          => ShR (m + n) i -> ShR n i-shrDrop (ShR ix) = ShR $ listrDrop ix--shrSplitAt :: (KnownNat m, KnownNat n)-           => ShR (m + n) i -> (ShR m i, ShR n i)-shrSplitAt ix = (shrTake ix, shrDrop ix)+-} -listrTake :: forall len n i. (KnownNat n, KnownNat len)-          => ListR (len + n) i -> ListR len i-listrTake ix = fromList $ take (valueOf @len) $ toList ix+shrTake :: forall m n. KnownNat m+        => IShR (m + n) -> IShR m+shrTake (ShR shx) =+  gcastWith (unsafeCoerceRefl+             :: Take m (Replicate (m + n) (Nothing @Nat))+                :~: Replicate m Nothing) $+  ShR $ shxTake @m shx -listrDrop :: forall len n i. (KnownNat len, KnownNat n)-          => ListR (len + n) i -> ListR n i-listrDrop ix = fromList $ drop (valueOf @len) $ toList ix+shrDrop :: forall m n. KnownNat m+        => IShR (m + n) -> IShR n+shrDrop (ShR shx) =+  gcastWith (unsafeCoerceRefl+             :: Drop m (Replicate (m + n) (Nothing @Nat))+                :~: Replicate n Nothing) $+  ShR $ shxDrop @m shx -listrSplitAt :: (KnownNat m, KnownNat n)-             => ListR (m + n) i -> (ListR m i, ListR n i)-listrSplitAt ix = (listrTake ix, listrDrop ix)+ixsTake :: forall len sh i. KnownNat len+        => IxS sh i -> IxS (Take len sh) i+ixsTake sh = go (SNat @len) sh+  where+    go :: SNat n -> IxS shn i -> IxS (Take n shn) i+    go SZ _ = ZIS+    go (SS (k :: SNat k)) ((:.$) @m @rest m rest) =+      gcastWith (unsafeCoerceRefl+                 :: Take (k + 1) (m : rest) :~: m : Take k rest) $+      m :.$ go k rest+    go (SS k) ZIS = error $ "ixsTake: list too short, missing "+                            ++ show (fromSNat' k + 1) -ixsDrop :: forall len sh i. (KnownShS sh, KnownNat len, KnownShS (Drop len sh))+ixsDrop :: forall len sh i. KnownNat len         => IxS sh i -> IxS (Drop len sh) i-ixsDrop (IxS ix) = IxS $ listsDrop ix+ixsDrop sh = go (SNat @len) sh+  where+    go :: SNat n -> IxS shn i -> IxS (Drop n shn) i+    go SZ shn = shn+    go (SS (k :: SNat k)) ((:.$) @m @rest _ rest) =+      gcastWith (unsafeCoerceRefl+                 :: Drop (k + 1) (m : rest) :~: Drop k rest) $+      go k rest+    go (SS k) ZIS = error $ "ixsDrop: list too short, missing "+                            ++ show (fromSNat' k + 1) --- TODO-shsTake :: forall len sh. (KnownNat len, KnownShS sh)+ixsSplitAt :: forall sh len i. KnownNat len+           => IxS sh i -> (IxS (Take len sh) i, IxS (Drop len sh) i)+ixsSplitAt ix = (ixsTake @len ix, ixsDrop @len ix)++shsTake :: forall len sh. KnownNat len         => ShS sh -> ShS (Take len sh)-shsTake sh0 = fromList2 $ take (valueOf @len) $ toList sh0- where-  fromList2 topl = ShS (go (knownShS @sh) topl)-    where  -- TODO: induction over (unary) SNat?-      go :: forall sh'. ShS sh' -> [Int] -> ListS (Take len sh') SNat-      go _ [] = gcastWith (unsafeCoerceRefl :: len :~: 0) $ gcastWith (unsafeCoerceRefl :: sh' :~: '[]) ZS-      go (sn :$$ sh) (i : is)-        | i == fromSNat' sn = unsafeCoerce $ sn ::$ go sh is-        | otherwise = error $ "shsTake: Value does not match typing (type says "-                                ++ show (fromSNat' sn) ++ ", list contains " ++ show i ++ ")"-      go _ _ = error $ "shsTake: Mismatched list length (type says "-                         ++ show (shsLength (knownShS @sh)) ++ ", list has length "-                         ++ show (length topl) ++ ")"+shsTake (ShS shx) =+  gcastWith (unsafeCoerceRefl+             :: Take len (MapJust sh) :~: MapJust (Take len sh)) $+  ShS $ shxTake @len shx --- TODO-shsDrop :: forall len sh. (KnownNat len, KnownShS sh)+shsDrop :: forall len sh. KnownNat len         => ShS sh -> ShS (Drop len sh)-shsDrop sh0 = fromList2 $ drop (valueOf @len) $ toList sh0- where-  fromList2 topl = ShS (go (knownShS @sh) $ replicate (valueOf @len) (-1) ++ topl)-    where  -- TODO: induction over (unary) SNat?-      go :: forall sh'. ShS sh' -> [Int] -> ListS (Drop len sh') SNat-      go _ [] = gcastWith (unsafeCoerceRefl :: len :~: 0) $ gcastWith (unsafeCoerceRefl :: sh' :~: '[]) ZS-      go (sn :$$ sh) (i : is)-        | i == -1 = unsafeCoerce $ go sh is-        | i == fromSNat' sn = unsafeCoerce $ sn ::$ go sh is-        | otherwise = error $ "shsDrop: Value does not match typing (type says "-                                ++ show (fromSNat' sn) ++ ", list contains " ++ show i ++ ")"-      go _ _ = error $ "shsDrop: Mismatched list length (type says "-                         ++ show (shsLength (knownShS @sh)) ++ ", list has length "-                         ++ show (length topl) ++ ")"--shxTake :: forall len sh. (KnownNat len, KnownShX sh, KnownShX (Take len sh))-        => IShX sh -> IShX (Take len sh)-shxTake sh0 = fromList $ take (valueOf @len) $ toList sh0--shxDrop :: forall len sh. (KnownNat len, KnownShX sh, KnownShX (Drop len sh))-        => IShX sh -> IShX (Drop len sh)-shxDrop sh0 = fromList $ drop (valueOf @len) $ toList sh0+shsDrop (ShS shx) =+  gcastWith (unsafeCoerceRefl+             :: Drop len (MapJust sh) :~: MapJust (Drop len sh)) $+  ShS $ shxDrop @len shx -ixxTake :: forall len sh i. (KnownNat len, KnownShX sh, KnownShX (Take len sh))+{- TODO:+ixxTake :: forall len sh i. (KnownNat len, KnownShX (Take len sh))         => IxX sh i -> IxX (Take len sh) i-ixxTake sh0 = fromList $ take (valueOf @len) $ toList sh0+ixxTake sh = -ixxDrop' :: forall len sh i. (KnownNat len, KnownShX sh, KnownShX (Drop len sh))+ixxDrop' :: forall len sh i. (KnownNat len, KnownShX (Drop len sh))          => IxX sh i -> IxX (Drop len sh) i-ixxDrop' sh0 = fromList $ drop (valueOf @len) $ toList sh0--listsTakeLen :: forall f g sh1 sh2.-                ListS sh1 f -> ListS sh2 g -> ListS (TakeLen sh1 sh2) g-listsTakeLen ZS _ = ZS-listsTakeLen (_ ::$ sh1) (n ::$ sh2) = n ::$ listsTakeLen sh1 sh2-listsTakeLen (_ ::$ _) ZS = error "listsTakeLen: list too short"--listsDropLen :: forall f g sh1 sh2.-                ListS sh1 f -> ListS sh2 g -> ListS (DropLen sh1 sh2) g-listsDropLen ZS sh = sh-listsDropLen (_ ::$ sh1) (_ ::$ sh2) = listsDropLen sh1 sh2-listsDropLen (_ ::$ _) ZS = error "listsDropLen: list too short"--shsDropLen :: Permutation.Perm is -> ShS sh -> ShS (DropLen is sh)-shsDropLen = coerce (listsDropLenPerm @SNat)---- This is only needed as a workaround for other ops not provided.+ixxDrop' sh =+-} -listxTake :: forall f g sh sh'. ListX (sh ++ sh') f -> ListX sh g -> ListX sh f-listxTake _ ZX = ZX-listxTake (i ::% long') ((::%) @_ @_ @sh2 _ short) =-  i ::% listxTake @f @g @sh2 @sh' long' short+shxTake :: forall len sh. KnownNat len+        => IShX sh -> IShX (Take len sh)+shxTake sh = go (SNat @len) sh+  where+    go :: SNat n -> IShX shn -> IShX (Take n shn)+    go SZ _ = ZSX+    go (SS k@(SNat @k)) ((:$%) @mn @rest mn rest) =+      -- This is needed probably because the type family can't decide+      -- which of the equations to use:+      gcastWith (unsafeCoerceRefl+                 :: Take (k + 1) (mn : rest) :~: mn : Take k rest) $+      mn :$% go k rest+    -- We error out where normal list take operation would complete fine+    -- in order to keep the type family simple:+    go (SS k) ZSX = error $ "shxTake: list too short, missing "+                            ++ show (fromSNat' k + 1) -ssxTakeIx :: forall sh sh' i. StaticShX (sh ++ sh') -> IxX sh i -> StaticShX sh-ssxTakeIx = coerce (listxTake @(Nested.SMayNat () SNat) @(Const i) @_ @sh')+shxDrop :: forall len sh. KnownNat len+        => IShX sh -> IShX (Drop len sh)+shxDrop sh = go (SNat @len) sh+  where+    go :: SNat n -> IShX shn -> IShX (Drop n shn)+    go SZ shn = shn+    go (SS k@(SNat @k)) ((:$%) @mn @rest _ rest) =+      gcastWith (unsafeCoerceRefl+                 :: Drop (k + 1) (mn : rest) :~: Drop k rest) $+      go k rest+    go (SS k) ZSX = error $ "shxDrop: list too short, missing "+                            ++ show (fromSNat' k + 1)
src/HordeAd/Core/Unwind.hs view
@@ -3,22 +3,24 @@ -- | Types and functions needed to define general tensor operations -- that work for any tensor kind, including nested (product) arrays -- and an assortment of such operations.+--+-- Large portions of this module are copied to HordeAd.Core.UnwindNum+-- in order to have more accurate typing and pattern exhaustiveness checks. module HordeAd.Core.Unwind-  ( addTarget, multTarget, sum0Target, dot0Target-  , replTarget, defTarget, concreteTarget+  ( concreteTarget   , toADTensorKindShared, fromADTensorKindShared   ) where  import Prelude  import Data.Default-import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Type.Equality ((:~:) (Refl)) import GHC.TypeLits (type (+))-import Type.Reflection (typeRep)  import Data.Array.Nested (MapJust, Replicate, type (++)) import Data.Array.Nested qualified as Nested import Data.Array.Nested.Convert+  (shxFromShR, shxFromShS, withShsFromShR, withShsFromShX) import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape@@ -29,7 +31,6 @@ import HordeAd.Core.Ops import HordeAd.Core.TensorKind import HordeAd.Core.Types-import HordeAd.OpsTensor  -- * Winding and unwinding @@ -67,162 +68,64 @@   WFTKProduct :: FullShapeTKW y -> FullShapeTKW z               -> FullShapeTKW (TKProduct y z) -addRepW :: forall y target. BaseTensor target-        => RepW target y -> RepW target y -> RepW target y-addRepW a b = case (a, b) of-  (WTKScalar ta, WTKScalar tb) -> WTKScalar $ ta + tb-  (WTKR ta, WTKR tb) -> WTKR $ ta + tb-  (WTKS ta, WTKS tb) -> WTKS $ ta + tb-  (WTKX ta, WTKX tb) -> WTKX $ ta + tb-  (WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->-    WTKProduct (addRepW ta1 tb1) (addRepW ta2 tb2)--multRepW :: forall y target. BaseTensor target-         => RepW target y -> RepW target y -> RepW target y-multRepW a b = case (a, b) of-  (WTKScalar ta, WTKScalar tb) -> WTKScalar $ ta * tb-  (WTKR ta, WTKR tb) -> WTKR $ ta * tb-  (WTKS ta, WTKS tb) -> WTKS $ ta * tb-  (WTKX ta, WTKX tb) -> WTKX $ ta * tb-  (WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->-    WTKProduct (multRepW ta1 tb1) (multRepW ta2 tb2)--sum0RepW :: forall y target. (BaseTensor target, ConvertTensor target)-         => FullShapeTKW y -> RepW target y-         -> target (TKScalar Double)-sum0RepW ftk a = case (ftk, a) of-  (_, WTKScalar @r ta) ->-    ifDifferentiable @r (kcast ta) 0-  (WFTKR sh, WTKR @r ta) | SNat <- shrRank sh ->-    ifDifferentiable @r (kcast $ kfromR $ rsum0 ta) 0-  (WFTKS sh, WTKS @r ta) ->-    withKnownShS sh $-    ifDifferentiable @r (kcast $ kfromS $ ssum0 ta) 0-  (WFTKX sh, WTKX @r ta) ->-    withKnownShX (ssxFromShX sh) $-    ifDifferentiable @r (kcast $ kfromX $ xsum0 ta) 0-  (WFTKProduct ftk1 ftk2, WTKProduct ta1 ta2) ->-    sum0RepW ftk1 ta1 + sum0RepW ftk2 ta2--dot0RepW :: forall y target. (BaseTensor target, ConvertTensor target)-         => FullShapeTKW y -> RepW target y -> RepW target y-         -> target (TKScalar Double)-dot0RepW ftk a b = case (ftk, a, b) of-  (_, WTKScalar @r ta, WTKScalar tb) ->-    ifDifferentiable @r (kcast $ ta * tb) 0-  (WFTKR sh, WTKR @r ta, WTKR tb) | SNat <- shrRank sh ->-    ifDifferentiable @r (kcast $ kfromR $ rdot0 ta tb) 0-  (WFTKS sh, WTKS @r ta, WTKS tb) ->-    withKnownShS sh $-    ifDifferentiable @r (kcast $ kfromS $ sdot0 ta tb) 0-  (WFTKX sh, WTKX @r ta, WTKX tb) ->-    withKnownShX (ssxFromShX sh) $-    ifDifferentiable @r (kcast $ kfromX $ xdot0 ta tb) 0-  (WFTKProduct ftk1 ftk2, WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->-    dot0RepW ftk1 ta1 tb1 + dot0RepW ftk2 ta2 tb2--replRepW :: forall y target. BaseTensor target-         => (forall r. GoodScalar r => r)-         -> FullShapeTKW y -> RepW target y-replRepW r = \case-  WFTKScalar -> WTKScalar $ kconcrete r-  WFTKR sh -> WTKR $ rrepl sh r-  WFTKS sh -> WTKS $ sconcrete $ Nested.sreplicateScal sh r-  WFTKX sh -> WTKX $ xrepl sh r-  WFTKProduct ftk1 ftk2 ->-    WTKProduct (replRepW r ftk1) (replRepW r ftk2)--defRepW :: forall y target. BaseTensor target-        => FullShapeTKW y -> RepW target y-defRepW = \case-  WFTKScalar -> WTKScalar $ kconcrete def-  WFTKR sh -> WTKR $ rrepl sh def-  WFTKS sh -> WTKS $ sconcrete $ Nested.sreplicateScal sh def-  WFTKX sh -> WTKX $ xrepl sh def-  WFTKProduct ftk1 ftk2 ->-    WTKProduct (defRepW ftk1) (defRepW ftk2)- concreteRepW   :: forall y target. (ConvertTensor Concrete, ConvertTensor target)   => (forall r. GoodScalar r => Concrete (TKScalar r) -> target (TKScalar r))   -> (forall r sh. GoodScalar r => Concrete (TKS sh r) -> target (TKS sh r))-  -> (forall x z. FullShapeTK z -> target x -> target z)+  -> (forall r sh. GoodScalar r+      => ShS sh -> target (TKS sh r) -> target (TKR (Rank sh) r))+  -> (forall r sh sh'. (GoodScalar r, Rank sh ~ Rank sh')+      => IShX sh' -> target (TKS sh r) -> target (TKX sh' r))   -> RepW Concrete y -> RepW target y {-# INLINE concreteRepW #-}-concreteRepW concreteK concreteS fromS w = case w of+concreteRepW concreteK concreteS toRfromS toXfromS w = case w of   WTKScalar v -> WTKScalar $ concreteK v   WTKR v -> WTKR $     let sh' = Nested.rshape $ unConcrete v     in withShsFromShR sh' $ \(sh :: ShS sh) ->       withKnownShS sh $-      fromS (FTKR sh' FTKScalar)-      $ concreteS (sfromR @_ @sh v)+      toRfromS sh $ concreteS (sfromR @_ @sh v)   WTKS v -> WTKS $ concreteS v   WTKX v -> WTKX $     let sh' = Nested.mshape $ unConcrete v     in withShsFromShX sh' $ \(sh :: ShS sh) ->       withKnownShS sh $-      fromS (FTKX sh' FTKScalar)-      $ concreteS (sfromX @_ @sh v)+      toXfromS sh' $ concreteS (sfromX @_ @sh v)   WTKProduct v1 v2 ->-    WTKProduct (concreteRepW concreteK concreteS fromS v1)-               (concreteRepW concreteK concreteS fromS v2)+    WTKProduct (concreteRepW concreteK concreteS toRfromS toXfromS v1)+               (concreteRepW concreteK concreteS toRfromS toXfromS v2)  toADTensorKindW   :: forall y target. BaseTensor target   => RepW target y -> FullShapeTKW y -> RepW target (ADTensorKind y) toADTensorKindW t = \case-  WFTKScalar @r -> case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> t-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> t-      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $-           WTKScalar $ kconcrete Z1-  WFTKR @r sh -> case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> t-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> t-      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $-           WTKR $ rrepl @_ @_ @target sh Z1-  WFTKS @r sh -> case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> t-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> t-      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $-           WTKS $ sconcrete $ Nested.sreplicateScal sh Z1-  WFTKX @r sh -> case testEquality (typeRep @r) (typeRep @Double) of-    Just Refl -> t-    _ -> case testEquality (typeRep @r) (typeRep @Float) of-      Just Refl -> t-      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $-           WTKX $ xrepl @_ @_ @target sh Z1+  WFTKScalar @r ->+    ifDifferentiable @r t (WTKScalar $ tkconcrete Z1)+  WFTKR @r sh ->+    ifDifferentiable @r t (WTKR $ rrepl @_ @_ @target sh Z1)+  WFTKS @r sh ->+    ifDifferentiable @r t (WTKS $ tsconcrete $ Nested.sreplicatePrim sh Z1)+  WFTKX @r sh ->+    ifDifferentiable @r t (WTKX $ xrepl @_ @_ @target sh Z1)   WFTKProduct ftk1 ftk2 -> case t of     WTKProduct t1 t2 ->       WTKProduct (toADTensorKindW t1 ftk1) (toADTensorKindW t2 ftk2)  fromADTensorKindW   :: forall y target. BaseTensor target-  => SingletonTK y -> RepW target (ADTensorKind y) -> RepW target y-fromADTensorKindW stk t = case (stk, t) of-  (STKScalar @r1, WTKScalar @r2 _) ->-    case testEquality (typeRep @r1) (typeRep @r2) of-      Just Refl -> t-      _ -> replRepW 0 WFTKScalar-  (STKR _ (STKScalar @r1), WTKR @r2 v) ->-    case testEquality (typeRep @r1) (typeRep @r2) of-      Just Refl -> t-      _ -> replRepW 0 (WFTKR (rshape v))-  (STKS sh (STKScalar @r1), WTKS @r2 _) ->-    case testEquality (typeRep @r1) (typeRep @r2) of-      Just Refl -> t-      _ -> replRepW 0 (WFTKS sh)-  (STKX _ (STKScalar @r1), WTKX @r2 v) ->-    case testEquality (typeRep @r1) (typeRep @r2) of-      Just Refl -> t-      _ -> replRepW 0 (WFTKX (xshape v))-  (STKProduct stk1 stk2, WTKProduct t1 t2) ->-    WTKProduct (fromADTensorKindW stk1 t1) (fromADTensorKindW stk2 t2)-  _ -> error "fromADTensorKindW: impossible SingletonTK"+  => RepW target (ADTensorKind y) -> FullShapeTKW y -> RepW target y+fromADTensorKindW t = \case+  WFTKScalar @r ->+    ifDifferentiable @r t (WTKScalar $ tkconcrete def)+  WFTKR @r sh ->+    ifDifferentiable @r t (WTKR $ rrepl sh def)+  WFTKS @r sh ->+    ifDifferentiable @r t (WTKS $ tsconcrete $ Nested.sreplicatePrim sh def)+  WFTKX @r sh ->+    ifDifferentiable @r t (WTKX $ xrepl sh def)+  WFTKProduct stk1 stk2 -> case t of+    WTKProduct t1 t2 ->+      WTKProduct (fromADTensorKindW t1 stk1) (fromADTensorKindW t2 stk2)  type family UnWind y where   UnWind (TKScalar r) =@@ -260,40 +163,6 @@   UnWind (TKProduct y z) =     TKProduct (UnWind y) (UnWind z) -unWindSTK :: SingletonTK y -> SingletonTK (UnWind y)-unWindSTK = \case-  stk@STKScalar -> stk-  stk@(STKR _ STKScalar) -> stk-  STKR (SNat @n) (STKR (SNat @m) stk2) ->-    unWindSTK $ STKR (SNat @(n + m)) stk2-  STKR n (STKS sh2 stk2) ->-    unWindSTK-    $ STKX (ssxReplicate n `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2-  STKR n (STKX sh2 stk2) ->-    unWindSTK $ STKX (ssxReplicate n `ssxAppend` sh2) stk2-  STKR n@SNat (STKProduct y z) ->-    unWindSTK $ STKProduct (STKR n y) (STKR n z)-  stk@(STKS _ STKScalar) -> stk-  STKS sh1 (STKR m stk2) ->-    unWindSTK-    $ STKX (ssxFromShX (shxFromShS sh1) `ssxAppend` ssxReplicate m) stk2-  STKS sh1 (STKS sh2 stk2) ->-    unWindSTK $ STKS (sh1 `shsAppend` sh2) stk2-  STKS sh1 (STKX sh2 stk2) ->-    unWindSTK $ STKX (ssxFromShX (shxFromShS sh1) `ssxAppend` sh2) stk2-  STKS sh1 (STKProduct y z)->-    unWindSTK $ STKProduct (STKS sh1 y) (STKS sh1 z)-  stk@(STKX _ STKScalar) -> stk-  STKX sh1 (STKR m stk2) ->-    unWindSTK $ STKX (sh1 `ssxAppend` ssxReplicate m) stk2-  STKX sh1 (STKS sh2 stk2) ->-    unWindSTK $ STKX (sh1 `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2-  STKX sh1 (STKX sh2 stk2) ->-    unWindSTK $ STKX (sh1 `ssxAppend` sh2) stk2-  STKX sh1 (STKProduct y z) ->-    unWindSTK $ STKProduct (STKX sh1 y) (STKX sh1 z)-  STKProduct y z -> STKProduct (unWindSTK y) (unWindSTK z)- unWindFTK :: FullShapeTK y -> FullShapeTKW (UnWind y) unWindFTK = \case   FTKScalar -> WFTKScalar@@ -328,7 +197,7 @@     unWindFTK $ FTKProduct (FTKX sh1 y) (FTKX sh1 z)   FTKProduct y z -> WFTKProduct (unWindFTK y) (unWindFTK z) --- This uses tunpairConv so to preserve sharing, @target@ either has+-- This uses tunpairConv, so to preserve sharing, @target@ either has -- to have a `ShareTensor` instance or the argument has to be duplicable. -- Only the argument of the first call, not of recursive calls, -- is assumed to be duplicable. In the AST case, this creates@@ -440,73 +309,27 @@  -- * Operations defined using unwinding --- | Add two (nested pairs of) tensors. Requires duplicable arguments--- or a `ShareTensor` instance.-addTarget :: (BaseTensor target, ConvertTensor target)-          => SingletonTK y -> target y -> target y -> target y-addTarget stk a b =-  let a2 = unWindTarget stk a-      b2 = unWindTarget stk b-  in windTarget stk $ addRepW a2 b2---- | Multiply two (nested pairs of) tensors. Requires duplicable arguments--- or a `ShareTensor` instance.-multTarget :: (BaseTensor target, ConvertTensor target)-           => SingletonTK y -> target y -> target y -> target y-multTarget stk a b =-  let a2 = unWindTarget stk a-      b2 = unWindTarget stk b-  in windTarget stk $ multRepW a2 b2---- | Sum all dimensions of each component and then sum it all.--- Requires duplicable arguments or a `ShareTensor` instance.-sum0Target :: (BaseTensor target, ConvertTensor target)-           => FullShapeTK y -> target y-           -> target (TKScalar Double)-sum0Target ftk a =-  let a2 = unWindTarget (ftkToSTK ftk) a-  in sum0RepW (unWindFTK ftk) a2---- | Dot product each component and then sum it all.--- Requires duplicable arguments or a `ShareTensor` instance.-dot0Target :: (BaseTensor target, ConvertTensor target)-           => FullShapeTK y -> target y -> target y-           -> target (TKScalar Double)-dot0Target ftk a b =-  let a2 = unWindTarget (ftkToSTK ftk) a-      b2 = unWindTarget (ftkToSTK ftk) b-  in dot0RepW (unWindFTK ftk) a2 b2---- | Replicate a scalar along the given full shape singleton.-replTarget :: forall y target. (BaseTensor target, ConvertTensor target)-           => (forall r. GoodScalar r => r)-           -> FullShapeTK y -> target y-replTarget r ftk =-  windTarget (ftkToSTK ftk) $ replRepW r (unWindFTK ftk)---- | Replicate the default value along the given full shape singleton.-defTarget :: forall y target. (BaseTensor target, ConvertTensor target)-          => FullShapeTK y -> target y-defTarget ftk =-  windTarget (ftkToSTK ftk) $ defRepW (unWindFTK ftk)- concreteTarget   :: forall y target. (ConvertTensor Concrete, ConvertTensor target)   => (forall r. GoodScalar r => Concrete (TKScalar r) -> target (TKScalar r))   -> (forall r sh. GoodScalar r => Concrete (TKS sh r) -> target (TKS sh r))-  -> (forall x z. FullShapeTK z -> target x -> target z)+  -> (forall r sh. GoodScalar r+      => ShS sh -> target (TKS sh r) -> target (TKR (Rank sh) r))+  -> (forall r sh sh'. (GoodScalar r, Rank sh ~ Rank sh')+      => IShX sh' -> target (TKS sh r) -> target (TKX sh' r))   -> SingletonTK y -> Concrete y   -> target y-concreteTarget concreteK concreteS fromS stk v =+{-# INLINE concreteTarget #-}+concreteTarget concreteK concreteS toRfromS toXfromS stk v =   windTarget stk-  $ concreteRepW concreteK concreteS fromS+  $ concreteRepW concreteK concreteS toRfromS toXfromS   $ unWindTarget stk v  lemUnWindOfAD :: SingletonTK y               -> UnWind (ADTensorKind y) :~: ADTensorKind (UnWind y) lemUnWindOfAD _ = unsafeCoerceRefl --- | Convert a tensor into a tensor with only trivial non-differentiable+-- | Convert a tensor into a tensor with only trivial (Z1) non-differentiable -- scalars. The `ShareTensor` constraint is needed, despite what GHC says, -- in order not to require duplicable arguments. toADTensorKindShared@@ -517,15 +340,15 @@   windTarget (adSTK $ ftkToSTK ftk)   $ toADTensorKindW (unWindTarget (ftkToSTK ftk) a) (unWindFTK ftk) --- | Convert a tensor with only trivial non-differentiable scalars+-- | Convert a tensor with only trivial (Z1) non-differentiable scalars -- into a tensor with the non-differentiable scalars given by the singleton -- and with zero values at the non-differentiable types. The `ShareTensor` -- constraint is needed, despite what GHC says, in order not to require -- duplicable arguments. fromADTensorKindShared   :: (BaseTensor target, ConvertTensor target, ShareTensor target)-  => SingletonTK y -> target (ADTensorKind y)+  => FullShapeTK y -> target (ADTensorKind y)   -> target y-fromADTensorKindShared stk a | Refl <- lemUnWindOfAD stk =-  windTarget stk-  $ fromADTensorKindW (unWindSTK stk) $ unWindTarget (adSTK stk) a+fromADTensorKindShared ftk a | Refl <- lemUnWindOfAD (ftkToSTK ftk) =+  windTarget (ftkToSTK ftk)+  $ fromADTensorKindW (unWindTarget (adSTK $ ftkToSTK ftk) a) (unWindFTK ftk)
+ src/HordeAd/Core/UnwindNum.hs view
@@ -0,0 +1,355 @@+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+-- | Types and functions needed to define general tensor operations+-- that work for any tensor kind, including nested (product) arrays+-- and an assortment of such operations.+--+-- Large portions of this module are a copy of HordeAd.Core.Unwind+-- with the addition of the @Num@ constraint on the underlying scalars.+module HordeAd.Core.UnwindNum+  ( addTarget, multTarget, sum0Target, dot0Target+  ) where++import Prelude++import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:))+import GHC.TypeLits (type (+))++import Data.Array.Nested (MapJust, Replicate, type (++))+import Data.Array.Nested.Convert+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Winding and unwinding++-- | This captures the normal form of type family UnWind and also+-- corresponds to the portion of ox-arrays that has Num defined.+type role RepW nominal nominal+data RepW target y where+  WTKScalar :: NumScalar r+            => target (TKScalar r)+            -> RepW target (TKScalar r)+  WTKR :: NumScalar r+       => target (TKR n r)+       -> RepW target (TKR n r)+  WTKS :: NumScalar r+       => target (TKS sh r)+       -> RepW target (TKS sh r)+  WTKX :: NumScalar r+       => target (TKX sh r)+       -> RepW target (TKX sh r)+  WTKProduct :: RepW target x -> RepW target z+             -> RepW target (TKProduct x z)++-- | This captures the normal form of type family UnWind for full shape+-- singletons.+type role FullShapeTKW nominal+data FullShapeTKW y where+  WFTKScalar :: NumScalar r+             => FullShapeTKW (TKScalar r)+  WFTKR :: NumScalar r+        => IShR n -> FullShapeTKW (TKR n r)+  WFTKS :: NumScalar r+        => ShS sh -> FullShapeTKW (TKS sh r)+  WFTKX :: NumScalar r+        => IShX sh -> FullShapeTKW (TKX sh r)+  WFTKProduct :: FullShapeTKW y -> FullShapeTKW z+              -> FullShapeTKW (TKProduct y z)++addRepW :: forall y target. BaseTensor target+        => RepW target y -> RepW target y -> RepW target y+addRepW a b = case (a, b) of+  (WTKScalar ta, WTKScalar tb) -> WTKScalar $ ta + tb+  (WTKR ta, WTKR tb) -> WTKR $ ta + tb+  (WTKS ta, WTKS tb) -> WTKS $ ta + tb+  (WTKX ta, WTKX tb) -> WTKX $ ta + tb+  (WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->+    WTKProduct (addRepW ta1 tb1) (addRepW ta2 tb2)++multRepW :: forall y target. BaseTensor target+         => RepW target y -> RepW target y -> RepW target y+multRepW a b = case (a, b) of+  (WTKScalar ta, WTKScalar tb) -> WTKScalar $ ta * tb+  (WTKR ta, WTKR tb) -> WTKR $ ta * tb+  (WTKS ta, WTKS tb) -> WTKS $ ta * tb+  (WTKX ta, WTKX tb) -> WTKX $ ta * tb+  (WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->+    WTKProduct (multRepW ta1 tb1) (multRepW ta2 tb2)++sum0RepW :: forall y target. (BaseTensor target, ConvertTensor target)+         => FullShapeTKW y -> RepW target y+         -> target (TKScalar Double)+sum0RepW ftk a = case (ftk, a) of+  (_, WTKScalar @r ta) ->+    ifDifferentiable @r (tkcast ta) 0+  (WFTKR sh, WTKR @r ta) | SNat <- shrRank sh ->+    ifDifferentiable @r (tkcast $ trsum0 ta) 0+  (WFTKS sh, WTKS @r ta) ->+    withKnownShS sh $+    ifDifferentiable @r (tkcast $ tssum0 ta) 0+  (WFTKX sh, WTKX @r ta) ->+    withKnownShX (ssxFromShX sh) $+    ifDifferentiable @r (tkcast $ txsum0 ta) 0+  (WFTKProduct ftk1 ftk2, WTKProduct ta1 ta2) ->+    sum0RepW ftk1 ta1 + sum0RepW ftk2 ta2++dot0RepW :: forall y target. (BaseTensor target, ConvertTensor target)+         => FullShapeTKW y -> RepW target y -> RepW target y+         -> target (TKScalar Double)+dot0RepW ftk a b = case (ftk, a, b) of+  (_, WTKScalar @r ta, WTKScalar tb) ->+    ifDifferentiable @r (tkcast $ ta * tb) 0+  (WFTKR sh, WTKR @r ta, WTKR tb) | SNat <- shrRank sh ->+    ifDifferentiable @r (tkcast $ trdot0 ta tb) 0+  (WFTKS sh, WTKS @r ta, WTKS tb) ->+    withKnownShS sh $+    ifDifferentiable @r (tkcast $ tsdot0 ta tb) 0+  (WFTKX sh, WTKX @r ta, WTKX tb) ->+    withKnownShX (ssxFromShX sh) $+    ifDifferentiable @r (tkcast $ txdot0 ta tb) 0+  (WFTKProduct ftk1 ftk2, WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->+    dot0RepW ftk1 ta1 tb1 + dot0RepW ftk2 ta2 tb2++type family UnWind y where+  UnWind (TKScalar r) =+    TKScalar r+  UnWind (TKR2 n (TKScalar r)) =+    TKR2 n (TKScalar r)+  UnWind (TKR2 n (TKR2 m x)) =+    UnWind (TKR2 (n + m) x)+  UnWind (TKR2 n (TKS2 sh2 x)) =+    UnWind (TKX2 (Replicate n Nothing ++ MapJust sh2) x)+  UnWind (TKR2 n (TKX2 sh2 x)) =+    UnWind (TKX2 (Replicate n Nothing ++ sh2) x)+  UnWind (TKR2 n (TKProduct y z)) =+    TKProduct (UnWind (TKR2 n y)) (UnWind (TKR2 n z))+  UnWind (TKS2 sh1 (TKScalar r)) =+    TKS2 sh1 (TKScalar r)+  UnWind (TKS2 sh1 (TKR2 m x)) =+    UnWind (TKX2 (MapJust sh1 ++ Replicate m Nothing) x)+  UnWind (TKS2 sh1 (TKS2 sh2 x)) =+    UnWind (TKS2 (sh1 ++ sh2) x)+  UnWind (TKS2 sh1 (TKX2 sh2 x)) =+    UnWind (TKX2 (MapJust sh1 ++ sh2) x)+  UnWind (TKS2 sh1 (TKProduct y z)) =+    TKProduct (UnWind (TKS2 sh1 y)) (UnWind (TKS2 sh1 z))+  UnWind (TKX2 sh1 (TKScalar r)) =+    TKX2 sh1 (TKScalar r)+  UnWind (TKX2 sh1 (TKR2 m x)) =+    UnWind (TKX2 (sh1 ++ Replicate m Nothing) x)+  UnWind (TKX2 sh1 (TKS2 sh2 x)) =+    UnWind (TKX2 (sh1 ++ MapJust sh2) x)+  UnWind (TKX2 sh1 (TKX2 sh2 x)) =+    UnWind (TKX2 (sh1 ++ sh2) x)+  UnWind (TKX2 sh1 (TKProduct y z)) =+    TKProduct (UnWind (TKX2 sh1 y)) (UnWind (TKX2 sh1 z))+  UnWind (TKProduct y z) =+    TKProduct (UnWind y) (UnWind z)++unWindFTK :: TKAllNum y+          => FullShapeTK y -> FullShapeTKW (UnWind y)+unWindFTK = \case+  FTKScalar @r | Dict0 <- numFromTKAllNum (Proxy @r) -> WFTKScalar+  FTKR sh (FTKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r) -> WFTKR sh+  FTKR sh1 (FTKR sh2 ftk2) ->+    unWindFTK $ FTKR (sh1 `shrAppend` sh2) ftk2+  FTKR sh1 (FTKS sh2 ftk2) ->+    unWindFTK+    $ FTKX (shxFromShR sh1 `shxAppend` shxFromShS sh2) ftk2+  FTKR sh1 (FTKX sh2 ftk2) ->+    unWindFTK $ FTKX (shxFromShR sh1 `shxAppend` sh2) ftk2+  FTKR sh1 (FTKProduct y z) ->+    unWindFTK $ FTKProduct (FTKR sh1 y) (FTKR sh1 z)+  FTKS sh (FTKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r) -> WFTKS sh+  FTKS sh1 (FTKR sh2 ftk2) ->+    unWindFTK+    $ FTKX (shxFromShS sh1 `shxAppend` shxFromShR sh2) ftk2+  FTKS sh1 (FTKS sh2 ftk2) ->+    unWindFTK $ FTKS (sh1 `shsAppend` sh2) ftk2+  FTKS sh1 (FTKX sh2 ftk2) ->+    unWindFTK $ FTKX (shxFromShS sh1 `shxAppend` sh2) ftk2+  FTKS sh1 (FTKProduct y z) ->+    unWindFTK $ FTKProduct (FTKS sh1 y) (FTKS sh1 z)+  FTKX sh (FTKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r) -> WFTKX sh+  FTKX sh1 (FTKR sh2 ftk2) ->+    unWindFTK $ FTKX (sh1 `shxAppend` shxFromShR sh2) ftk2+  FTKX sh1 (FTKS sh2 ftk2) ->+    unWindFTK $ FTKX (sh1 `shxAppend` shxFromShS sh2) ftk2+  FTKX sh1 (FTKX sh2 ftk2) ->+    unWindFTK $ FTKX (sh1 `shxAppend` sh2) ftk2+  FTKX sh1 (FTKProduct y z) ->+    unWindFTK $ FTKProduct (FTKX sh1 y) (FTKX sh1 z)+  FTKProduct y z -> WFTKProduct (unWindFTK y) (unWindFTK z)++-- This uses tunpairConv, so to preserve sharing, @target@ either has+-- to have a `ShareTensor` instance or the argument has to be duplicable.+-- Only the argument of the first call, not of recursive calls,+-- is assumed to be duplicable. In the AST case, this creates+-- a tower of projections for product, but if it's balanced,+-- that's of logarithmic length, so maybe even better than sharing+-- excessively, which is hard for technical typing reasons.+unWindTarget :: (TKAllNum y, ConvertTensor target)+             => SingletonTK y -> target y -> RepW target (UnWind y)+unWindTarget stk t = case stk of+  STKScalar @r | Dict0 <- numFromTKAllNum (Proxy @r) ->+    WTKScalar t+  STKR SNat (STKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r) ->+    WTKR t+  STKR (SNat @n) (STKR (SNat @m) stk2) | Dict <- lemKnownSTK stk2 ->+    unWindTarget (STKR (SNat @(n + m)) stk2) (runNest t)+  STKR n@SNat (STKS sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh2 $+    unWindTarget (STKX (ssxReplicate n+                        `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2)+                 (runNestS t)+  STKR n@SNat (STKX sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $+    unWindTarget (STKX (ssxReplicate n `ssxAppend` sh2) stk2)+                 (runNestX t)+  STKR n@SNat (STKProduct stk1 stk2) ->+    unWindTarget (STKProduct (STKR n stk1) (STKR n stk2)) (runzip t)+  STKS sh1 (STKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r) ->+    withKnownShS sh1 $ WTKS t+  STKS sh1 (STKR m@(SNat @m) stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh1 $+    unWindTarget (STKX (ssxFromShX (shxFromShS sh1)+                        `ssxAppend` ssxReplicate m) stk2) (sunNestR @_ @_ @m t)+  STKS sh1 (STKS sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh1 $ withKnownShS sh2 $+    unWindTarget (STKS (sh1 `shsAppend` sh2) stk2) (sunNest t)+  STKS sh1 (STKX sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $ withKnownShS sh1 $+    unWindTarget (STKX (ssxFromShX (shxFromShS sh1) `ssxAppend` sh2) stk2)+                 (sunNestX t)+  STKS sh1 (STKProduct stk1 stk2)->+    unWindTarget (STKProduct (STKS sh1 stk1) (STKS sh1 stk2)) (sunzip t)+  STKX sh1 (STKScalar @r) | Dict0 <- numFromTKAllNum (Proxy @r) ->+    withKnownShX sh1 $ WTKX t+  STKX sh1 (STKR m@(SNat @m) stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh1 $+    unWindTarget (STKX (sh1 `ssxAppend` ssxReplicate m) stk2)+                 (xunNestR @_ @_ @m t)+  STKX sh1 (STKS sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh1 $ withKnownShS sh2 $+    unWindTarget (STKX (sh1 `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2)+                 (xunNestS t)+  STKX sh1 (STKX sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh1 $ withKnownShX sh2 $+    unWindTarget (STKX (sh1 `ssxAppend` sh2) stk2) (xunNest t)+  STKX sh1 (STKProduct stk1 stk2) ->+    unWindTarget (STKProduct (STKX sh1 stk1) (STKX sh1 stk2)) (xunzip t)+  STKProduct stk1 stk2 ->+    let (t1, t2) = tunpairConv t+    in WTKProduct (unWindTarget stk1 t1) (unWindTarget stk2 t2)++windTarget :: ConvertTensor target+           => SingletonTK y -> RepW target (UnWind y) -> target y+windTarget stk t = case (stk, t) of+  (STKScalar, WTKScalar v) -> v+  (STKR _ STKScalar, WTKR v) -> v+  (STKR n@(SNat @n) (STKR (SNat @m) stk2), _)+   | Dict <- lemKnownSTK stk2 ->+    rnest n $ windTarget (STKR (SNat @(n + m)) stk2) t+  (STKR n (STKS sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh2 $+    rnestS n+    $ windTarget (STKX (ssxReplicate n+                        `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2) t+  (STKR n (STKX sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $+    rnestX n+    $ windTarget (STKX (ssxReplicate n `ssxAppend` sh2) stk2) t+  (STKR n@SNat (STKProduct stk1 stk2), _) | Dict <- lemKnownSTK stk1+                                          , Dict <- lemKnownSTK stk2 ->+    rzip $ windTarget (STKProduct (STKR n stk1) (STKR n stk2)) t+  (STKS _ STKScalar, WTKS v) -> v+  (STKS sh1 (STKR m@SNat stk2), _) | Dict <- lemKnownSTK stk2 ->+    snestR sh1+    $ windTarget (STKX (ssxFromShX (shxFromShS sh1)+                        `ssxAppend` ssxReplicate m) stk2) t+  (STKS sh1 (STKS sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh2 $+    snest sh1 $ windTarget (STKS (shsAppend sh1 sh2) stk2) t+  (STKS sh1 (STKX sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $+    snestX sh1 $ windTarget (STKX (ssxFromShX (shxFromShS sh1)+                                   `ssxAppend` sh2) stk2) t+  (STKS sh1 (STKProduct stk1 stk2), _) | Dict <- lemKnownSTK stk1+                                       , Dict <- lemKnownSTK stk2 ->+    szip $ windTarget (STKProduct (STKS sh1 stk1) (STKS sh1 stk2)) t+  (STKX _ STKScalar, WTKX v) -> v+  (STKX sh1 (STKR m@SNat stk2), _) | Dict <- lemKnownSTK stk2 ->+    xnestR sh1+    $ windTarget (STKX (sh1 `ssxAppend` ssxReplicate m) stk2) t+  (STKX sh1 (STKS sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh2 $+    xnestS sh1+    $ windTarget (STKX (sh1 `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2) t+  (STKX sh1 (STKX sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $+    xnest sh1 $ windTarget (STKX (ssxAppend sh1 sh2) stk2) t+  (STKX sh1 (STKProduct stk1 stk2), _) | Dict <- lemKnownSTK stk1+                                       , Dict <- lemKnownSTK stk2 ->+    xzip $ windTarget (STKProduct (STKX sh1 stk1) (STKX sh1 stk2)) t+  (STKProduct stk1 stk2, WTKProduct t1 t2) ->+    tpairConv (windTarget stk1 t1) (windTarget stk2 t2)+++-- * Operations defined using unwinding++-- | Add two (nested pairs of) tensors. Requires duplicable arguments+-- or a `ShareTensor` instance.+addTarget :: forall y target.+             (TKAllNum y, BaseTensor target, ConvertTensor target)+          => SingletonTK y -> target y -> target y -> target y+addTarget stk a b =+  let a2 = unWindTarget stk a+      b2 = unWindTarget stk b+  in gcastWith (unsafeCoerceRefl :: TKAllNum (UnWind y) :~: TKAllNum y)+     $ windTarget stk $ addRepW a2 b2++-- | Multiply two (nested pairs of) tensors. Requires duplicable arguments+-- or a `ShareTensor` instance.+multTarget :: forall y target.+              (TKAllNum y, BaseTensor target, ConvertTensor target)+           => SingletonTK y -> target y -> target y -> target y+multTarget stk a b =+  let a2 = unWindTarget stk a+      b2 = unWindTarget stk b+  in gcastWith (unsafeCoerceRefl :: TKAllNum (UnWind y) :~: TKAllNum y)+     $ windTarget stk $ multRepW a2 b2++-- | Sum all dimensions of each component and then sum it all. Ignore all+-- subtensors with non-differentiable elements.+--+-- Requires duplicable arguments or a 'ShareTensor' instance.+sum0Target :: forall y target.+              (TKAllNum y, BaseTensor target, ConvertTensor target)+           => FullShapeTK y -> target y+           -> target (TKScalar Double)+sum0Target ftk a =+  let a2 = unWindTarget (ftkToSTK ftk) a+  in gcastWith (unsafeCoerceRefl :: TKAllNum (UnWind y) :~: TKAllNum y)+     $ sum0RepW (unWindFTK ftk) a2++-- | Dot product each component and then sum it all. Ignore all+-- tensors with non-differentiable elements.+--+-- Requires duplicable arguments or a 'ShareTensor' instance.+dot0Target :: forall y target.+              (TKAllNum y, BaseTensor target, ConvertTensor target)+           => FullShapeTK y -> target y -> target y+           -> target (TKScalar Double)+dot0Target ftk a b =+  let a2 = unWindTarget (ftkToSTK ftk) a+      b2 = unWindTarget (ftkToSTK ftk) b+  in gcastWith (unsafeCoerceRefl :: TKAllNum (UnWind y) :~: TKAllNum y)+     $ dot0RepW (unWindFTK ftk) a2 b2
src/HordeAd/External/CommonRankedOps.hs view
@@ -1,4 +1,4 @@-{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+{-# LANGUAGE OverloadedLists #-} -- | Commonly used ranked operations on tensors. module HordeAd.External.CommonRankedOps   ( module HordeAd.External.CommonRankedOps@@ -7,18 +7,17 @@ import Prelude  import Control.Exception.Assert.Sugar-import Data.List.NonEmpty qualified as NonEmpty-import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, (:~:) (Refl))-import GHC.Exts (IsList (..))-import GHC.TypeLits (KnownNat, sameNat)+import Data.Type.Equality (gcastWith, (:~:))+import Data.Vector.Generic qualified as V+import GHC.TypeLits (KnownNat) -import Data.Array.Nested.Types (unsafeCoerceRefl)-import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (withShsFromShR) import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)  import HordeAd.Core.ConvertTensor-import HordeAd.Core.Ops+import HordeAd.Core.Ops (tD) import HordeAd.Core.TensorKind import HordeAd.Core.Types import HordeAd.OpsTensor@@ -29,115 +28,77 @@  rminimum :: forall target n r.             ( BaseTensor target, ConvertTensor target, LetTensor target-            , GoodScalar r )-         => target (TKR n r) -> target (TKR 0 r)-rminimum t =-  tlet (rflatten t) $ \tf ->-    rindex0 tf (tprimalPart (kfromR (rminIndex tf)) :.: ZIR)+            , NumScalar r )+         => target (TKR n r) -> target (TKScalar r)+rminimum t = tlet (rflatten t) $ \tf ->+               rindex0 tf (tplainPart (kfromR (rargMin tf)) :.: ZIR)  rmaximum :: forall target n r.             ( BaseTensor target, ConvertTensor target, LetTensor target-            , GoodScalar r )-         => target (TKR n r) -> target (TKR 0 r)-rmaximum t =-  tlet (rflatten t) $ \tf ->-    rindex0 tf (tprimalPart (kfromR (rmaxIndex tf)) :.: ZIR)--rfromIndex0 :: forall r target.-               (BaseTensor target, ConvertTensor target, GoodScalar r)-            => IntOf target -> target (TKR 0 r)-rfromIndex0 = rfromIntegral . rfromK . tfromPrimal STKScalar--rfromIndex1 :: forall n r target.-               ( KnownNat n , BaseTensor target-               , BaseTensor (PrimalOf target), ConvertTensor (PrimalOf target)-               , GoodScalar r )-            => IxROf target n -> target (TKR 1 r)-rfromIndex1 = case sameNat (Proxy @n) (Proxy @0) of-  Just Refl -> const $ rconcrete $ Nested.rfromListPrimLinear (0 :$: ZSR) []-  _ -> rfromIntegral . rfromPrimal . rfromList . NonEmpty.fromList-       . map rfromK . toList--{--rint64FromIndex1 :: forall n target.-                    ( KnownNat n-                    , BaseTensor target, BaseTensor (PrimalOf target) )-                 => IxROf target n -> target (TKR Int64 1)-rint64FromIndex1 = case sameNat (Proxy @n) (Proxy @0) of-  Just Refl -> const $ rconcrete $ Nested.rfromListPrimLinear (0 :$: ZSR) []-  _ -> rfromPrimal . rfromList . NonEmpty.fromList . indexToList--rint64ToIndex1 :: forall n target.-                  ( KnownNat n-                  , BaseTensor target, BaseTensor (PrimalOf target) )-               => target (TKR Int64 1) -> IxROf target n-rint64ToIndex1 v = listToIndex $ runravelToList $ rprimalPart v--tletIx :: ( KnownNat n, KnownNat m, GoodScalar r-          , BaseTensor target, BaseTensor (PrimalOf target)-          , LetTensor target )-       => IxROf target n -> (IxROf target n -> target (TKR m r))-       -> target (TKR m r)-tletIx ix0 f = tlet (rint64FromIndex1 ix0) $ \ixT -> f $ rint64ToIndex1 ixT--}+            , NumScalar r )+         => target (TKR n r) -> target (TKScalar r)+rmaximum t = tlet (rflatten t) $ \tf ->+               rindex0 tf (tplainPart (kfromR (rargMax tf)) :.: ZIR)  relu, reluLeaky   :: forall target n r.-     (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+     (ADReady target, NumScalar r, KnownNat n, Differentiable r)   => target (TKR n r) -> target (TKR n r) relu v0 = tlet v0 $ \v ->-  let oneIfGtZero =-        rmap0N (\x -> ifH (x <=. rscalar 0) (rscalar 0.0) (rscalar 1.0)) v+  let oneIfGtZero = rmap0N (\x -> ifH (x <=. 0) 0.0 1.0) v   in oneIfGtZero * v  reluLeaky v0 = tlet v0 $ \v ->-  let oneIfGtZero =-        rmap0N (\x -> ifH (x <=. rscalar 0) (rscalar 0.01) (rscalar 1.0)) v+  let oneIfGtZero = rmap0N (\x -> ifH (x <=. 0) 0.01 1.0) v   in oneIfGtZero * v  logistic :: forall target r n.             ( BaseTensor target, LetTensor target, BaseTensor (PrimalOf target)-            , KnownNat n, GoodScalar r, Differentiable r )+            , KnownNat n, NumScalar r, Differentiable r )          => target (TKR n r) -> target (TKR n r) logistic d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart   let one = rrepl (rshape d) 1       y0 = recip (one + exp (- rprimalPart d))-  in ttletPrimal y0 $ \y ->+  in tletPrimal y0 $ \y ->        rfromPrimal y + rfromDual (rScale @target (y * (one - y)) $ rdualPart d) +rsquare :: (NumScalar a, ADReady target)+        => target (TKR n a) -> target (TKR n a)+rsquare x' = tlet x' $ \x -> x * x+  -- slower even symbolically: rsquare x = x ** rrepl (rshape x) 2+ -- Optimized and more clearly written @u ** 2@. It's not clear if this is -- currently faster than @u ** 2@ and in which pipelines, but it's different, -- so useful as a test.-square :: forall target r n.-          ( BaseTensor target, LetTensor target-          , KnownNat n, Num (PrimalOf target (TKR n r)), GoodScalar r )-       => target (TKR n r) -> target (TKR n r)-square d = let u = rprimalPart @target d-               u' = rdualPart @target d-           in tD knownSTK (u * u) (rScale @target (2 * u) u')+squareR :: forall target r n.+           ( BaseTensor target, LetTensor target+           , KnownNat n, Num (PrimalOf target (TKR n r)), NumScalar r )+        => target (TKR n r) -> target (TKR n r)+squareR d' = tlet d' $ \d ->+  let u = rprimalPart @target d+      u' = rdualPart @target d+  in tD knownSTK (u * u) (rScale @target (2 * u) u')  squaredDifference   :: forall target n r.      ( BaseTensor target, LetTensor target-     , KnownNat n, Num (PrimalOf target (TKR n r)), GoodScalar r )+     , KnownNat n, Num (PrimalOf target (TKR n r)), NumScalar r )   => PrimalOf target (TKR n r) -> target (TKR n r) -> target (TKR n r)-squaredDifference targ res = square @target $ res - rfromPrimal @target targ+squaredDifference targ res = squareR @target $ res - rfromPrimal @target targ  lossCrossEntropyV-  :: ( BaseTensor target, ConvertTensor target-     , KnownNat n, GoodScalar r, Differentiable r )+  :: (BaseTensor target, KnownNat n, NumScalar r, Differentiable r)   => target (TKR n r) -> target (TKR n r) -> target (TKScalar r)-lossCrossEntropyV targ res = kfromR $ negate $ log res `rdot0` targ+lossCrossEntropyV targ res = negate $ log res `rdot0` targ  -- | Note that this is equivalent to a composition of softMax and cross entropy -- only when @expected@ is one-hot. Otherwise, results vary wildly. In our -- rendering of the MNIST data all labels are one-hot. lossSoftMaxCrossEntropyR   :: forall target n r.-     ( BaseTensor target, ConvertTensor target, LetTensor target+     ( BaseTensor target, LetTensor target      , BaseTensor (PrimalOf target), ConvertTensor (PrimalOf target)-     , LetTensor (PrimalOf target)-     , KnownNat n, GoodScalar r, Differentiable r )+     , LetTensor (PrimalOf target), KnownNat n, NumScalar r, Differentiable r )   => PrimalOf target (TKR n r) -> target (TKR n r) -> target (TKScalar r) lossSoftMaxCrossEntropyR expected d' = tlet d' $ \d ->   -- The following protects from underflows, overflows and exploding gradients@@ -147,65 +108,120 @@   -- and https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/xent_op.h   let u = rprimalPart d       expU' = exp (u - rreplicate0N (rshape u) (rminimum u))-  in ttletPrimal expU' $ \expU ->+  in tletPrimal expU' $ \expU ->     let softMaxU0 =           let sumExpU = rsum0 expU               recipSum = recip sumExpU           in rreplicate0N (rshape u) recipSum * expU-    in ttletPrimal softMaxU0 $ \softMaxU -> kfromR $-      tD knownSTK+    in tletPrimal softMaxU0 $ \softMaxU ->+      tD STKScalar          (negate $ log softMaxU `rdot0` expected)            -- TODO: avoid: log . exp-         (rdualPart $ rfromPrimal (softMaxU - expected) `rdot0` d)+         (kdualPart $ rfromPrimal (softMaxU - expected) `rdot0` d) +-- Fails for empty x'.+rlogsumexp :: forall n r target. (NumScalar r, Differentiable r, ADReady target)+           => target (TKR n r) -> target (TKScalar r)+rlogsumexp x' = tlet x' $ \x -> tlet (rmaximum x) $ \maxx ->+  let shiftedx = x - rreplicate0N (rshape x) maxx+      logged = log (rsum0 (exp shiftedx))+  in logged + maxx+ -- | No padding; remaining areas ignored. maxPool1 :: ( BaseTensor target, ConvertTensor target, LetTensor target-            , GoodScalar r )+            , NumScalar r )          => Int -> Int -> target (TKR 1 r) -> target (TKR 1 r) maxPool1 ksize stride v =   let slices = [rslice i ksize v | i <- [0, stride .. rwidth v - ksize]]-  in rfromList $ NonEmpty.fromList $ map rmaximum slices+      k = (rwidth v - ksize) `div` stride+  in rfromVectorLinear (k :$: ZSR) $ V.fromList $ map rmaximum slices  softMax1 :: ( BaseTensor target, LetTensor target-            , KnownNat n, GoodScalar r, Differentiable r )+            , NumScalar r, Differentiable r )          => target (TKR n r) -> target (TKR n r) softMax1 d =   let expU0 = exp d-  in tlet expU0 $ \expU -> rreplicate0N (rshape d) (recip $ rsum0 expU) * expU+  in tlet expU0 $ \expU ->+       rreplicate0N (rshape d) (recip $ rsum0 expU) * expU --- | Unpadded full convolution,---   where the output size is the same as the input size.+-- | Unpadded full convolution, where the output image size is the same+-- as the input size. -- -- BTW, the indexing lower bounds in the code are spurious, -- so they get simplified away in the resulting AST program.-conv2dUnpadded-  :: (ADReady target, GoodScalar r)+conv2dSame+  :: (ADReady target, NumScalar r)   => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded arrK arrA =+conv2dSame arrK arrA =   let [nImgs, nCinpA, nAh, nAw] = rshape arrA-      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      [nCout, nCinpK, nKh, nKw] = rshape arrK       nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA-      shB = [nImgs, nCoutK, nAh, nAw]+      shB = [nImgs, nCout, nAh, nAw]       shK1 = [1, nCinp, nKh, nKw]   in rbuild shB $ \case     [iImg, iCout, iBh, iBw] ->       let arrAt = slicez shK1 arrA [iImg, 0, iBh, iBw]           arrKt = slicez shK1 arrK [iCout, 0, 0, 0]-      in rdot0 arrAt arrKt-    _ -> error "conv2dUnpadded: impossible pattern needlessly required"+      in rfromK $ rdot0 arrAt arrKt+    _ -> error "conv2dSame: impossible pattern needlessly required" --- | Full convolution with custom padding,---   where the output size depends on the input size, kernel size and padding.+-- | Full convolution with only enough padding to ensure all output points+-- are affected by the same number of input points,+-- where the output size shrinks depending on the input size and kernel size.+-- Also no input points are ever ignored, though some are read less often.+--+-- This corresponds to+-- https://hackage.haskell.org/package/hmatrix-0.20.2/docs/Numeric-LinearAlgebra.html#v:corr2+conv2dShrinking+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dShrinking arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCout, nCinpK, nKh, nKw] = rshape arrK+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCout, nAh - nKh + 1, nAw - nKw + 1]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicez shK1 arrA [iImg, 0, iBh, iBw]+          arrKt = slicez shK1 arrK [iCout, 0, 0, 0]+      in rfromK $ rdot0 arrAt arrKt+    _ -> error "conv2dShrinking: impossible pattern needlessly required"++-- | Full convolution with enough padding to apply kernels at all+-- positions that give non-zero results. This corresponds to+-- https://hackage.haskell.org/package/hmatrix-0.20.2/docs/Numeric-LinearAlgebra.html#v:conv2+-- though it doesn't do the kernel flipping.+conv2dPadded+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPadded arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCout, nCinpK, nKh, nKw] = rshape arrK+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCout, nAh + nKh - 1, nAw + nKw - 1]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicez shK1 arrA [ iImg, 0+                                   , iBh - fromIntegral nKh + 1+                                   , iBw - fromIntegral nKw + 1 ]+          arrKt = slicez shK1 arrK [iCout, 0, 0, 0]+      in rfromK $ rdot0 arrAt arrKt+    _ -> error "conv2dPadded: impossible pattern needlessly required"++-- | Full convolution with custom padding, where the output image size+-- depends on the input size, kernel size and padding. conv2dCustomPadded-  :: (ADReady target, GoodScalar r)+  :: (ADReady target, NumScalar r)   => (Int, Int) -> target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r) conv2dCustomPadded (nPh, nPw) arrK arrA =   let [nImgs, nCinpA, nAh, nAw] = rshape arrA-      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      [nCout, nCinpK, nKh, nKw] = rshape arrK       nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA       nBh = nAh + 2 * nPh - nKh + 1       nBw = nAw + 2 * nPw - nKw + 1-      shB = [nImgs, nCoutK, nBh, nBw]+      shB = [nImgs, nCout, nBh, nBw]       shK1 = [1, nCinp, nKh, nKw]   in rbuild shB $ \case     [iImg, iCout, iBh, iBw] ->@@ -213,48 +229,33 @@           iFw = iBw - fromIntegral nPw           arrAt = slicez shK1 arrA [iImg, 0, iFh, iFw]           arrKt = slicez shK1 arrK [iCout, 0, 0, 0]-      in rdot0 arrAt arrKt+      in rfromK $ rdot0 arrAt arrKt     _ -> error "conv2dCustomPadded: impossible pattern needlessly required" --- | Full convolution with just enough padding to ensure all output points---   are affected by the same number of input points,---   where the output size shrinks depending on the input size and kernel size.---   Also no input points are ever ignored, though some are read less often.-conv2dShrinking-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)-conv2dShrinking arrK arrA =-  let [nImgs, nCinpA, nAh, nAw] = rshape arrA-      [nCoutK, nCinpK, nKh, nKw] = rshape arrK-      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA-      shB = [nImgs, nCoutK, nAh - nKh, nAw - nKw]-      shK1 = [1, nCinp, nKh, nKw]-  in rbuild shB $ \case-    [iImg, iCout, iBh, iBw] ->-      let arrAt = slicez shK1 arrA [iImg, 0, iBh, iBw]-          arrKt = slicez shK1 arrK [iCout, 0, 0, 0]-      in rdot0 arrAt arrKt-    _ -> error "conv2dShrinking: impossible pattern needlessly required"- -- | Slice a section out of a tensor, --   given a base offset and shape of the section. -- --   If the slice extends out side the source array then the corresponding --   elements are set to zero. slicez-  :: (ADReady target, GoodScalar r, KnownNat n)+  :: (ADReady target, NumScalar r, KnownNat n)   => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r) slicez shOut d ixBase =-  rbuild shOut $ \ixResult -> rindex0 d (ixrZipWith (+) ixBase ixResult)+  rbuild shOut $ \ixResult -> rindex @_ @0 d (ixrZipWith (+) ixBase ixResult)  maxPool2dUnpadded-  :: (ADReady target, GoodScalar r)+  :: (ADReady target, NumScalar r)   => Int -> Int -> target (TKR 4 r) -> target (TKR 4 r) maxPool2dUnpadded ksize stride arr =   let [batch_size, channels, h, w] = rshape arr-      shOut = [batch_size, channels, h `div` stride, w `div` stride]+      shOutR :: IShR 4+      shOutR = [batch_size, channels, h `div` stride, w `div` stride]+      shK1 :: IShR 4       shK1 = [1, 1, ksize, ksize]-  in rbuild shOut $ \case+  in+    withShsFromShR shOutR $ \(sh :: ShS shOut) ->+    withKnownShS sh $+    rfromS $ kbuild @shOut $ \case     [iImg, iChan, iBh, iBw] ->       rmaximum $ slicez shK1 arr [ iImg, iChan                                  , fromIntegral stride * iBh@@ -262,6 +263,6 @@     _ -> error "maxPool2dUnpadded: impossible pattern needlessly required"  xfromIndex0 :: forall r target.-               (BaseTensor target, ConvertTensor target, GoodScalar r)+               (BaseTensor target, ConvertTensor target, NumScalar r)             => IntOf target -> target (TKX '[] r)-xfromIndex0 = xfromIntegral . xfromK . tfromPrimal STKScalar+xfromIndex0 = xfromIntegral . xfromK . tfromPlain STKScalar
src/HordeAd/External/CommonShapedOps.hs view
@@ -1,4 +1,4 @@-{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+{-# LANGUAGE OverloadedLists #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} -- | Commonly used operations on shaped tensors.@@ -8,84 +8,53 @@  import Prelude -import Data.List.NonEmpty qualified as NonEmpty-import Data.Proxy (Proxy (Proxy)) import Data.Type.Equality (gcastWith, (:~:) (Refl)) import Data.Type.Ord (Compare)-import GHC.Exts (IsList (..))-import GHC.TypeLits-  (Div, KnownNat, SomeNat (..), sameNat, someNatVal, type (-), type (<=))+import Data.Vector.Generic qualified as V+import GHC.TypeLits (Div, KnownNat, type (+), type (-), type (<=)) -import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Convert (ixrFromIxS, ixsFromIxR')+import Data.Array.Nested.Convert (ixrFromIxS, ixsFromIxR)+import Data.Array.Nested.Lemmas import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape import Data.Array.Nested.Types (unsafeCoerceRefl) -import HordeAd.Core.ConvertTensor-import HordeAd.Core.Ops+import HordeAd.Core.Ops (tD) import HordeAd.Core.TensorKind import HordeAd.Core.Types import HordeAd.OpsTensor -sminimum :: forall r sh target. (ADReady target, GoodScalar r, KnownShS sh)-         => target (TKS sh r) -> target (TKS '[] r)-sminimum t | SNat <- shsProduct (knownShS @sh) =-  tlet (sflatten t) $ \tf ->-    sindex0 tf (tprimalPart (kfromS (sminIndex tf)) :.$ ZIS)--smaximum :: forall r sh target. (ADReady target, GoodScalar r, KnownShS sh)-         => target (TKS sh r) -> target (TKS '[] r)-smaximum t | SNat <- shsProduct (knownShS @sh) =-  tlet (sflatten t) $ \tf ->-    sindex0 tf (tprimalPart (kfromS (smaxIndex tf)) :.$ ZIS)--sfromIndex0 :: forall r target. (ADReady target, GoodScalar r)-            => IntOf target -> target (TKS '[] r)-sfromIndex0 = sfromR . rfromIntegral . rfromPrimal . rfromK--sfromIndex1 :: forall r sh target.-               (ADReady target, GoodScalar r, KnownShS sh)-            => IxSOf target sh -> target (TKS '[Rank sh] r)-sfromIndex1 | SNat <- shsRank (knownShS @sh) =-  case sameNat (Proxy @(Rank sh)) (Proxy @0) of-    Just Refl -> const $ sconcrete $ Nested.sfromListPrimLinear knownShS []-    _ -> sfromR . rfromIntegral . rfromPrimal . rfromList-         . NonEmpty.fromList . map rfromK . toList+sminimum :: forall r sh target. (ADReady target, NumScalar r, KnownShS sh)+         => target (TKS sh r) -> target (TKScalar r)+sminimum t = tlet (sflatten t) $ \tf ->+               sindex0 tf (tplainPart (kargMin tf) :.$ ZIS) -{--sletIx :: forall r sh n target.-          (ADReady target, GoodScalar r, KnownShS sh, KnownNat n)-       => IxROf target n -> (IxROf target n -> target (TKS sh r))-       -> target (TKS sh r)-sletIx ix0 f = tlet (sfromR @target @Int64 @'[n]-                     $ rint64FromIndex1 ix0) $ \ixT ->-                 f $ rint64ToIndex1 $ rfromS @target ixT--}+smaximum :: forall r sh target. (ADReady target, NumScalar r, KnownShS sh)+         => target (TKS sh r) -> target (TKScalar r)+smaximum t = tlet (sflatten t) $ \tf ->+               sindex0 tf (tplainPart (kargMax tf) :.$ ZIS)  reluS, reluLeakyS   :: forall target sh r.-     (KnownShS sh, ADReady target, GoodScalar r, Differentiable r)+     (KnownShS sh, ADReady target, NumScalar r, Differentiable r)   => target (TKS sh r) -> target (TKS sh r) reluS v0 = tlet v0 $ \v ->-  let oneIfGtZero =-        smap0N (\x -> ifH (x <=. sscalar 0) (sscalar 0.0) (sscalar 1.0)) v+  let oneIfGtZero = smap0N (\x -> ifH (x <=. 0) 0.0 1.0) v   in oneIfGtZero * v  reluLeakyS v0 = tlet v0 $ \v ->-  let oneIfGtZero =-        smap0N (\x -> ifH (x <=. sscalar 0) (sscalar 00.01) (sscalar 01.0)) v+  let oneIfGtZero = smap0N (\x -> ifH (x <=. 0) 00.01 01.0) v   in oneIfGtZero * v  logisticS :: forall target r sh.              ( BaseTensor target, LetTensor target, BaseTensor (PrimalOf target)-             , KnownShS sh, GoodScalar r, Differentiable r )+             , KnownShS sh, NumScalar r, Differentiable r )           => target (TKS sh r) -> target (TKS sh r) logisticS d0 = tlet d0 $ \d ->  -- used in rprimalPart and in sdualPart   let one = srepl 1       y0 = recip (one + exp (- sprimalPart d))-  in ttletPrimal y0 $ \y ->+  in tletPrimal y0 $ \y ->        sfromPrimal y + sfromDual (sScale @target (y * (one - y)) $ sdualPart d)  -- Optimized and more clearly written @u ** 2@. It's not clear if this is@@ -93,32 +62,33 @@ -- so useful as a test. squareS :: forall target r sh.            ( KnownShS sh, BaseTensor target, LetTensor target-           , Num (PrimalOf target (TKS sh r)), GoodScalar r )+           , Num (PrimalOf target (TKS sh r)), NumScalar r )         => target (TKS sh r) -> target (TKS sh r)-squareS d = let u = sprimalPart d-                u' = sdualPart d-            in tD knownSTK (u * u) (sScale @target (2 * u) u')+squareS d' = tlet d' $ \d ->+  let u = sprimalPart d+      u' = sdualPart d+  in tD knownSTK (u * u) (sScale @target (2 * u) u')  squaredDifferenceS   :: forall target sh r.      ( KnownShS sh, BaseTensor target, LetTensor target-     , Num (PrimalOf target (TKS sh r)), GoodScalar r )+     , Num (PrimalOf target (TKS sh r)), NumScalar r )   => PrimalOf target (TKS sh r) -> target (TKS sh r) -> target (TKS sh r) squaredDifferenceS targ res = squareS $ res - sfromPrimal targ -lossCrossEntropyVS :: ( KnownShS sh, GoodScalar r, Differentiable r-                      , BaseTensor target, ConvertTensor target )+lossCrossEntropyVS :: ( KnownShS sh, NumScalar r, Differentiable r+                      , BaseTensor target )                    => target (TKS sh r)                    -> target (TKS sh r)                    -> target (TKScalar r)-lossCrossEntropyVS targ res = kfromS $ negate $ log res `sdot0` targ+lossCrossEntropyVS targ res = negate $ log res `sdot0` targ  -- | Note that this is equivalent to a composition of softMax and cross entropy -- only when @expected@ is one-hot. Otherwise, results vary wildly. In our -- rendering of the MNIST data all labels are one-hot. lossSoftMaxCrossEntropyS   :: forall target sh r.-     ( ADReady target, ADReady (PrimalOf target), GoodScalar r, KnownShS sh+     ( ADReady target, ADReady (PrimalOf target), NumScalar r, KnownShS sh      , Differentiable r )   => PrimalOf target (TKS sh r) -> target (TKS sh r) -> target (TKScalar r) lossSoftMaxCrossEntropyS expected d' = tlet d' $ \d ->@@ -134,59 +104,119 @@           let sumExpU = ssum0 expU               recipSum = recip sumExpU           in sreplicate0N recipSum * expU-  in ttletPrimal softMaxU0 $ \softMaxU -> kfromS $-    tD knownSTK+  in tletPrimal softMaxU0 $ \softMaxU ->+    tD STKScalar        (negate $ log softMaxU `sdot0` expected)          -- TODO: avoid: log . exp-       (sdualPart $ sfromPrimal (softMaxU - expected) `sdot0` d)+       (kdualPart $ sfromPrimal (softMaxU - expected) `sdot0` d)  -- | No padding; remaining areas ignored. maxPool1S :: forall ksize stride m target r.-             ( ADReady target, GoodScalar r-             , KnownNat ksize, KnownNat stride, KnownNat m )-          => target (TKS '[m] r) -> target (TKS '[m] r)-maxPool1S v =+             (ADReady target, NumScalar r, KnownNat m)+          => SNat ksize -> SNat stride+          -> target (TKS '[m] r) -> target (TKS '[m] r)+maxPool1S SNat SNat v =   let l = [0, valueOf @stride .. swidth v - valueOf @ksize]-      maxOfSlice i =-        case someNatVal $ toInteger i of-          Just (SomeNat @i _proxy) ->-            gcastWith (unsafeCoerceRefl :: Compare i m :~: LT) $-            gcastWith (unsafeCoerceRefl :: Compare ksize (m - i) :~: LT) $-            smaximum $ sslice @i @(m - i - ksize) @ksize SNat SNat SNat v-          Nothing -> error "maxPool1S: impossible someNatVal error"-  in sfromList $ NonEmpty.fromList $ map maxOfSlice l+      maxOfSlice i = withSNat i $ \ (SNat @i) ->+        gcastWith (unsafeCoerceRefl :: Compare i m :~: LT) $+        gcastWith (unsafeCoerceRefl :: Compare ksize (m - i) :~: LT) $+        smaximum $ sslice @i @(m - i - ksize) @ksize SNat SNat SNat v+  in sfromVectorLinear (SNat @m :$$ ZSS) $ V.fromList $ map maxOfSlice l  softMax1S :: forall target sh r.-             ( KnownShS sh, BaseTensor target, LetTensor target-             , GoodScalar r, Differentiable r )+             ( KnownShS sh, NumScalar r, Differentiable r+             , BaseTensor target, LetTensor target )           => target (TKS sh r) -> target (TKS sh r) softMax1S d =   let expU0 = exp d   in tlet expU0 $ \expU -> sreplicate0N (recip $ ssum0 expU) * expU --- | Unpadded full convolution,---   where the output size is the same as the input size.-conv2dUnpaddedS-  :: forall nCoutK nCinpK nKh nKw nImgs nCinpA nAh nAw-            target r shB shK1.-     ( KnownNat nCoutK, KnownNat nCinpK, KnownNat nKh, KnownNat nKw-     , KnownNat nImgs, KnownNat nAh, KnownNat nAw-     , ADReady target, GoodScalar r-     , nCinpA ~ nCinpK-     , shB ~ '[nImgs, nCoutK, nAh, nAw]+-- | Full convolution, where the output image size is the same+-- as the input size.+conv2dSameS+  :: forall nImgs nCinp nCinpA nCout nAh nAw nKh nKw shK shA shB shK1+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh, KnownNat nAw, KnownNat nKh, KnownNat nKw+     , ADReady target, NumScalar r+     , nCinpA ~ nCinp+     , shK  ~ '[nCout, nCinp, nKh, nKw]+     , shA  ~ '[nImgs, nCinp, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh, nAw]      , shK1 ~ '[1, nCinpA, nKh, nKw]      )-  => target (TKS '[nCoutK, nCinpK, nKh, nKw] r)-  -> target (TKS '[nImgs, nCinpA, nAh, nAw] r)-  -> target (TKS shB r)-conv2dUnpaddedS arrK arrA =-  sbuild @(Rank shB) $ \case+  => target (TKS shK r) -> target (TKS shA r) -> target (TKS shB r)+conv2dSameS arrK arrA =+  kbuild $ \case     [iImg, iCout, iBh, iBw] ->-      let arrAt = slicezS @shK1 arrA [iImg, 0, iBh, iBw]-          arrKt = slicezS arrK [iCout, 0, 0, 0]+      let arrAt = slicezS @shK1 arrA+                          [iImg, 0, iBh, iBw]+          arrKt = slicezS arrK+                          [iCout, 0, 0, 0]       in sdot0 arrAt arrKt-    _ -> error "conv2dUnpaddedS: impossible pattern needlessly required"+    _ -> error "conv2dSameS: impossible pattern needlessly required" +-- | Full convolution with only enough padding to ensure all output points+-- are affected by the same number of input points,+-- where the output size shrinks depending on the input size and kernel size.+-- Also no input points are ever ignored, though some are read less often.+--+-- This corresponds to+-- https://hackage.haskell.org/package/hmatrix-0.20.2/docs/Numeric-LinearAlgebra.html#v:corr2+conv2dShrinkingS+  :: forall nImgs nCinp nCinpA nCout nAh_nKh1 nAw_nKw1 nKh1 nKw1+            shK shA shB shK1+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh_nKh1, KnownNat nAw_nKw1, KnownNat nKh1, KnownNat nKw1+     , ADReady target, NumScalar r+     , nCinpA ~ nCinp+     , shK  ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA  ~ '[nImgs, nCinpA, nAh_nKh1 + nKh1, nAw_nKw1 + nKw1]+     , shB  ~ '[nImgs, nCout, nAh_nKh1, nAw_nKw1]+     , shK1 ~ '[1, nCinpA, nKh1 + 1, nKw1 + 1]+     )+  => target (TKS shK r) -> target (TKS shA r) -> target (TKS shB r)+conv2dShrinkingS arrK arrA =+  kbuild $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezS @shK1 arrA+                          [iImg, 0, iBh, iBw]+          arrKt = slicezS arrK+                          [iCout, 0, 0, 0]+      in sdot0 arrAt arrKt+    _ -> error "conv2dShrinkingS: impossible pattern needlessly required"++-- | Full convolution with enough padding to apply kernels at all+-- positons that give non-zero results. This corresponds to+-- https://hackage.haskell.org/package/hmatrix-0.20.2/docs/Numeric-LinearAlgebra.html#v:conv2+-- though it doesn't do the kernel flipping.+conv2dPaddedS+  :: forall nImgs nCinp nCinpA nCout nAh nAw nKh1 nKw1+            shK shA shB shK1+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh, KnownNat nAw, KnownNat nKh1, KnownNat nKw1+     , ADReady target, NumScalar r+     , nCinpA ~ nCinp+     , shK  ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA  ~ '[nImgs, nCinpA, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh + nKh1, nAw + nKw1]+     , shK1 ~ '[1, nCinpA, nKh1 + 1, nKw1 + 1]+     )+  => target (TKS shK r) -> target (TKS shA r) -> target (TKS shB r)+conv2dPaddedS arrK arrA =+  kbuild $ \case+    [iImg, iCout, iBh, iBw] ->+      let nKh1 = valueOf @nKh1+          nKw1 = valueOf @nKw1+          arrAt = slicezS @shK1 arrA+                          [iImg, 0, iBh - nKh1, iBw - nKw1]+          arrKt = slicezS arrK+                          [iCout, 0, 0, 0]+      in sdot0 arrAt arrKt+    _ -> error "conv2dPaddedS: impossible pattern needlessly required"+ -- | Slice a section out of a tensor, --   given a base offset and shape of the section. --@@ -194,34 +224,35 @@ --   elements are set to zero. slicezS   :: forall shOut sh target r.-     ( KnownShS sh, KnownShS shOut, KnownShS (Take (Rank sh) shOut)-     , KnownNat (Rank sh)-     , Rank shOut ~ Rank sh, ADReady target, GoodScalar r )+     (KnownShS shOut, Rank shOut ~ Rank sh, ADReady target, NumScalar r)   => target (TKS sh r) -> IxSOf target sh -> target (TKS shOut r)-slicezS d ixBase =+slicezS d ixBase | Refl <- lemAppNil @sh =   gcastWith (unsafeCoerceRefl              :: Rank (Take (Rank shOut) shOut) :~: Rank shOut) $   gcastWith (unsafeCoerceRefl :: Drop (Rank sh) shOut :~: '[]) $-  sbuild @(Rank shOut)+  kbuild   $ \ixResult ->-      sindex0 d (ixsFromIxR' knownShS $ ixrZipWith (+) (ixrFromIxS ixBase) (ixrFromIxS ixResult))-      -- TODO: this doesn't work, because ixsZipWith has too strict a type:-      -- sbuild @(Rank shOut) $ \ixResult -> sindex0 d (ixsZipWith (+) ixBase ixResult)+      sindex0 @sh d+              (ixsFromIxR+               $ ixrZipWith (+) (ixrFromIxS ixBase) (ixrFromIxS ixResult))+  -- TODO: this doesn't work, because ixsZipWith has too strict a type:+  -- sbuild @(Rank shOut) $ \ixResult -> sindex d (ixsZipWith (+) ixBase ixResult)  maxPool2dUnpaddedS   :: forall ksize stride batch_size channels h w target r shOut shK1.-     ( KnownNat ksize, KnownNat stride, KnownNat batch_size, KnownNat channels+     ( KnownNat batch_size, KnownNat channels      , KnownNat h, KnownNat w      , 1 <= stride  -- wrongly reported as redundant due to plugins-     , ADReady target, GoodScalar r+     , ADReady target, NumScalar r      , shOut ~ '[batch_size, channels, h `Div` stride, w `Div` stride]      , shK1 ~ '[1, 1, ksize, ksize]      )-  => target (TKS '[batch_size, channels, h, w] r)+  => SNat ksize -> SNat stride+  -> target (TKS '[batch_size, channels, h, w] r)   -> target (TKS shOut r)-maxPool2dUnpaddedS arr =+maxPool2dUnpaddedS SNat SNat arr =   let stride = valueOf @stride :: Int-  in sbuild @(Rank shOut) $ \case+  in kbuild @shOut $ \case     [iImg, iChan, iBh, iBw] ->       smaximum $ slicezS @shK1 arr [ iImg, iChan                                    , fromIntegral stride * iBh
src/HordeAd/External/Optimizer.hs view
@@ -13,11 +13,12 @@ import HordeAd.Core.OpsADVal import HordeAd.Core.OpsConcrete () import HordeAd.Core.TensorKind+import HordeAd.Core.Types import HordeAd.External.OptimizerTools  -- | Stochastic Gradient Descent.-sgdSTK :: forall a x z.-          SingletonTK x+sgdSTK :: forall a x z. TKAllNum (ADTensorKind z)+       => SingletonTK x        -> Double  -- ^ gamma (learning_rate?)        -> (a -> ADVal Concrete x -> ADVal Concrete z)        -> [a]  -- ^ training data@@ -32,18 +33,19 @@   go (a : rest) !parameters =     let inputs :: ADVal Concrete x         inputs = dDnotShared parameters deltaInputs-        (gradients, valueNew) = crevOnADInputs Nothing (f a) zftk inputs+        (valueNew, gradients) = crevOnADInputs Nothing (f a) zftk inputs         parametersNew = updateWithGradient gamma stk parameters gradients     in if null rest        then (parametersNew, valueNew)        else go rest parametersNew -sgd :: forall a x z. KnownSTK x+sgd :: forall a x z. (TKAllNum (ADTensorKind z), KnownSTK x)     => Double  -- ^ gamma (learning_rate?)     -> (a -> ADVal Concrete x -> ADVal Concrete z)     -> [a]  -- ^ training data     -> Concrete x  -- ^ initial parameters     -> (Concrete x, Concrete z)+{-# INLINE sgd #-} sgd = sgdSTK knownSTK  -- We inline (possibly causing a binary blowup) until we are able to work around@@ -51,7 +53,7 @@ -- and specialize. -- | An implementation of the Adam gradient descent. sgdAdam-  :: forall a x z . KnownSTK x+  :: forall a x z . (TKAllNum (ADTensorKind z), KnownSTK x)   => (a -> ADVal Concrete x -> ADVal Concrete z)   -> [a]   -> Concrete x@@ -61,14 +63,13 @@ sgdAdam = sgdAdamArgs defaultArgsAdam  sgdAdamArgs-  :: forall a x z. KnownSTK x+  :: forall a x z. (TKAllNum (ADTensorKind z), KnownSTK x)   => ArgsAdam   -> (a -> ADVal Concrete x -> ADVal Concrete z)   -> [a]   -> Concrete x   -> StateAdam x   -> (Concrete x, StateAdam x)-{-# INLINE sgdAdamArgs #-} sgdAdamArgs argsAdam f trainingData !parameters0 !stateAdam0 =   go trainingData parameters0 stateAdam0  where@@ -80,7 +81,7 @@   go (a : rest) !parameters !stateAdam =     let inputs :: ADVal Concrete x         inputs = dDnotShared parameters deltaInputs-        gradients = fst $ crevOnADInputs Nothing (f a) zftk inputs+        gradients = snd $ crevOnADInputs Nothing (f a) zftk inputs         (parametersNew, stateAdamNew) =           updateWithGradientAdam             argsAdam stateAdam knownSTK parameters gradients
src/HordeAd/External/OptimizerTools.hs view
@@ -9,9 +9,6 @@  import Prelude -import Data.Type.Equality (gcastWith, (:~:))--import Data.Array.Nested.Types (unsafeCoerceRefl) import Data.Array.Nested qualified as Nested  import HordeAd.Core.CarriersConcrete@@ -28,27 +25,23 @@                              g@(Concrete gradient) = case stk of   STKScalar @r -> Concrete $     ifDifferentiable @r-      (gcastWith (unsafeCoerceRefl :: y :~: ADTensorKind y) $-       params - realToFrac gamma * gradient)+      (params - realToFrac gamma * gradient)       params   STKR _ (STKScalar @r) -> Concrete $     ifDifferentiable @r-      (gcastWith (unsafeCoerceRefl :: y :~: ADTensorKind y) $-       params - Nested.rreplicateScal (Nested.rshape params)+      (params - Nested.rreplicatePrim (Nested.rshape params)                                       (realToFrac gamma)                 * gradient)       params   STKS _ (STKScalar @r) -> Concrete $     ifDifferentiable @r-      (gcastWith (unsafeCoerceRefl :: y :~: ADTensorKind y) $-       params - Nested.sreplicateScal (Nested.sshape params)+      (params - Nested.sreplicatePrim (Nested.sshape params)                                       (realToFrac gamma)                 * gradient)       params   STKX _ (STKScalar @r) -> Concrete $     ifDifferentiable @r-      (gcastWith (unsafeCoerceRefl :: y :~: ADTensorKind y) $-       params - Nested.mreplicateScal (Nested.mshape params)+      (params - Nested.mreplicatePrim (Nested.mshape params)                                       (realToFrac gamma)                 * gradient)       params@@ -154,18 +147,18 @@               -> (Nested.Ranked n r, Nested.Ranked n r, Nested.Ranked n r)       updateR mA vA p g =         let sh = Nested.rshape g-            mANew = Nested.rreplicateScal sh (realToFrac betaOne) * mA-                    + Nested.rreplicateScal sh (realToFrac oneMinusBeta1) * g-            vANew = Nested.rreplicateScal sh (realToFrac betaTwo) * vA-                    + Nested.rreplicateScal sh (realToFrac oneMinusBeta2)+            mANew = Nested.rreplicatePrim sh (realToFrac betaOne) * mA+                    + Nested.rreplicatePrim sh (realToFrac oneMinusBeta1) * g+            vANew = Nested.rreplicatePrim sh (realToFrac betaTwo) * vA+                    + Nested.rreplicatePrim sh (realToFrac oneMinusBeta2)                       * (g * g)             alphat = alpha * sqrt (1 - betaTwo ^ tAdamNew)                              / (1 - betaOne ^ tAdamNew)         in ( mANew            , vANew-           , p - (Nested.rreplicateScal sh (realToFrac alphat) * mANew)+           , p - (Nested.rreplicatePrim sh (realToFrac alphat) * mANew)                  / (sqrt vANew-                    + Nested.rreplicateScal sh (realToFrac epsilon)) )+                    + Nested.rreplicatePrim sh (realToFrac epsilon)) )       updateProd :: forall y2.                     SingletonTK y2                  -> Concrete y2 -> Concrete y2@@ -175,8 +168,7 @@                      (Concrete p) (Concrete g) = case stk of         STKScalar @r ->           ifDifferentiable @r-            (gcastWith (unsafeCoerceRefl :: y2 :~: ADTensorKind y2) $-             let !(!mAN, !vAN, !pN) =+            (let !(!mAN, !vAN, !pN) =                    updateR (Nested.rscalar mA)                            (Nested.rscalar vA)                            (Nested.rscalar p)@@ -188,14 +180,12 @@             (Concrete ((mA, vA), p))         STKR SNat (STKScalar @r) ->           ifDifferentiable @r-            (gcastWith (unsafeCoerceRefl :: y2 :~: ADTensorKind y2) $-             let !(!mAN, !vAN, !pN) = updateR mA vA p g+            (let !(!mAN, !vAN, !pN) = updateR mA vA p g              in Concrete ((mAN, vAN), pN))             (Concrete ((mA, vA), p))         STKS sh (STKScalar @r) ->           ifDifferentiable @r-            (gcastWith (unsafeCoerceRefl :: y2 :~: ADTensorKind y2) $-             let !(!mAN, !vAN, !pN) =+            (let !(!mAN, !vAN, !pN) =                    updateR (Nested.stoRanked mA)                            (Nested.stoRanked vA)                            (Nested.stoRanked p)@@ -207,8 +197,7 @@             (Concrete ((mA, vA), p))         STKX _ (STKScalar @r) ->           ifDifferentiable @r-            (gcastWith (unsafeCoerceRefl :: y2 :~: ADTensorKind y2) $-             let !(!mAN, !vAN, !pN) =+            (let !(!mAN, !vAN, !pN) =                    updateR (Nested.mtoRanked mA)                            (Nested.mtoRanked vA)                            (Nested.mtoRanked p)
src/HordeAd/OpsTensor.hs view
@@ -19,18 +19,25 @@   , xshape, xlength, xsize, xwidth   , tsize, tftk     -- * Constructing arrays from concrete values, lists and vectors+  , kconcrete   , rconcrete, rscalar, rrepl, ringestData, rfromListLinear   , sconcrete, sscalar, srepl, singestData, sfromListLinear   , xconcrete, xscalar, xrepl, xingestData, xfromListLinear-  , kconcrete-  , rfromList, rfromVector, rfromVector0N, rfromList0N, runravelToList-  , sfromList, sfromVector, sfromVector0N, sfromList0N, sunravelToList-  , xfromList, xfromVector, xfromVector0N, xfromList0N, xunravelToList+  , rfromList, rfromVector, rfromVectorN, rfromVectorLinear+  , runravelToList, runravelToListN, rtoListLinear+  , sfromList, sfromVector, sfromVectorN, sfromVectorLinear+  , sunravelToList, sunravelToListN, stoListLinear+  , xfromList, xfromVector, xfromVectorN, xfromVectorLinear+  , xunravelToList, xunravelToListN, xtoListLinear     -- * Main array operations-  , tunit, tlet, ifH, minH, maxH, tpair, tproject1, tproject2-  , rsum, rsum0, rdot0, rdot1In, rmatvecmul, rmatmul2, rreplicate, rreplicate0N-  , ssum, ssum0, sdot0, sdot1In, smatvecmul, smatmul2, sreplicate, sreplicate0N-  , xsum, xsum0, xdot0, xdot1In, xmatvecmul, xmatmul2, xreplicate, xreplicate0N+  , tunit, tlet, tletPrimal, tletPlain, ifH, minH, maxH+  , tpair, tproject1, tproject2+  , rsum, rsumN, rsum0, rdot0, rdot1In, rmatvecmul, rmatmul2+  , rreplicate, rreplicateN, rreplicate0N+  , ssum, ssumN, ssum0, sdot0, sdot1In, smatvecmul, smatmul2+  , sreplicate, sreplicateN, sreplicate0N+  , xsum, xsumN, xsum0, xdot0, xdot1In, xmatvecmul, xmatmul2+  , xreplicate, xreplicateN, xreplicate0N   , rindex, (!), rindex0, roneHot, rscatter, rscatter1, rgather, rgather1   , sindex, (!$), sindex0, soneHot, sscatter, sscatter1, sgather, sgather1   , xindex, xindex0, xoneHot, xscatter, xscatter1, xgather, xgather1@@ -38,42 +45,45 @@   , str, stranspose, sflatten, sreshape   , xtr, xtranspose, xflatten, xreshape    -- * Auxiliary array operations-  , rfloor, rfromIntegral, rcast, rminIndex, rmaxIndex, riota-  , sfloor, sfromIntegral, scast, sminIndex, smaxIndex, siota-  , xfloor, xfromIntegral, xcast, xminIndex, xmaxIndex, xiota-  , kfloor, kfromIntegral, kcast+  , kfloor, kfromIntegral, kcast, kargMin, kargMax+  , rfloor, rfromIntegral, rcast, rargMin, rargMax, riota+  , sfloor, sfromIntegral, scast, sargMin, sargMax, siota+  , xfloor, xfromIntegral, xcast, xargMin, xargMax, xiota   , rappend, rconcat, rslice, runcons, rreverse   , sappend, sslice, suncons, sreverse-  , xappend, xappend0, xconcat, xslice, xuncons, xreverse+  , xappend, xconcat, xslice, xuncons, xreverse     -- * Array operations derived from @build@+  , kbuild, kbuild1   , rbuild, rbuild1, rmap, rmap1, rmap0N, rzipWith, rzipWith1, rzipWith0N   , rzipWith3, rzipWith31, rzipWith30N, rzipWith4, rzipWith41, rzipWith40N   , sbuild, sbuild1, smap, smap1, smap0N, szipWith, szipWith1, szipWith0N   , szipWith3, szipWith31, szipWith30N, szipWith4, szipWith41, szipWith40N   , xbuild, xbuild1     -- * Array operations derived from @mapAccum@-  , rfold, rscan, sfold, sscan, xfold, xscan, tmapAccumR, tmapAccumL+  , rfold, rscan, sfold, sscan, xfold, xscan+  , tfold, tscan, tmapAccumR, tmapAccumL     -- * Array operations producing derivatives   , kgrad, rvjp, rjvp, svjp, sjvp     -- * Operations dealing with dual numbers+  , kprimalPart, kdualPart, kplainPart, kfromPrimal, kfromDual, kfromPlain+  , kScale   , rprimalPart, rdualPart, rfromPrimal, rfromDual, rScale   , sprimalPart, sdualPart, sfromPrimal, sfromDual, sScale   , xprimalPart, xdualPart, xfromPrimal, xfromDual, xScale-  , kprimalPart, kdualPart, kfromPrimal, kfromDual, kScale+  , tplainPart, tfromPlain, tprimalPart, tfromPrimal     -- * Array operations that utilize unwinding of nested arrays   , treplTarget, tdefTarget, taddTarget, tmultTarget, tsum0Target, tdot0Target     -- * Minimal re-exports to make this module a higher level replacement for "HordeAd.Core.Ops"-  , ADReady+  , ADReady, ADReadyNoLet, ShareTensor   , LetTensor, BaseTensor   ) where  import Prelude -import Data.Int (Int64) import Data.List.NonEmpty (NonEmpty) import Data.List.NonEmpty qualified as NonEmpty import Data.Proxy (Proxy (Proxy))-import Data.Type.Equality (gcastWith, (:~:))+import Data.Type.Equality (gcastWith, (:~:) (Refl)) import Data.Vector.Generic qualified as V import Data.Vector.Strict qualified as Data.Vector import GHC.TypeLits@@ -81,6 +91,7 @@  import Data.Array.Nested (type (++)) import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Lemmas import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Permutation qualified as Permutation import Data.Array.Nested.Ranked.Shape@@ -99,10 +110,6 @@ rscalar :: (GoodScalar r, BaseTensor target)         => r -> target (TKR 0 r) rscalar r = rconcrete $ Nested.rscalar r-rrepl :: forall n r target. (GoodScalar r, BaseTensor target)-      => IShR n -> r -> target (TKR n r)-rrepl sh a = tconcrete (FTKR sh FTKScalar)-                       (Concrete $ Nested.rreplicateScal sh a) ringestData :: forall n r target. (GoodScalar r, BaseTensor target)             => IShR n -> [r] -> target (TKR n r) ringestData sh l =@@ -110,8 +117,12 @@ -- | Create a tensor from a list of individual underlying scalar values. -- -- An operation with the same name is used in @ox-arrays@ to represent--- pretty-printed concrete arrays. The types agree, so the same representation--- results in analogous horde-ad array when this operation is used instead.+-- pretty-printed concrete arrays. The types mostly agree,+-- so the same representation results in analogous horde-ad array+-- when this operation is used instead.+--+-- Note that 'rfromVectorLinear' has very different typing that coincides+-- only for the concrete instance of @target@. rfromListLinear :: forall n r target. (GoodScalar r, BaseTensor target)                 => IShR n -> NonEmpty r -> target (TKR n r) rfromListLinear sh = ringestData sh . NonEmpty.toList@@ -122,17 +133,18 @@ sscalar :: (GoodScalar r, BaseTensor target)         => r -> target (TKS '[] r) sscalar r = sconcrete $ Nested.sscalar r-srepl :: (KnownShS sh, GoodScalar r, BaseTensor target)-      => r -> target (TKS sh r)-srepl = sconcrete . Nested.sreplicateScal knownShS singestData :: (KnownShS sh, GoodScalar r, BaseTensor target)             => [r] -> target (TKS sh r) singestData l = sconcrete $ Nested.sfromListPrimLinear knownShS l -- | Create a tensor from a list of individual underlying scalar values. -- -- An operation with the same name is used in @ox-arrays@ to represent--- pretty-printed concrete arrays. The types agree, so the same representation--- results in analogous horde-ad array when this operation is used instead.+-- pretty-printed concrete arrays. The types mostly agree,+-- so the same representation results in analogous horde-ad array+-- when this operation is used instead.+--+-- Note that 'sfromVectorLinear' has very different typing that coincides+-- only for the concrete instance of @target@. sfromListLinear :: forall sh r target. (GoodScalar r, BaseTensor target)                 => ShS sh -> NonEmpty r -> target (TKS sh r) sfromListLinear sh = sconcrete . Nested.sfromListPrimLinear sh . NonEmpty.toList@@ -143,14 +155,19 @@ xscalar :: (GoodScalar r, BaseTensor target)         => r -> target (TKX '[] r) xscalar r = xconcrete $ Nested.mscalar r-xrepl :: forall sh r target. (GoodScalar r, BaseTensor target)-      => IShX sh -> r -> target (TKX sh r)-xrepl sh a = tconcrete (FTKX sh FTKScalar)-                       (Concrete $ Nested.mreplicateScal sh a) xingestData :: forall sh r target. (GoodScalar r, BaseTensor target)             => IShX sh -> [r] -> target (TKX sh r) xingestData sh l =   tconcrete (FTKX sh FTKScalar) (Concrete $ Nested.mfromListPrimLinear sh l)+-- | Create a tensor from a list of individual underlying scalar values.+--+-- An operation with the same name is used in @ox-arrays@ to represent+-- pretty-printed concrete arrays. The types mostly agree,+-- so the same representation results in analogous horde-ad array+-- when this operation is used instead.+--+-- Note that 'xfromVectorLinear' has very different typing that coincides+-- only for the concrete instance of @target@. xfromListLinear :: forall sh r target. (GoodScalar r, BaseTensor target)                 => IShX sh -> NonEmpty r -> target (TKX sh r) xfromListLinear sh = xingestData sh . NonEmpty.toList@@ -171,14 +188,18 @@             => Data.Vector.Vector (target (TKR2 n x))             -> target (TKR2 (1 + n) x) rfromVector = trfromVector-rfromVector0N :: forall n x target. (KnownSTK x, BaseTensor target)-              => IShR n -> Data.Vector.Vector (target (TKR2 0 x))-              -> target (TKR2 n x)-rfromVector0N = trfromVector0N-rfromList0N :: forall n x target. (KnownSTK x, BaseTensor target)-            => IShR n -> [target (TKR2 0 x)]-            -> target (TKR2 n x)-rfromList0N sh = trfromVector0N sh . V.fromList+-- | Create a tensor from a non-empty strict boxed vector with the given+-- outermost part of the shape of the result.+rfromVectorN :: forall m n x target. (KnownNat n, KnownSTK x, BaseTensor target)+             => IShR m -> Data.Vector.Vector (target (TKR2 n x))+             -> target (TKR2 (m + n) x)+rfromVectorN = trfromVectorN+-- | Create a tensor with the given shape from a non-empty strict boxed+-- vector of the elements.+rfromVectorLinear :: forall m r target. (GoodScalar r, BaseTensor target)+                  => IShR m -> Data.Vector.Vector (target (TKScalar r))+                  -> target (TKR m r)+rfromVectorLinear = trfromVectorLinear -- | Unravel a tensor into a list of its immediate subtensors. -- -- Warning: during computation, sharing between the elements@@ -188,64 +209,99 @@                   (KnownSTK x, KnownNat n, BaseTensor target)                => target (TKR2 (1 + n) x) -> [target (TKR2 n x)] runravelToList = trunravelToList+runravelToListN :: forall m n x target.+                   (KnownSTK x, KnownNat n, BaseTensor target)+                => IShR m -> target (TKR2 (m + n) x)+                -> [target (TKR2 n x)]+runravelToListN = trunravelToListN+rtoListLinear :: forall m r target. (GoodScalar r, BaseTensor target)+              => target (TKR m r) -> [target (TKScalar r)]+rtoListLinear = trtoListLinear  -- | Create a tensor from a list treated as the outermost dimension, -- going through strict boxed vectors, because laziness is risky with -- impurity, e.g., it easily perturbs results of fragile tests.-sfromList :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)-          => NonEmpty (target (TKS2 sh x)) -> target (TKS2 (n ': sh) x)-sfromList = tsfromVector . V.fromList . NonEmpty.toList+sfromList :: forall k shn x target.+             (KnownNat k, KnownShS shn, KnownSTK x, BaseTensor target)+          => NonEmpty (target (TKS2 shn x)) -> target (TKS2 (k ': shn) x)+sfromList = tsfromVector . V.fromListN (valueOf @k) . NonEmpty.toList -- | Create a tensor from a non-empty strict boxed vector treated -- as the outermost dimension.-sfromVector :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)-            => Data.Vector.Vector (target (TKS2 sh x))-            -> target (TKS2 (n ': sh) x)+sfromVector :: (KnownNat k, KnownShS shn, KnownSTK x, BaseTensor target)+            => Data.Vector.Vector (target (TKS2 shn x))+            -> target (TKS2 (k ': shn) x) sfromVector = tsfromVector-sfromVector0N :: (KnownShS sh, KnownSTK x, BaseTensor target)-              => Data.Vector.Vector (target (TKS2 '[] x))-              -> target (TKS2 sh x)-sfromVector0N = tsfromVector0N-sfromList0N :: (KnownShS sh, KnownSTK x, BaseTensor target)-            => [target (TKS2 '[] x)]-            -> target (TKS2 sh x)-sfromList0N = tsfromVector0N . V.fromList+-- | Create a tensor from a non-empty strict boxed vector with the given+-- outermost part of the shape of the result.+sfromVectorN :: forall shm shn x target.+                (KnownShS shn, KnownSTK x, BaseTensor target)+             => ShS shm -> Data.Vector.Vector (target (TKS2 shn x))+             -> target (TKS2 (shm ++ shn) x)+sfromVectorN = tsfromVectorN+-- | Create a tensor with the given shape from a non-empty strict boxed+-- vector of the elements.+sfromVectorLinear :: forall shm r target. (GoodScalar r, BaseTensor target)+                  => ShS shm -> Data.Vector.Vector (target (TKScalar r))+                  -> target (TKS shm r)+sfromVectorLinear = tsfromVectorLinear -- | Unravel a tensor into a list of its immediate subtensors. -- -- Warning: during computation, sharing between the elements -- of the resulting list is likely to be lost, so it needs to be ensured -- by explicit sharing, e.g., 'ttlet'.-sunravelToList :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)-               => target (TKS2 (n ': sh) x) -> [target (TKS2 sh x)]+sunravelToList :: (KnownNat k, KnownShS shn, KnownSTK x, BaseTensor target)+               => target (TKS2 (k ': shn) x) -> [target (TKS2 shn x)] sunravelToList = tsunravelToList+sunravelToListN :: forall shm shn x target.+                   (KnownShS shn, KnownSTK x, BaseTensor target)+                => ShS shm -> target (TKS2 (shm ++ shn) x)+                -> [target (TKS2 shn x)]+sunravelToListN = tsunravelToListN+stoListLinear :: forall shm r target. (GoodScalar r, BaseTensor target)+              => target (TKS shm r) -> [target (TKScalar r)]+stoListLinear = tstoListLinear  -- | Create a tensor from a list treated as the outermost dimension, -- going through strict boxed vectors, because laziness is risky with -- impurity, e.g., it easily perturbs results of fragile tests.-xfromList :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)-          => NonEmpty (target (TKX2 sh x)) -> target (TKX2 (Just n ': sh) x)+xfromList :: (KnownNat k, KnownShX shn, KnownSTK x, BaseTensor target)+          => NonEmpty (target (TKX2 shn x)) -> target (TKX2 (Just k ': shn) x) xfromList = txfromVector . V.fromList . NonEmpty.toList -- | Create a tensor from a non-empty strict boxed vector treated -- as the outermost dimension.-xfromVector :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)-            => Data.Vector.Vector (target (TKX2 sh x))-            -> target (TKX2 (Just n ': sh) x)+xfromVector :: (KnownNat k, KnownShX shn, KnownSTK x, BaseTensor target)+            => Data.Vector.Vector (target (TKX2 shn x))+            -> target (TKX2 (Just k ': shn) x) xfromVector = txfromVector-xfromVector0N :: forall sh x target. (KnownSTK x, BaseTensor target)-              => IShX sh -> Data.Vector.Vector (target (TKX2 '[] x))-              -> target (TKX2 sh x)-xfromVector0N = txfromVector0N-xfromList0N :: forall sh x target. (KnownSTK x, BaseTensor target)-            => IShX sh -> [target (TKX2 '[] x)]-            -> target (TKX2 sh x)-xfromList0N sh = txfromVector0N sh . V.fromList+-- | Create a tensor from a non-empty strict boxed vector with the given+-- outermost part of the shape of the result.+xfromVectorN :: forall shm shn x target.+                (KnownShX shn, KnownSTK x, BaseTensor target)+             => IShX shm -> Data.Vector.Vector (target (TKX2 shn x))+             -> target (TKX2 (shm ++ shn) x)+xfromVectorN = txfromVectorN+-- | Create a tensor with the given shape from a non-empty strict boxed+-- vector of the elements.+xfromVectorLinear :: forall shm r target. (GoodScalar r, BaseTensor target)+                  => IShX shm -> Data.Vector.Vector (target (TKScalar r))+                  -> target (TKX shm r)+xfromVectorLinear = txfromVectorLinear -- | Unravel a tensor into a list of its immediate subtensors. -- -- Warning: during computation, sharing between the elements -- of the resulting list is likely to be lost, so it needs to be ensured -- by explicit sharing, e.g., 'ttlet'.-xunravelToList :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)-               => target (TKX2 (Just n ': sh) x) -> [target (TKX2 sh x)]+xunravelToList :: (KnownNat k, KnownShX shn, KnownSTK x, BaseTensor target)+               => target (TKX2 (Just k ': shn) x) -> [target (TKX2 shn x)] xunravelToList = txunravelToList+xunravelToListN :: forall shm shn x target.+                   (KnownShX shn, KnownSTK x, BaseTensor target)+                => IShX shm -> target (TKX2 (shm ++ shn) x)+                -> [target (TKX2 shn x)]+xunravelToListN = txunravelToListN+xtoListLinear :: forall shm r target. (GoodScalar r, BaseTensor target)+              => target (TKX shm r) -> [target (TKScalar r)]+xtoListLinear = txtoListLinear  tunit :: BaseTensor target       => target TKUnit@@ -254,144 +310,181 @@ tlet :: forall x z target. LetTensor target      => target x -> (target x -> target z) -> target z tlet = ttlet+tletPrimal :: forall x z target. LetTensor target+           => PrimalOf target x -> (PrimalOf target x -> target z) -> target z+tletPrimal = ttletPrimal+tletPlain :: forall x z target. LetTensor target+          => PlainOf target x -> (PlainOf target x -> target z) -> target z+tletPlain = ttletPlain  -- | The operation is potentially strict in all arguments. ifH :: (KnownSTK y, Boolean (BoolOf target), BaseTensor target)     => BoolOf target -> target y -> target y -> target y ifH = tcond knownSTK-minH :: (KnownSTK y, OrdH target y, BaseTensor target)+minH :: (KnownSTK y, OrdH target y, BaseTensor target, LetTensor target)      => target y -> target y -> target y-minH u v = ifH (u <=. v) u v-maxH :: (KnownSTK y, OrdH target y, BaseTensor target)+minH u' v' = ttlet u' $ \u -> ttlet v' $ \v -> ifH (u <=. v) u v+maxH :: (KnownSTK y, OrdH target y, BaseTensor target, LetTensor target)      => target y -> target y -> target y-maxH u v = ifH (u >=. v) u v+maxH u' v' = ttlet u' $ \u -> ttlet v' $ \v -> ifH (u >=. v) u v -rsum :: (KnownNat n, KnownSTK x, BaseTensor target)+rsum :: forall n x target.+        (KnownNat n, TKAllNum x, KnownSTK x, BaseTensor target)      => target (TKR2 (1 + n) x) -> target (TKR2 n x) rsum = trsum-rsum0 :: (KnownNat n, KnownSTK x, BaseTensor target)-      => target (TKR2 n x) -> target (TKR2 0 x)+rsumN :: forall m n x target.+         (KnownNat m, KnownNat n, TKAllNum x, KnownSTK x, BaseTensor target)+      => target (TKR2 (m + n) x) -> target (TKR2 n x)+rsumN = trsumN+rsum0 :: forall m r target. (NumScalar r, BaseTensor target)+      => target (TKR m r) -> target (TKScalar r) rsum0 = trsum0-rdot0 :: ( KnownNat n, GoodScalar r, BaseTensor target)-      => target (TKR n r) -> target (TKR n r) -> target (TKR 0 r)+rdot0 :: forall n r target. (KnownNat n, NumScalar r, BaseTensor target)+      => target (TKR n r) -> target (TKR n r) -> target (TKScalar r) rdot0 = trdot0-rdot1In :: (KnownNat n, GoodScalar r, BaseTensor target)+rdot1In :: forall n r target. (KnownNat n, NumScalar r, BaseTensor target)         => target (TKR (1 + n) r) -> target (TKR (1 + n) r)         -> target (TKR n r) rdot1In = trdot1In-rmatvecmul :: (GoodScalar r, BaseTensor target)+rmatvecmul :: (NumScalar r, BaseTensor target, ConvertTensor target)            => target (TKR 2 r) -> target (TKR 1 r) -> target (TKR 1 r) rmatvecmul = trmatvecmul-rmatmul2 :: (GoodScalar r, BaseTensor target)+rmatmul2 :: (NumScalar r, BaseTensor target, ConvertTensor target)          => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r) rmatmul2 = trmatmul2 -- | Copy the given tensor along the new, outermost dimension.-rreplicate :: (KnownNat n, KnownSTK x, BaseTensor target)+rreplicate :: forall n x target. (KnownSTK x, BaseTensor target)            => Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x) rreplicate = trreplicate-rreplicate0N :: (KnownNat n, KnownSTK x, BaseTensor target)-             => IShR n -> target (TKR2 0 x) -> target (TKR2 n x)+rreplicateN :: forall m n x target. (KnownSTK x, BaseTensor target)+            => IShR m -> target (TKR2 n x) -> target (TKR2 (m + n) x)+rreplicateN = trreplicateN+rreplicate0N :: forall m r target. (GoodScalar r, BaseTensor target)+             => IShR m -> target (TKScalar r) -> target (TKR m r) rreplicate0N = trreplicate0N -ssum :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)+ssum :: forall n sh x target.+        (KnownNat n, KnownShS sh, TKAllNum x, KnownSTK x, BaseTensor target)      => target (TKS2 (n ': sh) x) -> target (TKS2 sh x) ssum = tssum-ssum0 :: (KnownShS sh, KnownSTK x, BaseTensor target)-      => target (TKS2 sh x) -> target (TKS2 '[] x)+ssumN :: forall shm shn x target.+         (KnownShS shm, KnownShS shn, TKAllNum x, KnownSTK x, BaseTensor target)+      => target (TKS2 (shm ++ shn) x) -> target (TKS2 shn x)+ssumN = tssumN @_ @shm+ssum0 :: forall shm r target. (NumScalar r, BaseTensor target)+      => target (TKS shm r) -> target (TKScalar r) ssum0 = tssum0-sdot0 :: (KnownShS sh, GoodScalar r, BaseTensor target)-      => target (TKS sh r) -> target (TKS sh r) -> target (TKS '[] r)+sdot0 :: forall sh r target. (KnownShS sh, NumScalar r, BaseTensor target)+      => target (TKS sh r) -> target (TKS sh r) -> target (TKScalar r) sdot0 = tsdot0-sdot1In :: (KnownShS sh, KnownNat n, GoodScalar r, BaseTensor target)+sdot1In :: forall sh n r target.+           (KnownShS sh, KnownNat n, NumScalar r, BaseTensor target)         => target (TKS (sh ++ '[n]) r) -> target (TKS (sh ++ '[n]) r)         -> target (TKS sh r)-sdot1In @sh @n = tsdot1In @_ @sh (SNat @n)-smatvecmul :: (KnownNat m, KnownNat n, GoodScalar r, BaseTensor target)+sdot1In = tsdot1In @_ @sh (SNat @n)+smatvecmul :: (KnownNat m, KnownNat n, NumScalar r, BaseTensor target)            => target (TKS '[m, n] r) -> target (TKS '[n] r)            -> target (TKS '[m] r) smatvecmul = tsmatvecmul-smatmul2 :: ( KnownNat m, KnownNat n, KnownNat p-            , GoodScalar r, BaseTensor target )+smatmul2 :: (KnownNat m, KnownNat n, KnownNat p, NumScalar r, BaseTensor target)          => target (TKS '[m, n] r) -> target (TKS '[n, p] r)          -> target (TKS '[m, p] r) smatmul2 = tsmatmul2-sreplicate :: (KnownNat k, KnownShS sh, KnownSTK x, BaseTensor target)+sreplicate :: forall k sh x target. (KnownNat k, KnownSTK x, BaseTensor target)            => target (TKS2 sh x) -> target (TKS2 (k ': sh) x)-sreplicate = tsreplicate SNat knownShS-sreplicate0N :: (KnownShS sh, KnownSTK x, BaseTensor target)-             => target (TKS2 '[] x) -> target (TKS2 sh x)+sreplicate = tsreplicate SNat+sreplicateN :: forall shm shn x target.+               (KnownShS shm, KnownSTK x, BaseTensor target)+             => target (TKS2 shn x) -> target (TKS2 (shm ++ shn) x)+sreplicateN = tsreplicateN (knownShS @shm)+sreplicate0N :: (KnownShS shm, GoodScalar r, BaseTensor target)+             => target (TKScalar r) -> target (TKS shm r) sreplicate0N = tsreplicate0N knownShS -xsum :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)+xsum :: forall n sh x target.+        (KnownNat n, KnownShX sh, TKAllNum x, KnownSTK x, BaseTensor target)      => target (TKX2 (Just n ': sh) x) -> target (TKX2 sh x) xsum = txsum-xsum0 :: (KnownShX sh, KnownSTK x, BaseTensor target, ConvertTensor target)-      => target (TKX2 sh x) -> target (TKX2 '[] x)+xsumN :: forall shm shn x target.+         (KnownShX shm, KnownShX shn, TKAllNum x, KnownSTK x, BaseTensor target)+      => target (TKX2 (shm ++ shn) x) -> target (TKX2 shn x)+xsumN = txsumN @_ @shm+xsum0 :: forall shm r target. (NumScalar r, BaseTensor target)+      => target (TKX shm r) -> target (TKScalar r) xsum0 = txsum0-xdot0 :: ( KnownShX sh, GoodScalar r-         , BaseTensor target, ConvertTensor target )-      => target (TKX sh r) -> target (TKX sh r) -> target (TKX '[] r)+xdot0 :: forall sh r target. (KnownShX sh, NumScalar r, BaseTensor target)+      => target (TKX sh r) -> target (TKX sh r) -> target (TKScalar r) xdot0 = txdot0-xdot1In :: (KnownShX sh, KnownNat n, GoodScalar r, BaseTensor target)+xdot1In :: forall sh n r target.+           (KnownShX sh, KnownNat n, NumScalar r, BaseTensor target)         => target (TKX (sh ++ '[Just n]) r)         -> target (TKX (sh ++ '[Just n]) r)         -> target (TKX sh r)-xdot1In @sh @n = txdot1In @_ @sh (SNat @n)+xdot1In = txdot1In @_ @sh (SNat @n) xmatvecmul :: forall mm mn r target.-              (GoodScalar r, BaseTensor target, ConvertTensor target)-           => Nested.SMayNat Int SNat mm -> Nested.SMayNat Int SNat mn+              (NumScalar r, BaseTensor target, ConvertTensor target)+           => Nested.SMayNat Int mm -> Nested.SMayNat Int mn            -> target (TKX '[mm, mn] r) -> target (TKX '[mn] r)            -> target (TKX '[mm] r) xmatvecmul = txmatvecmul xmatmul2 :: ( KnownNat m, KnownNat n, KnownNat p-            , GoodScalar r, BaseTensor target, ConvertTensor target )+            , NumScalar r, BaseTensor target, ConvertTensor target )          => target (TKX '[Just m, Just n] r)          -> target (TKX '[Just n, Just p] r)          -> target (TKX '[Just m, Just p] r) xmatmul2 = txmatmul2-xreplicate :: (KnownNat k, KnownShX sh, KnownSTK x, BaseTensor target)-           => target (TKX2 sh x) -> target (TKX2 (Just k ': sh) x)-xreplicate = txreplicate SNat knownShX-xreplicate0N :: (KnownShX sh, KnownSTK x, BaseTensor target)-             => IShX sh -> target (TKX2 '[] x) -> target (TKX2 sh x)+xreplicate :: forall k sh x target. (KnownSTK x, BaseTensor target)+           => SNat k -> target (TKX2 sh x) -> target (TKX2 (Just k ': sh) x)+xreplicate = txreplicate+xreplicateN :: forall shm shn x target.+               (KnownSTK x, BaseTensor target)+            => IShX shm -> target (TKX2 shn x) -> target (TKX2 (shm ++ shn) x)+xreplicateN = txreplicateN+xreplicate0N :: forall shm r target. (GoodScalar r, BaseTensor target)+             => IShX shm -> target (TKScalar r) -> target (TKX shm r) xreplicate0N = txreplicate0N  -- | First index is for outermost dimension; empty index means identity, -- if index is out of bounds, the result is defined and is @def@, which is @0@. -- The same holds for all operations with indexes.-rindex, (!) :: (KnownNat m, KnownNat n, KnownSTK x, BaseTensor target)+rindex, (!) :: forall m n x target. (KnownNat n, KnownSTK x, BaseTensor target)             => target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x) rindex = trindex infixl 9 ! (!) = rindex  -- prefix form better when type applications are necessary-rindex0 :: (KnownNat m, KnownSTK x, BaseTensor target)-        => target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)+rindex0 :: forall m r target. (GoodScalar r, BaseTensor target)+        => target (TKR m r) -> IxROf target m -> target (TKScalar r) rindex0 = trindex0-roneHot :: ( KnownNat m, KnownNat n, KnownSTK x-           , BoolOf (PrimalOf target) ~ BoolOf target-           , EqH (PrimalOf target) (TKScalar Int64), BaseTensor target)+roneHot :: ( KnownNat m, KnownNat n, TKAllNum x, KnownSTK x+           , PlainOf (PlainOf target) ~ PlainOf target+           , EqH (PlainOf target) (TKScalar Int), BaseTensor target)         => IShR m -> target (TKR2 n x) -> IxROf target m         -> target (TKR2 (m + n) x) roneHot = troneHot-rscatter :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x, BaseTensor target)-         => IShR (p + n) -> target (TKR2 (m + n) x)+rscatter :: ( KnownNat m, KnownNat n, KnownNat p, TKAllNum x, KnownSTK x+            , BaseTensor target )+         => IShR p -> target (TKR2 (m + n) x)          -> (IxROf target m -> IxROf target p)          -> target (TKR2 (p + n) x)+{-# INLINE rscatter #-} rscatter = trscatter -- | Build a tensor by adding up tensors of rank @n@ taken from -- the second argument and inserted in a zero tensor -- at indexes of length @p@. -- The semantics of the operation permits index out of bounds -- and then no tensor is added at such an index.-rscatter1 :: (KnownNat n, KnownNat p, KnownSTK x, BaseTensor target)-          => IShR (p + n) -> target (TKR2 (1 + n) x)+rscatter1 :: ( KnownNat n, KnownNat p, TKAllNum x, KnownSTK x+             , BaseTensor target )+          => IShR p -> target (TKR2 (1 + n) x)           -> (IntOf target -> IxROf target p)           -> target (TKR2 (p + n) x)+{-# INLINE rscatter1 #-} rscatter1 = trscatter1 rgather :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x, BaseTensor target)-        => IShR (m + n) -> target (TKR2 (p + n) x)+        => IShR m -> target (TKR2 (p + n) x)         -> (IxROf target m -> IxROf target p)         -> target (TKR2 (m + n) x)+{-# INLINE rgather #-} rgather = trgather -- | Build a tensor by collecting tensors of rank @n@ obtained by indexing -- in the second argument at the given indexes of length @p@.@@ -401,88 +494,97 @@          => Int -> target (TKR2 (p + n) x)          -> (IntOf target -> IxROf target p)          -> target (TKR2 (1 + n) x)+{-# INLINE rgather1 #-} rgather1 = trgather1 -sindex, (!$) :: (KnownShS shm, KnownShS shn, KnownSTK x, BaseTensor target)+sindex, (!$) :: forall shm shn x target.+                (KnownShS shn, KnownSTK x, BaseTensor target)              => target (TKS2 (shm ++ shn) x) -> IxSOf target shm              -> target (TKS2 shn x) sindex = tsindex infixl 9 !$ (!$) = sindex  -- prefix form better when type applications are necessary-sindex0 :: (KnownShS sh1, KnownSTK x, BaseTensor target)-        => target (TKS2 sh1 x) -> IxSOf target sh1-        -> target (TKS2 '[] x)+sindex0 :: forall shm r target. (GoodScalar r, BaseTensor target)+        => target (TKS shm r) -> IxSOf target shm -> target (TKScalar r) sindex0 = tsindex0-soneHot :: ( KnownShS sh1, KnownShS sh2, KnownSTK x-           , BoolOf (PrimalOf target) ~ BoolOf target-           , EqH (PrimalOf target) (TKScalar Int64), BaseTensor target )+soneHot :: ( KnownShS sh1, KnownShS sh2, TKAllNum x, KnownSTK x+           , PlainOf (PlainOf target) ~ PlainOf target+           , EqH (PlainOf target) (TKScalar Int), BaseTensor target )         => target (TKS2 sh2 x) -> IxSOf target sh1         -> target (TKS2 (sh1 ++ sh2) x) soneHot = tsoneHot sscatter-  :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x, BaseTensor target)+  :: ( KnownShS shm, KnownShS shn, KnownShS shp, TKAllNum x, KnownSTK x+     , BaseTensor target )   => target (TKS2 (shm ++ shn) x)   -> (IxSOf target shm -> IxSOf target shp)   -> target (TKS2 (shp ++ shn) x)+{-# INLINE sscatter #-} sscatter @shm @shn @shp = tsscatter @_ @shm @shn @shp sscatter1-  :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x, BaseTensor target)+  :: ( KnownNat n2, KnownShS shn, KnownShS shp, TKAllNum x, KnownSTK x+     , BaseTensor target )   => target (TKS2 (n2 ': shn) x)   -> (IntOf target -> IxSOf target shp)   -> target (TKS2 (shp ++ shn) x)+{-# INLINE sscatter1 #-} sscatter1 = tsscatter1 sgather   :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x, BaseTensor target)   => target (TKS2 (shp ++ shn) x)   -> (IxSOf target shm -> IxSOf target shp)   -> target (TKS2 (shm ++ shn) x)+{-# INLINE sgather #-} sgather @shm @shn @shp = tsgather @_ @shm @shn @shp sgather1   :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x, BaseTensor target)   => target (TKS2 (shp ++ shn) x)   -> (IntOf target -> IxSOf target shp)   -> target (TKS2 (n2 ': shn) x)+{-# INLINE sgather1 #-} sgather1 = tsgather1 -xindex :: (KnownShX sh1, KnownShX sh2, KnownSTK x, BaseTensor target)-       => target (TKX2 (sh1 ++ sh2) x) -> IxXOf target sh1-       -> target (TKX2 sh2 x)+xindex :: forall shm shn x target. (KnownShX shn, KnownSTK x, BaseTensor target)+       => target (TKX2 (shm ++ shn) x) -> IxXOf target shm+       -> target (TKX2 shn x) xindex = txindex-xindex0 :: (KnownShX sh1, KnownSTK x, BaseTensor target)-        => target (TKX2 sh1 x) -> IxXOf target sh1-        -> target (TKX2 '[] x)+xindex0 :: forall shm r target. (GoodScalar r, BaseTensor target)+        => target (TKX shm r) -> IxXOf target shm -> target (TKScalar r) xindex0 = txindex0-xoneHot :: ( KnownShX sh1, KnownShX sh2, KnownSTK x-           , BoolOf (PrimalOf target) ~ BoolOf target-           , EqH (PrimalOf target) (TKScalar Int64)-           , BaseTensor target, ConvertTensor target )+xoneHot :: ( KnownShX sh1, KnownShX sh2, TKAllNum x, KnownSTK x+           , PlainOf (PlainOf target) ~ PlainOf target+           , EqH (PlainOf target) (TKScalar Int), BaseTensor target )         => IShX sh1 -> target (TKX2 sh2 x) -> IxXOf target sh1         -> target (TKX2 (sh1 ++ sh2) x) xoneHot = txoneHot-xscatter :: ( KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x+xscatter :: ( KnownShX shm, KnownShX shn, KnownShX shp, TKAllNum x, KnownSTK x             , BaseTensor target )-         => IShX (shp ++ shn) -> target (TKX2 (shm ++ shn) x)+         => IShX shp -> target (TKX2 (shm ++ shn) x)          -> (IxXOf target shm -> IxXOf target shp)          -> target (TKX2 (shp ++ shn) x)+{-# INLINE xscatter #-} xscatter @shm @shn @shp = txscatter @_ @shm @shn @shp-xscatter1 :: ( KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x+xscatter1 :: ( KnownNat n2, KnownShX shn, KnownShX shp, TKAllNum x, KnownSTK x              , BaseTensor target )-          => IShX (shp ++ shn) -> target (TKX2 (Just n2 ': shn) x)+          => IShX shp -> target (TKX2 (Just n2 ': shn) x)           -> (IntOf target -> IxXOf target shp)           -> target (TKX2 (shp ++ shn) x)+{-# INLINE xscatter1 #-} xscatter1 = txscatter1 xgather :: ( KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x            , BaseTensor target )-        => IShX (shm ++ shn)+        => IShX shm         -> target (TKX2 (shp ++ shn) x)         -> (IxXOf target shm -> IxXOf target shp)         -> target (TKX2 (shm ++ shn) x)+{-# INLINE xgather #-} xgather @shm @shn @shp = txgather @_ @shm @shn @shp xgather1 :: ( KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x             , BaseTensor target )          => SNat n2 -> target (TKX2 (shp ++ shn) x)          -> (IntOf target -> IxXOf target shp)          -> target (TKX2 (Just n2 ': shn) x)+{-# INLINE xgather1 #-} xgather1 = txgather1  -- | Transpose according to the permutation.@@ -511,87 +613,98 @@          => IShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x) xreshape = txreshape -rfloor :: ( GoodScalar r, RealFrac r, GoodScalar r2, Integral r2+kfloor :: ( NumScalar r, Differentiable r, NumScalar r2, Integral r2           , BaseTensor target )+       => target (TKScalar r) -> target (TKScalar r2)+kfloor = tkfloor+kfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2, BaseTensor target)+              => target (TKScalar r1) -> target (TKScalar r2)+kfromIntegral = tkfromIntegral+kcast :: ( Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2+         , BaseTensor target )+      => target (TKScalar r1) -> target (TKScalar r2)+kcast = tkcast+-- | Throws if the array is empty.+kargMin, kargMax+  :: forall n r target. (NumScalar r, BaseTensor target)+  => target (TKS '[n] r) -> target (TKScalar Int)+kargMin = tkargMin+kargMax = tkargMax++rfloor :: ( NumScalar r, Differentiable r, NumScalar r2, Integral r2+          , BaseTensor target )        => target (TKR n r) -> target (TKR n r2) rfloor = trfloor-rfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target)+rfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2, BaseTensor target)               => target (TKR n r1) -> target (TKR n r2) rfromIntegral = trfromIntegral-rcast :: ( RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2+rcast :: ( Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2          , BaseTensor target )       => target (TKR n r1) -> target (TKR n r2) rcast = trcast-rminIndex, rmaxIndex  -- partial-  :: forall n r r2 target. (GoodScalar r, GoodScalar r2, BaseTensor target)-  => target (TKR (1 + n) r) -> target (TKR n r2)-rminIndex = trminIndex-rmaxIndex = trmaxIndex-riota :: (GoodScalar r, BaseTensor target)+-- | Throws if the array is empty.+rargMin, rargMax+  :: forall n r target. (NumScalar r, BaseTensor target)+  => target (TKR (1 + n) r) -> target (TKR n Int)+rargMin = trargMin+rargMax = trargMax+riota :: (NumScalar r, BaseTensor target)       => Int -> target (TKR 1 r)  -- from 0 to n - 1 riota = triota -sfloor :: ( GoodScalar r, RealFrac r, GoodScalar r2, Integral r2+sfloor :: ( NumScalar r, Differentiable r, NumScalar r2, Integral r2           , BaseTensor target )        => target (TKS sh r) -> target (TKS sh r2) sfloor = tsfloor-sfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target)+sfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2, BaseTensor target)               => target (TKS sh r1) -> target (TKS sh r2) sfromIntegral = tsfromIntegral-scast :: ( RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2+scast :: ( Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2          , BaseTensor target )       => target (TKS sh r1) -> target (TKS sh r2) scast = tscast-sminIndex, smaxIndex  -- partial-  :: forall n sh r r2 target. (GoodScalar r, GoodScalar r2, BaseTensor target)-  => target (TKS (n ': sh) r) -> target (TKS (Init (n ': sh)) r2)-sminIndex = tsminIndex-smaxIndex = tsmaxIndex-siota :: (KnownNat n, GoodScalar r, BaseTensor target)+-- | Throws if the array is empty.+sargMin, sargMax+  :: forall n sh r target. (NumScalar r, BaseTensor target)+  => target (TKS (n ': sh) r) -> target (TKS (Init (n ': sh)) Int)+sargMin = tsargMin+sargMax = tsargMax+siota :: (KnownNat n, NumScalar r, BaseTensor target)       => target (TKS '[n] r)  -- from 0 to n - 1 siota = tsiota -xfloor :: ( GoodScalar r, RealFrac r, GoodScalar r2, Integral r2+xfloor :: ( NumScalar r, Differentiable r, NumScalar r2, Integral r2           , BaseTensor target )        => target (TKX sh r) -> target (TKX sh r2) xfloor = txfloor-xfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target)+xfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2, BaseTensor target)               => target (TKX sh r1) -> target (TKX sh r2) xfromIntegral = txfromIntegral-xcast :: ( RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2+xcast :: ( Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2          , BaseTensor target )       => target (TKX sh r1) -> target (TKX sh r2) xcast = txcast-xminIndex, xmaxIndex  -- partial-  :: forall mn sh r r2 target. (GoodScalar r, GoodScalar r2, BaseTensor target)-  => target (TKX (mn ': sh) r) -> target (TKX (Init (mn ': sh)) r2)-xminIndex = txminIndex-xmaxIndex = txmaxIndex-xiota :: (KnownNat n, GoodScalar r, BaseTensor target)+-- | Throws if the array is empty.+xargMin, xargMax+  :: forall mn sh r target. (NumScalar r, BaseTensor target)+  => target (TKX (mn ': sh) r) -> target (TKX (Init (mn ': sh)) Int)+xargMin = txargMin+xargMax = txargMax+xiota :: (KnownNat n, NumScalar r, BaseTensor target)       => target (TKX '[Just n] r)  -- from 0 to n - 1 xiota = txiota -kfloor :: ( GoodScalar r, RealFrac r, GoodScalar r2, Integral r2-          , BaseTensor target )-       => target (TKScalar r) -> target (TKScalar r2)-kfloor = tkfloor-kfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target)-              => target (TKScalar r1) -> target (TKScalar r2)-kfromIntegral = tkfromIntegral-kcast :: ( RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2-         , BaseTensor target )-      => target (TKScalar r1) -> target (TKScalar r2)-kcast = tkcast- -- | Append two arrays along the outermost dimension. -- All dimensions, except the outermost, must be the same. rappend :: forall n x target. (KnownSTK x, BaseTensor target)         => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)         -> target (TKR2 (1 + n) x) rappend = trappend+-- | Append a list of arrays that agree on all dimensions, including+-- in arrays in their elements, except for the outermost dimension. rconcat :: forall n x target. (KnownSTK x, BaseTensor target)         => NonEmpty (target (TKR2 (1 + n) x)) -> target (TKR2 (1 + n) x)-rconcat = foldr1 rappend+rconcat = trconcat -- | Extract a slice of an array along the outermost dimension. -- The extracted slice must fall within the dimension. rslice :: forall n x target. (KnownSTK x, BaseTensor target)@@ -629,70 +742,84 @@ sreverse = tsreverse  xappend :: forall m n sh x target. (KnownSTK x, BaseTensor target)-        => target (TKX2 (Just m ': sh) x) -> target (TKX2 (Just n ': sh) x)-        -> target (TKX2 (Just (m + n) ': sh) x)+        => target (TKX2 (m ': sh) x) -> target (TKX2 (n ': sh) x)+        -> target (TKX2 (AddMaybe m n ': sh) x) xappend = txappend-xappend0 :: forall sh x target.-            (KnownSTK x, BaseTensor target, ConvertTensor target)-         => target (TKX2 (Nothing ': sh) x) -> target (TKX2 (Nothing ': sh) x)-         -> target (TKX2 (Nothing ': sh) x)-xappend0 a b = case xshape a of-  mmsnat :$% sh ->-    withSNat (fromSMayNat' mmsnat) $ \msnat ->-    withSNat (shxLength $ xshape b) $ \nsnat ->-    let sh0 = Nested.SUnknown () :!% ssxFromShX sh-        sha = Nested.SKnown msnat :!% ssxFromShX sh-        shb = Nested.SKnown nsnat :!% ssxFromShX sh-    in withKnownShX (ssxFromShX sh) $-       xmcast sh0 $ xappend (xmcast sha a) (xmcast shb b)-xconcat :: forall sh x target.-           (KnownSTK x, BaseTensor target, ConvertTensor target)+-- | Append a list of arrays that agree on all dimensions, including+-- in arrays in their elements, except for the outermost dimension.+xconcat :: forall sh x target. (KnownSTK x, BaseTensor target)         => NonEmpty (target (TKX2 (Nothing ': sh) x))         -> target (TKX2 (Nothing ': sh) x)-xconcat = foldr1 xappend0+xconcat = txconcat xslice :: forall i n k sh x target. (KnownSTK x, BaseTensor target)-       => SNat i -> SNat n -> SNat k-       -> target (TKX2 (Just (i + n + k) ': sh) x)-       -> target (TKX2 (Just n ': sh) x)+       => SMayNat Int i -> SMayNat Int n -> SMayNat Int k+       -> target (TKX2 (AddMaybe (AddMaybe i n) k ': sh) x)+       -> target (TKX2 (n ': sh) x) xslice = txslice xuncons :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)         => target (TKX2 (Just n ': sh) x)         -> Maybe (target (TKX2 sh x), target (TKX2 (Just (n - 1) ': sh) x)) xuncons @n v = case cmpNat (Proxy @1) (Proxy @n) of   EQI -> Just ( v `xindex` (0 :.% ZIX)-              , xslice @1 @(n - 1) @0 SNat SNat SNat v )+              , xslice @(Just 1) @(Just (n - 1)) @(Just 0)+                       (SKnown SNat) (SKnown SNat) (SKnown SNat) v )   LTI -> Just ( v `xindex` (0 :.% ZIX)-              , xslice @1 @(n - 1) @0 SNat SNat SNat v )+              , xslice @(Just 1) @(Just (n - 1)) @(Just 0)+                       (SKnown SNat) (SKnown SNat) (SKnown SNat) v )   _ -> Nothing xreverse :: forall mn sh x target. (KnownSTK x, BaseTensor target)          => target (TKX2 (mn ': sh) x) -> target (TKX2 (mn ': sh) x) xreverse = txreverse +kbuild1 :: (KnownNat k, GoodScalar r, BaseTensor target)+        => (IntOf target -> target (TKScalar r))+        -> target (TKS '[k] r)+{-# INLINE kbuild1 #-}+kbuild1 = tkbuild1+kbuild :: (KnownShS sh, GoodScalar r, BaseTensor target, ConvertTensor target)+       => (IxSOf target sh -> target (TKScalar r))+       -> target (TKS sh r)+{-# INLINE kbuild #-}+kbuild = tkbuild+ rbuild1 :: (KnownNat n, KnownSTK x, BaseTensor target)         => Int  -- ^ width of the outermost dimension of the created tensor         -> (IntOf target -> target (TKR2 n x))  -- ^ the function to build with         -> target (TKR2 (1 + n) x)+{-# INLINE rbuild1 #-} rbuild1 = trbuild1+-- | Building a tensor (also known as @generate@ and @tabulate@).+-- The resulting tensor should have no zero dimensions.+-- See https://futhark-lang.org/blog/2025-09-26-the-biggest-semantic-mess.html+-- for why trying to handle zero dimensions complicates the system greatly.+rbuild :: (KnownNat m, KnownNat n, KnownSTK x, BaseTensor target)+       => IShR m+       -> (IxROf target m -> target (TKR2 n x))+            -- ^ the function to build with+       -> target (TKR2 (m + n) x)+{-# INLINE rbuild #-}+rbuild = trbuild rmap :: (KnownNat m, KnownNat n, KnownSTK x, KnownSTK x2, BaseTensor target)      => (target (TKR2 n x) -> target (TKR2 n x2))  -- ^ the function to map with      -> target (TKR2 (m + n) x)  -- ^ the tensor to map over      -> target (TKR2 (m + n) x2)-rmap f v = rbuild (rshape v) (\ix -> f (v ! ix))+rmap @m f v = rbuild (shrTake @m $ rshape v) (\ix -> f (v ! ix)) rmap1 :: (KnownNat n, KnownSTK x, KnownSTK x2, BaseTensor target)       => (target (TKR2 n x) -> target (TKR2 n x2))             -- ^ the function to map with       -> target (TKR2 (1 + n) x)  -- ^ the tensor to map over       -> target (TKR2 (1 + n) x2) rmap1 f u = rbuild1 (rwidth u) (\i -> f (u ! [i]))-rmap0N :: (KnownNat n, KnownSTK x, KnownSTK x1, BaseTensor target)-       => (target (TKR2 0 x1) -> target (TKR2 0 x))+rmap0N :: ( KnownNat n, GoodScalar r1, GoodScalar r+          , BaseTensor target, ConvertTensor target )+       => (target (TKScalar r1) -> target (TKScalar r))             -- ^ the function to map with-       -> target (TKR2 n x1)  -- ^ the tensor to map over-       -> target (TKR2 n x)+       -> target (TKR n r1)  -- ^ the tensor to map over+       -> target (TKR n r) rmap0N = trmap0N rzipWith :: ( KnownNat m, KnownNat n1, KnownNat n2, KnownNat n, KnownSTK x             , KnownSTK x1, KnownSTK x2, BaseTensor target )-         => IShR (m + n)  -- ^ the shape of the resulting tensor+         => IShR m          -> (target (TKR2 n1 x1) -> target (TKR2 n2 x2) -> target (TKR2 n x))               -- ^ the function to zip with          -> target (TKR2 (m + n1) x1)  -- ^ the first tensor to zip over@@ -707,18 +834,19 @@           -> target (TKR2 (1 + n2) x2)  -- ^ the second tensor to zip over           -> target (TKR2 (1 + n) x) rzipWith1 f u v = rbuild1 (rwidth u) (\i -> f (u ! [i]) (v ! [i]))-rzipWith0N :: ( KnownNat n, KnownSTK x, KnownSTK x1, KnownSTK x2-              , BaseTensor target )-           => (target (TKR2 0 x1) -> target (TKR2 0 x2) -> target (TKR2 0 x))+rzipWith0N :: ( KnownNat n, GoodScalar r, GoodScalar r1, GoodScalar r2+              , BaseTensor target, ConvertTensor target )+           => (target (TKScalar r1) -> target (TKScalar r2)+               -> target (TKScalar r))                 -- ^ the function to zip with-           -> target (TKR2 n x1)  -- ^ the first tensor to zip over-           -> target (TKR2 n x2)  -- ^ the second tensor to zip over-           -> target (TKR2 n x)-rzipWith0N  = trzipWith0N+           -> target (TKR n r1)  -- ^ the first tensor to zip over+           -> target (TKR n r2)  -- ^ the second tensor to zip over+           -> target (TKR n r)+rzipWith0N = trzipWith0N rzipWith3 :: ( KnownNat m, KnownNat n1, KnownNat n2, KnownNat n3              , KnownNat n, KnownSTK x              , KnownSTK x1, KnownSTK x2, KnownSTK x3, BaseTensor target )-          => IShR (m + n)  -- ^ the shape of the resulting tensor+          => IShR m           -> (target (TKR2 n1 x1) -> target (TKR2 n2 x2) -> target (TKR2 n3 x3)               -> target (TKR2 n x))  -- ^ the function to zip with           -> target (TKR2 (m + n1) x1)  -- ^ the first tensor to zip over@@ -736,6 +864,8 @@            -> target (TKR2 (1 + n) x) rzipWith31 f u v w =   rbuild1 (rwidth u) (\i -> f (u ! [i]) (v ! [i]) (w ! [i]))+-- TODO: change the rank 0 ranked tensors to scalars unless many cases+-- such as szipWithNested emerge and then even rzipWith0N gets changed back rzipWith30N :: ( KnownNat n, KnownSTK x                , KnownSTK x1, KnownSTK x2, KnownSTK x3, BaseTensor target )             => (target (TKR2 0 x1) -> target (TKR2 0 x2) -> target (TKR2 0 x3)@@ -745,13 +875,13 @@             -> target (TKR2 n x3)  -- ^ the third tensor to zip over             -> target (TKR2 n x) rzipWith30N f u v w =-  rbuild (rshape v) (\ix -> f (rindex0 u ix) (rindex0 v ix) (rindex0 w ix))+  rbuild (rshape v) (\ix -> f (rindex u ix) (rindex v ix) (rindex w ix)) rzipWith4 :: ( KnownNat m              , KnownNat n1, KnownNat n2, KnownNat n3, KnownNat n4              , KnownNat n, KnownSTK x              , KnownSTK x1, KnownSTK x2, KnownSTK x3, KnownSTK x4              , BaseTensor target )-          => IShR (m + n)  -- ^ the shape of the resulting tensor+          => IShR m           -> (target (TKR2 n1 x1) -> target (TKR2 n2 x2)               -> target (TKR2 n3 x3) -> target (TKR2 n4 x4)               -> target (TKR2 n x))  -- ^ the function to zip with@@ -788,37 +918,46 @@             -> target (TKR2 n x4)  -- ^ the fourth tensor to zip over             -> target (TKR2 n x) rzipWith40N f u v w x =-  rbuild (rshape v) (\ix -> f (rindex0 u ix) (rindex0 v ix) (rindex0 w ix)-                              (rindex0 x ix))+  rbuild (rshape v) (\ix -> f (rindex u ix) (rindex v ix) (rindex w ix)+                              (rindex x ix))  sbuild1 :: (KnownNat k, KnownShS sh, KnownSTK x, BaseTensor target)         => (IntOf target -> target (TKS2 sh x))  -- ^ the function to build with         -> target (TKS2 (k ': sh) x)+{-# INLINE sbuild1 #-} sbuild1 = tsbuild1-smap :: ( KnownShS (Take m sh), KnownShS (Drop m sh), KnownShS sh-        , KnownSTK x, KnownSTK x2-        , BaseTensor target )+sbuild :: (KnownShS shm, KnownShS shn, KnownSTK x, BaseTensor target)+       => (IxSOf target shm -> target (TKS2 shn x))+            -- ^ the function to build with+       -> target (TKS2 (shm ++ shn) x)+{-# INLINE sbuild #-}+sbuild = tsbuild+smap :: ( KnownShS (Take m sh), KnownShS (Drop m sh)+        , KnownSTK x, KnownSTK x2, BaseTensor target )      => (target (TKS2 (Drop m sh) x) -> target (TKS2 (Drop m sh) x2))           -- ^ the function to map with      -> target (TKS2 sh x)  -- ^ the tensor to map over      -> target (TKS2 sh x2) smap @m @sh f v = gcastWith (unsafeCoerceRefl                              :: sh :~: Take m sh ++ Drop m sh)-                  $ sbuild (\ix -> f (v !$ ix))+                  $ sbuild @(Take m sh) (\ix -> f (v !$ ix)) smap1 :: (KnownNat n, KnownShS sh, KnownSTK x, KnownSTK x2, BaseTensor target)       => (target (TKS2 sh x) -> target (TKS2 sh x2))            -- ^ the function to map with       -> target (TKS2 (n ': sh) x)  -- ^ the tensor to map over       -> target (TKS2 (n ': sh) x2) smap1 f u = sbuild1 (\i -> f (u !$ (i :.$ ZIS)))-smap0N :: (KnownShS sh, KnownSTK x1, KnownSTK x, BaseTensor target)-       => (target (TKS2 '[] x1) -> target (TKS2 '[] x))+smap0N :: ( KnownShS sh, GoodScalar r1, GoodScalar r+          , BaseTensor target, ConvertTensor target )+       => (target (TKScalar r1) -> target (TKScalar r))             -- ^ the function to map with-       -> target (TKS2 sh x1)  -- ^ the tensor to map over-       -> target (TKS2 sh x)+       -> target (TKS sh r1)  -- ^ the tensor to map over+       -> target (TKS sh r) smap0N = tsmap0N-szipWith :: ( KnownShS (Drop m sh1), KnownShS (Drop m sh2), KnownShS (Take m sh)-            , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownShS sh+szipWith :: ( KnownShS (Drop m sh1), KnownShS (Drop m sh2)+            , KnownShS (Take m sh), KnownShS (Drop m sh)+            , KnownSTK x, KnownSTK x1, KnownSTK x2+            , sh ~ Take m sh ++ Drop m sh             , sh1 ~ Take m sh ++ Drop m sh1             , sh2 ~ Take m sh ++ Drop m sh2, BaseTensor target )          => (target (TKS2 (Drop m sh1) x1) -> target (TKS2 (Drop m sh2) x2)@@ -826,7 +965,7 @@          -> target (TKS2 sh1 x1)  -- ^ the first tensor to zip over          -> target (TKS2 sh2 x2)  -- ^ the second tensor to zip over          -> target (TKS2 sh x)-szipWith f u v = sbuild (\ix -> f (u !$ ix) (v !$ ix))+szipWith @m @_ @_ @sh f u v = sbuild @(Take m sh) (\ix -> f (u !$ ix) (v !$ ix)) szipWith1 :: ( KnownNat n, KnownShS sh1, KnownShS sh2, KnownShS sh              , KnownSTK x, KnownSTK x1, KnownSTK x2, BaseTensor target )           => (target (TKS2 sh1 x1) -> target (TKS2 sh2 x2)@@ -836,17 +975,18 @@           -> target (TKS2 (n ': sh) x) szipWith1 f u v = sbuild1 (\i -> f (u !$ (i :.$ ZIS))                                    (v !$ (i :.$ ZIS)))-szipWith0N :: ( KnownShS sh, KnownSTK x, KnownSTK x1, KnownSTK x2-              , BaseTensor target )-           => (target (TKS2 '[] x1) -> target (TKS2 '[] x2)-               -> target (TKS2 '[] x))  -- ^ the function to zip with-           -> target (TKS2 sh x1)  -- ^ the first tensor to zip over-           -> target (TKS2 sh x2)  -- ^ the second tensor to zip over-           -> target (TKS2 sh x)+szipWith0N :: ( KnownShS sh, GoodScalar r, GoodScalar r1, GoodScalar r2+              , BaseTensor target, ConvertTensor target )+           => (target (TKScalar r1) -> target (TKScalar r2)+               -> target (TKScalar r))  -- ^ the function to zip with+           -> target (TKS sh r1)  -- ^ the first tensor to zip over+           -> target (TKS sh r2)  -- ^ the second tensor to zip over+           -> target (TKS sh r) szipWith0N = tszipWith0N szipWith3 :: ( KnownShS (Drop m sh1), KnownShS (Drop m sh2)-             , KnownShS (Drop m sh3), KnownShS (Take m sh), KnownShS sh+             , KnownShS (Drop m sh3), KnownShS (Take m sh), KnownShS (Drop m sh)              , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownSTK x3+             , sh ~ Take m sh ++ Drop m sh              , sh1 ~ Take m sh ++ Drop m sh1              , sh2 ~ Take m sh ++ Drop m sh2              , sh3 ~ Take m sh ++ Drop m sh3, BaseTensor target )@@ -857,7 +997,8 @@           -> target (TKS2 sh2 x2)  -- ^ the second tensor to zip over           -> target (TKS2 sh3 x3)  -- ^ the third tensor to zip over           -> target (TKS2 sh x)-szipWith3 f u v w = sbuild (\ix -> f (u !$ ix) (v !$ ix) (w !$ ix))+szipWith3 @m @_ @_ @_ @sh f u v w =+  sbuild @(Take m sh) (\ix -> f (u !$ ix) (v !$ ix) (w !$ ix)) szipWith31 :: ( KnownNat n               , KnownShS sh1, KnownShS sh2, KnownShS sh3, KnownShS sh               , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownSTK x3@@ -881,16 +1022,13 @@             -> target (TKS2 sh x2)  -- ^ the second tensor to zip over             -> target (TKS2 sh x3)  -- ^ the third tensor to zip over             -> target (TKS2 sh x)-szipWith30N @sh f u v w =-  gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])-  $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)-  $ sbuild @(Rank sh) (\ix -> f (sindex0 u ix)-                                (sindex0 v ix)-                                (sindex0 w ix))+szipWith30N @sh f u v w | Refl <- lemAppNil @sh =+  sbuild @sh (\ix -> f (sindex u ix) (sindex v ix) (sindex w ix)) szipWith4 :: ( KnownShS (Drop m sh1), KnownShS (Drop m sh2)              , KnownShS (Drop m sh3), KnownShS (Drop m sh4)-             , KnownShS (Take m sh), KnownShS sh+             , KnownShS (Take m sh), KnownShS (Drop m sh)              , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownSTK x3, KnownSTK x4+             , sh ~ Take m sh ++ Drop m sh              , sh1 ~ Take m sh ++ Drop m sh1              , sh2 ~ Take m sh ++ Drop m sh2              , sh3 ~ Take m sh ++ Drop m sh3@@ -903,8 +1041,8 @@           -> target (TKS2 sh3 x3)  -- ^ the third tensor to zip over           -> target (TKS2 sh4 x4)  -- ^ the fourth tensor to zip over           -> target (TKS2 sh x)-szipWith4 f u v w x =-  sbuild (\ix -> f (u !$ ix) (v !$ ix) (w !$ ix) (x !$ ix))+szipWith4 @m @_ @_ @_ @_ @sh f u v w x =+  sbuild @(Take m sh) (\ix -> f (u !$ ix) (v !$ ix) (w !$ ix) (x !$ ix)) szipWith41 :: ( KnownNat n               , KnownShS sh1, KnownShS sh2, KnownShS sh3, KnownShS sh4               , KnownShS sh@@ -933,18 +1071,21 @@             -> target (TKS2 sh x3)  -- ^ the third tensor to zip over             -> target (TKS2 sh x4)  -- ^ the fourth tensor to zip over             -> target (TKS2 sh x)-szipWith40N @sh f u v w x =-  gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])-  $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)-  $ sbuild @(Rank sh) (\ix -> f (sindex0 u ix)-                                (sindex0 v ix)-                                (sindex0 w ix)-                                (sindex0 x ix))+szipWith40N @sh f u v w x | Refl <- lemAppNil @sh =+  sbuild @sh (\ix -> f (sindex u ix) (sindex v ix) (sindex w ix) (sindex x ix))  xbuild1 :: (KnownNat k, KnownShX sh, KnownSTK x, BaseTensor target)         => (IntOf target -> target (TKX2 sh x))  -- ^ the function to build with         -> target (TKX2 (Just k ': sh) x)+{-# INLINE xbuild1 #-} xbuild1 = txbuild1+xbuild :: forall shm shn x target. (KnownShX shn, KnownSTK x, BaseTensor target)+       => IShX shm+       -> (IxXOf target shm -> target (TKX2 shn x))+            -- ^ the function to build with+       -> target (TKX2 (shm ++ shn) x)+{-# INLINE xbuild #-}+xbuild = txbuild -- xmap and other special cases of build can be defined by the user.  rfold@@ -960,7 +1101,7 @@   withSNat (rwidth es) $ \k -> tfold k knownSTK knownSTK f acc0 es rscan   :: ( KnownNat n, KnownNat m, KnownSTK xn, KnownSTK xm-     , BaseTensor target, LetTensor target )+     , BaseTensor target, LetTensor target, ConvertTensor target )   => (forall f. ADReady f => f (TKR2 n xn) -> f (TKR2 m xm) -> f (TKR2 n xn))        -- ^ the function to fold with   -> target (TKR2 n xn)  -- ^ the initial accumulator@@ -982,7 +1123,7 @@ sfold @k = tfold (SNat @k) knownSTK knownSTK sscan   :: ( KnownNat k, KnownShS sh, KnownShS shm, KnownSTK xn, KnownSTK xm-     , BaseTensor target, LetTensor target )+     , BaseTensor target, LetTensor target, ConvertTensor target )   => (forall f. ADReady f       => f (TKS2 sh xn) -> f (TKS2 shm xm) -> f (TKS2 sh xn))        -- ^ the function to scan with@@ -1004,7 +1145,7 @@ xfold @k = tfold (SNat @k) knownSTK knownSTK xscan   :: ( KnownNat k, KnownShX sh, KnownShX shm, KnownSTK xn, KnownSTK xm-     , BaseTensor target, LetTensor target )+     , BaseTensor target, LetTensor target, ConvertTensor target )   => (forall f. ADReady f       => f (TKX2 sh xn) -> f (TKX2 shm xm) -> f (TKX2 sh xn))        -- ^ the function to scan with@@ -1019,46 +1160,75 @@ -- The function argument needs to be quantified, -- because otherwise in the ADVal instance one could put an illegal -- @DeltaInput@ there, confusing different levels of contangents.-kgrad :: forall x r target. BaseTensor target+kgrad :: forall x r target. (GoodScalar r, BaseTensor target)       => (forall f. ADReady f => f x -> f (TKScalar r))  -- ^ x |-> TKScalar r       -> FullShapeTK x  -- ^ shape of x and dx       -> target x  -- ^ input x       -> target (ADTensorKind x)  -- ^ gradient dx+{-# INLINE kgrad #-} kgrad f xftk =-  \ !es -> tApply (tgrad @target xftk (HFun f)) es+  \ !es -> tapply (tgrad @target xftk (HFun f)) es rvjp :: forall n x r target. BaseTensor target      => (forall f. ADReady f => f x -> f (TKR2 n r))  -- ^ x |-> z      -> FullShapeTK x  -- ^ shape of x and dx      -> target x  -- ^ input x      -> target (ADTensorKind (TKR2 n r))  -- ^ incoming cotangent dz      -> target (ADTensorKind x)  -- ^ gradient dx+{-# INLINE rvjp #-} rvjp f xftk =-  \ !es !dt -> tApply (tvjp @target xftk $ HFun f) (tpair dt es)+  \ !es !dt -> tapply (tvjp @target xftk $ HFun f) (tpair dt es) rjvp :: forall n x r target. BaseTensor target      => (forall f. ADReady f => f x -> f (TKR2 n r))  -- ^ x |-> z      -> FullShapeTK x  -- ^ shape of x and dx      -> target x  -- ^ input x      -> target (ADTensorKind x)  -- ^ incoming tangent dx      -> target (ADTensorKind (TKR2 n r))  -- ^ derivative dz+{-# INLINE rjvp #-} rjvp f xftk =-  \ !es !ds -> tApply (tjvp @target xftk $ HFun f) (tpair ds es)+  \ !es !ds -> tapply (tjvp @target xftk $ HFun f) (tpair ds es) svjp :: forall sh x r target. BaseTensor target      => (forall f. ADReady f => f x -> f (TKS2 sh r))  -- ^ x |-> z      -> FullShapeTK x  -- ^ shape of x and dx      -> target x  -- ^ input x      -> target (ADTensorKind (TKS2 sh r))  -- ^ incoming cotangent dz      -> target (ADTensorKind x)  -- ^ gradient dx+{-# INLINE svjp #-} svjp f xftk =-  \ !es !dt -> tApply (tvjp @target xftk $ HFun f) (tpair dt es)+  \ !es !dt -> tapply (tvjp @target xftk $ HFun f) (tpair dt es) sjvp :: forall sh x r target. BaseTensor target      => (forall f. ADReady f => f x -> f (TKS2 sh r))      -> FullShapeTK x  -- ^ shape of x and dx      -> target x  -- ^ input x      -> target (ADTensorKind x)  -- ^ incoming tangent dx      -> target (ADTensorKind (TKS2 sh r))  -- ^ derivative dz+{-# INLINE sjvp #-} sjvp f xftk =-  \ !es !ds -> tApply (tjvp @target xftk $ HFun f) (tpair ds es)+  \ !es !ds -> tapply (tjvp @target xftk $ HFun f) (tpair ds es) +kprimalPart :: BaseTensor target+            => target (TKScalar r) -> PrimalOf target (TKScalar r)+kprimalPart = tprimalPart+kdualPart :: (BaseTensor target, GoodScalar r)+          => target (TKScalar r) -> DualOf target (TKScalar r)+kdualPart = tdualPart knownSTK+kplainPart :: BaseTensor target+           => target (TKScalar r) -> PlainOf target (TKScalar r)+kplainPart = tplainPart+kfromPrimal :: (BaseTensor target, GoodScalar r)+            => PrimalOf target (TKScalar r) -> target (TKScalar r)+kfromPrimal = tfromPrimal knownSTK+kfromDual :: BaseTensor target+          => DualOf target (TKScalar r) -> target (TKScalar r)+kfromDual = tfromDual+kfromPlain :: (BaseTensor target, GoodScalar r)+            => PlainOf target (TKScalar r) -> target (TKScalar r)+kfromPlain = tfromPlain knownSTK+kScale :: ( BaseTensor target, GoodScalar r+          , Num (target (TKScalar r)), Num (PrimalOf target (TKScalar r)) )+       => PrimalOf target (TKScalar r) -> DualOf target (TKScalar r)+       -> DualOf target (TKScalar r)+kScale @target = tScale @target knownSTK+ -- These take @target@ first, because they change the target. rprimalPart :: BaseTensor target             => target (TKR2 n x) -> PrimalOf target (TKR2 n x)@@ -1113,21 +1283,3 @@        => PrimalOf target (TKX sh r) -> DualOf target (TKX sh r)        -> DualOf target (TKX sh r) xScale @target = tScale @target knownSTK--kprimalPart :: BaseTensor target-            => target (TKScalar r) -> PrimalOf target (TKScalar r)-kprimalPart = tprimalPart-kdualPart :: (BaseTensor target, GoodScalar r)-          => target (TKScalar r) -> DualOf target (TKScalar r)-kdualPart = tdualPart knownSTK-kfromPrimal :: (BaseTensor target, GoodScalar r)-            => PrimalOf target (TKScalar r) -> target (TKScalar r)-kfromPrimal = tfromPrimal knownSTK-kfromDual :: BaseTensor target-          => DualOf target (TKScalar r) -> target (TKScalar r)-kfromDual = tfromDual-kScale :: ( BaseTensor target, GoodScalar r-          , Num (target (TKScalar r)), Num (PrimalOf target (TKScalar r)) )-       => PrimalOf target (TKScalar r) -> DualOf target (TKScalar r)-       -> DualOf target (TKScalar r)-kScale @target = tScale @target knownSTK
src/HordeAd/OpsTensorMixed.hs view
@@ -1,4 +1,3 @@-{-# LANGUAGE AllowAmbiguousTypes #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} -- | The mixed tensor operations intended for the casual library user.@@ -19,31 +18,33 @@     xshape, xlength, xsize, xwidth   , tsize, tftk     -- * Constructing arrays from concrete values, lists and vectors-  , xconcrete, xscalar, xrepl, xingestData, xfromListLinear-  , kconcrete-  , xfromList, xfromVector, xfromVector0N, xfromList0N, xunravelToList+  , kconcrete, xconcrete, xscalar, xrepl, xingestData, xfromListLinear+  , xfromList, xfromVector, xfromVectorN, xfromVectorLinear+  , xunravelToList, xunravelToListN, xtoListLinear     -- * Main array operations-  , tunit, tlet, ifH, minH, maxH, tpair, tproject1, tproject2-  , xsum, xsum0, xdot0, xdot1In, xmatvecmul, xmatmul2, xreplicate, xreplicate0N+  , tunit, tlet, tletPrimal, tletPlain, ifH, minH, maxH+  , tpair, tproject1, tproject2+  , xsum, xsumN, xsum0, xdot0, xdot1In, xmatvecmul, xmatmul2+  , xreplicate, xreplicate0N, xreplicateN   , xindex, xindex0, xoneHot, xscatter, xscatter1, xgather, xgather1   , xtr, xtranspose, xflatten, xreshape    -- * Auxiliary array operations-  , xfloor, xfromIntegral, xcast, xminIndex, xmaxIndex, xiota   , kfloor, kfromIntegral, kcast-  , xappend, xappend0, xconcat, xslice, xuncons, xreverse+  , xfloor, xfromIntegral, xcast, xargMin, xargMax, xiota+  , xappend, xconcat, xslice, xuncons, xreverse     -- * Array operations derived from @build@   , xbuild, xbuild1     -- * Array operations derived from @mapAccum@-  , xfold, xscan, tmapAccumR, tmapAccumL+  , xfold, xscan, tfold, tscan, tmapAccumR, tmapAccumL     -- * Array operations producing derivatives   , kgrad     -- * Operations dealing with dual numbers-  , xprimalPart, xdualPart, xfromPrimal, xfromDual, xScale   , kprimalPart, kdualPart, kfromPrimal, kfromDual, kScale+  , xprimalPart, xdualPart, xfromPrimal, xfromDual, xScale     -- * Array operations that utilize unwinding of nested arrays   , treplTarget, tdefTarget, taddTarget, tmultTarget, tsum0Target, tdot0Target     -- * Minimal re-exports to make this module a higher level replacement for the mixed part of "HordeAd.Core.Ops"-  , ADReady+  , ADReady, ADReadyNoLet, ShareTensor   , LetTensor, BaseTensor   ) where
src/HordeAd/OpsTensorRanked.hs view
@@ -1,4 +1,3 @@-{-# LANGUAGE AllowAmbiguousTypes #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} -- | The ranked tensor operations intended for the casual library user.@@ -19,32 +18,34 @@     rshape, rlength, rsize, rwidth   , tsize, tftk     -- * Constructing arrays from concrete values, lists and vectors-  , rconcrete, rscalar, rrepl, ringestData, rfromListLinear-  , kconcrete-  , rfromList, rfromVector, rfromVector0N, rfromList0N, runravelToList+  , kconcrete, rconcrete, rscalar, rrepl, ringestData, rfromListLinear+  , rfromList, rfromVector, rfromVectorN, rfromVectorLinear+  , runravelToList, runravelToListN, rtoListLinear     -- * Main array operations-  , tunit, tlet, ifH, minH, maxH, tpair, tproject1, tproject2-  , rsum, rsum0, rdot0, rdot1In, rmatvecmul, rmatmul2, rreplicate, rreplicate0N+  , tunit, tlet, tletPrimal, tletPlain, ifH, minH, maxH+  , tpair, tproject1, tproject2+  , rsum, rsumN, rsum0, rdot0, rdot1In, rmatvecmul, rmatmul2+  , rreplicate, rreplicate0N, rreplicateN   , rindex, (!), rindex0, roneHot, rscatter, rscatter1, rgather, rgather1   , rtr, rtranspose, rflatten, rreshape    -- * Auxiliary array operations-  , rfloor, rfromIntegral, rcast, rminIndex, rmaxIndex, riota   , kfloor, kfromIntegral, kcast+  , rfloor, rfromIntegral, rcast, rargMin, rargMax, riota   , rappend, rconcat, rslice, runcons, rreverse     -- * Array operations derived from @build@   , rbuild, rbuild1, rmap, rmap1, rmap0N, rzipWith, rzipWith1, rzipWith0N   , rzipWith3, rzipWith31, rzipWith30N, rzipWith4, rzipWith41, rzipWith40N     -- * Array operations derived from @mapAccum@-  , rfold, rscan, tmapAccumR, tmapAccumL+  , rfold, rscan, tfold, tscan, tmapAccumR, tmapAccumL     -- * Array operations producing derivatives   , kgrad, rvjp, rjvp     -- * Operations dealing with dual numbers-  , rprimalPart, rdualPart, rfromPrimal, rfromDual, rScale   , kprimalPart, kdualPart, kfromPrimal, kfromDual, kScale+  , rprimalPart, rdualPart, rfromPrimal, rfromDual, rScale     -- * Array operations that utilize unwinding of nested arrays   , treplTarget, tdefTarget, taddTarget, tmultTarget, tsum0Target, tdot0Target     -- * Minimal re-exports to make this module a higher level replacement for the ranked part of "HordeAd.Core.Ops"-  , ADReady+  , ADReady, ADReadyNoLet, ShareTensor   , LetTensor, BaseTensor   ) where
src/HordeAd/OpsTensorShaped.hs view
@@ -1,4 +1,3 @@-{-# LANGUAGE AllowAmbiguousTypes #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} -- | The shaped tensor operations intended for the casual library user.@@ -19,32 +18,34 @@     sshape, slength, ssize, swidth   , tsize, tftk     -- * Constructing arrays from concrete values, lists and vectors-  , sconcrete, sscalar, srepl, singestData, sfromListLinear-  , kconcrete-  , sfromList, sfromVector, sfromVector0N, sfromList0N, sunravelToList+  , kconcrete, sconcrete, sscalar, srepl, singestData, sfromListLinear+  , sfromList, sfromVector, sfromVectorN, sfromVectorLinear+  , sunravelToList, sunravelToListN, stoListLinear     -- * Main array operations-  , tunit, tlet, ifH, minH, maxH, tpair, tproject1, tproject2-  , ssum, ssum0, sdot0, sdot1In, smatvecmul, smatmul2, sreplicate, sreplicate0N+  , tunit, tlet, tletPrimal, tletPlain, ifH, minH, maxH+  , tpair, tproject1, tproject2+  , ssum, ssumN, ssum0, sdot0, sdot1In, smatvecmul, smatmul2+  , sreplicate, sreplicate0N, sreplicateN   , sindex, (!$), sindex0, soneHot, sscatter, sscatter1, sgather, sgather1   , str, stranspose, sflatten, sreshape    -- * Auxiliary array operations-  , sfloor, sfromIntegral, scast, sminIndex, smaxIndex, siota   , kfloor, kfromIntegral, kcast+  , sfloor, sfromIntegral, scast, sargMin, sargMax, siota   , sappend, sslice, suncons, sreverse     -- * Array operations derived from @build@   , sbuild, sbuild1, smap, smap1, smap0N, szipWith, szipWith1, szipWith0N   , szipWith3, szipWith31, szipWith30N, szipWith4, szipWith41, szipWith40N     -- * Array operations derived from @mapAccum@-  , sfold, sscan, tmapAccumR, tmapAccumL+  , sfold, sscan, tfold, tscan, tmapAccumR, tmapAccumL     -- * Array operations producing derivatives   , kgrad, svjp, sjvp     -- * Operations dealing with dual numbers-  , sprimalPart, sdualPart, sfromPrimal, sfromDual, sScale   , kprimalPart, kdualPart, kfromPrimal, kfromDual, kScale+  , sprimalPart, sdualPart, sfromPrimal, sfromDual, sScale     -- * Array operations that utilize unwinding of nested arrays   , treplTarget, tdefTarget, taddTarget, tmultTarget, tsum0Target, tdot0Target     -- * Minimal re-exports to make this module a higher level replacement for the shaped part of "HordeAd.Core.Ops"-  , ADReady+  , ADReady, ADReadyNoLet, ShareTensor   , LetTensor, BaseTensor   ) where
test/CAFlessTest.hs view
@@ -29,8 +29,8 @@  tests :: TestTree tests =-  testGroup "Tests for simplified horde-ad that don't create big CAFs"-    [ testGroup "Short_tests"+  inOrderTestGroup "Tests for simplified horde-ad that don't create big CAFs"+    [ inOrderTestGroup "Short_tests"         (TestAdaptorSimplified.testTrees          ++ TestConvSimplified.testTrees          ++ TestGatherSimplified.testTrees
test/FullTest.hs view
@@ -10,6 +10,7 @@  import EqEpsilon import TestAdaptorSimplified qualified+import TestConvQuickCheck qualified import TestConvSimplified qualified import TestGatherSimplified qualified import TestHighRankSimplified qualified@@ -35,7 +36,7 @@ tests :: TestTree tests =   testGroup "Full set of tests for simplified horde-ad"-    [ testGroup "Short_tests"+    [ testGroup "Short_tests"  -- TODO: inOrderTestGroup needed here?         (TestAdaptorSimplified.testTrees          ++ TestConvSimplified.testTrees          ++ TestGatherSimplified.testTrees@@ -47,5 +48,6 @@          ++ TestMnistCNNS.testTrees          ++ TestMnistFCNNR.testTrees          ++ TestMnistRNNR.testTrees-         ++ TestMnistRNNS.testTrees)+         ++ TestMnistRNNS.testTrees+         ++ TestConvQuickCheck.testTrees)     ]
test/MinimalTest.hs view
@@ -24,5 +24,5 @@  tests :: TestTree tests =-  testGroup "Minimal test that doesn't require any dataset"+  inOrderTestGroup "Minimal test that doesn't require any dataset"     TestGatherSimplified.testTrees
test/ParallelTest.hs view
@@ -9,6 +9,7 @@ import Test.Tasty.Runners  import EqEpsilon+import TestConvQuickCheck qualified import TestMnistCNNR qualified import TestMnistCNNS qualified import TestMnistFCNNR qualified@@ -28,11 +29,15 @@  tests :: TestTree tests =-  testGroup "The set of tests for simplified horde-ad that can run in parallel"+  testGroup "The set of tests for horde-ad that can be run in parallel"     [ testGroup "Long_tests"-        (TestMnistCNNR.testTrees+        (TestConvQuickCheck.testTrees+         ++ TestMnistCNNR.testTrees          ++ TestMnistCNNS.testTrees+         ++ TestConvQuickCheck.testTrees  -- saturates cores to prevent OOM          ++ TestMnistFCNNR.testTrees+         ++ TestConvQuickCheck.testTrees          ++ TestMnistRNNR.testTrees+         ++ TestConvQuickCheck.testTrees          ++ TestMnistRNNS.testTrees)     ]
+ test/SequentialMnistTest.hs view
@@ -0,0 +1,38 @@+module Main (main) where++import Prelude++import Data.Proxy+import System.IO qualified as SIO+import Test.Tasty+import Test.Tasty.Options+import Test.Tasty.Runners++import EqEpsilon+import TestMnistCNNR qualified+import TestMnistCNNS qualified+import TestMnistFCNNR qualified+import TestMnistRNNR qualified+import TestMnistRNNS qualified++main :: IO ()+main = do+  -- Limit interleaving of characters in parallel tests.+  SIO.hSetBuffering SIO.stdout SIO.LineBuffering+  SIO.hSetBuffering SIO.stderr SIO.LineBuffering+  opts <- parseOptions (ingredients : defaultIngredients) tests+  setEpsilonEq (lookupOption opts :: EqEpsilon)+  defaultMainWithIngredients (ingredients : defaultIngredients) tests+ where+  ingredients = includingOptions [Option (Proxy :: Proxy EqEpsilon)]++tests :: TestTree+tests =+  inOrderTestGroup "The set of MNIST tests for horde-ad that can serve as a benchmark when run sequentially"+    [ inOrderTestGroup "Long_tests"+        (TestMnistCNNR.testTrees+         ++ TestMnistCNNS.testTrees+         ++ TestMnistFCNNR.testTrees+         ++ TestMnistRNNR.testTrees+         ++ TestMnistRNNS.testTrees)+    ]
test/simplified/TestAdaptorSimplified.hs view
@@ -7,2456 +7,2664 @@  import Prelude -import Data.Int (Int64)-import Data.Kind (Type)-import Data.List.NonEmpty qualified as NonEmpty-import Foreign.C (CInt)-import GHC.Exts (IsList (..))-import GHC.TypeLits (KnownNat)-import Test.Tasty-import Test.Tasty.HUnit hiding (assert)--import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Mixed.Shape-import Data.Array.Nested.Ranked.Shape-import Data.Array.Nested.Shaped.Shape--import HordeAd-import HordeAd.Core.Adaptor-import HordeAd.Core.AstEnv-import HordeAd.Core.AstFreshId (funToAst, resetVarCounter)-import HordeAd.Core.AstInterpret-import HordeAd.Core.DeltaFreshId (resetIdCounter)--import CrossTesting-import EqEpsilon--testTrees :: [TestTree]-testTrees =-  [ testCase "2zeroZ" testZeroZ-  , testCase "2zeroS" testZeroS-  , testCase "2CFwdZeroS" testCFwdZeroS-  , testCase "2FwdZeroS" testFwdZeroS-  , testCase "2zero2S" testZero2S-  , testCase "2CFwdZero2S" testCFwdZero2S-  , testCase "2FwdZero2S" testFwdZero2S-  , testCase "2zero3S" testZero3S-  , testCase "2CFwdZero3S" testCFwdZero3S-  , testCase "2FwdZero3S" testFwdZero3S-  , testCase "2zero4s" testZero4s-  , testCase "2zero5S" testZero5S-  , testCase "2zero6S" testZero6S-  , testCase "2zero7S" testZero7S-  , testCase "2zero8" testZero8-  , testCase "2zero9S" testZero9S-  , testCase "2CFwdZero9S" testCFwdZero9S-  , testCase "2FwdZero9S" testFwdZero9S-  , testCase "2zero10S" testZero10S-  , testCase "2CFwdZero10S" testCFwdZero10S-  , testCase "2FwdZero10S" testFwdZero10S-  , testCase "2zero11S" testZero11S-  , testCase "2CFwdZero11S" testCFwdZero11S-  , testCase "2FwdZero11S" testFwdZero11S-  , testCase "2piecewiseLinear" testPiecewiseLinear-  , testCase "2piecewiseLinearPP" testPiecewiseLinearPP-  , testCase "2piecewiseLinear2" testPiecewiseLinear2-  , testCase "2piecewiseLinear2PP" testPiecewiseLinear2PP-  , testCase "2overleaf" testOverleaf-  , testCase "2overleafInt64n" testOverleafInt64n-  , testCase "2overleafCIntn" testOverleafCIntn-  , testCase "2overleafCIntToFloatn" testOverleafCIntToFloatn-  , testCase "2overleafInt64p" testOverleafInt64p-  , testCase "2overleafCIntp" testOverleafCIntp-  , testCase "2overleafCIntToFloatp" testOverleafCIntToFloatp-  , testCase "2overleafPP" testOverleafPP-  , testCase "2foo" testFoo-  , testCase "2fooGradDouble" testGradFooDouble-  , testCase "2fooMatrix" testFooMatrix-  , testCase "2fooGradMatrix" testGradFooMatrix-  , testCase "2fooLetGradMatrixPP" testGradFooLetMatrixPP-  , testCase "2fooGradMatrixVjp" testGradFooMatrixVjp-  , testCase "2fooGradMatrixRev" testGradFooMatrixRev-  , testCase "2fooLetGradMatrixSimpPP" testGradFooLetMatrixSimpPP-  , testCase "2fooLetGradMatrixSimpRPP" testGradFooLetMatrixSimpRPP-  , testCase "2fooSumMatrix" testfooSumMatrix-  , testCase "2fooGradMatrix2" testGradFooMatrix2-  , testCase "2fooGradMatrixPP" testGradFooMatrixPP-  , testCase "2fooGradMatrixSimpPP" testGradFooMatrixSimpPP-  , testCase "2fooGradScalar" testGradFooScalar-  , testCase "2fooGradCScalar" testGradCFooScalar-  , testCase "2fooS" testFooS-  , testCase "2fooSToFloat" testFooSToFloat-  , testCase "2fooSBoth" testFooSBoth-  , testCase "2fooBoth" testFooBoth-  , testCase "2trustVstackConcatRepl10" testTrustVstackConcatRepl10-  , testCase "2trustVstackConcatIota10" testTrustVstackConcatIota10-  , testCase "2trustVstackConcatReplIota10" testTrustVstackConcatReplIota10-  , testCase "2vstackWarmup" testVstackWarmup-  , testCase "2vstackConcatConcrete" testVstackConcatConcrete-  , testCase "2vstackBuildConcrete" testVstackBuildConcrete-  , testCase "2vstackConcatAst" testVstackConcatAst-  , testCase "2vstackBuildAst" testVstackBuildAst-  , testCase "2vstackBuildAstSimp" testVstackBuildAstSimp-  , testCase "2vstackBuildAstPP" testVstackBuildAstPP-  , testCase "2vstackConcatConcrete2" testVstackConcatConcrete2-  , testCase "2vstackBuildConcrete2" testVstackBuildConcrete2-  , testCase "2vstackConcatAst2" testVstackConcatAst2-  , testCase "2vstackBuildAst2" testVstackBuildAst2-  , testCase "2vstackBuildAstSimp2" testVstackBuildAstSimp2-  , testCase "2vstackBuildAstPP2" testVstackBuildAstPP2-  , testCase "2fooPP" testFooPP-  , testCase "2fooLet" testFooLet-  , testCase "2fooLetPP" testFooLetPP-  , testCase "2listProdPP" testListProdPP-  , testCase "2listProdrPP" testListProdrPP-  , testCase "2listProdrLongPP" testListProdrLongPP-  , testCase "2listProd" testListProd-  , testCase "2listProdr" testListProdr-  , testCase "2listSumrPP" testListSumrPP-  , testCase "2listSum2rPP" testListSum2rPP-  , testCase "2listSum22rPP" testListSum22rPP-  , testCase "2listSumk22rPP" testListSumk22rPP-  , testCase "2listSum2xpyrPP" testListSum2xpyrPP-  , testCase "2listSum2xyrPP" testListSum2xyrPP-  , testCase "2list2xyPP" test2xyPP-  , testCase "2listSum23rPP" testListSum23rPP-  , testCase "2list23PP" test23PP-  , testCase "2reluPP" testReluPP-  , testCase "2reluPP2" testReluPP2-  , testCase "2reluSimpler" testReluSimpler-  , testCase "2reluSimplerPP" testReluSimplerPP-  , testCase "2reluSimplerPP2" testReluSimplerPP2-  , testCase "2reluSimplerPP3" testReluSimplerPP3-  , testCase "2reluSimpler3" testReluSimpler3-  , testCase "2reluSimplerPP4" testReluSimplerPP4-  , testCase "2reluSimpler4" testReluSimpler4-  , testCase "2reluSimplerPP4s" testReluSimplerPP4s-  , testCase "2reluSimplerPP4s2" testReluSimplerPP4s2-  , testCase "2reluSimpler4s" testReluSimpler4s-  , testCase "2reluSimplerPP7s2" testReluSimplerPP7s2-  , testCase "2reluSimpler7s" testReluSimpler7s-  , testCase "2reluMax" testReluMax-  , testCase "2reluMaxPP" testReluMaxPP-  , testCase "2reluMaxPP2" testReluMaxPP2-  , testCase "2reluMax3" testReluMax3-  , testCase "2dot1PP" testDot1PP-  , testCase "2dot2PP" testDot2PP-  , testCase "2matvecmulPP" testMatvecmulPP-  , testCase "2matmul2PP" testMatmul2PP-  , testCase "2matmul2FromMatvecmulPP" testMatmul2FromMatvecmulPP-  , testCase "2matmul2PaperPP" testMatmul2PaperPP-  , testCase "2matmul2PPS" testMatmul2PPS-  , testCase "2addSpeedBig" testAddSpeedBig-  , testCase "2matmul2SpeedBig" testMatmul2SpeedBig-  , testCase "2bar" testBar-  , testCase "2barS" testBarS-  , testCase "2bar2S" testBar2S-  , testCase "2barCFwd" testBarCFwd-  , testCase "2barFwd" testBarFwd-  , testCase "2baz old to force fooFromPrimal" testBaz-  , testCase "2baz if repetition breaks things" testBaz-  , testCase "2baz again with renumbered terms" testBazRenumbered-  , testCase "2fooD T Double [1.1, 2.2, 3.3]" testFooD-  , testCase "2fooBuildDt" testFooBuildDt-  , testCase "2fooBuildCFwd" testFooBuildCFwd-  , testCase "2fooBuildFwd" testFooBuildFwd-  , testCase "2fooBuild" testFooBuild-  , testCase "2fooNoGo0" testFooNoGo0-  , testCase "2nestedBuildMap1" testNestedBuildMap1-  , testCase "2nestedSumBuild" testNestedSumBuild-  , testCase "2nestedBuildIndex" testNestedBuildIndex-  , testCase "2barReluDt" testBarReluDt-  , testCase "2barRelu" testBarRelu-  , testCase "2barRelu3" testBarRelu3-  , testCase "2barReluMaxDt" testBarReluMaxDt-  , testCase "2barReluMax" testBarReluMax-  , testCase "2barReluMax30" testBarReluMax30-  , testCase "2barReluMax31" testBarReluMax31-  , testCase "2barReluMax3CFwd" testBarReluMax3CFwd-  , testCase "2barReluMax3FwdS" testBarReluMax3FwdS-  , testCase "2barReluMax3FwdFrom" testBarReluMax3FwdFrom-  , testCase "2barReluMax3FwdR" testBarReluMax3FwdR-  , testCase "2F1" testF1-  , testCase "2F11" testF11-  , testCase "2F2" testF2-  , testCase "2F21" testF21-  , testCase "2F2CFwd" testF2CFwd-  , testCase "2F2Fwd" testF2Fwd-  , testCase "2braidedBuilds1" testBraidedBuilds1-  , testCase "2recycled1" testRecycled1-  , testCase "2concatBuild1" testConcatBuild1-  , testCase "2concatBuild2" testConcatBuild2-  , testCase "2concatBuild3" testConcatBuild3-  , testCase "2concatBuild4" testConcatBuild4-  , testCase "2concatBuild5" testConcatBuild5-  , testCase "2concatBuild6" testConcatBuild6-  , testCase "2concatBuild7" testConcatBuild7---  , testCase "2concatBuild8" testConcatBuild8---  , testCase "2concatBuild9" testConcatBuild9-  , testCase "2concatBuild10" testConcatBuild10-  , testCase "2concatBuild11" testConcatBuild11-  , testCase "2concatBuild12" testConcatBuild12-  , testCase "2emptyArgs0" testEmptyArgs0-  , testCase "2emptyArgs1" testEmptyArgs1-  , testCase "2emptyArgs4" testEmptyArgs4-  , testCase "2filterPositiveFail" testFilterPositiveFail-  , testCase "2blowupPP" fblowupPP-  , testCase "2blowup2LetPP" fblowupLetPP-  , testCase "2blowup2LetPP23" fblowupLetPP23-  , testCase "2blowup2LetPP10" fblowupLetPP10-  , blowupTests-  , testCase "22concatBuild3PP" testConcatBuild3PP-  , testCase "22concatBuild3PP2" testConcatBuild3PP2-  ]--testZeroZ :: Assertion-testZeroZ =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 0)-    (rev' @Double @0 (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)-                          f = const (rscalar 3)-                      in f) (rscalar 42))--testZeroS :: Assertion-testZeroS =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])-    (cgrad (kfromS . ssum0 .-            let f :: ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)-                  -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)-                f = const (srepl 3)-            in f) (srepl 42))--testCFwdZeroS :: Assertion-testCFwdZeroS =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])-    (cjvp (let f :: ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)-                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)-               f = const (srepl 3)-           in f) 42 41)--testFwdZeroS :: Assertion-testFwdZeroS =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])-    (jvp (let f :: AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)-                -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)-              f = const (srepl 3)-          in f) (srepl 42) (srepl 41))--testZero2S :: Assertion-testZero2S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [1])-    (cgrad (kfromS @_ @Double .-           let f :: a -> a-               f = id-           in f) (srepl 42))--testCFwdZero2S :: Assertion-testCFwdZero2S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [41])-    (cjvp @_ @(TKS '[] Double)-          (let f :: a -> a-               f = id-           in f) (srepl 42) (srepl 41))--testFwdZero2S :: Assertion-testFwdZero2S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [41])-    (jvp @_ @(TKS '[] Double)-          (let f :: a -> a-               f = id-           in f) (srepl 42) (srepl 41))--testZero3S :: Assertion-testZero3S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[33, 2] knownShS (replicate 66 3.6174114266850617))-    (cgrad (kfromS . ssum0 . (\x -> barF @(ADVal Concrete (TKS '[33, 2] Double)) (x, x))) (srepl 1))--testCFwdZero3S :: Assertion-testCFwdZero3S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[33, 2] knownShS (replicate 66 3.9791525693535674))-    (cjvp (\x -> barF @(ADVal Concrete (TKS '[33, 2] Double)) (x, x)) (srepl 1) (srepl 1.1))--testFwdZero3S :: Assertion-testFwdZero3S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[33, 2] knownShS (replicate 66 3.9791525693535674))-    (jvp (\x -> barF @(AstTensor AstMethodLet FullSpan (TKS '[33, 2] Double)) (x, x)) (srepl 1) (srepl 1.1))--testZero4s :: Assertion-testZero4s =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [0])-    (grad @(AstTensor AstMethodLet FullSpan (TKS '[] Double))-          (kfromS @_ @Double .-          let f = const (srepl 3)-          in f) (srepl 42))--testZero5S :: Assertion-testZero5S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[44] knownShS (replicate 44 1))-    (grad (kfromS . ssum0 .-           let f :: a -> a-               f = id-           in f @(AstTensor AstMethodLet FullSpan (TKS '[44] Double))) (srepl 42))--testZero6S :: Assertion-testZero6S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] knownShS (replicate (product ([2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] :: [Int])) 3.6174114266850617))-    (grad (kfromS . ssum0 @'[2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] @(TKScalar Double) . (\x -> barF (x, x))) (srepl 1))--testZero7S :: Assertion-testZero7S =-  assertEqualUpToEpsilon 1e-10-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [0])-    (grad (kfromR . (const (rscalar 3) :: AstTensor AstMethodLet FullSpan (TKS '[] Double) -> AstTensor AstMethodLet FullSpan (TKR 0 Double))) (srepl 42))--testZero8 :: Assertion-testZero8 =-  assertEqualUpToEpsilon 1e-10-    (rfromList0N [] [rscalar 0])-    (grad (kfromS . (const (sscalar 3) :: AstTensor AstMethodLet FullSpan (TKR 0 Double) -> AstTensor AstMethodLet FullSpan (TKS '[] Double))) (rscalar 42))--testZero9S :: Assertion-testZero9S =-  assertEqualUpToEpsilon 1e-9-    (rfromList0N [0, 2, 4, 0, 1] [])-    (cgrad (kfromS . ssum0 .-            let f :: ADVal Concrete (TKR 5 Double)-                  -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)-                f = const (srepl 3)-            in f)-           (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42)))--testCFwdZero9S :: Assertion-testCFwdZero9S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])-    (cjvp (let f :: ADVal Concrete (TKR 5 Double)-                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)-               f = const (srepl 3)-           in f)-          (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42)) (rreplicate0N [0, 2, 4, 0, 1] (rscalar 41)))--testFwdZero9S :: Assertion-testFwdZero9S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])-    (jvp (let f :: AstTensor AstMethodLet FullSpan (TKR 5 Double)-                -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)-              f = const (srepl 3)-          in f)-         (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42)) (rreplicate0N [0, 2, 4, 0, 1] (rscalar 41)))--testZero10S :: Assertion-testZero10S =-  assertEqualUpToEpsilon 1e-9-    ( rfromList0N [0, 2, 4, 0, 1] []-    , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )-    (cgrad (kfromS . ssum0 .-            let f = const (srepl 3) . snd-            in f :: ( ADVal Concrete (TKR 5 Double)-                    , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )-                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double))-          (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42), srepl 21))--testCFwdZero10S :: Assertion-testCFwdZero10S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])-    (cjvp (let f = const (srepl 3) . snd-           in f :: ( ADVal Concrete (TKR 5 Double)-                   , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )-                   -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double))-          ( rfromList0N [0, 2, 4, 0, 1] []-          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )-          ( rfromList0N [0, 2, 4, 0, 1] []-          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] ))--testFwdZero10S :: Assertion-testFwdZero10S =-  assertEqualUpToEpsilon 1e-9-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])-    (jvp  (let f = const (srepl 3) . snd-           in f :: ( AstTensor AstMethodLet FullSpan (TKR 5 Double)-                   , AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double) )-                   -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double))-          ( rfromList0N [0, 2, 4, 0, 1] []-          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )-          ( rfromList0N [0, 2, 4, 0, 1] []-          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] ))--testZero11S :: Assertion-testZero11S =-  assertEqualUpToEpsilon 1e-9-    ( rfromList0N [0, 2, 4, 0, 1] []-    , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )-    (cgrad (kfromR . rsum0 .-            let f = const (rreplicate0N [0, 2, 4, 0, 1] (rscalar 3)) . snd-            in f :: ( ADVal Concrete (TKR 5 Double)-                    , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )-                 -> ADVal Concrete (TKR 5 Double))-          (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42), srepl 21))--testCFwdZero11S :: Assertion-testCFwdZero11S =-  assertEqualUpToEpsilon 1e-9-    (rfromList0N [0, 2, 4, 0, 1] [])-    (cjvp (let f = const (rreplicate0N [0, 2, 4, 0, 1] (rscalar 3)) . snd-           in f :: ( ADVal Concrete (TKR 5 Double)-                   , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )-                   -> ADVal Concrete (TKR 5 Double))-          ( rfromList0N [0, 2, 4, 0, 1] []-          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )-          ( rfromList0N [0, 2, 4, 0, 1] []-          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] ))--testFwdZero11S :: Assertion-testFwdZero11S =-  assertEqualUpToEpsilon 1e-9-    (rfromList0N [0, 2, 4, 0, 1] [])-    (jvp  (let f = const (rreplicate0N [0, 2, 4, 0, 1] (rscalar 3)) . snd-           in f :: ( AstTensor AstMethodLet FullSpan (TKR 5 Double)-                   , AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double) )-                   -> AstTensor AstMethodLet FullSpan (TKR 5 Double))-          ( rfromList0N [0, 2, 4, 0, 1] []-          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )-          ( rfromList0N [0, 2, 4, 0, 1] []-          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] ))--testPiecewiseLinear :: Assertion-testPiecewiseLinear =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 2)-    (let fT :: ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)-         fT x = ifH (x >. rscalar 0) (rscalar 2 * x) (rscalar 5 * x)-     in rev' @Double @0 fT (rscalar 42))--testPiecewiseLinearPP :: Assertion-testPiecewiseLinearPP = do-  resetVarCounter-  let fT :: AstTensor AstMethodLet FullSpan (TKR 0 Double)-         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT x = ifH (x >. rscalar 0) (rscalar 2 * x) (rscalar 5 * x)-      (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR ZSR FTKScalar)-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> rfromS (let v4 = soneHot (sfromR dret) [ifH (sscalar -0.0 <=. negate (sfromR x1)) 0 1] in sscalar 2.0 * v4 !$ [1] + sscalar 5.0 * v4 !$ [0])"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (ifH (sscalar -0.0 <=. negate (sfromR x1)) (sscalar 5.0 * sfromR x1) (sscalar 2.0 * sfromR x1))"--testPiecewiseLinear2 :: Assertion-testPiecewiseLinear2 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 5)-    (let fT :: ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)-         fT x = ifH (x >. rscalar 0) (rscalar 2) (rscalar 5) * x-     in rev' @Double @0 fT (rscalar (-42)))--testPiecewiseLinear2PP :: Assertion-testPiecewiseLinear2PP = do-  resetVarCounter-  let fT :: AstTensor AstMethodLet FullSpan (TKR 0 Double)-         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT x = ifH (x >. rscalar 0) (rscalar 2) (rscalar 5) * x-      (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR ZSR FTKScalar)-  printArtifactPretty artifactRev-    @?= "\\dret x1 -> let x2 = ifH (sscalar -0.0 <=. negate (sfromR x1)) (sscalar 5.0) (sscalar 2.0) in rfromS (x2 * sfromR dret)"-  printArtifactPrimalPretty artifactRev-    @?= "\\x1 -> let x2 = ifH (sscalar -0.0 <=. negate (sfromR x1)) (sscalar 5.0) (sscalar 2.0) in rfromS (x2 * sfromR x1)"--overleaf :: forall r target. (BaseTensor target, GoodScalar r)-         => target (TKR 1 r) -> target (TKR 0 r)-overleaf v = let wrap i = i `remH` fromIntegral (rwidth v)-             in rsum (rbuild @1 [50] (\[i] -> rindex v [wrap i]))--testOverleaf :: Assertion-testOverleaf =-  assertEqualUpToEpsilon' 1e-10-    (ringestData @_ @Double [28] [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0])-    (rev' @Double @0 overleaf (ringestData [28] [0 .. 27]))--testOverleafInt64n :: Assertion-testOverleafInt64n =-  assertEqualUpToEpsilon 1e-10-    (ringestData [28] (map round [0 :: Double,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]))-    (cgrad (kfromR @_ @Int64 . overleaf) (ringestData [28] [0 .. 27]))--testOverleafCIntn :: Assertion-testOverleafCIntn =-  assertEqualUpToEpsilon 1e-10-    (ringestData [28] (map round [0 :: Double,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]))-    (grad (kfromR @_ @CInt . overleaf) (ringestData [28] [0 .. 27]))--testOverleafCIntToFloatn :: Assertion-testOverleafCIntToFloatn =-  assertEqualUpToEpsilon 1e-10-    (rfromList0N [28] (replicate 28 (rscalar 0.0)))-    (grad (kfromR @_ @CInt . rfromIntegral . overleaf @CInt . rfloor) (ringestData @_ @Float [28] [0 .. 27]))--testOverleafInt64p :: Assertion-testOverleafInt64p =-  assertEqualUpToEpsilon' 1e-10-    (ringestData @_ @Int64 [28] (map round [2.0 :: Double,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0]))-    (rev' @Int64 @0 overleaf (ringestData [28] [0 .. 27]))--testOverleafCIntp :: Assertion-testOverleafCIntp =-  assertEqualUpToEpsilon' 1e-10-    (ringestData @_ @CInt [28] (map round [2.0 :: Double,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0]))-    (rev' @CInt @0 overleaf (ringestData [28] [0 .. 27]))--testOverleafCIntToFloatp :: Assertion-testOverleafCIntToFloatp =-  assertEqualUpToEpsilon' 1e-10-    (ringestData @1 @Float [28] (replicate 28 0.0))-    (let f :: forall f. ADReady f => f (TKR 1 Float) -> f (TKR 0 Float)-         f = rfromIntegral . overleaf @CInt . rfloor-     in rev' @Float @0 f (ringestData [28] [0 .. 27]))--testOverleafPP :: Assertion-testOverleafPP = do-  resetVarCounter >> resetIdCounter-  let fT :: AstTensor AstMethodLet FullSpan (TKR 1 Double)-         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = overleaf-      (var3, ast3) = funToAst (FTKR [28] FTKScalar) Nothing fT-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\v1 -> rfromS (ssum @50 (sgather [] (sfromR v1) (\\[i2] -> [remH i2 28])))"-  resetVarCounter-  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR [28] FTKScalar)-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret v1 -> rfromS (sscatter (sreplicate @50 (sfromR dret)) (\\[i5] -> [remH i5 28]))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\v1 -> rfromS (ssum0 (sgather (sfromR v1) (\\[i3] -> [remH i3 28])))"-  show deltas-    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000002 (DeltaSum (SNat @50) (STKS [] STKScalar) (DeltaShare 100000001 (DeltaGatherS [50] [] [28] (DeltaConvert (ConvCmp (ConvXS' (FTKS [28] FTKScalar)) ConvRX) (DeltaInput (InputId 0))) <function>))))"--foo :: RealFloat a => (a, a, a) -> a-foo (x, y, z) =-  let w = x * sin y-  in atan2 z w + z * w  -- note that w appears twice--testFoo :: Assertion-testFoo = do-  assertEqualUpToEpsilon 1e-10-    (rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627)-    (cgrad (kfromR @_ @Double . foo) (rscalar 1.1, rscalar 2.2, rscalar 3.3))--gradFooDouble :: (Double, Double, Double) -> (Double, Double, Double)-gradFooDouble = fromDValue . cgrad foo . fromValue--testGradFooDouble :: Assertion-testGradFooDouble =-  assertEqualUpToEpsilon 1e-10-    (2.4396285219055063, -1.953374825727421, 0.9654825811012627)-    (gradFooDouble (1.1, 2.2, 3.3))--type Matrix2x2 :: Type -> Type-type Matrix2x2 r = Concrete (TKS '[2, 2] r)-type ThreeMatrices r = (Matrix2x2 r, Matrix2x2 r, Matrix2x2 r)-threeSimpleMatrices :: ThreeMatrices Double-threeSimpleMatrices = (srepl 1.1, srepl 2.2, srepl 3.3)-fooMatrixValue :: Matrix2x2 Double-fooMatrixValue = foo threeSimpleMatrices-gradSumFooMatrix :: ThreeMatrices Double -> ThreeMatrices Double-gradSumFooMatrix = cgrad (kfromS . ssum0 . foo)--testFooMatrix :: Assertion-testFooMatrix =-  assertEqualUpToEpsilon 1e-10-    (sfromListLinear [2,2] [4.242393641025528,4.242393641025528,4.242393641025528,4.242393641025528])-    fooMatrixValue--testGradFooMatrix :: Assertion-testGradFooMatrix =-  assertEqualUpToEpsilon 1e-10-    (sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])-    (gradSumFooMatrix threeSimpleMatrices)--fooLet :: (RealFloatH (f r), ADReady f)-       => (f r, f r, f r) -> f r-fooLet (x, y, z) =-  tlet (x * sin y) $ \w ->-    atan2H z w + z * w  -- note that w still appears twice--artifact :: AstArtifactRev (X (ThreeMatrices Double)) (TKS '[2, 2] Double)-artifact = vjpArtifact fooLet threeSimpleMatrices--testGradFooLetMatrixPP :: Assertion-testGradFooLetMatrixPP = do-  resetVarCounter-  printArtifactPretty artifact-    @?= "\\dret m1 -> let m3 = sin (tproject2 (tproject1 m1)) ; m4 = tproject1 (tproject1 m1) * m3 ; m5 = recip (tproject2 m1 * tproject2 m1 + m4 * m4) ; m7 = (negate (tproject2 m1) * m5) * dret + tproject2 m1 * dret in tpair (tpair (m3 * m7) (cos (tproject2 (tproject1 m1)) * (tproject1 (tproject1 m1) * m7))) ((m4 * m5) * dret + m4 * dret)"--testGradFooMatrixVjp :: Assertion-testGradFooMatrixVjp =-  assertEqualUpToEpsilon 1e-10-    ((sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627]) :: ThreeMatrices Double)-    (vjpInterpretArtifact artifact (toTarget threeSimpleMatrices) (srepl 1))--testGradFooMatrixRev :: Assertion-testGradFooMatrixRev =-  assertEqualUpToEpsilon 1e-10-    (sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])-    (grad (kfromS . ssum0 . fooLet) threeSimpleMatrices)--testGradFooLetMatrixSimpPP :: Assertion-testGradFooLetMatrixSimpPP = do-  resetVarCounter-  (let ftk = FTKS @'[2, 2] [2, 2] (FTKScalar @Double)-   in printArtifactPretty-        (simplifyArtifact $ revArtifactAdapt UseIncomingCotangent fooLet (FTKProduct (FTKProduct ftk ftk) ftk)))-      @?= "\\dret m1 -> let m3 = sin (tproject2 (tproject1 m1)) ; m4 = tproject1 (tproject1 m1) * m3 ; m5 = recip (tproject2 m1 * tproject2 m1 + m4 * m4) ; m7 = (negate (tproject2 m1) * m5) * dret + tproject2 m1 * dret in tpair (tpair (m3 * m7) (cos (tproject2 (tproject1 m1)) * (tproject1 (tproject1 m1) * m7))) ((m4 * m5) * dret + m4 * dret)"--testGradFooLetMatrixSimpRPP :: Assertion-testGradFooLetMatrixSimpRPP = do-  resetVarCounter-  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)-    in printArtifactPretty (simplifyArtifact $ revArtifactAdapt UseIncomingCotangent fooLet (FTKProduct (FTKProduct ftk ftk) ftk)))-       @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2] FTKScalar)) ConvSX))) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2,2] STKScalar)) (let m3 = sin (sfromR (tproject2 (tproject1 m1))) ; m4 = sfromR (tproject1 (tproject1 m1)) * m3 ; m5 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m4 * m4) ; m7 = (negate (sfromR (tproject2 m1)) * m5) * sfromR dret + sfromR (tproject2 m1) * sfromR dret in tpair (tpair (m3 * m7) (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m7))) ((m4 * m5) * sfromR dret + m4 * sfromR dret))"--type Matrix2x2f :: Target -> Type -> Type-type Matrix2x2f f r = f (TKS '[2, 2] r)--sumFooMatrix :: (ADReady f, RealFloat (Matrix2x2f f r), GoodScalar r)-             => (Matrix2x2f f r, Matrix2x2f f r, Matrix2x2f f r)-             -> f (TKScalar r)-sumFooMatrix = kfromS . ssum0 . foo--testfooSumMatrix :: Assertion-testfooSumMatrix =-  assertEqualUpToEpsilon 1e-10-    16.96957456410211-    (sumFooMatrix threeSimpleMatrices)--foo2 :: RealFloatH a => (a, a, a) -> a-foo2 (x, y, z) =-  let w = x * sin y-  in atan2H z w + z * w--gradFooMatrix2 :: (Differentiable r, GoodScalar r)-               => (Concrete (TKR 2 r), Concrete (TKR 2 r), Concrete (TKR 2 r))-               -> (Concrete (TKR 2 r), Concrete (TKR 2 r), Concrete (TKR 2 r))-gradFooMatrix2 = grad (kfromR . rsum0 . foo2)--testGradFooMatrix2 :: Assertion-testGradFooMatrix2 =-  assertEqualUpToEpsilon 1e-10-    (ringestData [2, 2] [2.4396285219055063 :: Double,2.4396285219055063,2.4396285219055063,2.4396285219055063], ringestData [2, 2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421], ringestData [2, 2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])-    (gradFooMatrix2 (rreplicate0N [2, 2] (rscalar 1.1), rreplicate0N [2, 2] (rscalar 2.2), rreplicate0N [2, 2] (rscalar 3.3)))--testGradFooMatrixPP :: Assertion-testGradFooMatrixPP = do-  resetVarCounter-  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)-    in printArtifactPretty (revArtifactAdapt UseIncomingCotangent foo2 (FTKProduct (FTKProduct ftk ftk) ftk)))-      @?= "\\dret m1 -> let m2 = sin (sfromR (tproject2 (tproject1 m1))) ; m3 = sfromR (tproject1 (tproject1 m1)) * m2 ; m4 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m3 * m3) ; m5 = sin (sfromR (tproject2 (tproject1 m1))) ; m6 = sfromR (tproject1 (tproject1 m1)) * m5 ; m8 = sfromR (tproject2 m1) * sfromR dret ; m9 = (negate (sfromR (tproject2 m1)) * m4) * sfromR dret in tpair (tpair (rfromS (m2 * m9 + m5 * m8)) (rfromS (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m9) + cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m8)))) (rfromS ((m3 * m4) * sfromR dret + m6 * sfromR dret))"--testGradFooMatrixSimpPP :: Assertion-testGradFooMatrixSimpPP = do-  resetVarCounter-  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)-    in printArtifactPretty (simplifyArtifact $ revArtifactAdapt UseIncomingCotangent foo2 (FTKProduct (FTKProduct ftk ftk) ftk)))-      @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2,2] STKScalar)) (let m2 = sin (sfromR (tproject2 (tproject1 m1))) ; m3 = sfromR (tproject1 (tproject1 m1)) * m2 ; m4 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m3 * m3) ; m5 = sin (sfromR (tproject2 (tproject1 m1))) ; m8 = sfromR (tproject2 m1) * sfromR dret ; m9 = (negate (sfromR (tproject2 m1)) * m4) * sfromR dret in tpair (tpair (m2 * m9 + m5 * m8) (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m9) + cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m8))) ((m3 * m4) * sfromR dret + (sfromR (tproject1 (tproject1 m1)) * m5) * sfromR dret))"--gradFooScalar :: forall r. r ~ Double-              => (r, r, r) -> (r, r, r)-gradFooScalar = fromDValue . grad foo2 . fromValue--testGradFooScalar :: Assertion-testGradFooScalar =-  assertEqualUpToEpsilon 1e-10-    (2.4396285219055063, -1.953374825727421, 0.9654825811012627)-    (gradFooScalar (1.1, 2.2, 3.3))--gradCFooScalar :: forall r. r ~ Float-               => (r, r, r) -> (r, r, r)-gradCFooScalar = fromDValue . cgrad foo2 . fromValue--testGradCFooScalar :: Assertion-testGradCFooScalar =-  assertEqualUpToEpsilon 1e-10-    (2.4396284,-1.9533751,0.96548253)-    (gradCFooScalar (1.1, 2.2, 3.3))--testFooS :: Assertion-testFooS = do-  assertEqualUpToEpsilon 1e-10-    (srepl 2.4396285219055063, srepl (-1.953374825727421), srepl 0.9654825811012627)-    (grad (kfromS @_ @Double . ssum0 @'[3, 534, 3] @(TKScalar Double) . foo2) (srepl 1.1, srepl 2.2, srepl 3.3))--testFooSToFloat :: Assertion-testFooSToFloat = do-  assertEqualUpToEpsilon 1e-5-    (srepl 2.4396285219055063, srepl (-1.953374825727421), srepl 0.9654825811012627)-    (grad (kfromS @_ @Float . ssum0 . scast . foo2)-         (srepl 1.1 :: Concrete (TKS '[3, 534, 3] Double), srepl 2.2, srepl 3.3))--testFooSBoth :: Assertion-testFooSBoth = do-  assertEqualUpToEpsilon 1e-5-    (srepl 2.439628436155373, srepl (-1.9533749), srepl 0.9654825479484146)-    (grad (kfromS @_ @Float . ssum0 . scast . foo2 . (\(d, f, d2) -> (d, scast f, d2)))-         ( srepl 1.1 :: Concrete (TKS '[3, 534, 3] Double)-         , srepl 2.2 :: Concrete (TKS '[3, 534, 3] Float)-         , srepl 3.3 ))--testFooBoth :: Assertion-testFooBoth = do-  assertEqualUpToEpsilon 1e-5-    (rscalar 2.439628436155373, rscalar (-1.9533749), rscalar 0.9654825479484146)-    (grad (kfromR @_ @Float . rcast . foo2 . (\(d, f, d2) -> (d, rcast f, d2)))-         ( rscalar 1.1 :: Concrete (TKR 0 Double)-         , rscalar 2.2 :: Concrete (TKR 0 Float)-         , rscalar 3.3 ))---- Add arrays a,b,c, but shifting b and c one to left/right--- and then remove the first and last element.------ In PyTorch--- vstack( a[0] + b[1]---       , a[1:N-1] + b[2:N] + c[:N-2]---       , a[N-1] + c[N-2] )-vstackABC :: (ADReady target, GoodScalar r)-          => (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))-          -> target (TKR 1 r)-vstackABC (a, b, c) =-  let n = rwidth a-  in rconcat [ rreplicate 1 (a ! [0] + b ! [1])-             , rslice 1 (n - 2) a + rslice 2 (n - 2) b + rslice 0 (n - 2) c-             , rreplicate 1 (a ! [fromIntegral n - 1]-                             + c ! [fromIntegral n - 2]) ]--vstackBuild :: (ADReady target, GoodScalar r)-            => (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))-            -> target (TKR 1 r)-vstackBuild (a, b, c) =-  let n = rwidth a-  in rbuild1 n (\i ->-       ifH (i ==. 0)-           (a ! [0] + b ! [1])-           (ifH (i ==. fromIntegral n - 1)-                (a ! [fromIntegral n - 1] + c ! [fromIntegral n - 2])-                (a ! [i] + b ! [i + 1] + c ! [i - 1])))--testTrustVstackConcatRepl10 :: Assertion-testTrustVstackConcatRepl10 = do-  vstackABC @Concrete @Double (rrepl [10] 1, rrepl [10] 2, rrepl [10] 3)-  @?= rfromListLinear [10] [3.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,4.0]--testTrustVstackConcatIota10 :: Assertion-testTrustVstackConcatIota10 = do-  vstackABC @Concrete @Double (riota 10, riota 10, riota 10)-  @?= rfromListLinear [10] [1.0,3.0,6.0,9.0,12.0,15.0,18.0,21.0,24.0,17.0]--replIota :: (ADReady target, GoodScalar r)-         => Int -> (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))-replIota n =-  ( rconcrete (unConcrete $ rrepl [n] 1 * riota n)-  , rconcrete (unConcrete $ rrepl [n] 2 * riota n)-  , rconcrete (unConcrete $ rrepl [n] 3 * riota n) )--testTrustVstackConcatReplIota10 :: Assertion-testTrustVstackConcatReplIota10 = do-  vstackABC @Concrete @Double (replIota 10)-  @?= rfromListLinear [10] [2.0,5.0,11.0,17.0,23.0,29.0,35.0,41.0,47.0,33.0]--nN :: Int-nN = (round :: Double -> Int) 1e6  -- 1e6--trustedResult :: Concrete (TKR 1 Double)-trustedResult =-  rcast (vstackABC @Concrete @Float (replIota nN))-    -- the cast prevents computation sharing with the first test below--testVstackWarmup :: Assertion-testVstackWarmup = do-  trustedResult-  @?= trustedResult--testVstackConcatConcrete :: Assertion-testVstackConcatConcrete = do-  vstackABC @Concrete @Double (replIota nN)-  @?= trustedResult--testVstackBuildConcrete :: Assertion-testVstackBuildConcrete = do-  vstackBuild @Concrete @Double (replIota nN)-  @?= trustedResult--testVstackConcatAst :: Assertion-testVstackConcatAst = do-  interpretAstFull @Concrete-    emptyEnv-    (vstackABC @(AstTensor AstMethodLet FullSpan) @Double-               (replIota nN))-  @?= trustedResult--testVstackBuildAst :: Assertion-testVstackBuildAst = do-  interpretAstFull @Concrete-    emptyEnv-    (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double-                 (replIota nN))-  @?= trustedResult--testVstackBuildAstSimp :: Assertion-testVstackBuildAstSimp = do-  interpretAstFull @Concrete-    emptyEnv-      (simplifyInlineContract-         (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double-                      (replIota nN)))-  @?= trustedResult--testVstackBuildAstPP :: Assertion-testVstackBuildAstPP = do-  resetVarCounter-  let (var3, ast3) =-        funToAst (FTKProduct (FTKProduct (FTKR [10] FTKScalar) (FTKR [10] FTKScalar)) (FTKR [10] FTKScalar))-          Nothing (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double . fromTarget)-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstPretty ast3-    @?= "\\v1 -> rfromS (let x9 = sfromR (tproject1 (tproject1 v1)) !$ [0] + sfromR (tproject2 (tproject1 v1)) !$ [1] ; v10 = let x6 = sfromR (tproject1 (tproject1 v1)) !$ [9] + sfromR (tproject2 v1) !$ [8] ; v7 = (sfromR (tproject1 (tproject1 v1)) + sappend (sslice (SNat @1) (SNat @9) (sfromR (tproject2 (tproject1 v1)))) (sconcrete (sfromListLinear [1] [0.0]))) + sappend (sconcrete (sfromListLinear [1] [0.0])) (sslice (SNat @0) (SNat @9) (sfromR (tproject2 v1))) in sappend (sslice (SNat @0) (SNat @9) v7) (sreplicate @1 x6) in sappend (sreplicate @1 x9) (sslice (SNat @1) (SNat @9) v10))"-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstPretty (simplifyInlineContract ast3)-    @?= "\\v1 -> rfromS (sappend (sreplicate @1 (sfromR (tproject1 (tproject1 v1)) !$ [0] + sfromR (tproject2 (tproject1 v1)) !$ [1])) (sappend ((sslice (SNat @1) (SNat @8) (sfromR (tproject1 (tproject1 v1))) + sslice (SNat @2) (SNat @8) (sfromR (tproject2 (tproject1 v1)))) + sslice (SNat @0) (SNat @8) (sfromR (tproject2 v1))) (sreplicate @1 (sfromR (tproject1 (tproject1 v1)) !$ [9] + sfromR (tproject2 v1) !$ [8]))))"--{- The above is:-\v1 ->-  rfromS-    (sappend-       (sreplicate @1-          (sfromR (tproject1 (tproject1 v1)) !$ [0] +-           sfromR (tproject2 (tproject1 v1)) !$ [1]))-       (sappend-          ((sslice (SNat @1) (SNat @8) (sfromR (tproject1 (tproject1 v1))) +-            sslice (SNat @2) (SNat @8) (sfromR (tproject2 (tproject1 v1)))) +-           sslice (SNat @0) (SNat @8) (sfromR (tproject2 v1)))-          (sreplicate @1-             (sfromR (tproject1 (tproject1 v1)) !$ [9] +-              sfromR (tproject2 v1) !$ [8]))))--}--replIota2 :: (ADReady target, GoodScalar r)-          => Int -> (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))-replIota2 n =-  (rrepl [n] 1 * riota n, rrepl [n] 2 * riota n, rrepl [n] 3 * riota n)--testVstackConcatConcrete2 :: Assertion-testVstackConcatConcrete2 = do-  vstackABC @Concrete @Double (replIota nN)-  @?= trustedResult--testVstackBuildConcrete2 :: Assertion-testVstackBuildConcrete2 = do-  vstackBuild @Concrete @Double (replIota2 nN)-  @?= trustedResult--testVstackConcatAst2 :: Assertion-testVstackConcatAst2 = do-  interpretAstFull @Concrete-    emptyEnv-    (vstackABC @(AstTensor AstMethodLet FullSpan) @Double-               (replIota2 nN))-  @?= trustedResult--testVstackBuildAst2 :: Assertion-testVstackBuildAst2 = do-  interpretAstFull @Concrete-    emptyEnv-    (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double-                 (replIota2 nN))-  @?= trustedResult--testVstackBuildAstSimp2 :: Assertion-testVstackBuildAstSimp2 = do-  interpretAstFull @Concrete-    emptyEnv-      (simplifyInlineContract-         (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double-                      (replIota2 nN)))-  @?= trustedResult--testVstackBuildAstPP2 :: Assertion-testVstackBuildAstPP2 = do-  resetVarCounter-  (printAstPretty-     (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double-                  (replIota2 10)))-    @?= "rfromS (let x14 = sfromIntegral (sscalar 0) + sscalar 2.0 * sfromIntegral (sscalar 1) ; v15 = let x11 = sfromIntegral (sscalar 9) + sscalar 3.0 * sfromIntegral (sscalar 8) ; v12 = (siota (SNat @10) + (let v2 = sconcrete (sreplicate [10] 1) + siota (SNat @10) in sgather (sconcrete (sreplicate [10] 2.0)) (\\[i8] -> [kfromS (v2 !$ [i8])]) * sgather (sfromVector (fromList [sfromIntegral v2, sconcrete (sreplicate [10] 0.0)])) (\\[i9] -> [ifH (sscalar (-9) <=. negate (v2 !$ [i9])) 0 1, i9]))) + (let v3 = sconcrete (sreplicate [10] (-1)) + siota (SNat @10) in sgather (sconcrete (sreplicate [10] 3.0)) (\\[i6] -> [kfromS (v3 !$ [i6])]) * sgather (sfromVector (fromList [sfromIntegral v3, sconcrete (sreplicate [10] 0.0)])) (\\[i7] -> [ifH (sscalar 0 <=. v3 !$ [i7]) 0 1, i7])) in sappend (sslice (SNat @0) (SNat @9) v12) (sreplicate @1 x11) in sappend (sreplicate @1 x14) (sslice (SNat @1) (SNat @9) v15))"-  (printAstPretty-     (simplifyInlineContract-        (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double-                     (replIota2 10))))-    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [2.0])) (sappend (sconcrete (sfromListLinear [8] [1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0]) + (sgather (sconcrete (sreplicate [10] 2.0)) (\\[i23] -> [kfromS (sconcrete (sfromListLinear [10] [1,2,3,4,5,6,7,8,9,10]) !$ [1 + i23])]) * sgather (sfromVector (fromList [sfromIntegral (sconcrete (sfromListLinear [8] [2,3,4,5,6,7,8,9])), sconcrete (sreplicate [8] 0.0)])) (\\[i24] -> [ifH (sscalar (-9) <=. negate (sconcrete (sfromListLinear [10] [1,2,3,4,5,6,7,8,9,10]) !$ [1 + i24])) 0 1, i24]) + sgather (sconcrete (sreplicate [10] 3.0)) (\\[i21] -> [kfromS (sconcrete (sfromListLinear [10] [-1,0,1,2,3,4,5,6,7,8]) !$ [1 + i21])]) * sgather (sfromVector (fromList [sfromIntegral (sconcrete (sfromListLinear [8] [0,1,2,3,4,5,6,7])), sconcrete (sreplicate [8] 0.0)])) (\\[i22] -> [ifH (sscalar 0 <=. sconcrete (sfromListLinear [10] [-1,0,1,2,3,4,5,6,7,8]) !$ [1 + i22]) 0 1, i22]))) (sconcrete (sfromListLinear [1] [33.0]))))"-      -- TODO: this reduced to "rfromS (sconcrete (sfromListLinear [10] [2.0,5.0,11.0,17.0,23.0,29.0,35.0,41.0,47.0,33.0]))" when tshare/tlet were inserted less often--testFooPP :: Assertion-testFooPP = do-  resetVarCounter-  let fooT = foo2 @(AstTensor AstMethodLet FullSpan (TKR 0 Double))-      foo3 x = fooT (x, x, x)-      (var3, ast3) = funToAst (FTKR ZSR FTKScalar) Nothing foo3-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\x1 -> rfromS (atan2H (sfromR x1) (sfromR x1 * sin (sfromR x1)) + sfromR x1 * (sfromR x1 * sin (sfromR x1)))"-  resetVarCounter-  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fooT (FTKProduct (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar)) (FTKR ZSR FTKScalar))-  printArtifactSimple artifactRev-    @?= "\\dret x1 -> tlet (sin (sfromR (tproject2 (tproject1 x1)))) (\\x2 -> tlet (sfromR (tproject1 (tproject1 x1)) * x2) (\\x3 -> tlet (recip (sfromR (tproject2 x1) * sfromR (tproject2 x1) + x3 * x3)) (\\x4 -> tlet (sin (sfromR (tproject2 (tproject1 x1)))) (\\x5 -> tlet (sfromR (tproject1 (tproject1 x1)) * x5) (\\x6 -> tlet (sfromR (tproject2 x1) * sfromR dret) (\\x8 -> tlet ((negate (sfromR (tproject2 x1)) * x4) * sfromR dret) (\\x9 -> tpair (tpair (rfromS (x2 * x9 + x5 * x8)) (rfromS (cos (sfromR (tproject2 (tproject1 x1))) * (sfromR (tproject1 (tproject1 x1)) * x9) + cos (sfromR (tproject2 (tproject1 x1))) * (sfromR (tproject1 (tproject1 x1)) * x8)))) (rfromS ((x3 * x4) * sfromR dret + x6 * sfromR dret)))))))))"-  printArtifactPrimalSimple artifactRev-    @?= "\\x1 -> tlet (sin (sfromR (tproject2 (tproject1 x1)))) (\\x2 -> tlet (sfromR (tproject1 (tproject1 x1)) * x2) (\\x3 -> tlet (sin (sfromR (tproject2 (tproject1 x1)))) (\\x5 -> tlet (sfromR (tproject1 (tproject1 x1)) * x5) (\\x6 -> rfromS (atan2H (sfromR (tproject2 x1)) x3 + sfromR (tproject2 x1) * x6)))))"--fooLetOld :: forall target r n.-          (RealFloatH (target (TKR n r)), LetTensor target)-       => (target (TKR n r), target (TKR n r), target (TKR n r)) -> target (TKR n r)-fooLetOld (x, y, z) =-  let w0 = x * sin y-  in tlet w0 $ \w ->-     atan2H z w + z * w--testFooLet :: Assertion-testFooLet = do-  assertEqualUpToEpsilon 1e-10-    (rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627)-    (grad (kfromR @_ @Double . fooLetOld) (rscalar 1.1, rscalar 2.2, rscalar 3.3))--testFooLetPP :: Assertion-testFooLetPP = do-  resetVarCounter-  let fooLetT = fooLetOld @(AstTensor AstMethodLet FullSpan) @Double-      fooLet3 x = fooLetT (x, x, x)-      (var3, ast3) = funToAst (FTKR ZSR FTKScalar) Nothing fooLet3-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\x1 -> rfromS (tlet (sfromR x1 * sin (sfromR x1)) (\\x2 -> atan2H (sfromR x1) x2 + sfromR x1 * x2))"-  resetVarCounter-  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fooLetT (FTKProduct (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar)) (FTKR ZSR FTKScalar))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX))) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (STKS [] STKScalar)) (let x3 = sin (sfromR (tproject2 (tproject1 x1))) ; x4 = sfromR (tproject1 (tproject1 x1)) * x3 ; x5 = recip (sfromR (tproject2 x1) * sfromR (tproject2 x1) + x4 * x4) ; x7 = (negate (sfromR (tproject2 x1)) * x5) * sfromR dret + sfromR (tproject2 x1) * sfromR dret in tpair (tpair (x3 * x7) (cos (sfromR (tproject2 (tproject1 x1))) * (sfromR (tproject1 (tproject1 x1)) * x7))) ((x4 * x5) * sfromR dret + x4 * sfromR dret))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (let x4 = sfromR (tproject1 (tproject1 x1)) * sin (sfromR (tproject2 (tproject1 x1))) in atan2H (sfromR (tproject2 x1)) x4 + sfromR (tproject2 x1) * x4)"--shapedListProd :: forall k target r. (BaseTensor target, GoodScalar r)-               => ListR k (target (TKS '[] r)) -> target (TKS '[] r)-shapedListProd = foldr1 (*)--testListProdPP :: Assertion-testListProdPP = do-  resetVarCounter-  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKS '[] Double)) -> AstTensor AstMethodLet FullSpan (TKS '[] Double)-      fT = shapedListProd-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) ftkUnit))))-  printArtifactSimple artifactRev-    @?= "\\dret x1 -> tlet (tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1)))) (\\x2 -> tlet (tproject1 (tproject2 x1) * x2) (\\x3 -> tlet (tproject1 x1 * dret) (\\x5 -> tlet (tproject1 (tproject2 x1) * x5) (\\x6 -> tpair (x3 * dret) (tpair (x2 * x5) (tpair (tproject1 (tproject2 (tproject2 (tproject2 x1))) * x6) (tpair (tproject1 (tproject2 (tproject2 x1)) * x6) Z1)))))))"-  printArtifactPrimalSimple artifactRev-    @?= "\\x1 -> tlet (tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1)))) (\\x2 -> tlet (tproject1 (tproject2 x1) * x2) (\\x3 -> tproject1 x1 * x3))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> let x2 = tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1))) ; x5 = tproject1 x1 * dret ; x6 = tproject1 (tproject2 x1) * x5 in tpair ((tproject1 (tproject2 x1) * x2) * dret) (tpair (x2 * x5) (tpair (tproject1 (tproject2 (tproject2 (tproject2 x1))) * x6) (tpair (tproject1 (tproject2 (tproject2 x1)) * x6) Z1)))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> tproject1 x1 * (tproject1 (tproject2 x1) * (tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1)))))"--rankedListProdr :: forall k target r. (BaseTensor target, GoodScalar r)-                => ListR k (target (TKR 0 r)) -> target (TKR 0 r)-rankedListProdr = foldr1 (*)--testListProdrPP :: Assertion-testListProdrPP = do-  resetVarCounter-  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListProdr-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x2 = sfromR (tproject1 (tproject2 (tproject2 x1))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) ; x5 = sfromR (tproject1 x1) * sfromR dret ; x6 = sfromR (tproject1 (tproject2 x1)) * x5 in tpair ((sfromR (tproject1 (tproject2 x1)) * x2) * sfromR dret) (tpair (x2 * x5) (tpair (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x6) (tpair (sfromR (tproject1 (tproject2 (tproject2 x1))) * x6) Z1))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sfromR (tproject1 x1) * (sfromR (tproject1 (tproject2 x1)) * (sfromR (tproject1 (tproject2 (tproject2 x1))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"--testListProdrLongPP :: Assertion-testListProdrLongPP = do-  resetVarCounter-  let fT :: ListR 13 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListProdr-  let (artifactRev, _) =-        revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit)))))))))))))-  printArtifactSimple artifactRev-    @?= "\\dret x1 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))))) (\\x2 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x2) (\\x3 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x3) (\\x4 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x4) (\\x5 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x5) (\\x6 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x6) (\\x7 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x7) (\\x8 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x8) (\\x9 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x9) (\\x10 -> tlet (sfromR (tproject1 (tproject2 (tproject2 x1))) * x10) (\\x11 -> tlet (sfromR (tproject1 (tproject2 x1)) * x11) (\\x12 -> tlet (sfromR (tproject1 x1) * sfromR dret) (\\x14 -> tlet (sfromR (tproject1 (tproject2 x1)) * x14) (\\x15 -> tlet (sfromR (tproject1 (tproject2 (tproject2 x1))) * x15) (\\x16 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x16) (\\x17 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x17) (\\x18 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x18) (\\x19 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x19) (\\x20 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x20) (\\x21 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x21) (\\x22 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x22) (\\x23 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x23) (\\x24 -> tpair (rfromS (x12 * sfromR dret)) (tpair (rfromS (x11 * x14)) (tpair (rfromS (x10 * x15)) (tpair (rfromS (x9 * x16)) (tpair (rfromS (x8 * x17)) (tpair (rfromS (x7 * x18)) (tpair (rfromS (x6 * x19)) (tpair (rfromS (x5 * x20)) (tpair (rfromS (x4 * x21)) (tpair (rfromS (x3 * x22)) (tpair (rfromS (x2 * x23)) (tpair (rfromS (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))))) * x24)) (tpair (rfromS (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * x24)) Z1))))))))))))))))))))))))))))))))))"-  printArtifactPrimalSimple artifactRev-    @?= "\\x1 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))))) (\\x2 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x2) (\\x3 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x3) (\\x4 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x4) (\\x5 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x5) (\\x6 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x6) (\\x7 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x7) (\\x8 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x8) (\\x9 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x9) (\\x10 -> tlet (sfromR (tproject1 (tproject2 (tproject2 x1))) * x10) (\\x11 -> tlet (sfromR (tproject1 (tproject2 x1)) * x11) (\\x12 -> rfromS (sfromR (tproject1 x1) * x12))))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId))))))))))))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar))))))))))))) (let x2 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))))) ; x3 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x2 ; x4 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x3 ; x5 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x4 ; x6 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x5 ; x7 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x6 ; x8 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x7 ; x9 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x8 ; x10 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x9 ; x11 = sfromR (tproject1 (tproject2 (tproject2 x1))) * x10 ; x14 = sfromR (tproject1 x1) * sfromR dret ; x15 = sfromR (tproject1 (tproject2 x1)) * x14 ; x16 = sfromR (tproject1 (tproject2 (tproject2 x1))) * x15 ; x17 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x16 ; x18 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x17 ; x19 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x18 ; x20 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x19 ; x21 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x20 ; x22 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x21 ; x23 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x22 ; x24 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x23 in tpair ((sfromR (tproject1 (tproject2 x1)) * x11) * sfromR dret) (tpair (x11 * x14) (tpair (x10 * x15) (tpair (x9 * x16) (tpair (x8 * x17) (tpair (x7 * x18) (tpair (x6 * x19) (tpair (x5 * x20) (tpair (x4 * x21) (tpair (x3 * x22) (tpair (x2 * x23) (tpair (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))))) * x24) (tpair (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * x24) Z1)))))))))))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sfromR (tproject1 x1) * (sfromR (tproject1 (tproject2 x1)) * (sfromR (tproject1 (tproject2 (tproject2 x1))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))))))))))))))))"--testListProd :: Assertion-testListProd = do-  assertEqualUpToEpsilon 1e-10-    [srepl 24, srepl 12, srepl 8, srepl 6]-    (grad (kfromS @_ @Double . shapedListProd @4) [srepl 1, srepl 2, srepl 3, srepl 4])--testListProdr :: Assertion-testListProdr = do-  assertEqualUpToEpsilon 1e-10-    [rscalar 24, rscalar 12, rscalar 8, rscalar 6]-    (grad (kfromR @_ @Double . rankedListProdr @4) [rscalar 1, rscalar 2, rscalar 3, rscalar 4])--rankedListSumr :: (BaseTensor target, GoodScalar r)-               => ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)-rankedListSumr = foldr1 (+)--testListSumrPP :: Assertion-testListSumrPP = do-  resetVarCounter >> resetIdCounter-  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListSumr-  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tpair dret (tpair dret (tpair dret (tpair dret Z1)))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sfromR (tproject1 x1) + (sfromR (tproject1 (tproject2 x1)) + (sfromR (tproject1 (tproject2 (tproject2 x1))) + sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"-  show deltas-    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000003 (DeltaAdd (DeltaConvert (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (DeltaInput (InputId 0))) (DeltaShare 100000002 (DeltaAdd (DeltaConvert (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (DeltaInput (InputId 1))) (DeltaShare 100000001 (DeltaAdd (DeltaConvert (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (DeltaInput (InputId 2))) (DeltaConvert (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (DeltaInput (InputId 3)))))))))"---- Note that the function is not associative, so foldr vs foldl matters.-rankedListSum2r :: (BaseTensor target, GoodScalar r)-                => ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)-rankedListSum2r = foldr1 (\x y -> x + rscalar 2 * y)--testListSum2rPP :: Assertion-testListSum2rPP = do-  resetVarCounter-  let fT ::  ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListSum2r-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKR (SNat @0) STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x5 = sscalar 2.0 * sfromR dret ; x6 = sscalar 2.0 * x5 in tpair dret (tpair x5 (tpair x6 (tpair (sscalar 2.0 * x6) Z1))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sfromR (tproject1 x1) + (sscalar 2.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 4.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"--rankedListSum22r :: (BaseTensor target, GoodScalar r)-                 => ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)-rankedListSum22r = foldr1 (\x y -> rscalar 2 * x + rscalar 2 * y)--testListSum22rPP :: Assertion-testListSum22rPP = do-  resetVarCounter-  let fT ::  ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListSum22r-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x5 = sscalar 2.0 * sfromR dret ; x6 = sscalar 2.0 * x5 in tpair (sscalar 2.0 * sfromR dret) (tpair (sscalar 2.0 * x5) (tpair (sscalar 2.0 * x6) (tpair (sscalar 2.0 * x6) Z1))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + (sscalar 4.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"---- Note how this tlet did not change anything, in particular the sharing.-rankedListSumk22r :: ( BaseTensor target, LetTensor target-                     , GoodScalar r )-                 =>  ListR k (target (TKR 0 r)) -> target (TKR 0 r)-rankedListSumk22r = foldr1 (\x y -> tlet (rscalar 2) (\k -> k * x + k * y))--testListSumk22rPP :: Assertion-testListSumk22rPP = do-  resetVarCounter-  let fT ::  ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListSumk22r-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x5 = sscalar 2.0 * sfromR dret ; x6 = sscalar 2.0 * x5 in tpair (sscalar 2.0 * sfromR dret) (tpair (sscalar 2.0 * x5) (tpair (sscalar 2.0 * x6) (tpair (sscalar 2.0 * x6) Z1))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + (sscalar 4.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"--rankedListSum2xpyr :: (BaseTensor target, GoodScalar r)-                   =>  ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)-rankedListSum2xpyr = foldr1 (\x y -> rscalar 2 * (x + y))--testListSum2xpyrPP :: Assertion-testListSum2xpyrPP = do-  resetVarCounter-  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListSum2xpyr-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x6 = sscalar 2.0 * sfromR dret ; x7 = sscalar 2.0 * x6 ; x8 = sscalar 2.0 * x7 in tpair x6 (tpair x7 (tpair x8 (tpair x8 Z1))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + (sscalar 4.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"--rankedListSum2xyr :: (BaseTensor target, GoodScalar r)-                  => ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)-rankedListSum2xyr = foldr1 (\x y -> rscalar 2 * (x * y))--testListSum2xyrPP :: Assertion-testListSum2xyrPP = do-  resetVarCounter-  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListSum2xyr-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x3 = sscalar 2.0 * (sfromR (tproject1 (tproject2 (tproject2 x1))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1))))) ; x8 = sscalar 2.0 * sfromR dret ; x9 = sscalar 2.0 * (sfromR (tproject1 x1) * x8) ; x10 = sscalar 2.0 * (sfromR (tproject1 (tproject2 x1)) * x9) in tpair (sscalar 2.0 * ((sfromR (tproject1 (tproject2 x1)) * x3) * x8)) (tpair (x3 * x9) (tpair (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x10) (tpair (sfromR (tproject1 (tproject2 (tproject2 x1))) * x10) Z1))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sscalar 8.0 * (sfromR (tproject1 x1) * (sfromR (tproject1 (tproject2 x1)) * (sfromR (tproject1 (tproject2 (tproject2 x1))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1))))))))"--ranked2xy :: (BaseTensor target, GoodScalar r)-          => (target (TKR 0 r), target (TKR 0 r)) -> target (TKR 0 r)-ranked2xy = \(x, y) -> rscalar 2 * x * y--test2xyPP :: Assertion-test2xyPP = do-  resetVarCounter-  let fT :: (AstTensor AstMethodLet FullSpan (TKR 0 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = ranked2xy-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair (sscalar 2.0 * (sfromR (tproject2 x1) * sfromR dret)) (sscalar 2.0 * (sfromR (tproject1 x1) * sfromR dret)))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sscalar 2.0 * (sfromR (tproject1 x1) * sfromR (tproject2 x1)))"---- Note that the function is not associative, so foldr vs foldl matters.-rankedListSum23r :: forall k target r. (BaseTensor target, GoodScalar r)-                 => ListR k (target (TKR 0 r)) -> target (TKR 0 r)-rankedListSum23r = foldr1 (\x y -> rscalar 2 * x + rscalar 3 * y)--testListSum23rPP :: Assertion-testListSum23rPP = do-  resetVarCounter-  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = rankedListSum23r-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x5 = sscalar 3.0 * sfromR dret ; x6 = sscalar 3.0 * x5 in tpair (sscalar 2.0 * sfromR dret) (tpair (sscalar 2.0 * x5) (tpair (sscalar 2.0 * x6) (tpair (sscalar 3.0 * x6) Z1))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + (sscalar 6.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 18.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 27.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"--ranked23 :: (BaseTensor target, GoodScalar r)-         => (target (TKR 0 r), target (TKR 0 r)) -> target (TKR 0 r)-ranked23 = \(x, y) -> rscalar 2 * x + rscalar 3 * y--test23PP :: Assertion-test23PP = do-  resetVarCounter-  let fT :: (AstTensor AstMethodLet FullSpan (TKR 0 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)-      fT = ranked23-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair (sscalar 2.0 * sfromR dret) (sscalar 3.0 * sfromR dret))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + sscalar 3.0 * sfromR (tproject2 x1))"--reluPrimal-  :: forall target n r.-     (ADReady target, GoodScalar r, KnownNat n, Differentiable r)-  => target (TKR n r) -> target (TKR n r)-reluPrimal v =-  let oneIfGtZero = rmap0N (\x -> ifH (x <=. rscalar 0) (rscalar 0.0) (rscalar 1.0))-                           (rprimalPart v)-  in scale2 oneIfGtZero v--scale2 :: forall target r n.-          (ADReady target, GoodScalar r, KnownNat n)-       => PrimalOf target (TKR n r) -> target (TKR n r) -> target (TKR n r)-scale2 a d = rfromPrimal @target a * d--testReluPP :: Assertion-testReluPP = do-  resetVarCounter >> resetIdCounter-  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)-      reluT = reluPrimal-      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing reluT-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\m1 -> rfromS (tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i4, i3] -> [ifH (sscalar -0.0 <=. negate (tprimalPart (sfromR m1) !$ [i4, i3])) 0 1])) * sfromR m1)"-  resetVarCounter-  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i5, i6])) 0 1]) * sfromR dret)"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i5, i6])) 0 1]) * sfromR m1)"-  show deltas-    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000004 (DeltaScale <primal> (DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0)))))"--testReluPP2 :: Assertion-testReluPP2 = do-  resetVarCounter-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)-      reluT2 (t, r) = reluPrimal (t * rreplicate 5 r)-      (var3, ast3) = funToAst (FTKR [5] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\v1 -> rfromS (tfromPrimal (STKS [5] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i2] -> [ifH (sscalar -0.0 <=. sscalar (-7.0) * tprimalPart (sfromR v1) !$ [i2]) 0 1])) * (sfromR v1 * tfromPrimal (STKS [5] STKScalar) (sconcrete (sreplicate [5] 7.0))))"-  resetVarCounter-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [5] STKScalar) (STKS [] STKScalar)) (let v9 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 v1) !$ [i3]) * sfromR (tproject2 v1)) 0 1]) * sfromR dret in tpair (sreplicate @5 (sfromR (tproject2 v1)) * v9) (sdot0 (sfromR (tproject1 v1)) v9))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\v1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 v1) !$ [i3]) * sfromR (tproject2 v1)) 0 1]) * (sfromR (tproject1 v1) * sreplicate @5 (sfromR (tproject2 v1))))"--testReluSimpler :: Assertion-testReluSimpler = do-  assertEqualUpToEpsilon' 1e-10-    (ringestData [3, 4] [1.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,0.0,0.0,1.0,1.0])-    (rev' @Double @2 relu (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12]))--testReluSimplerPP :: Assertion-testReluSimplerPP = do-  resetVarCounter >> resetIdCounter-  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)-      reluT = relu-      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing reluT-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\m1 -> rfromS (tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i4, i3] -> [ifH (sscalar -0.0 <=. negate (tprimalPart (sfromR m1) !$ [i4, i3])) 0 1])) * sfromR m1)"-  resetVarCounter-  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i5, i6])) 0 1]) * sfromR dret)"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i5, i6])) 0 1]) * sfromR m1)"-  show deltas-    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000004 (DeltaScale <primal> (DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0)))))"--testReluSimplerPP2 :: Assertion-testReluSimplerPP2 = do-  resetVarCounter-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)-      reluT2 (t, r) = relu (t * rreplicate 5 r)-      (var3, ast3) = funToAst (FTKR [5] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\v1 -> rfromS (tlet (sfromR v1 * tfromPrimal (STKS [5] STKScalar) (sconcrete (sreplicate [5] 7.0))) (\\v2 -> tfromPrimal (STKS [5] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3] -> [ifH (sscalar -0.0 <=. negate (sfromR (tprimalPart (rfromS (v2 !$ [i3]))))) 0 1])) * v2))"-  resetVarCounter-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [5] STKScalar) (STKS [] STKScalar)) (let v8 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 v1) !$ [i5]) * sfromR (tproject2 v1)) 0 1]) * sfromR dret in tpair (sreplicate @5 (sfromR (tproject2 v1)) * v8) (sdot0 (sfromR (tproject1 v1)) v8))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\v1 -> rfromS (let v4 = sfromR (tproject1 v1) * sreplicate @5 (sfromR (tproject2 v1)) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5] -> [ifH (sscalar -0.0 <=. negate (v4 !$ [i5])) 0 1]) * v4)"--testReluSimplerPP3 :: Assertion-testReluSimplerPP3 = do-  resetVarCounter-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)-      reluT2 (t, r) = relu (t * rreplicate 3 (rreplicate 4 r))-      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\m1 -> rfromS (tlet (sfromR m1 * tfromPrimal (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0))) (\\m2 -> tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (sscalar -0.0 <=. negate (tprimalPart m2 !$ [i5, i4])) 0 1])) * m2))"-  resetVarCounter-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [3, 4] FTKScalar) (FTKR ZSR FTKScalar))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [3,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [3,4] STKScalar) (STKS [] STKScalar)) (let m11 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 m1) !$ [i7, i8]) * sfromR (tproject2 m1)) 0 1]) * sfromR dret in tpair (sreplicate @3 (sreplicate @4 (sfromR (tproject2 m1))) * m11) (sdot0 (sfromR (tproject1 m1)) m11))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (let m6 = sfromR (tproject1 m1) * sreplicate @3 (sreplicate @4 (sfromR (tproject2 m1))) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) * m6)"--testReluSimpler3 :: Assertion-testReluSimpler3 = do-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)-      reluT2 (t, r) = relu (t * rreplicate 3 (rreplicate 4 r))-  assertEqualUpToEpsilon 1e-10-    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]-    , rscalar 57.1 )-    (grad (kfromR . rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))--testReluSimplerPP4 :: Assertion-testReluSimplerPP4 = do-  resetVarCounter-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)-      reluT2 (t, r) = relu (t * rreplicate0N [3, 4] r)-      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\m1 -> rfromS (tlet (sfromR m1 * tfromPrimal (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0))) (\\m2 -> tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (sscalar -0.0 <=. negate (tprimalPart m2 !$ [i5, i4])) 0 1])) * m2))"-  resetVarCounter-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [3, 4] FTKScalar) (FTKR ZSR FTKScalar))-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [3,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX))) (STKProduct (STKS [3,4] STKScalar) (STKS [] STKScalar)) (let m11 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 m1) !$ [i7, i8]) * sfromR (tproject2 m1)) 0 1]) * sfromR dret in tpair (sreplicate @3 (sreplicate @4 (sfromR (tproject2 m1))) * m11) (sdot0 (sfromR (tproject1 m1)) m11))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (let m6 = sfromR (tproject1 m1) * sreplicate @3 (sreplicate @4 (sfromR (tproject2 m1))) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) * m6)"--testReluSimpler4 :: Assertion-testReluSimpler4 = do-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)-      reluT2 (t, r) = relu (t * rreplicate0N [3, 4] r)-  assertEqualUpToEpsilon 1e-10-    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]-    , rscalar 57.1 )-    (grad (kfromR . rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))--testReluSimplerPP4s :: Assertion-testReluSimplerPP4s = do-  resetVarCounter-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Float), AstTensor AstMethodLet FullSpan (TKS '[] Float))-             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Float)-      reluT2 (t, r) = reluS (t * sreplicate0N r)-      (var3, ast3) = funToAst (FTKS knownShS FTKScalar) Nothing (\t -> reluT2 (t, srepl 7))-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\m1 -> tlet (m1 * tfromPrimal (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0))) (\\m2 -> tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (sscalar -0.0 <=. negate (tprimalPart m2 !$ [i5, i4])) 0 1])) * m2)"--testReluSimplerPP4s2 :: Assertion-testReluSimplerPP4s2 = do-  resetVarCounter-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))-             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)-      -- This is tweaked compared to above to avoid test artifacts coming-      -- from counter resets, which are inherently unsafe (cse, etc.).-      reluT2 (t, r) = reluS (t * sreplicate0N r)-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKS [3, 4] FTKScalar) (FTKS ZSS FTKScalar))-  printArtifactPretty artifactRev-    @?= "\\dret m1 -> let m6 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) ; m9 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) ; m11 = m9 * dret in tpair (sreplicate @3 (sreplicate @4 (tproject2 m1)) * m11) (ssum @4 (ssum @3 (tproject1 m1 * m11)))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> let m6 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) ; m9 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) in m9 * m6"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> let m11 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (tproject1 m1 !$ [i7, i8]) * tproject2 m1) 0 1]) * dret in tpair (sreplicate @3 (sreplicate @4 (tproject2 m1)) * m11) (sdot0 (tproject1 m1) m11)"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> let m6 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) * m6"--testReluSimpler4s :: Assertion-testReluSimpler4s = do-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))-             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)-      reluT2 (t, r) = reluS (t * sreplicate0N r)-  assertEqualUpToEpsilon 1e-10-    ( sconcrete-      $ Nested.sfromListPrimLinear @_ @'[3, 4] knownShS [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]-    , srepl 57.1 )-    (grad (kfromS . ssum0 . reluT2) (sconcrete $ Nested.sfromListPrimLinear @_ @'[3, 4] knownShS [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], srepl 7))--testReluSimplerPP7s2 :: Assertion-testReluSimplerPP7s2 = do-  resetVarCounter-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))-             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)-      reluT2 (t, r) = smap0N (flip tlet $ \x -> ifH (x <=. sscalar 0) (sscalar 0.0) x) (t * sreplicate0N r)-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKS [3, 4] FTKScalar) (FTKS ZSS FTKScalar))-  printArtifactPretty artifactRev-    @?= "\\dret m1 -> let m9 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) ; t15 = sscatter dret (\\[i13, i14] -> [ifH (sscalar -0.0 <=. negate (m9 !$ [i13, i14])) 0 1, i13, i14]) ; m16 = t15 !$ [1] in tpair (sreplicate @3 (sreplicate @4 (tproject2 m1)) * m16) (ssum @4 (ssum @3 (tproject1 m1 * m16)))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> let m9 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) in sgather (sfromVector (fromList [sconcrete (sreplicate [3,4] 0.0), m9])) (\\[i10, i11] -> [ifH (sscalar -0.0 <=. negate (m9 !$ [i10, i11])) 0 1, i10, i11])"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> let m16 = sscatter dret (\\[i13, i14] -> [ifH (sscalar -0.0 <=. negate (tproject1 m1 !$ [i13, i14]) * tproject2 m1) 0 1, i13, i14]) !$ [1] in tpair (sreplicate @3 (sreplicate @4 (tproject2 m1)) * m16) (sdot0 (tproject1 m1) m16)"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> let m9 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) in sgather (sfromVector (fromList [sconcrete (sreplicate [3,4] 0.0), m9])) (\\[i10, i11] -> [ifH (sscalar -0.0 <=. negate (m9 !$ [i10, i11])) 0 1, i10, i11])"--testReluSimpler7s :: Assertion-testReluSimpler7s = do-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))-             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)-      reluT2 (t, r) = smap0N (flip tlet $ \x -> ifH (x <=. sscalar 0) (sscalar 0.0) x) (t * sreplicate0N r)-  assertEqualUpToEpsilon 1e-10-    ( sconcrete-      $ Nested.sfromListPrimLinear @_ @'[3, 4] knownShS [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]-    , srepl 57.1 )-    (grad (kfromS . ssum0 . reluT2) (sconcrete $ Nested.sfromListPrimLinear @_ @'[3, 4] knownShS [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], srepl 7))--reluMax :: forall target n r. (ADReady target, GoodScalar r, KnownNat n)-        => target (TKR n r) -> target (TKR n r)-reluMax = rmap0N (maxH (rscalar 0))--testReluMax :: Assertion-testReluMax = do-  assertEqualUpToEpsilon' 1e-10-    (ringestData [3, 4] [1.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,0.0,0.0,1.0,1.0])-    (rev' @Double @2 reluMax (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12]))--testReluMaxPP :: Assertion-testReluMaxPP = do-  resetVarCounter >> resetIdCounter-  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)-      reluT = reluMax-      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing reluT-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\m1 -> rfromS (sgather [] (tfromVector (SNat @2) (STKS [3,4] STKScalar) (fromList [tfromPrimal (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 0.0)), sfromR m1])) (\\[i5, i4] -> [ifH (sscalar -0.0 <=. negate (tprimalPart (sfromR m1) !$ [i5, i4])) 0 1, i5, i4]))"-  resetVarCounter-  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> rfromS (sscatter (sfromR dret) (\\[i9, i10] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i9, i10])) 0 1, i9, i10]) !$ [1])"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (sgather (sfromVector (fromList [sconcrete (sreplicate [3,4] 0.0), sfromR m1])) (\\[i6, i7] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i6, i7])) 0 1, i6, i7]))"-  show deltas-    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000004 (DeltaGatherS [3,4] [] [2,3,4] (DeltaShare 100000001 (DeltaFromVector (SNat @2) (STKS [3,4] STKScalar) [DeltaZero (FTKS [3,4] FTKScalar),DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0))])) <function>))"--testReluMaxPP2 :: Assertion-testReluMaxPP2 = do-  resetVarCounter-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)-      reluT2 (t, r) = reluMax (t * rreplicate 5 r)-      (var3, ast3) = funToAst (FTKR [5] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\v1 -> rfromS (sgather [] (tfromVector (SNat @2) (STKS [5] STKScalar) (fromList [tfromPrimal (STKS [5] STKScalar) (sconcrete (sreplicate [5] 0.0)), sfromR v1 * tfromPrimal (STKS [5] STKScalar) (sconcrete (sreplicate [5] 7.0))])) (\\[i3] -> [ifH (sscalar -0.0 <=. sscalar (-7.0) * tprimalPart (sfromR v1) !$ [i3]) 0 1, i3]))"-  resetVarCounter-  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))-  printArtifactPretty artifactRev-    @?= "\\dret v1 -> let m11 = sscatter (sfromR dret) (\\[i8] -> [let x9 = negate (sfromR (tproject1 v1) !$ [i8]) ; x10 = sfromR (rreplicate 5 (tproject2 v1)) !$ [i8] in ifH (sscalar -0.0 <=. x9 * x10) 0 1, i8]) ; v12 = m11 !$ [1] in tpair (rfromS (sfromR (rreplicate 5 (tproject2 v1)) * v12)) (rsum (rfromS (sfromR (tproject1 v1) * v12)))"-  printArtifactPrimalPretty artifactRev-    @?= "\\v1 -> rfromS (sgather (sfromVector (fromList [sconcrete (sreplicate [5] 0.0), sfromR (tproject1 v1) * sfromR (rreplicate 5 (tproject2 v1))])) (\\[i4] -> [let x5 = negate (sfromR (tproject1 v1) !$ [i4]) ; x6 = sfromR (rreplicate 5 (tproject2 v1)) !$ [i4] in ifH (sscalar -0.0 <=. x5 * x6) 0 1, i4]))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [5] STKScalar) (STKS [] STKScalar)) (let v12 = sscatter (sfromR dret) (\\[i8] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 v1) !$ [i8]) * sfromR (tproject2 v1)) 0 1, i8]) !$ [1] in tpair (sreplicate @5 (sfromR (tproject2 v1)) * v12) (sdot0 (sfromR (tproject1 v1)) v12))"--testReluMax3 :: Assertion-testReluMax3 = do-  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))-             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)-      reluT2 (t, r) = reluMax (t * rreplicate 3 (rreplicate 4 r))-  assertEqualUpToEpsilon 1e-10-    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]-    , rscalar 57.1 )-    (grad (kfromR . rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))--testDot1PP :: Assertion-testDot1PP = do-  resetVarCounter-  let (artifactRev, _) =-        revArtifactDelta UseIncomingCotangent (uncurry (rdot0 @1 @Double))-                 (FTKProduct (FTKR [3] FTKScalar) (FTKR [3] FTKScalar))-  printArtifactPretty artifactRev-    @?= "\\dret v1 -> tpair (rfromS (sfromR (tproject2 v1) * sreplicate @3 (sfromR dret))) (rfromS (sfromR (tproject1 v1) * sreplicate @3 (sfromR dret)))"-  printArtifactPrimalPretty artifactRev-    @?= "\\v1 -> rfromS (ssum @3 (sfromR (tproject1 v1) * sfromR (tproject2 v1)))"--testDot2PP :: Assertion-testDot2PP = do-  resetVarCounter-  let (artifactRev, _deltas) =-        revArtifactDelta UseIncomingCotangent (uncurry (rdot0 @2 @Double))-                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [2, 3] FTKScalar))-  printArtifactPretty artifactRev-    @?= "\\dret m1 -> let m3 = sreshape @[2,3] (sreplicate @6 (sfromR dret)) in tpair (rfromS (sfromR (tproject2 m1) * m3)) (rfromS (sfromR (tproject1 m1) * m3))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> rfromS (ssum @6 (sreshape @[6] (sfromR (tproject1 m1) * sfromR (tproject2 m1))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,3] FTKScalar)) ConvSX))) (STKProduct (STKS [2,3] STKScalar) (STKS [2,3] STKScalar)) (tpair (sfromR (tproject2 m1) * sreplicate @2 (sreplicate @3 (sfromR dret))) (sfromR (tproject1 m1) * sreplicate @2 (sreplicate @3 (sfromR dret))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (sdot0 (sfromR (tproject1 m1)) (sfromR (tproject2 m1)))"--testMatvecmulPP :: Assertion-testMatvecmulPP = do-  resetVarCounter-  let (artifactRev, _) =-        revArtifactDelta-                 UseIncomingCotangent (uncurry rmatvecmul)-                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3] FTKScalar))-  printArtifactPretty @_ @(TKR 1 Double) artifactRev-    @?= "\\dret m1 -> tpair (rfromS (str (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret)))) (rfromS (ssum @2 (str (str (sfromR (tproject1 m1)) * sreplicate @3 (sfromR dret)))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> rfromS (ssum @3 (str (sreplicate @2 (sfromR (tproject2 m1))) * str (sfromR (tproject1 m1))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [3] FTKScalar)) ConvSX))) (STKProduct (STKS [2,3] STKScalar) (STKS [3] STKScalar)) (tpair (sreplicate @2 (sfromR (tproject2 m1)) * str (sreplicate @3 (sfromR dret))) (sdot1In (str (sfromR (tproject1 m1))) (sreplicate @3 (sfromR dret))))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (sdot1In (sreplicate @2 (sfromR (tproject2 m1))) (sfromR (tproject1 m1)))"--testMatmul2PP :: Assertion-testMatmul2PP = do-  resetVarCounter-  let (artifactRev, _) =-        revArtifactDelta-                 UseIncomingCotangent (uncurry rmatmul2)-                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))-  printArtifactPretty @_ @(TKR 2 Double) artifactRev-    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2,1,0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [3,4] FTKScalar)) ConvSX))) (STKProduct (STKS [2,3] STKScalar) (STKS [3,4] STKScalar)) (tpair (smatmul2 (sfromR dret) (str (sfromR (tproject2 m1)))) (smatmul2 (str (sfromR (tproject1 m1))) (sfromR dret)))"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (smatmul2 (sfromR (tproject1 m1)) (sfromR (tproject2 m1)))"--testMatmul2FromMatvecmulPP :: Assertion-testMatmul2FromMatvecmulPP = do-  resetVarCounter-  let rmatmul2F :: (BaseTensor target, GoodScalar r)-                => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)-      rmatmul2F m1 m2 =-        rbuild1 (rwidth m1) (\i -> rmatvecmul (rtr m2) (m1 ! [i]))-      (artifactRev, _) =-        revArtifactDelta-                 UseIncomingCotangent (uncurry rmatmul2F)-                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))-  printArtifactPretty @_ @(TKR 2 Double) artifactRev-    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2,1,0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"--testMatmul2PaperPP :: Assertion-testMatmul2PaperPP = do-  resetVarCounter-  let rmatmul2P :: (BaseTensor target, GoodScalar r)-                => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)-      rmatmul2P a b =-        let k :$: m :$: _ = rshape a-            _ :$: n :$: _ = rshape b-        in rbuild1 k (\i ->-             rbuild1 n (\j ->-               rsum (rbuild1 m (\p -> a ! [i, p] * b ! [p, j]))))-      (artifactRev, _) =-        revArtifactDelta-                 UseIncomingCotangent (uncurry rmatmul2P)-                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))-  printArtifactPretty @_ @(TKR 2 Double) artifactRev-    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2,1,0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"--testMatmul2PPS :: Assertion-testMatmul2PPS = do-  resetVarCounter-  let (artifactRev, _) =-        revArtifactDelta-                 UseIncomingCotangent (uncurry smatmul2)-                 (FTKProduct (FTKS (SNat @2 :$$ SNat @3 :$$ ZSS) (FTKScalar @Float)) (FTKS (SNat @3 :$$ SNat @4 :$$ ZSS) (FTKScalar @Float)))-  printArtifactPretty artifactRev-    @?= "\\dret m1 -> tpair (ssum @4 (stranspose @[2,1,0] (str (sreplicate @2 (tproject2 m1)) * sreplicate @3 dret))) (ssum @2 (str (stranspose @[2,1,0] (sreplicate @4 (tproject1 m1)) * sreplicate @3 dret)))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> ssum @3 (stranspose @[2,1,0] (sreplicate @4 (tproject1 m1)) * str (sreplicate @2 (tproject2 m1)))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tpair (smatmul2 dret (str (tproject2 m1))) (smatmul2 (str (tproject1 m1)) dret)"-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> smatmul2 (tproject1 m1) (tproject2 m1)"--testAddSpeedBig :: Assertion-testAddSpeedBig =-  let m1 :: Concrete (TKR 2 Double)-      m1 = ringestData [1000,1000] [1 .. 1000000]-      m2 = m1 + m1-  in assertEqualUpToEpsilon 1e-10 m2 m2--testMatmul2SpeedBig :: Assertion-testMatmul2SpeedBig =-  let m1 :: Concrete (TKR 2 Double)-      m1 = ringestData [1000,1000] [1 .. 1000000]-      m2 = rmatmul2 m1 m1-  in assertEqualUpToEpsilon 1e-10 m2 m2--bar :: forall a. RealFloatH a => (a, a) -> a-bar (x, y) =-  let w = foo2 (x, y, x) * sin y-  in atan2H x w + y * w--barF :: forall a. RealFloatH a => (a, a) -> a-barF (x, y) =-  let w = foo2 (x, y, x) * sin y-  in atan2H x w + y * w--testBar :: Assertion-testBar =-  assertEqualUpToEpsilon 1e-9-    (rscalar 3.1435239435581166,rscalar (-1.1053869545195814))-    (cgrad (kfromR . bar @(ADVal Concrete (TKR 0 Double))) (rscalar 1.1, rscalar 2.2))--testBarS :: Assertion-testBarS =-  assertEqualUpToEpsilon 1e-9-    (srepl 3.1435239435581166, srepl (-1.1053869545195814))-    (cgrad (kfromS . barF @(ADVal Concrete (TKS '[] Double))) (srepl 1.1, srepl 2.2))--testBar2S :: Assertion-testBar2S =-  assertEqualUpToEpsilon 1e-9-    (srepl 3.1435239435581166, srepl (-1.1053869545195814))-    (grad (kfromS . ssum0 . barF @(AstTensor AstMethodLet FullSpan (TKS '[52, 2, 2, 1, 1, 3] Double))) (srepl 1.1, srepl 2.2))--testBarCFwd :: Assertion-testBarCFwd =-  assertEqualUpToEpsilon 1e-9-    (rscalar 9.327500345189534e-2)-    (cjvp (bar @(ADVal Concrete (TKR 0 Double))) (rscalar 1.1, rscalar 2.2) (rscalar 0.1, rscalar 0.2))--testBarFwd :: Assertion-testBarFwd =-  assertEqualUpToEpsilon 1e-9-    (rscalar 9.327500345189534e-2)-    (jvp @_ @(TKR 0 Double)-         bar (rscalar 1.1, rscalar 2.2) (rscalar 0.1, rscalar 0.2))--barADVal2 :: forall a. RealFloatH a-          => (a, a, a) -> a-barADVal2 (x,y,z) =-  let w = foo2 (x,y,z) * sin y-  in atan2H z w + z * w---- A check if gradient computation is re-entrant.--- This test fails (not on first run, but on repetition) if old terms,--- from before current iteration of gradient descent, are not soundly--- handled in new computations (due to naive impurity, TH, plugins,--- or just following the papers that assume a single isolated run).--- This example requires monomorphic types and is contrived,--- but GHC does optimize and factor out some accidentally constant--- subterms in real examples (presumably after it monomorphizes them)--- causing exactly the same danger.--- This example also tests unused parameters (x), another common cause--- of crashes in naive gradient computing code.-baz :: ( ADVal Concrete (TKR 0 Double)-       , ADVal Concrete (TKR 0 Double)-       , ADVal Concrete (TKR 0 Double) )-    -> ADVal Concrete (TKR 0 Double)-baz (_x,y,z) =-  let w = fooFromPrimal * barADVal2 (y,y,z) * sin y-  in atan2H z w + z * w---- An "old term", computed once, stored at top level.-fooFromPrimal :: ADVal Concrete (TKR 0 Double)-fooFromPrimal = foo2 (rscalar 7, rscalar 8, rscalar 9)--testBaz :: Assertion-testBaz =-  assertEqualUpToEpsilon 1e-9-    (rscalar 0, rscalar (-5219.20995030263), rscalar 2782.276274462047)-    (cgrad (kfromR . baz) (rscalar 1.1, rscalar 2.2, rscalar 3.3))---- If terms are numbered and @z@ is, wrongly, decorated with number 0,--- here @fooFromPrimal@ is likely to clash with @z@, since it was numbered--- starting from 0, too.--- The test fails if term numbering is reset to 0 between gradient computation--- runs (verified) as well as when delta-expression evaluation assumes--- it can only encounter terms numbered in the current run and,--- e.g., allocates an array with only that much space (not verified).--- Actually, with term counter zeroed just before @cgrad@ application,--- even a single @testBaz@ execution produces wrong results, but this test--- is likely to fail in less naive implementations, as well.-testBazRenumbered :: Assertion-testBazRenumbered =-  assertEqualUpToEpsilon 1e-9-    (rscalar 0, rscalar (-5219.20995030263), rscalar 2783.276274462047)-    (cgrad (kfromR . (\(x,y,z) -> z + baz (x,y,z))) (rscalar 1.1, rscalar 2.2, rscalar 3.3))---- A dual-number and list-based version of a function that goes--- from `R^3` to `R`.-fooD :: forall r. r ~ Double-     => ListR 3 (ADVal Concrete (TKR 0 r)) -> ADVal Concrete (TKR 0 r)-fooD (x ::: y ::: z ::: ZR) =-  let w = x * sin y-  in atan2H z w + z * w--testFooD :: Assertion-testFooD =-  assertEqualUpToEpsilon 1e-10-    (fromList [rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627])-    (cgrad (kfromR . fooD) (fromList [rscalar 1.1, rscalar 2.2, rscalar 3.3]))--fooBuild1 :: (ADReady target, GoodScalar r, Differentiable r)-          => target (TKR 1 r) -> target (TKR 1 r)-fooBuild1 v =-  let r = rsum0 v-      v' = rminimum v-  in rbuild1 3 $ \ix ->-       r * foo2 ( rscalar 3-               , rscalar 5 * r-               , r * rminimum v * v')-       + bar (r, rindex v [ix + 1])--testFooBuildDt :: Assertion-testFooBuildDt =-  assertEqualUpToEpsilon 1e-5-    (rconcrete $ Nested.rfromListPrimLinear [4] [-189890.46351219364,-233886.08744601303,-222532.22669716467,-206108.68889329425])-    (vjp @_ @(TKR 1 Double)-           fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [3] (rscalar 42)))--testFooBuildCFwd :: Assertion-testFooBuildCFwd =-  assertEqualUpToEpsilon 1e-5-    (rconcrete $ Nested.rfromListPrimLinear [3] [-296584.8166864211,-290062.472288043,-265770.1775742018])-    (cjvp @_ @(TKR 1 Double)-          fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [4] (rscalar 42)))--testFooBuildFwd :: Assertion-testFooBuildFwd =-  assertEqualUpToEpsilon 1e-5-    (rconcrete $ Nested.rfromListPrimLinear [3] [-296584.8166864211,-290062.472288043,-265770.1775742018])-    (jvp @_ @(TKR 1 Double)-         fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [4] (rscalar 42)))--testFooBuild :: Assertion-testFooBuild =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [4] [-4521.201512195087,-5568.7163677622175,-5298.386349932494,-4907.349735554627])-    (rev' @Double @1 fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]))--fooNoGo :: forall target r. (ADReady target, GoodScalar r, Differentiable r)-        => target (TKR 1 r) -> target (TKR 1 r)-fooNoGo v =-  let r = rsum0 v-  in rbuild1 3 (\ix' -> let ix :: PrimalOf target (TKS '[] Int64)-                            ix = sfromR $ rfromK ix' in-       bar (rscalar 3.14, bar (rscalar 3.14, rindex v [kfromS $ (ix + (sprimalPart . sfloor . sfromR) r) `minH` sscalar 2 `maxH` sscalar 0]))-       + ifH ((&&*) (rindex @1 v [ix' * 2] <=. rscalar 0)-                    (sscalar 6 >. abs ix))-               r (rscalar 5 * r))-     / rslice 1 3 (rmap0N (\x -> ifH (x >. r) r x) v)-     * rbuild1 3 (const (rscalar 1))--testFooNoGo0 :: Assertion-testFooNoGo0 =-  assertEqualUpToEpsilon' 1e-6-   (ringestData [5] [5.037878787878788,-14.394255484765257,43.23648655081373,-0.8403418295960368,5.037878787878788])-   (rev' @Double @1 fooNoGo-         (ringestData [5] [1.1 :: Double, 2.2, 3.3, 4, 5]))--nestedBuildMap :: forall target r.-                  (ADReady target, GoodScalar r, Differentiable r)-               => target (TKR 0 r) -> target (TKR 1 r)-nestedBuildMap r =-  let w = rreplicate0N [4]-      v0' = rreplicate0N [177] r :: target (TKR 1 r)-  in tlet v0' $ \v' ->-    let nestedMap x0 = tlet x0 $ \x -> rmap0N (x /) (w x)-        variableLengthBuild iy = rbuild1 7 (\ix -> rindex v' (ix + iy :.: ZIR))-        doublyBuild = rbuild1 5 (rminimum . variableLengthBuild)-        fooMap1 :: target (TKR 0 r) -> target (TKR 1 r)-        fooMap1 r2 =-          let v = fooBuild1 $ rreplicate0N [130] r2-          in rmap0N (\x -> x * r2 + rscalar 5) v-    in rmap0N (\x0 -> tlet x0 $ \x -> x * rsum0-                           (rbuild1 3 (\ix -> bar (x, rindex v' [ix]))-                            + (tlet (nestedMap x) $ \x3 -> fooBuild1 x3)-                            / (tlet x $ \x4 -> fooMap1 x4))-              ) doublyBuild--testNestedBuildMap1 :: Assertion-testNestedBuildMap1 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 107.25984443006627)-    (rev' @Double @1 nestedBuildMap (rscalar 1.1))--nestedSumBuild :: (ADReady target, GoodScalar r, Differentiable r)-               => target (TKR 1 r) -> target (TKR 1 r)-nestedSumBuild v0 = tlet v0 $ \v ->- tlet (rsum (rbuild1 9 rfromIndex0)) (\tbtf ->-  (rbuild1 13 (\ix ->-    rsum (rbuild1 4 (\ix2 ->-      flip rindex [ix2]-        (tlet (rsum v) $ \tsumv -> rbuild1 5 (const tsumv)-         * rfromList-             [ rfromIndex0 ix-             , tbtf-             , rsum (rbuild1 6 (\_ -> rsum v))-             , rindex v [ix2]-             , rsum (rbuild1 3 (\ix7 ->-                 rsum (rreplicate 5 (rfromIndex0 ix7))))-             ]))))))- + tlet (nestedBuildMap (rsum0 v)) (\nbmt -> (rbuild1 13 (\ix ->-     nbmt `rindex` [minH ix 4])))--testNestedSumBuild :: Assertion-testNestedSumBuild =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [5] [-14084.715065313612,-14084.715065313612,-14084.715065313612,-14014.775065313623,-14084.715065313612])-    (rev' @Double @1 nestedSumBuild (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))--nestedBuildIndex :: forall target r. (ADReady target, GoodScalar r)-                 => target (TKR 1 r) -> target (TKR 1 r)-nestedBuildIndex v =-  rbuild1 2 $ \ix2 -> rindex (rbuild1 3 $ \ix3 -> rindex (rbuild1 4 $ \ix4 -> rindex v (ix4 :.: ZIR)) [ix3]) (ix2 :.: ZIR)--testNestedBuildIndex :: Assertion-testNestedBuildIndex =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [5] [1,1,0,0,0])-    (rev' @Double @1 nestedBuildIndex (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))--barRelu-  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )-  => target (TKR n r) -> target (TKR n r)-barRelu x = relu $ bar (x, relu x)--testBarReluDt :: Assertion-testBarReluDt =-  assertEqualUpToEpsilon 1e-10-    (rconcrete $ Nested.rfromListPrimLinear [] [191.20462646925841])-    (vjp @_ @(TKR 0 Double)-           barRelu (rscalar 1.1) (rscalar 42.2))--testBarRelu :: Assertion-testBarRelu =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [] [4.5309153191767395])-    (rev' @Double @0 barRelu (ringestData [] [1.1]))--testBarRelu3 :: Assertion-testBarRelu3 =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [2, 1, 2] [4.5309153191767395,4.5302138998556,-9.39547533946234,95.29759282497125])-    (rev' @Double @3 barRelu (ringestData [2, 1, 2] [1.1, 2, 3, 4.2]))--barReluMax0-  :: ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )-  => target (TKR n r) -> target (TKR n r)-barReluMax0 x = reluMax $ bar (x, x)--barReluMax-  :: ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )-  => target (TKR n r) -> target (TKR n r)-barReluMax x = reluMax $ bar (x, reluMax x)--testBarReluMaxDt :: Assertion-testBarReluMaxDt =-  assertEqualUpToEpsilon 1e-10-    (rconcrete $ Nested.rfromListPrimLinear [] [191.20462646925841])-    (vjp @_ @(TKR 0 Double)-           barReluMax (rfromList0N [] [rscalar 1.1]) (rscalar 42.2))--testBarReluMax :: Assertion-testBarReluMax =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [] [4.5309153191767395])-    (rev' @Double @0 barReluMax (ringestData [] [1.1]))--testBarReluMax30 :: Assertion-testBarReluMax30 =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1] [4.5309153191767395])-    (rev' @Double @1 barReluMax0 (ringestData [1] [1.1]))--testBarReluMax31 :: Assertion-testBarReluMax31 =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [2, 1, 2] [4.5309153191767395,4.5302138998556,-9.39547533946234,95.29759282497125])-    (rev' @Double @3 barReluMax (ringestData [2, 1, 2] [1.1, 2, 3, 4.2]))--testBarReluMax3CFwd :: Assertion-testBarReluMax3CFwd =-  assertEqualUpToEpsilon 1e-10-    (rconcrete $ Nested.rfromListPrimLinear [2, 1, 2] [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])-    (cjvp @_ @(TKR 3 Double)-          barReluMax-                     (rconcrete $ Nested.rfromListPrimLinear (fromList [2, 1, 2]) [1.1, 2, 3, 4.2])-                     (ringestData [2, 1, 2] [0.1, 0.2, 0.3, 0.42]))--reluMaxS :: forall target sh r. (ADReady target, GoodScalar r, KnownShS sh)-         => target (TKS sh r) -> target (TKS sh r)-reluMaxS = smap0N (maxH (srepl 0))--barReluMaxS-  :: ( ADReady target, GoodScalar r, KnownShS sh-     , RealFloatH (target (TKS sh r)) )-  => target (TKS sh r) -> target (TKS sh r)-barReluMaxS x = reluMaxS $ barF (x, reluMaxS x)---- Previously the shape of FromListR[DeltaZero] couldn't be determined--- in buildDerivative, so this was needed. See below that it now works fine.-testBarReluMax3FwdS :: Assertion-testBarReluMax3FwdS =-  assertEqualUpToEpsilon 1e-10-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])-    (jvp @_ @(TKS '[2, 1, 2] Double)-         barReluMaxS-         (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [1.1, 2, 3, 4.2])-         (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [0.1, 0.2, 0.3, 0.42]))--testBarReluMax3FwdFrom :: Assertion-testBarReluMax3FwdFrom =-  assertEqualUpToEpsilon 1e-10-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])-    (jvp @_ @(TKS '[2, 1, 2] Double)-         (sfromR . barReluMax . rfromS)-         (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [1.1, 2, 3, 4.2])-         (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [0.1, 0.2, 0.3, 0.42]))--testBarReluMax3FwdR :: Assertion-testBarReluMax3FwdR =-  assertEqualUpToEpsilon 1e-10-    (rconcrete $ Nested.rfromListPrimLinear [2, 1, 2] [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])-    (jvp @_ @(TKR 3 Double)-         barReluMax-         (ringestData [2, 1, 2] [1.1, 2, 3, 4.2])-         (ringestData [2, 1, 2] [0.1, 0.2, 0.3, 0.42]))--f1 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 0 r)-f1 = \arg -> rsum0 (rbuild1 10 (\i -> arg * rfromIndex0 i))--testF1 :: Assertion-testF1 =-  assertEqualUpToEpsilon 1e-10-    (rscalar 45.0)-    (grad (kfromR @_ @Double . f1) (rscalar 1.1))--testF11 :: Assertion-testF11 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 45.0)-    (rev' @Double @0 f1 (rscalar 1.1))--f2 :: forall target r. (ADReady target, GoodScalar r)-   => target (TKR 0 r) -> target (TKR 0 r)-f2 = \arg ->-  let fun1 i = arg * rfromIndex0 i-      v1a = rsum0 (rbuild1 10 fun1)-      v1b = rsum0 (rbuild1 20 fun1)-      fun2 y i = y * rfromIndex0 @r i-      v2a = rsum0 (rbuild1 10 (fun2 arg))-      v2b = rsum0 (rbuild1 20 (fun2 (arg + rscalar 1)))-  in v1a + v1b + v2a + v2b--testF2 :: Assertion-testF2 =-  assertEqualUpToEpsilon 1e-10-    (rscalar 470)-    (grad (kfromR @_ @Double . f2) (rscalar 1.1))--testF21 :: Assertion-testF21 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 470)-    (rev' @Double @0 f2 (rscalar 1.1))--testF2CFwd :: Assertion-testF2CFwd =-  assertEqualUpToEpsilon 1e-10-    (rscalar 47)-    (cjvp @_ @(TKR 0 Double)-          f2 (rscalar 1.1) (rscalar 0.1))--testF2Fwd :: Assertion-testF2Fwd =-  assertEqualUpToEpsilon 1e-10-    (rscalar 47)-    (jvp @_ @(TKR 0 Double)-         f2 (rscalar 1.1) (rscalar 0.1))--braidedBuilds :: forall target r.-                 (ADReady target, GoodScalar r, Differentiable r)-              => target (TKR 0 r) -> target (TKR 2 r)-braidedBuilds r =-  rbuild1 3 (\ix1 ->-    rbuild1 4 (\ix2 -> rindex (rfromList0N [4]-      [rfromIndex0 ix2, rscalar 7, r, rscalar (-0.2)]) (ix1 :.: ZIR)))--testBraidedBuilds1 :: Assertion-testBraidedBuilds1 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 4.0)-    (rev' @Double @2 braidedBuilds (rscalar 3.4))--recycled :: (ADReady target, GoodScalar r, Differentiable r)-         => target (TKR 0 r) -> target (TKR 5 r)-recycled r =-  tlet (nestedSumBuild (rreplicate0N [7] r)) $ \nsb ->-    rbuild1 2 $ \_ -> rbuild1 4 $ \_ -> rbuild1 2 $ \_ -> rbuild1 3 $ const nsb--testRecycled1 :: Assertion-testRecycled1 =-  assertEqualUpToEpsilon 1e-6-    (rscalar 348356.9278600814)-    (grad (kfromR @_ @Double . rsum0 . recycled) (rscalar 0.0000001))--concatBuild :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 2 r)-concatBuild r =-  rbuild1 7 (\i ->-    rappend (rappend (rbuild1 5 (const r))-                     (rreplicate 1 (rfromIndex0 i)))-            (rbuild1 13 (const r)))--testConcatBuild1 :: Assertion-testConcatBuild1 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 126.0)-    (rev' @Double @2 concatBuild (rscalar 3.4))--concatBuild2 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)-concatBuild2 r =-  tlet (rfromList [r, rscalar 1, rscalar 2, rscalar 3, rscalar 4]) $ \a ->-    rbuild1 10 (\i -> ifH (i <. 5) (rindex a [i]) (rindex a [i - 5]))--testConcatBuild2 :: Assertion-testConcatBuild2 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 2.0)-    (rev' @Double @1 concatBuild2 (rscalar 3.4))--concatBuild3 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)-concatBuild3 r =-  tlet (rfromList [r, rscalar 1, rscalar 2, rscalar 3, rscalar 4]) $ \a ->-    rbuild1 10 (\i -> ifH (i <. 5) (rindex a [i]) (rindex a [i - 5 + (1 `quotH` maxH 1 (i - 5))]))--testConcatBuild3 :: Assertion-testConcatBuild3 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 1)-    (rev' @Double @1 concatBuild3 (rscalar 3.4))--concatBuild4 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)-concatBuild4 r =-  tlet (rgather1 5 (rreplicate 1 r)-                   (\i -> (1 `quotH` (4 + i)) :.: ZIR)) $ \a ->-    rappend a a--testConcatBuild4 :: Assertion-testConcatBuild4 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 10)-    (rev' @Double @1 concatBuild4 (rscalar 3.4))--concatBuild5 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)-concatBuild5 r =-  tlet (rgather1 5 (rreplicate 1 r)-                   (\i -> (1 `quotH` (5 + i)) :.: ZIR)) $ \a ->-    (rappend a (rgather1 5 (rreplicate 1 r)-                           (\i -> (1 `quotH` (5 + i)) :.: ZIR)))--testConcatBuild5 :: Assertion-testConcatBuild5 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 10)-    (rev' @Double @1 concatBuild5 (rscalar 3.4))--concatBuild6 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 2 r)-concatBuild6 r =-  rbuild1 7 (\j ->-    rappend (rappend-             (tlet (rgather1 5 (rreplicate 1 r)-                   (\i -> (1 `quotH` (4 + i)) :.: ZIR)) $ \a ->-    (rappend (rgather1 5 (rreplicate 1 r)-                         (\i -> (1 `quotH` (100 * maxH 1 (i - j))) :.: ZIR)) a))-                     (rreplicate 1 (rfromIndex0 j)))-            (rbuild1 13 (const r)))--testConcatBuild6 :: Assertion-testConcatBuild6 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 161)-    (rev' @Double @2 concatBuild6 (rscalar 3.4))--concatBuild7 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)-concatBuild7 r =-  rbuild1 10 $ \j ->-    (rappend (rreplicate 5 r) (rgather1 5 (rreplicate 1 r)-                                 (\i -> (1 `quotH` maxH 1 (j - i)) :.: ZIR)))-     ! (j :.: ZIR)--testConcatBuild7 :: Assertion-testConcatBuild7 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 10)-    (rev' @Double @1 concatBuild7 (rscalar 3.4))---- These tests show that term rewriting changes the value--- of out-of-bounds indexing, e.g., of gather(replicate)--- that reduces to a non-gather. Vectorization is not needed for that.-_concatBuild8 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)-_concatBuild8 r =-  tlet (rgather1 5 (rreplicate 1 r)-                   (\i -> (1 `quotH` (5 - i)) :.: ZIR)) $ \a ->-    (rappend a (rgather1 5 (rreplicate 1 r)-                           (\i -> (1 `quotH` (5 - i)) :.: ZIR)))--_testConcatBuild8 :: Assertion-_testConcatBuild8 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 10)-    (rev' @Double @1 _concatBuild8 (rscalar 3.4))--_concatBuild9 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 2 r)-_concatBuild9 r =-  rbuild1 7 (\j ->-    rappend (rappend-             (tlet (rgather1 5 (rreplicate 1 r)-                   (\i -> (1 `quotH` (4 - i)) :.: ZIR)) $ \a ->-    (rappend (rgather1 5 (rreplicate 1 r)-                         (\i -> (1 `quotH` (100 * maxH 0 (i - j))) :.: ZIR)) a))-                     (rreplicate 1 (rfromIndex0 j)))-            (rbuild1 13 (const r)))--_testConcatBuild9 :: Assertion-_testConcatBuild9 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 161)-    (rev' @Double @2 _concatBuild9 (rscalar 3.4))--concatBuild10 :: (ADReady target, GoodScalar r)-              => target (TKR 0 r) -> target (TKR 2 r)-concatBuild10 r =-  rbuild1 7 (\j ->-    rappend (rappend-               (tlet (rgather1 5 (rreplicate 1 r)-                       (\_i -> 10000 :.: ZIR)) $ \a ->-      (rappend (rgather1 5 (rreplicate 1 r)-                 (\_i -> (-1) :.: ZIR)) a))-                   (rreplicate 1 (rfromIndex0 j)))-            (rbuild1 13 (const r)))--testConcatBuild10 :: Assertion-testConcatBuild10 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 91)-    (rev' @Double @2 concatBuild10 (rscalar 3.4))--concatBuild11 :: (ADReady target, GoodScalar r)-              => target (TKR 0 r) -> target (TKR 1 r)-concatBuild11 r =-  rgather1 5 (rreplicate 1 r) (\_i -> (-1) :.: ZIR)--testConcatBuild11 :: Assertion-testConcatBuild11 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 0)-    (rev' @Double @1 concatBuild11 (rscalar 3.4))--concatBuild12 :: (ADReady target, GoodScalar r)-              => target (TKR 0 r) -> target (TKR 0 r)-concatBuild12 r =-  rindex (rreplicate 1 r) ((-1) :.: ZIR)--testConcatBuild12 :: Assertion-testConcatBuild12 =-  assertEqualUpToEpsilon' 1e-10-    (rscalar 0)-    (rev' @Double @0 concatBuild12 (rscalar 3.4))---- TODO: copy-paste a variant of emptyArgs with r not GoodScalar--- or maybe generalize emptyArgs and then in half of the tests do TKScalar-emptyArgs :: forall target r. (ADReady target, GoodScalar r, Differentiable r)-          => target (TKR 1 r) -> target (TKR 1 r)-emptyArgs t =-  emptyTensor-  - rfromList0N (rshape @target @_ @(TKScalar r) emptyTensor) []-  + rconcrete Nested.remptyArray-  - rsum (rfromList0N (0 :$: rshape @target @_ @(TKScalar r) emptyTensor) [])-  * rsum (rconcrete $ Nested.rreshape (0 :$: 0 :$: ZSR) Nested.remptyArray)-  * rconcrete (Nested.rsumOuter1 $ Nested.rfromListOuter-               $ NonEmpty.fromList [Nested.remptyArray])-  * rsum (rconcrete $ Nested.rfromListOuter-          $ NonEmpty.fromList [Nested.remptyArray])-  * rsum (rfromList [rconcrete Nested.remptyArray])-  * rsum (rfromList [emptyTensor])-  * rsum (rfromList [rsum (rfromList0N-                             (0 :$: rshape @target @_ @(TKScalar r)-                                      emptyTensor) [])])-  - rindex (rfromList0N (0 :$: 0 :$: ZSR) []) (42 :.: ZIR)-  - rindex (rfromList0N (0 :$: rshape @target @_ @(TKScalar r)-                                     emptyTensor) []) (42 :.: ZIR)-  - rreshape @1 [0] emptyTensor-  - rsum (rreshape @1 [0, 0] emptyTensor)--- TODO:---  - rgather1 0 emptyTensor (:.: ZIR)---  - rsum (rgather1 0 emptyTensor (const ZIR))---  - rsum (rgather @target @(TKScalar r) @2 (0 :$: 0 :$: ZSR) emptyTensor (const (0 :.: ZIR)))---  - rsum (rgather @target @(TKScalar r) @2 @0 @1 [0, 0] emptyTensor (const [0]))---  * rflatten (rtr (rgather1 0 t (const ZIR)))-  + rbuild1 0 (\i -> t ! (i :.: ZIR))-  + rbuild1 0 (\i -> t ! [fromIntegral (rlength t) `quotH` i] / rfromIndex0 i)-  + rbuild @1 (0 :$: ZSR) (const $ rscalar 73)-  - rsum (rbuild @0 (0 :$: 0 :$: ZSR)-                 (const (rreplicate 1 emptyTensor)))-  + rfromS-      (sfromList0N []-       + sconcrete (Nested.semptyArray ZSS)-       - ssum @0 (sfromList0N [])-       * ssum @0 (sconcrete $ Nested.semptyArray (SNat @0 :$$ ZSS))-       * sconcrete (Nested.ssumOuter1 $ Nested.sfromListOuter (SNat @1)-                    $ NonEmpty.fromList [Nested.semptyArray ZSS])-       * ssum @1 (sconcrete $ Nested.sfromListOuter SNat-                        $ NonEmpty.fromList [Nested.semptyArray ZSS])-       * ssum @1 (sfromList [sconcrete $ Nested.semptyArray ZSS])-       * ssum @1 (sfromList [sfromR @_ @'[0] emptyTensor])-       * ssum @1 (sfromList [ssum @0 (sfromList0N [])])-       - sindex @'[0] (sfromList0N []) (42 :.$ ZIS)-       - sindex @'[0] (sfromList0N []) (42 :.$ ZIS)-       - sreshape @_ @'[0] (sfromR @_ @'[0] emptyTensor)-       - ssum (sreshape @_ @'[0, 0] (sfromR @_ @'[0] emptyTensor))-       * sbuild1 @0 (\i -> sfromR @_ @'[0] (rslice 0 0 t) !$ (i :.$ ZIS))-       + sbuild1 @0 (\i -> sfromR @_ @'[0] (rslice 0 0 t)-                              !$ (fromIntegral (rlength t) `quotH` i :.$ ZIS)-                              / sfromIndex0 i)-       + sbuild @1 (const $ sscalar 73)-       - ssum (sbuild @0-                      (const (sreplicate @1 (sfromR emptyTensor)))))-  + rfromX-      (xfromList0N (SKnown (SNat @0) :$% ZSX) []-       + xconcrete (Nested.memptyArray ZSX)-       - xsum @0 (xfromList0N-                          (SKnown (SNat @0) :$% SKnown (SNat @0) :$% ZSX) [])-       * xsum @0 (xconcrete-                        $ Nested.memptyArray (SKnown (SNat @0) :$% ZSX))-       * xconcrete (Nested.msumOuter1 $ Nested.mfromListOuter-                    $ NonEmpty.fromList [Nested.memptyArray ZSX])-       * xsum @1 (xconcrete-                        $ Nested.mcast-                            (SKnown (SNat @1) :!% SKnown (SNat @0) :!% ZKX)-                        $ Nested.mfromListOuter-                        $ NonEmpty.fromList [Nested.memptyArray ZSX])-       * xsum @1 (xfromList [xconcrete $ Nested.memptyArray ZSX])-       * xsum @1 (xfromList [xfromR @_ @'[Just 0] emptyTensor])-       * xsum @1 (xfromList [xsum @0-                                     (xfromList0N-                                        (SKnown (SNat @0)-                                         :$% xshape @target @_ @(TKScalar r)-                                               (xfromR @_ @'[Just 0]-                                                  emptyTensor)) [])])-       - xindex @_ @'[Just 0] (xfromList0N-                                    (SKnown (SNat @0)-                                     :$% SKnown (SNat @0)-                                     :$% ZSX) []) (42 :.% ZIX)-       - xindex @_ @'[Just 0] (xfromList0N-                                    (SKnown (SNat @0)-                                     :$% xshape @target @_ @(TKScalar r)-                                             (xfromR @_ @'[Just 0]-                                                emptyTensor)) []) (42 :.% ZIX)-       - xreshape (SKnown (SNat @0) :$% ZSX) (xfromR @_ @'[Just 0] emptyTensor)-       - xsum (xreshape (SKnown (SNat @0) :$% SKnown (SNat @0) :$% ZSX) (xfromR @_ @'[Just 0] emptyTensor))-       * xbuild1 @0 (\i -> xfromR @_ @'[Nothing] (rslice 0 0 t)-                            `xindex` (i :.% ZIX))-       + xbuild1 @0 (\i -> xfromR @_ @'[Nothing] (rslice 0 0 t)-                              `xindex`-                              (fromIntegral (rlength t) `quotH` i :.% ZIX)-                              / xfromIndex0 i)-       + xbuild @1 (SKnown (SNat @0) :$% ZSX)-                (const $ xscalar 73)-       - xsum (xbuild @0 (SKnown (SNat @0)-                                                :$% SKnown (SNat @0)-                                                :$% ZSX)-                      (const (xreplicate (xfromR emptyTensor)))))-  -- - rsum (rbuild @2 (0 :$: 0 :$: ZSR) (const 73))-  -- - rsum (rbuild @1 (0 :$: 0 :$: ZSR) (const emptyTensor))-       -- these two fail and rightly so; TODO: make them fail earlier- where-  emptyTensor :: target (TKR 1 r)-  emptyTensor = rconcrete $ Nested.rfromListPrimLinear (fromList [0]) []--testEmptyArgs0 :: Assertion-testEmptyArgs0 =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [0] [])-    (rev' @Double @1 emptyArgs (ringestData [0] []))--testEmptyArgs1 :: Assertion-testEmptyArgs1 =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1] [0])-    (rev' @Float @1 emptyArgs (ringestData [1] [0.24]))--testEmptyArgs4 :: Assertion-testEmptyArgs4 =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1] [0])-    (rev' @Double @1-          (\t -> rbuild [2, 5, 11, 0] (const $ emptyArgs t))-          (ringestData [1] [0.24]))--filterPositiveFail :: ADReady target-                   => target (TKR 1 Double) -> target (TKR 1 Double)-filterPositiveFail v =-  let l = runravelToList v-      -- l2 = filter (\x -> x >= 0) l-      -- Could not deduce ‘Ord (target Double 0)’-      -- l2 = filter (\x -> x >=. 0) l-      -- Could not deduce ‘BoolOf target ~ Bool’-      l2 = take 3 l  -- the most I can do-  in rfromList $ NonEmpty.fromList l2--testFilterPositiveFail :: Assertion-testFilterPositiveFail =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [5] [1.0,1.0,1.0,0.0,0.0])-    (rev' @Double @1-          filterPositiveFail-          (ringestData [5] [0.24, 52, -0.5, 0.33, 0.1]))---- Catastrophic loss of sharing prevented.-fblowup :: forall target r. (ADReady target, GoodScalar r, Differentiable r)-        => Int -> target (TKR 1 r) -> target (TKR 0 r)-fblowup k inputs =-  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)-      blowup 0 y = y - rfromIndex0 0-      blowup n y =-        let ysum = y + y - rfromIndex0 0-            yscaled = rscalar 0.499999985 * ysum-              -- without the scaling we'd get NaN at once-        in blowup (pred n) yscaled-      y0 = (inputs ! [0]) / (inputs ! [1])-  in blowup k y0--fblowupLet :: forall target r. (ADReady target, GoodScalar r, Differentiable r)-           => IntOf target -> Int -> target (TKR 1 r) -> target (TKR 0 r)-fblowupLet i k inputs =-  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)-      blowup 0 y = y - rfromIndex0 i-      blowup n y1 = tlet y1 $ \y ->-        let ysum = y + y - rfromIndex0 i-            yscaled = rscalar 0.499999985 * ysum-              -- without the scaling we'd get NaN at once-        in blowup (pred n) yscaled-      y0 = (inputs ! [0]) / (inputs ! [1])-  in blowup k y0---- Catastrophic loss of sharing prevented also with non-trivial multiplication.-fblowupMult :: forall target r. (ADReady target, GoodScalar r, Differentiable r)-            => Int -> target (TKR 1 r) -> target (TKR 0 r)-fblowupMult k inputs =-  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)-      blowup 0 y = y-      blowup n y =-        let ysum = y + y * y / (y - rscalar 0.000000001)-            yscaled = sqrt $ rscalar 0.499999985 * rscalar 0.499999985 * ysum * ysum-              -- without the scaling we'd get NaN at once-        in blowup (pred n) yscaled - rfromIndex0 0-      y0 = (inputs ! [0 `remH` 2]) * (inputs ! [1])-  in blowup k y0--fblowupMultLet :: forall target r.-                  (ADReady target, GoodScalar r, Differentiable r)-               => IntOf target -> Int -> target (TKR 1 r) -> target (TKR 0 r)-fblowupMultLet i k inputs =-  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)-      blowup 0 y = y-      blowup n y1 = tlet y1 $ \y ->-        let ysum0 = y + y * y / (y - rscalar 0.000001)-            yscaled = tlet ysum0 $ \ysum ->-                        sqrt $ rscalar 0.499999985 * rscalar 0.499999985 * ysum * ysum-              -- without the scaling we'd get NaN at once-        in blowup (pred n) yscaled - rfromIndex0 i-      y0 = (inputs ! [i `remH` 2]) * (inputs ! [1])-  in blowup k y0--fblowupPP :: Assertion-fblowupPP = do-  resetVarCounter-  let fblowupT = fblowup @(AstTensor AstMethodLet FullSpan) @Double 1-  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupT (FTKR [4] FTKScalar)-  printArtifactSimple artifactRev-    @?= "\\dret v1 -> tlet (sfromR v1 !$ [0]) (\\x2 -> tlet (sfromR v1 !$ [1]) (\\x3 -> tlet (sfromR v1 !$ [0]) (\\x4 -> tlet (sfromR v1 !$ [1]) (\\x5 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x8 -> rfromS (soneHot (recip x3 * x8) [0] + (soneHot ((negate x2 / (x3 * x3)) * x8) [1] + (soneHot (recip x5 * x8) [0] + soneHot ((negate x4 / (x5 * x5)) * x8) [1]))))))))"-  printArtifactPrimalSimple artifactRev-    @?= "\\v1 -> tlet (sfromR v1 !$ [0]) (\\x2 -> tlet (sfromR v1 !$ [1]) (\\x3 -> tlet (sfromR v1 !$ [0]) (\\x4 -> tlet (sfromR v1 !$ [1]) (\\x5 -> tlet ((x2 / x3 + x4 / x5) + negate (sfromIntegral (sscalar 0))) (\\x6 -> rfromS (sscalar 0.499999985 * x6 + negate (sfromIntegral (sscalar 0))))))))"--fblowupLetPP :: Assertion-fblowupLetPP = do-  resetVarCounter-  let fblowupLetT = fblowupLet @(AstTensor AstMethodLet FullSpan) @Double 0 1-  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupLetT (FTKR [4] FTKScalar)-  printArtifactSimple artifactRev-    @?= "\\dret v1 -> tlet (sfromR v1 !$ [0]) (\\x3 -> tlet (sfromR v1 !$ [1]) (\\x4 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x8 -> tlet (x8 + x8) (\\x9 -> rfromS (soneHot (recip x4 * x9) [0] + soneHot ((negate x3 / (x4 * x4)) * x9) [1])))))"-  printArtifactPrimalSimple artifactRev-    @?= "\\v1 -> tlet (sfromR v1 !$ [0]) (\\x3 -> tlet (sfromR v1 !$ [1]) (\\x4 -> tlet (x3 / x4) (\\x5 -> tlet ((x5 + x5) + negate (sfromIntegral (sscalar 0))) (\\x6 -> rfromS (sscalar 0.499999985 * x6 + negate (sfromIntegral (sscalar 0)))))))"--fblowupLetPP23 :: Assertion-fblowupLetPP23 = do-  resetVarCounter-  let fblowupLetT = fblowupLet @(AstTensor AstMethodLet FullSpan) @Double 2 3-  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupLetT (FTKR [4] FTKScalar)-  printArtifactSimple artifactRev-    @?= "\\dret v1 -> tlet (sfromR v1 !$ [0]) (\\x5 -> tlet (sfromR v1 !$ [1]) (\\x6 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x14 -> tlet (sscalar 0.499999985 * x14 + sscalar 0.499999985 * x14) (\\x15 -> tlet (sscalar 0.499999985 * x15 + sscalar 0.499999985 * x15) (\\x16 -> tlet (x16 + x16) (\\x17 -> rfromS (soneHot (recip x6 * x17) [0] + soneHot ((negate x5 / (x6 * x6)) * x17) [1])))))))"-  printArtifactPrimalSimple artifactRev-    @?= "\\v1 -> tlet (sfromR v1 !$ [0]) (\\x5 -> tlet (sfromR v1 !$ [1]) (\\x6 -> tlet (x5 / x6) (\\x7 -> tlet ((x7 + x7) + negate (sfromIntegral (sscalar 2))) (\\x8 -> tlet (sscalar 0.499999985 * x8) (\\x9 -> tlet ((x9 + x9) + negate (sfromIntegral (sscalar 2))) (\\x10 -> tlet (sscalar 0.499999985 * x10) (\\x11 -> tlet ((x11 + x11) + negate (sfromIntegral (sscalar 2))) (\\x12 -> rfromS (sscalar 0.499999985 * x12 + negate (sfromIntegral (sscalar 2)))))))))))"-  printArtifactSimple (simplifyArtifact artifactRev)-    @?= "\\dret v1 -> rfromS (tlet (sfromR v1 !$ [1]) (\\x6 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x14 -> tlet (sscalar 0.499999985 * x14 + sscalar 0.499999985 * x14) (\\x15 -> tlet (sscalar 0.499999985 * x15 + sscalar 0.499999985 * x15) (\\x16 -> tlet (x16 + x16) (\\x17 -> soneHot (recip x6 * x17) [0] + soneHot ((negate (sfromR v1 !$ [0]) / (x6 * x6)) * x17) [1]))))))"--fblowupLetPP10 :: Assertion-fblowupLetPP10 = do-  resetVarCounter-  let fblowupLetT = fblowupLet @(AstTensor AstMethodLet FullSpan) @Double 0 6-  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupLetT (FTKR [2] FTKScalar)-  printArtifactSimple artifactRev-    @?= "\\dret v1 -> tlet (sfromR v1 !$ [0]) (\\x8 -> tlet (sfromR v1 !$ [1]) (\\x9 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x23 -> tlet (sscalar 0.499999985 * x23 + sscalar 0.499999985 * x23) (\\x24 -> tlet (sscalar 0.499999985 * x24 + sscalar 0.499999985 * x24) (\\x25 -> tlet (sscalar 0.499999985 * x25 + sscalar 0.499999985 * x25) (\\x26 -> tlet (sscalar 0.499999985 * x26 + sscalar 0.499999985 * x26) (\\x27 -> tlet (sscalar 0.499999985 * x27 + sscalar 0.499999985 * x27) (\\x28 -> tlet (x28 + x28) (\\x29 -> rfromS (soneHot (recip x9 * x29) [0] + soneHot ((negate x8 / (x9 * x9)) * x29) [1]))))))))))"-  printArtifactPrimalSimple artifactRev-    @?= "\\v1 -> tlet (sfromR v1 !$ [0]) (\\x8 -> tlet (sfromR v1 !$ [1]) (\\x9 -> tlet (x8 / x9) (\\x10 -> tlet ((x10 + x10) + negate (sfromIntegral (sscalar 0))) (\\x11 -> tlet (sscalar 0.499999985 * x11) (\\x12 -> tlet ((x12 + x12) + negate (sfromIntegral (sscalar 0))) (\\x13 -> tlet (sscalar 0.499999985 * x13) (\\x14 -> tlet ((x14 + x14) + negate (sfromIntegral (sscalar 0))) (\\x15 -> tlet (sscalar 0.499999985 * x15) (\\x16 -> tlet ((x16 + x16) + negate (sfromIntegral (sscalar 0))) (\\x17 -> tlet (sscalar 0.499999985 * x17) (\\x18 -> tlet ((x18 + x18) + negate (sfromIntegral (sscalar 0))) (\\x19 -> tlet (sscalar 0.499999985 * x19) (\\x20 -> tlet ((x20 + x20) + negate (sfromIntegral (sscalar 0))) (\\x21 -> rfromS (sscalar 0.499999985 * x21 + negate (sfromIntegral (sscalar 0)))))))))))))))))"-  printArtifactSimple (simplifyArtifact artifactRev)-    @?= "\\dret v1 -> rfromS (tlet (sfromR v1 !$ [1]) (\\x9 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x23 -> tlet (sscalar 0.499999985 * x23 + sscalar 0.499999985 * x23) (\\x24 -> tlet (sscalar 0.499999985 * x24 + sscalar 0.499999985 * x24) (\\x25 -> tlet (sscalar 0.499999985 * x25 + sscalar 0.499999985 * x25) (\\x26 -> tlet (sscalar 0.499999985 * x26 + sscalar 0.499999985 * x26) (\\x27 -> tlet (sscalar 0.499999985 * x27 + sscalar 0.499999985 * x27) (\\x28 -> tlet (x28 + x28) (\\x29 -> soneHot (recip x9 * x29) [0] + soneHot ((negate (sfromR v1 !$ [0]) / (x9 * x9)) * x29) [1])))))))))"---- TODO: should do 1000000 in a few seconds-blowupTests :: TestTree-blowupTests = testGroup "Catastrophic blowup avoidance tests"-  [ testCase "blowup prim 7" $ do-      assertEqualUpToEpsilon' 1e-5-        (ringestData [2] [0.3333332333333467,-0.22222215555556446])-        (rev' @Double @0 (fblowup 7) (ringestData [2] [2, 3]))-  , testCase "blowupLet prim 2000" $ do-      assertEqualUpToEpsilon' 1e-10-        (ringestData [2] [0.3333133339329949,-0.22220888928866325])-        (rev' @Double @0 (fblowupLet 1 2000) (ringestData [2] [2, 3]))-  , testCase "blowupLet 7000" $ do-      assertEqualUpToEpsilon 1e-10-        (ringestData [2] [0.3332633406816766,-0.22217556045445108])-        (grad (kfromR @_ @Double . rsum0 . fblowupLet 0 7000) (ringestData [2] [2, 3]))-  , testCase "blowupLet tbuild0" $ do-      assertEqualUpToEpsilon 1e-10-        (ringestData [2] [333.2633406816765,-222.175560454451])-        (grad (kfromR @_ @Double . rsum0 . (\intputs -> rbuild1 1000 (\_ -> fblowupLet 0 7000 intputs)))-              (ringestData [2] [2, 3]))-  , testCase "blowupLet tbuild2" $ do-      assertEqualUpToEpsilon 1e-10-        (ringestData [2] [333.2633406816765,-222.175560454451])-        (grad (kfromR @_ @Double . rsum0 . (\intputs -> rbuild1 1000 (\_ -> fblowupLet 2 7000 intputs)))-              (ringestData [2] [2, 3]))-  , testCase "blowupLet tbuildi" $ do-      assertEqualUpToEpsilon 1e-10-        (ringestData [2] [333.2983351701977,-222.19889011346513])-        (grad (kfromR @_ @Double . rsum0-               . (\intputs -> rbuild1 1000 (\i -> fblowupLet i 3500 intputs)))-              (ringestData [2] [2, 3]))-  , testCase "blowupLet tbuildc" $ do-      assertEqualUpToEpsilon 1e-7-        (ringestData [2] [333.326333406717,-222.21755560448116])-        (cgrad (kfromR @_ @Double . rsum0-                . (\intputs -> rbuild1 1000 (\i -> fblowupLet i 700 intputs)))-              (ringestData [2] [2, 3]))-  , testCase "blowupLet prim tbuild" $ do-      assertEqualUpToEpsilon 1e-7-        (ringestData [2] [33.33263334067178,-22.221755560447928])-        (grad (kfromR @_ @Double . rsum0-               . (\intputs -> rbuild1 100 (\i -> fblowupLet i 700 intputs)))-              (ringestData [2] [2, 3]))-  , testCase "blowupMult 3" $ do-      assertEqualUpToEpsilon' 1e-5-        (ringestData [2] [2.999999730000007,1.9999998200000046])-        (rev' @Double @0 (fblowupMult 3) (ringestData [2] [2, 3]))-  , testCase "blowupMultLet 5" $ do-      assertEqualUpToEpsilon' 1e-10-        (ringestData [2] [2.9999995500000267,1.9999997000000178])-        (rev' @Double @0 (fblowupMultLet 0 5)-                                   (ringestData [2] [2, 3]))-  , testCase "blowupMultLet 50" $ do-      assertEqualUpToEpsilon' 1e-10-        (ringestData [2] [2.999995500001215,1.99999700000081])-        (rev' @Double @0 (fblowupMultLet 0 50)-                                   (ringestData [2] [2, 3]))-  , testCase "blowupMultLet tbuild1" $ do-      assertEqualUpToEpsilon 1e-10-        (ringestData [2] [14.9999773958889,39.9999398380561])-        (grad (kfromR @_ @Double . rsum0-               . (\intputs -> rbuild1 100 (\i -> fblowupMultLet i 50 intputs)))-              (ringestData [2] [0.2, 0.3]))-  ]--concatBuild33 :: (ADReady target, GoodScalar r)-             => target (TKR 1 r) -> target (TKR 2 r)-concatBuild33 _r =-  rbuild1 5 (\i ->-    rbuild1 2 (\j -> rfromIndex0 (maxH j (i `quotH` (j + 1)))))--testConcatBuild3PP :: Assertion-testConcatBuild3PP = do-  resetVarCounter-  let t = concatBuild33 @(AstTensor AstMethodLet FullSpan) @Float-      (var3, ast3) = funToAst (FTKR [3] FTKScalar) Nothing t-  "\\" ++ printAstVarName var3-       ++ " -> " ++ printAstSimple ast3-    @?= "\\v1 -> tfromPrimal (STKR (SNat @2) STKScalar) (rfromS (sgather [] (sfromIntegral (tfromVector (SNat @2) (STKS [5,2] STKScalar) (fromList [sreplicate @5 (siota (SNat @2)), quotH (str (sreplicate @2 (siota (SNat @5)))) (sreplicate @5 (sconcrete (sreplicate [2] 1) + siota (SNat @2)))]))) (\\[i5, i4] -> [ifH (0 <=. i4 + quotH (negate i5) (1 + i4)) 0 1, i5, i4])))"--testConcatBuild3PP2 :: Assertion-testConcatBuild3PP2 = do-  resetVarCounter-  let t = concatBuild33 @(AstTensor AstMethodLet FullSpan) @Double-  let (artifactRev, _) =-        revArtifactDelta UseIncomingCotangent t (FTKR [3] FTKScalar)-  printArtifactSimple artifactRev-    @?= "\\dret v1 -> rfromS (sconcrete (sreplicate [3] 0.0))"-  printArtifactPrimalSimple artifactRev-    @?= "\\v1 -> rfromS (sgather [] (sfromIntegral (tfromVector (SNat @2) (STKS [5,2] STKScalar) (fromList [sreplicate @5 (siota (SNat @2)), quotH (str (sreplicate @2 (siota (SNat @5)))) (sreplicate @5 (sconcrete (sreplicate [2] 1) + siota (SNat @2)))]))) (\\[i6, i7] -> [ifH (0 <=. i7 + quotH (negate i6) (1 + i7)) 0 1, i6, i7]))"-  printArtifactPrimalSimple (simplifyArtifact artifactRev)-    @?= "\\v1 -> rfromS (sgather [] (sconcrete (sfromListLinear [2,5,2] [0.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0,0.0,0.0,1.0,0.0,2.0,1.0,3.0,1.0,4.0,2.0])) (\\[i6, i7] -> [ifH (0 <=. i7 + quotH (negate i6) (1 + i7)) 0 1, i6, i7]))"+import Data.Int (Int16, Int64, Int8)+import Data.Kind (Type)+import Data.List.NonEmpty qualified as NonEmpty+import Data.Vector.Generic qualified as V+import Foreign.C (CInt)+import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat)+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId (funToAst, resetVarCounter)+import HordeAd.Core.AstInterpret+import HordeAd.Core.DeltaFreshId (resetIdCounter)++import CrossTesting+import EqEpsilon++testTrees :: [TestTree]+testTrees =+  [ testCase "2zeroZ" testZeroZ+  , testCase "2zeroS" testZeroS+  , testCase "2CFwdZeroS" testCFwdZeroS+  , testCase "2FwdZeroS" testFwdZeroS+  , testCase "2zero2S" testZero2S+  , testCase "2CFwdZero2S" testCFwdZero2S+  , testCase "2FwdZero2S" testFwdZero2S+  , testCase "2zero3S" testZero3S+  , testCase "2CFwdZero3S" testCFwdZero3S+  , testCase "2FwdZero3S" testFwdZero3S+  , testCase "2zero4s" testZero4s+  , testCase "2zero5S" testZero5S+  , testCase "2zero6S" testZero6S+  , testCase "2zero7S" testZero7S+  , testCase "2zero8" testZero8+  , testCase "2zero9S" testZero9S+  , testCase "2CFwdZero9S" testCFwdZero9S+  , testCase "2FwdZero9S" testFwdZero9S+  , testCase "2zero10S" testZero10S+  , testCase "2CFwdZero10S" testCFwdZero10S+  , testCase "2FwdZero10S" testFwdZero10S+  , testCase "2zero11S" testZero11S+  , testCase "2CFwdZero11S" testCFwdZero11S+  , testCase "2FwdZero11S" testFwdZero11S+  , testCase "2piecewiseLinear" testPiecewiseLinear+  , testCase "2piecewiseLinearPP" testPiecewiseLinearPP+  , testCase "2piecewiseLinear2" testPiecewiseLinear2+  , testCase "2piecewiseLinear2PP" testPiecewiseLinear2PP+  , testCase "2overleaf" testOverleaf+  , testCase "2overleafGrad1" testOverleafGrad1+  , testCase "2overleafInt64n" testOverleafInt64n+  , testCase "2overleafCIntn" testOverleafCIntn+  , testCase "2overleafCIntToFloatn" testOverleafCIntToFloatn+  , testCase "2overleafInt64p" testOverleafInt64p+  , testCase "2overleafCIntp" testOverleafCIntp+  , testCase "2overleafCIntToFloatp" testOverleafCIntToFloatp+  , testCase "2overleafInt16p" testOverleafInt16p+  , testCase "2overleafInt16ToFloatp" testOverleafInt16ToFloatp+  , testCase "2overleafInt8p" testOverleafInt8p+  , testCase "2overleafInt8ToFloatp" testOverleafInt8ToFloatp+  , testCase "2overleafPP" testOverleafPP+  , testCase "2foo" testFoo+  , testCase "2fooGradDouble" testGradFooDouble+  , testCase "2fooMatrix" testFooMatrix+  , testCase "2fooGradMatrix" testGradFooMatrix+  , testCase "2fooLetGradMatrixPP" testGradFooLetMatrixPP+  , testCase "2fooGradMatrixVjp" testGradFooMatrixVjp+  , testCase "2fooGradMatrixRev" testGradFooMatrixRev+  , testCase "2fooLetGradMatrixSimpPP" testGradFooLetMatrixSimpPP+  , testCase "2fooLetGradMatrixSimpRPP" testGradFooLetMatrixSimpRPP+  , testCase "2fooSumMatrix" testfooSumMatrix+  , testCase "2fooGradMatrix2" testGradFooMatrix2+  , testCase "2fooGradMatrixPP" testGradFooMatrixPP+  , testCase "2fooGradMatrixSimpPP" testGradFooMatrixSimpPP+  , testCase "2fooGradScalar" testGradFooScalar+  , testCase "2fooGradCScalar" testGradCFooScalar+  , testCase "2fooS" testFooS+  , testCase "2fooSToFloat" testFooSToFloat+  , testCase "2fooSBoth" testFooSBoth+  , testCase "2fooBoth" testFooBoth+  , testCase "2trustVstackConcatRepl10" testTrustVstackConcatRepl10+  , testCase "2trustVstackConcatIota10" testTrustVstackConcatIota10+  , testCase "2trustVstackConcatReplIota10" testTrustVstackConcatReplIota10+  , testCase "2vstackWarmup" testVstackWarmup+  , testCase "2vstackConcatConcrete" testVstackConcatConcrete+  , testCase "2vstackBuildConcrete" testVstackBuildConcrete+  , testCase "2vstackConcatAst" testVstackConcatAst+  , testCase "2vstackBuildAst" testVstackBuildAst+  , testCase "2vstackBuildAstSimp" testVstackBuildAstSimp+  , testCase "2vstackBuildAstPP" testVstackBuildAstPP+  , testCase "2vstackConcatConcrete2" testVstackConcatConcrete2+  , testCase "2vstackBuildConcrete2" testVstackBuildConcrete2+  , testCase "2vstackConcatAst2" testVstackConcatAst2+  , testCase "2vstackBuildAst2" testVstackBuildAst2+  , testCase "2vstackBuildAstSimp2" testVstackBuildAstSimp2+  , testCase "2vstackBuildAstPP2" testVstackBuildAstPP2+  , testCase "2fooPP" testFooPP+  , testCase "2fooLet" testFooLet+  , testCase "2fooLetPP" testFooLetPP+  , testCase "2listProdPP" testListProdPP+  , testCase "2listProdrPP" testListProdrPP+  , testCase "2listProdrLongPP" testListProdrLongPP+  , testCase "2listProd" testListProd+  , testCase "2listProdr" testListProdr+  , testCase "2listSumrPP" testListSumrPP+  , testCase "2listSum2rPP" testListSum2rPP+  , testCase "2listSum22rPP" testListSum22rPP+  , testCase "2listSumk22rPP" testListSumk22rPP+  , testCase "2listSum2xpyrPP" testListSum2xpyrPP+  , testCase "2listSum2xyrPP" testListSum2xyrPP+  , testCase "2list2xyPP" test2xyPP+  , testCase "2listSum23rPP" testListSum23rPP+  , testCase "2list23PP" test23PP+  , testCase "2reluPP" testReluPP+  , testCase "2reluPP2" testReluPP2+  , testCase "2reluSimpler" testReluSimpler+  , testCase "2reluSimplerPP" testReluSimplerPP+  , testCase "2reluSimplerPP2" testReluSimplerPP2+  , testCase "2reluSimplerPP3" testReluSimplerPP3+  , testCase "2reluSimpler3" testReluSimpler3+  , testCase "2reluSimplerPP4" testReluSimplerPP4+  , testCase "2reluSimpler4" testReluSimpler4+  , testCase "2reluSimplerPP4s" testReluSimplerPP4s+  , testCase "2reluSimplerPP4s2" testReluSimplerPP4s2+  , testCase "2reluSimpler4s" testReluSimpler4s+  , testCase "2reluSimplerPP7s2" testReluSimplerPP7s2+  , testCase "2reluSimpler7s" testReluSimpler7s+  , testCase "2reluMax" testReluMax+  , testCase "2reluMaxPP" testReluMaxPP+  , testCase "2reluMaxPP2" testReluMaxPP2+  , testCase "2reluMax3" testReluMax3+  , testCase "2dot1PP" testDot1PP+  , testCase "2dot2PP" testDot2PP+  , testCase "2matvecmulPP" testMatvecmulPP+  , testCase "2matmul2PP" testMatmul2PP+  , testCase "2matmul2FromMatvecmulPP" testMatmul2FromMatvecmulPP+  , testCase "2matmul2PaperPP" testMatmul2PaperPP+  , testCase "2matmul2PPS" testMatmul2PPS+  , testCase "2addSpeedBig" testAddSpeedBig+  , testCase "2matmul2SpeedBig" testMatmul2SpeedBig+  , testCase "2bar" testBar+  , testCase "2barS" testBarS+  , testCase "2bar2S" testBar2S+  , testCase "2barCFwd" testBarCFwd+  , testCase "2barFwd" testBarFwd+  , testCase "2baz old to force fooFromPrimal" testBaz+  , testCase "2baz if repetition breaks things" testBaz+  , testCase "2baz again with renumbered terms" testBazRenumbered+  , testCase "2fooD T Double [1.1, 2.2, 3.3]" testFooD+  , testCase "2fooBuildDt" testFooBuildDt+  , testCase "2fooBuildCFwd" testFooBuildCFwd+  , testCase "2fooBuildFwd" testFooBuildFwd+  , testCase "2fooBuild" testFooBuild+  , testCase "2fooNoGo0" testFooNoGo0+  , testCase "2nestedBuildMap1" testNestedBuildMap1+  , testCase "2nestedSumBuildm1" testNestedSumBuildm1+  , testCase "2nestedSumBuild0" testNestedSumBuild0+  , testCase "2nestedSumBuild" testNestedSumBuild+  , testCase "2nestedBuildIndex" testNestedBuildIndex+  , testCase "2barReluDt" testBarReluDt+  , testCase "2barRelu" testBarRelu+  , testCase "2barRelu3" testBarRelu3+  , testCase "2barReluMaxDt" testBarReluMaxDt+  , testCase "2barReluMax" testBarReluMax+  , testCase "2barReluMax30" testBarReluMax30+  , testCase "2barReluMax31" testBarReluMax31+  , testCase "2barReluMax3CFwd" testBarReluMax3CFwd+  , testCase "2barReluMax3FwdS" testBarReluMax3FwdS+  , testCase "2barReluMax3FwdFrom" testBarReluMax3FwdFrom+  , testCase "2barReluMax3FwdR" testBarReluMax3FwdR+  , testCase "2F1" testF1+  , testCase "2F11" testF11+  , testCase "2F2" testF2+  , testCase "2F21" testF21+  , testCase "2F2CFwd" testF2CFwd+  , testCase "2F2Fwd" testF2Fwd+  , testCase "2braidedBuilds1" testBraidedBuilds1+  , testCase "2recycled1" testRecycled1+  , testCase "2concatBuild1" testConcatBuild1+  , testCase "2concatBuild2" testConcatBuild2+  , testCase "2concatBuild3" testConcatBuild3+  , testCase "2concatBuild4" testConcatBuild4+  , testCase "2concatBuild5" testConcatBuild5+  , testCase "2concatBuild6" testConcatBuild6+  , testCase "2concatBuild7" testConcatBuild7+  , testCase "2concatBuild8" testConcatBuild8+  , testCase "2concatBuild9" testConcatBuild9+  , testCase "2concatBuild10" testConcatBuild10+  , testCase "2concatBuild11" testConcatBuild11+  , testCase "2concatBuild12" testConcatBuild12+  , testCase "2concatBuild13" testConcatBuild13+  , testCase "2concatBuild14" testConcatBuild14+  , testCase "2concatBuild15" testConcatBuild15+  , testCase "2concatBuild16" testConcatBuild16+  , testCase "2concatBuild17" testConcatBuild17+  , testCase "2emptyArgs0" testEmptyArgs0+  , testCase "2emptyArgs1" testEmptyArgs1+  , testCase "2emptyArgs2" testEmptyArgs2+  , testCase "2emptyArgs3" testEmptyArgs3+  , testCase "2emptyArgs4" testEmptyArgs4+  , testCase "2emptyArgs5" testEmptyArgs5+  , testCase "2emptyArgs6" testEmptyArgs6+  , testCase "2filterPositiveFail0" testFilterPositiveFail0+  , testCase "2filterPositiveFail0PP" testFilterPositiveFail0PP+  , testCase "2filterPositiveFail1" testFilterPositiveFail1+  , testCase "2filterPositiveFail1PP" testFilterPositiveFail1PP+  , testCase "2filterPositiveFail" testFilterPositiveFail+  , testCase "2filterPositiveFailPP" testFilterPositiveFailPP+  , testCase "2blowupPP" fblowupPP+  , testCase "2blowup2LetPP" fblowupLetPP+  , testCase "2blowup2LetPP23" fblowupLetPP23+  , blowupTests+  , testCase "22concatBuild3PP" testConcatBuild3PP+  , testCase "22concatBuild3PP2" testConcatBuild3PP2+  ]++rfromIndex0 :: forall r target.+               (BaseTensor target, ConvertTensor target, NumScalar r)+            => IntOf target -> target (TKR 0 r)+rfromIndex0 = rfromIntegral . rfromK . tfromPlain STKScalar++testZeroZ :: Assertion+testZeroZ =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @0 (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+                          f = const (rscalar 3)+                      in f) (rscalar 42))++testZeroS :: Assertion+testZeroS =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [])+    (cgrad (ssum0 .+            let f :: ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+                  -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+                f = const (srepl 3)+            in f) (srepl 42))++testCFwdZeroS :: Assertion+testCFwdZeroS =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [])+    (cjvp (let f :: ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+               f = const (srepl 3)+           in f) (srepl 42) (srepl 41))++testFwdZeroS :: Assertion+testFwdZeroS =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [])+    (jvp (let f :: AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)+                -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)+              f = const (srepl 3)+          in f) (srepl 42) (srepl 41))++testZero2S :: Assertion+testZero2S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[] knownShS [1])+    (cgrad @_ @_ @_ @Concrete+           (kfromS @_ @Double .+           let f :: a -> a+               f = id+           in f) (srepl 42))++testCFwdZero2S :: Assertion+testCFwdZero2S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[] knownShS [41])+    (cjvp @_ @(TKS '[] Double) @_ @Concrete+          (let f :: a -> a+               f = id+           in f) (srepl 42) (srepl 41))++testFwdZero2S :: Assertion+testFwdZero2S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[] knownShS [41])+    (jvp @_ @(TKS '[] Double)+          (let f :: a -> a+               f = id+           in f) (srepl 42) (srepl 41))++testZero3S :: Assertion+testZero3S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[33, 2] knownShS (replicate 66 3.6174114266850617))+    (cgrad (ssum0 . (\x -> barF @(ADVal Concrete (TKS '[33, 2] Double)) (x, x))) (srepl 1))++testCFwdZero3S :: Assertion+testCFwdZero3S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[33, 2] knownShS (replicate 66 3.9791525693535674))+    (cjvp (\x -> barF @(ADVal Concrete (TKS '[33, 2] Double)) (x, x)) (srepl 1) (srepl 1.1))++testFwdZero3S :: Assertion+testFwdZero3S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[33, 2] knownShS (replicate 66 3.9791525693535674))+    (jvp (\x -> barF @(AstTensor AstMethodLet FullSpan (TKS '[33, 2] Double)) (x, x)) (srepl 1) (srepl 1.1))++testZero4s :: Assertion+testZero4s =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[] knownShS [0])+    (grad @(AstTensor AstMethodLet FullSpan (TKS '[] Double))+          (kfromS @_ @Double .+          let f = const (srepl 3)+          in f) (srepl 42))++testZero5S :: Assertion+testZero5S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[44] knownShS (replicate 44 1))+    (grad (ssum0 .+           let f :: a -> a+               f = id+           in f @(AstTensor AstMethodLet FullSpan (TKS '[44] Double))) (srepl 42))++testZero6S :: Assertion+testZero6S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] knownShS (replicate (product ([2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] :: [Int])) 3.6174114266850617))+    (grad (ssum0 @'[2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] @Double . (\x -> barF (x, x))) (srepl 1))++testZero7S :: Assertion+testZero7S =+  assertEqualUpToEpsilon 1e-10+    (sconcrete $ Nested.sfromListPrimLinear @'[] knownShS [0])+    (grad (kfromR . (const (rscalar 3) :: AstTensor AstMethodLet FullSpan (TKS '[] Double) -> AstTensor AstMethodLet FullSpan (TKR 0 Double))) (srepl 42))++testZero8 :: Assertion+testZero8 =+  assertEqualUpToEpsilon 1e-10+    (rrepl [] 0)+    (grad (kfromS . (const (sscalar 3) :: AstTensor AstMethodLet FullSpan (TKR 0 Double) -> AstTensor AstMethodLet FullSpan (TKS '[] Double))) (rscalar 42))++testZero9S :: Assertion+testZero9S =+  assertEqualUpToEpsilon 1e-9+    (ringestData [0, 2, 4, 0, 1] [])+    (cgrad (ssum0 .+            let f :: ADVal Concrete (TKR 5 Double)+                  -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+                f = const (srepl 3)+            in f)+           (rreplicate0N [0, 2, 4, 0, 1] 42))++testCFwdZero9S :: Assertion+testCFwdZero9S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [])+    (cjvp (let f :: ADVal Concrete (TKR 5 Double)+                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+               f = const (srepl 3)+           in f)+          (rreplicate0N [0, 2, 4, 0, 1] 42) (rreplicate0N [0, 2, 4, 0, 1] 41))++testFwdZero9S :: Assertion+testFwdZero9S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [])+    (jvp (let f :: AstTensor AstMethodLet FullSpan (TKR 5 Double)+                -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)+              f = const (srepl 3)+          in f)+         (rreplicate0N [0, 2, 4, 0, 1] 42) (rreplicate0N [0, 2, 4, 0, 1] 41))++testZero10S :: Assertion+testZero10S =+  assertEqualUpToEpsilon 1e-9+    ( ringestData [0, 2, 4, 0, 1] []+    , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] )+    (cgrad (ssum0 .+            let f = const (srepl 3) . snd+            in f :: ( ADVal Concrete (TKR 5 Double)+                    , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )+                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double))+          (rreplicate0N [0, 2, 4, 0, 1] 42, srepl 21))++testCFwdZero10S :: Assertion+testCFwdZero10S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [])+    (cjvp (let f = const (srepl 3) . snd+           in f :: ( ADVal Concrete (TKR 5 Double)+                   , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )+                   -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double))+          ( ringestData [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] )+          ( ringestData [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] ))++testFwdZero10S :: Assertion+testFwdZero10S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [])+    (jvp  (let f = const (srepl 3) . snd+           in f :: ( AstTensor AstMethodLet FullSpan (TKR 5 Double)+                   , AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double) )+                   -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double))+          ( ringestData [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] )+          ( ringestData [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] ))++testZero11S :: Assertion+testZero11S =+  assertEqualUpToEpsilon 1e-9+    ( ringestData [0, 2, 4, 0, 1] []+    , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] )+    (cgrad (rsum0 .+            let f = const (rreplicate0N [0, 2, 4, 0, 1] 3) . snd+            in f :: ( ADVal Concrete (TKR 5 Double)+                    , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )+                 -> ADVal Concrete (TKR 5 Double))+          (rreplicate0N [0, 2, 4, 0, 1] 42, srepl 21))++testCFwdZero11S :: Assertion+testCFwdZero11S =+  assertEqualUpToEpsilon 1e-9+    (ringestData [0, 2, 4, 0, 1] [])+    (cjvp (let f = const (rreplicate0N [0, 2, 4, 0, 1] 3) . snd+           in f :: ( ADVal Concrete (TKR 5 Double)+                   , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )+                   -> ADVal Concrete (TKR 5 Double))+          ( ringestData [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] )+          ( ringestData [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] ))++testFwdZero11S :: Assertion+testFwdZero11S =+  assertEqualUpToEpsilon 1e-9+    (ringestData [0, 2, 4, 0, 1] [])+    (jvp  (let f = const (rreplicate0N [0, 2, 4, 0, 1] 3) . snd+           in f :: ( AstTensor AstMethodLet FullSpan (TKR 5 Double)+                   , AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double) )+                   -> AstTensor AstMethodLet FullSpan (TKR 5 Double))+          ( ringestData [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] )+          ( ringestData [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @'[0, 2, 4, 0, 1] knownShS [] ))++testPiecewiseLinear :: Assertion+testPiecewiseLinear =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2)+    (let fT :: ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+         fT x = ifH (x >. rscalar 0) (rscalar 2 * x) (rscalar 5 * x)+     in rev' @Double @0 fT (rscalar 42))++testPiecewiseLinearPP :: Assertion+testPiecewiseLinearPP = do+  resetVarCounter+  let fT :: AstTensor AstMethodLet FullSpan (TKR 0 Double)+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT x = ifH (x >. rscalar 0) (rscalar 2 * x) (rscalar 5 * x)+      (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR ZSR FTKScalar)+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> rfromK (kfromPlain (ifH (0.0 <=. kplainPart (negate (kfromR x1))) 5.0 2.0) * kfromR dret)"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (kfromPlain (ifH (0.0 <=. kplainPart (negate (kfromR x1))) 5.0 2.0) * kfromR x1)"++testPiecewiseLinear2 :: Assertion+testPiecewiseLinear2 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 5)+    (let fT :: ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+         fT x = ifH (x >. rscalar 0) (rscalar 2) (rscalar 5) * x+     in rev' @Double @0 fT (rscalar (-42)))++testPiecewiseLinear2PP :: Assertion+testPiecewiseLinear2PP = do+  resetVarCounter+  let fT :: AstTensor AstMethodLet FullSpan (TKR 0 Double)+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT x = ifH (x >. rscalar 0) (rscalar 2) (rscalar 5) * x+      (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR ZSR FTKScalar)+  printArtifactPretty artifactRev+    @?= "\\dret x1 -> rfromK (kfromPlain (ifH (0.0 <=. kplainPart (negate (kfromR x1))) 5.0 2.0) * kfromR dret)"+  printArtifactPrimalPretty artifactRev+    @?= "\\x1 -> rfromK (kfromPlain (ifH (0.0 <=. kplainPart (negate (kfromR x1))) 5.0 2.0) * kfromR x1)"++overleaf :: forall r target. (ADReady target, NumScalar r)+         => target (TKR 1 r) -> target (TKR 0 r)+overleaf v = let wrap i = i `remH` fromIntegral (rwidth v)+             in rsum (rbuild @1 [50] (\[i] -> rindex v [wrap i]))++testOverleaf :: Assertion+testOverleaf =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @_ @Double [28] [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @0 overleaf (ringestData [28] [0 .. 27]))++testOverleafGrad1 :: Assertion+testOverleafGrad1 =+  assertEqualUpToEpsilon 1e-10+    (ringestData @_ @Double [1] [50])+    (grad (kfromR @_ @Double . overleaf) (ringestData [1] [27]))++testOverleafInt64n :: Assertion+testOverleafInt64n =+  assertEqualUpToEpsilon 1e-10+    (ringestData [28] (map round [0 :: Double,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]))+    (cgrad @_ @_ @_ @Concrete+           (kfromR @_ @Int64 . overleaf) (ringestData [28] [0 .. 27]))++testOverleafCIntn :: Assertion+testOverleafCIntn =+  assertEqualUpToEpsilon 1e-10+    (ringestData [28] (map round [0 :: Double,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]))+    (grad (kfromR @_ @CInt . overleaf) (ringestData [28] [0 .. 27]))++testOverleafCIntToFloatn :: Assertion+testOverleafCIntToFloatn =+  assertEqualUpToEpsilon 1e-10+    (ringestData [28] (replicate 28 0.0))+    (grad (kfromR @_ @CInt . rfromIntegral . overleaf @CInt . rfloor) (ringestData @_ @Float [28] [0 .. 27]))++testOverleafInt64p :: Assertion+testOverleafInt64p =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @_ @Int64 [28] (map round [2.0 :: Double,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0]))+    (rev' @Int64 @0 overleaf (ringestData [28] [0 .. 27]))++testOverleafCIntp :: Assertion+testOverleafCIntp =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @_ @CInt [28] (map round [2.0 :: Double,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0]))+    (rev' @CInt @0 overleaf (ringestData [28] [0 .. 27]))++testOverleafCIntToFloatp :: Assertion+testOverleafCIntToFloatp =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @1 @Float [28] (replicate 28 0.0))+    (let f :: forall f. ADReady f => f (TKR 1 Float) -> f (TKR 0 Float)+         f = rfromIntegral . overleaf @CInt . rfloor+     in rev' @Float @0 f (ringestData [28] [0 .. 27]))++testOverleafInt16p :: Assertion+testOverleafInt16p =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @_ @Int16 [28] (map round [2.0 :: Double,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0]))+    (rev' @Int16 @0 overleaf (ringestData [28] [0 .. 27]))++testOverleafInt16ToFloatp :: Assertion+testOverleafInt16ToFloatp =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @1 @Float [28] (replicate 28 0.0))+    (let f :: forall f. ADReady f => f (TKR 1 Float) -> f (TKR 0 Float)+         f = rfromIntegral . overleaf @Int16 . rfloor+     in rev' @Float @0 f (ringestData [28] [0 .. 27]))++testOverleafInt8p :: Assertion+testOverleafInt8p =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @_ @Int8 [28] (map round [2.0 :: Double,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0]))+    (rev' @Int8 @0 overleaf (ringestData [28] [0 .. 27]))++testOverleafInt8ToFloatp :: Assertion+testOverleafInt8ToFloatp =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @1 @Float [28] (replicate 28 0.0))+    (let f :: forall f. ADReady f => f (TKR 1 Float) -> f (TKR 0 Float)+         f = rfromIntegral . overleaf @Int8 . rfloor+     in rev' @Float @0 f (ringestData [28] [0 .. 27]))++testOverleafPP :: Assertion+testOverleafPP = do+  resetVarCounter >> resetIdCounter+  let fT :: AstTensor AstMethodLet FullSpan (TKR 1 Double)+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = overleaf+      (var3, ast3) = funToAst (FTKR [28] FTKScalar) fT+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> rfromK (ssum0 @[50] (sgather1 @50 (sfromR v1) (\\i2 -> [remH i2 28])))"+  resetVarCounter+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR [28] FTKScalar)+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> rfromS (sscatter1 @50 (sreplicate0N @[50] (kfromR dret)) (\\i5 -> [remH i5 28]))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromK (ssum0 @[50] (sgather1 @50 (sfromR v1) (\\i3 -> [remH i3 28])))"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (DeltaShare 100000002 (DeltaSum0S (DeltaShare 100000001 (DeltaGatherS [50] [] [28] (DeltaConvert (ConvCmp (ConvXS' (FTKS [28] FTKScalar)) ConvRX) (DeltaInput (InputId 0))) <function>))))"++foo :: RealFloat a => (a, a, a) -> a+foo (x, y, z) =+  let w = x * sin y+  in atan2 z w + z * w  -- note that w appears twice++testFoo :: Assertion+testFoo = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627)+    (cgrad @_ @_ @_ @Concrete+           (kfromR @_ @Double . foo) (rscalar 1.1, rscalar 2.2, rscalar 3.3))++gradFooDouble :: (Double, Double, Double) -> (Double, Double, Double)+gradFooDouble = fromDValue . cgrad @_ @_ @_ @Concrete foo . fromValue++testGradFooDouble :: Assertion+testGradFooDouble =+  assertEqualUpToEpsilon 1e-10+    (2.4396285219055063, -1.953374825727421, 0.9654825811012627)+    (gradFooDouble (1.1, 2.2, 3.3))++type Matrix2x2 :: Type -> Type+type Matrix2x2 r = Concrete (TKS '[2, 2] r)+type ThreeMatrices r = (Matrix2x2 r, Matrix2x2 r, Matrix2x2 r)+threeSimpleMatrices :: ThreeMatrices Double+threeSimpleMatrices = (srepl 1.1, srepl 2.2, srepl 3.3)+fooMatrixValue :: Matrix2x2 Double+fooMatrixValue = foo threeSimpleMatrices+gradSumFooMatrix :: ThreeMatrices Double -> ThreeMatrices Double+gradSumFooMatrix = cgrad @_ @_ @_ @Concrete (ssum0 . foo)++testFooMatrix :: Assertion+testFooMatrix =+  assertEqualUpToEpsilon 1e-10+    (sfromListLinear [2,2] [4.242393641025528,4.242393641025528,4.242393641025528,4.242393641025528])+    fooMatrixValue++testGradFooMatrix :: Assertion+testGradFooMatrix =+  assertEqualUpToEpsilon 1e-10+    (sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])+    (gradSumFooMatrix threeSimpleMatrices)++fooLet :: (RealFloatH (f r), ADReady f)+       => (f r, f r, f r) -> f r+fooLet (x, y, z) =+  tlet (x * sin y) $ \w ->+    atan2H z w + z * w  -- note that w still appears twice++artifact :: AstArtifactRev (X (ThreeMatrices Double)) (TKS '[2, 2] Double)+artifact = vjpArtifact fooLet threeSimpleMatrices++testGradFooLetMatrixPP :: Assertion+testGradFooLetMatrixPP = do+  resetVarCounter+  printArtifactPretty artifact+    @?= "\\dret m1 -> let m3 = sin (tproject2 (tproject1 m1)) ; m4 = tproject1 (tproject1 m1) * m3 ; m5 = recip (tproject2 m1 * tproject2 m1 + m4 * m4) in tpair (let m7 = (negate (tproject2 m1) * m5) * dret + tproject2 m1 * dret in tpair (m3 * m7) (cos (tproject2 (tproject1 m1)) * (tproject1 (tproject1 m1) * m7))) ((m4 * m5) * dret + m4 * dret)"++testGradFooMatrixVjp :: Assertion+testGradFooMatrixVjp =+  assertEqualUpToEpsilon 1e-10+    ((sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627]) :: ThreeMatrices Double)+    (vjpInterpretArtifact artifact (toTarget threeSimpleMatrices) (srepl 1))++testGradFooMatrixRev :: Assertion+testGradFooMatrixRev =+  assertEqualUpToEpsilon 1e-10+    (sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])+    (grad (ssum0 . fooLet) threeSimpleMatrices)++testGradFooLetMatrixSimpPP :: Assertion+testGradFooLetMatrixSimpPP = do+  resetVarCounter+  (let ftk = FTKS @'[2, 2] [2, 2] (FTKScalar @Double)+   in printArtifactPretty+        (simplifyArtifactRev $ revArtifactAdapt UseIncomingCotangent fooLet (FTKProduct (FTKProduct ftk ftk) ftk)))+      @?= "\\dret m1 -> let m3 = sin (tproject2 (tproject1 m1)) ; m4 = tproject1 (tproject1 m1) * m3 ; m5 = recip (tproject2 m1 * tproject2 m1 + m4 * m4) in tpair (let m7 = negate (tproject2 m1) * (m5 * dret) + tproject2 m1 * dret in tpair (m3 * m7) (cos (tproject2 (tproject1 m1)) * (tproject1 (tproject1 m1) * m7))) (m4 * (m5 * dret) + m4 * dret)"++testGradFooLetMatrixSimpRPP :: Assertion+testGradFooLetMatrixSimpRPP = do+  resetVarCounter+  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)+    in printArtifactPretty (simplifyArtifactRev $ revArtifactAdapt UseIncomingCotangent fooLet (FTKProduct (FTKProduct ftk ftk) ftk)))+       @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2,2] STKScalar)) (let m3 = sin (sfromR (tproject2 (tproject1 m1))) ; m4 = sfromR (tproject1 (tproject1 m1)) * m3 ; m5 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m4 * m4) in tpair (let m7 = negate (sfromR (tproject2 m1)) * (m5 * sfromR dret) + sfromR (tproject2 m1) * sfromR dret in tpair (m3 * m7) (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m7))) (m4 * (m5 * sfromR dret) + m4 * sfromR dret))"++type Matrix2x2f :: Target -> Type -> Type+type Matrix2x2f f r = f (TKS '[2, 2] r)++sumFooMatrix :: (ADReady f, RealFloat (Matrix2x2f f r), NumScalar r)+             => (Matrix2x2f f r, Matrix2x2f f r, Matrix2x2f f r)+             -> f (TKScalar r)+sumFooMatrix = ssum0 . foo++testfooSumMatrix :: Assertion+testfooSumMatrix =+  assertEqualUpToEpsilon 1e-10+    16.96957456410211+    (sumFooMatrix threeSimpleMatrices)++foo2 :: RealFloatH a => (a, a, a) -> a+foo2 (x, y, z) =+  let w = x * sin y+  in atan2H z w + z * w++gradFooMatrix2 :: (Differentiable r, NumScalar r)+               => (Concrete (TKR 2 r), Concrete (TKR 2 r), Concrete (TKR 2 r))+               -> (Concrete (TKR 2 r), Concrete (TKR 2 r), Concrete (TKR 2 r))+gradFooMatrix2 = grad (rsum0 . foo2)++testGradFooMatrix2 :: Assertion+testGradFooMatrix2 =+  assertEqualUpToEpsilon 1e-10+    (ringestData [2, 2] [2.4396285219055063 :: Double,2.4396285219055063,2.4396285219055063,2.4396285219055063], ringestData [2, 2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421], ringestData [2, 2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])+    (gradFooMatrix2 (rreplicate0N [2, 2] 1.1, rreplicate0N [2, 2] 2.2, rreplicate0N [2, 2] 3.3))++testGradFooMatrixPP :: Assertion+testGradFooMatrixPP = do+  resetVarCounter+  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)+    in printArtifactPretty (revArtifactAdapt UseIncomingCotangent foo2 (FTKProduct (FTKProduct ftk ftk) ftk)))+      @?= "\\dret m1 -> let m2 = sin (sfromR (tproject2 (tproject1 m1))) ; m3 = sfromR (tproject1 (tproject1 m1)) * m2 ; m4 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m3 * m3) ; m5 = sin (sfromR (tproject2 (tproject1 m1))) in tpair (let m8 = sfromR (tproject2 m1) * sfromR dret ; m9 = (negate (sfromR (tproject2 m1)) * m4) * sfromR dret in tpair (rfromS (m2 * m9 + m5 * m8)) (rfromS (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m9) + cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m8)))) (rfromS (let m6 = sfromR (tproject1 (tproject1 m1)) * m5 in (m3 * m4) * sfromR dret + m6 * sfromR dret))"++testGradFooMatrixSimpPP :: Assertion+testGradFooMatrixSimpPP = do+  resetVarCounter+  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)+    in printArtifactPretty (simplifyArtifactRev $ revArtifactAdapt UseIncomingCotangent foo2 (FTKProduct (FTKProduct ftk ftk) ftk)))+      @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2,2] STKScalar)) (let m2 = sin (sfromR (tproject2 (tproject1 m1))) ; m3 = sfromR (tproject1 (tproject1 m1)) * m2 ; m4 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m3 * m3) ; m5 = sin (sfromR (tproject2 (tproject1 m1))) in tpair (let m8 = sfromR (tproject2 m1) * sfromR dret ; m9 = negate (sfromR (tproject2 m1)) * (m4 * sfromR dret) in tpair (m2 * m9 + m5 * m8) (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m9) + cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m8))) (m3 * (m4 * sfromR dret) + sfromR (tproject1 (tproject1 m1)) * (m5 * sfromR dret)))"++gradFooScalar :: forall r. r ~ Double+              => (r, r, r) -> (r, r, r)+gradFooScalar = fromDValue . grad foo2 . fromValue++testGradFooScalar :: Assertion+testGradFooScalar =+  assertEqualUpToEpsilon 1e-10+    (2.4396285219055063, -1.953374825727421, 0.9654825811012627)+    (gradFooScalar (1.1, 2.2, 3.3))++gradCFooScalar :: forall r. r ~ Float+               => (r, r, r) -> (r, r, r)+gradCFooScalar = fromDValue . cgrad @_ @_ @_ @Concrete foo2 . fromValue++testGradCFooScalar :: Assertion+testGradCFooScalar =+  assertEqualUpToEpsilon 1e-10+    (2.4396284,-1.9533751,0.96548253)+    (gradCFooScalar (1.1, 2.2, 3.3))++testFooS :: Assertion+testFooS = do+  assertEqualUpToEpsilon 1e-10+    (srepl 2.4396285219055063, srepl (-1.953374825727421), srepl 0.9654825811012627)+    (grad (ssum0 @'[3, 534, 3] @Double . foo2) (srepl 1.1, srepl 2.2, srepl 3.3))++testFooSToFloat :: Assertion+testFooSToFloat = do+  assertEqualUpToEpsilon 1e-5+    (srepl 2.4396285219055063, srepl (-1.953374825727421), srepl 0.9654825811012627)+    (grad (ssum0 @_ @Float . scast . foo2)+         (srepl 1.1 :: Concrete (TKS '[3, 534, 3] Double), srepl 2.2, srepl 3.3))++testFooSBoth :: Assertion+testFooSBoth = do+  assertEqualUpToEpsilon 1e-5+    (srepl 2.439628436155373, srepl (-1.9533749), srepl 0.9654825479484146)+    (grad (ssum0 @_ @Float . scast . foo2 . (\(d, f, d2) -> (d, scast f, d2)))+         ( srepl 1.1 :: Concrete (TKS '[3, 534, 3] Double)+         , srepl 2.2 :: Concrete (TKS '[3, 534, 3] Float)+         , srepl 3.3 ))++testFooBoth :: Assertion+testFooBoth = do+  assertEqualUpToEpsilon 1e-5+    (rscalar 2.439628436155373, rscalar (-1.9533749), rscalar 0.9654825479484146)+    (grad (kfromR @_ @Float . rcast . foo2 . (\(d, f, d2) -> (d, rcast f, d2)))+         ( rscalar 1.1 :: Concrete (TKR 0 Double)+         , rscalar 2.2 :: Concrete (TKR 0 Float)+         , rscalar 3.3 ))++-- Add arrays a,b,c, but shifting b and c one to left/right+-- and then remove the first and last element.+--+-- In PyTorch+-- vstack( a[0] + b[1]+--       , a[1:N-1] + b[2:N] + c[:N-2]+--       , a[N-1] + c[N-2] )+vstackABC :: (ADReady target, NumScalar r)+          => (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))+          -> target (TKR 1 r)+vstackABC (a, b, c) =+  let n = rwidth a+  in rconcat [ rreplicate 1 (a ! [0] + b ! [1])+             , rslice 1 (n - 2) a + rslice 2 (n - 2) b + rslice 0 (n - 2) c+             , rreplicate 1 (a ! [fromIntegral n - 1]+                             + c ! [fromIntegral n - 2]) ]++vstackBuild :: (ADReady target, NumScalar r)+            => (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))+            -> target (TKR 1 r)+vstackBuild (a, b, c) =+  let n = rwidth a+  in rbuild1 n (\i ->+       ifH (i ==. 0)+           (a ! [0] + b ! [1])+           (ifH (i ==. fromIntegral n - 1)+                (a ! [fromIntegral n - 1] + c ! [fromIntegral n - 2])+                (a ! [i] + b ! [i + 1] + c ! [i - 1])))++testTrustVstackConcatRepl10 :: Assertion+testTrustVstackConcatRepl10 = do+  vstackABC @Concrete @Double (rrepl [10] 1, rrepl [10] 2, rrepl [10] 3)+  @?= rfromListLinear [10] [3.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,4.0]++testTrustVstackConcatIota10 :: Assertion+testTrustVstackConcatIota10 = do+  vstackABC @Concrete @Double (riota 10, riota 10, riota 10)+  @?= rfromListLinear [10] [1.0,3.0,6.0,9.0,12.0,15.0,18.0,21.0,24.0,17.0]++replIota :: (ADReady target, NumScalar r)+         => Int -> (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))+replIota n =+  ( rconcrete (unConcrete $ rrepl [n] 1 * riota n)+  , rconcrete (unConcrete $ rrepl [n] 2 * riota n)+  , rconcrete (unConcrete $ rrepl [n] 3 * riota n) )++testTrustVstackConcatReplIota10 :: Assertion+testTrustVstackConcatReplIota10 = do+  vstackABC @Concrete @Double (replIota 10)+  @?= rfromListLinear [10] [2.0,5.0,11.0,17.0,23.0,29.0,35.0,41.0,47.0,33.0]++nN :: Int+nN = (round :: Double -> Int) 1e6  -- 1e6++trustedResult :: Concrete (TKR 1 Double)+trustedResult =+  rcast (vstackABC @Concrete @Float (replIota nN))+    -- the cast prevents computation sharing with the first test below++testVstackWarmup :: Assertion+testVstackWarmup = do+  trustedResult+  @?= trustedResult++testVstackConcatConcrete :: Assertion+testVstackConcatConcrete = do+  vstackABC @Concrete @Double (replIota nN)+  @?= trustedResult++testVstackBuildConcrete :: Assertion+testVstackBuildConcrete = do+  vstackBuild @Concrete @Double (replIota nN)+  @?= trustedResult++testVstackConcatAst :: Assertion+testVstackConcatAst = do+  interpretAstFull @Concrete+    emptyEnv+    (vstackABC @(AstTensor AstMethodLet FullSpan) @Double+               (replIota nN))+  @?= trustedResult++testVstackBuildAst :: Assertion+testVstackBuildAst = do+  interpretAstFull @Concrete+    emptyEnv+    (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                 (replIota nN))+  @?= trustedResult++testVstackBuildAstSimp :: Assertion+testVstackBuildAstSimp = do+  interpretAstFull @Concrete+    emptyEnv+      (simplifyInlineContract+         (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                      (replIota nN)))+  @?= trustedResult++testVstackBuildAstPP :: Assertion+testVstackBuildAstPP = do+  resetVarCounter+  let (var3, ast3) =+        funToAst (FTKProduct (FTKProduct (FTKR [10] FTKScalar) (FTKR [10] FTKScalar)) (FTKR [10] FTKScalar))+          (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double . fromTarget)+  "\\" ++ printAstVarName @FullSpan var3+       ++ " -> " ++ printAstPretty ast3+    @?= "\\v1 -> rfromS (let v9 = let x6 = sfromR (tproject1 (tproject1 v1)) `sindex0` [9] + sfromR (tproject2 v1) `sindex0` [8] ; v7 = sfromR (tproject1 (tproject1 v1)) + (sappend (sslice (SNat @1) (SNat @9) (sfromR (tproject2 (tproject1 v1)))) (sconcrete (sfromListLinear [1] [0.0])) + sappend (sconcrete (sfromListLinear [1] [0.0])) (sslice (SNat @0) (SNat @9) (sfromR (tproject2 v1)))) in sappend (sslice (SNat @0) (SNat @9) v7) (sreplicate0N @[1] x6) ; x10 = sfromR (tproject1 (tproject1 v1)) `sindex0` [0] + sfromR (tproject2 (tproject1 v1)) `sindex0` [1] in sappend (sreplicate0N @[1] x10) (sslice (SNat @1) (SNat @9) v9))"+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstPretty (simplifyInlineContract ast3)+    @?= "\\v1 -> rfromS (sappend (sreplicate0N @[1] (sfromR (tproject1 (tproject1 v1)) `sindex0` [0] + sfromR (tproject2 (tproject1 v1)) `sindex0` [1])) (sappend (sslice (SNat @1) (SNat @8) (sfromR (tproject1 (tproject1 v1))) + (sslice (SNat @2) (SNat @8) (sfromR (tproject2 (tproject1 v1))) + sslice (SNat @0) (SNat @8) (sfromR (tproject2 v1)))) (sreplicate0N @[1] (sfromR (tproject1 (tproject1 v1)) `sindex0` [9] + sfromR (tproject2 v1) `sindex0` [8]))))"++replIota2 :: (ADReady target, NumScalar r)+          => Int -> (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))+replIota2 n =+  (rrepl [n] 1 * riota n, rrepl [n] 2 * riota n, rrepl [n] 3 * riota n)++testVstackConcatConcrete2 :: Assertion+testVstackConcatConcrete2 = do+  vstackABC @Concrete @Double (replIota nN)+  @?= trustedResult++testVstackBuildConcrete2 :: Assertion+testVstackBuildConcrete2 = do+  vstackBuild @Concrete @Double (replIota2 nN)+  @?= trustedResult++testVstackConcatAst2 :: Assertion+testVstackConcatAst2 = do+  interpretAstFull @Concrete+    emptyEnv+    (vstackABC @(AstTensor AstMethodLet FullSpan) @Double+               (replIota2 nN))+  @?= trustedResult++testVstackBuildAst2 :: Assertion+testVstackBuildAst2 = do+  interpretAstFull @Concrete+    emptyEnv+    (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                 (replIota2 nN))+  @?= trustedResult++testVstackBuildAstSimp2 :: Assertion+testVstackBuildAstSimp2 = do+  interpretAstFull @Concrete+    emptyEnv+      (simplifyInlineContract+         (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                      (replIota2 nN)))+  @?= trustedResult++testVstackBuildAstPP2 :: Assertion+testVstackBuildAstPP2 = do+  resetVarCounter+  (printAstPretty+     (vstackBuild @(AstTensor AstMethodLet PlainSpan) @Double+                  (replIota2 10)))+    @?= "rfromS (let v7 = let v5 = siota (SNat @10) + (sconcrete (sreplicate [10] 2.0) * sfromIntegral (sconcrete (sreplicate [10] 1) + siota (SNat @10)) + sconcrete (sreplicate [10] 3.0) * sfromIntegral (sconcrete (sreplicate [10] (-1)) + siota (SNat @10))) in sappend (sslice (SNat @0) (SNat @9) v5) (sconcrete (sfromListLinear [1] [33.0])) in sappend (sconcrete (sfromListLinear [1] [2.0])) (sslice (SNat @1) (SNat @9) v7))"+  (printAstPretty+     (simplifyInlineContract+        (vstackBuild @(AstTensor AstMethodLet PlainSpan) @Double+                     (replIota2 10))))+    @?= "rfromS (sconcrete (sfromListLinear [10] [2.0,5.0,11.0,17.0,23.0,29.0,35.0,41.0,47.0,33.0]))"++testFooPP :: Assertion+testFooPP = do+  resetVarCounter+  let fooT = foo2 @(AstTensor AstMethodLet FullSpan (TKR 0 Double))+      foo3 x = fooT (x, x, x)+      (var3, ast3) = funToAst (FTKR ZSR FTKScalar) foo3+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\x1 -> rfromK (atan2H (kfromR x1) (kfromR x1 * sin (kfromR x1)) + kfromR x1 * (kfromR x1 * sin (kfromR x1)))"+  resetVarCounter+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fooT (FTKProduct (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar)) (FTKR ZSR FTKScalar))+  printArtifactSimple (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct (STKProduct STKScalar STKScalar) STKScalar) (tlet (sin (kfromR (tproject2 (tproject1 x1)))) (\\x2 -> tlet (kfromR (tproject1 (tproject1 x1)) * x2) (\\x3 -> tlet (recip (kfromR (tproject2 x1) * kfromR (tproject2 x1) + x3 * x3)) (\\x4 -> tlet (sin (kfromR (tproject2 (tproject1 x1)))) (\\x5 -> tpair (tlet (kfromR (tproject2 x1) * kfromR dret) (\\x8 -> tlet (negate (kfromR (tproject2 x1)) * (x4 * kfromR dret)) (\\x9 -> tpair (x2 * x9 + x5 * x8) (cos (kfromR (tproject2 (tproject1 x1))) * (kfromR (tproject1 (tproject1 x1)) * x9) + cos (kfromR (tproject2 (tproject1 x1))) * (kfromR (tproject1 (tproject1 x1)) * x8))))) (x3 * (x4 * kfromR dret) + kfromR (tproject1 (tproject1 x1)) * (x5 * kfromR dret)))))))"+  printArtifactPrimalSimple (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (atan2H (kfromR (tproject2 x1)) (kfromR (tproject1 (tproject1 x1)) * sin (kfromR (tproject2 (tproject1 x1)))) + kfromR (tproject2 x1) * (kfromR (tproject1 (tproject1 x1)) * sin (kfromR (tproject2 (tproject1 x1)))))"++fooLetOld :: forall target r n.+          (RealFloatH (target (TKR n r)), LetTensor target)+       => (target (TKR n r), target (TKR n r), target (TKR n r)) -> target (TKR n r)+fooLetOld (x, y, z) =+  let w0 = x * sin y+  in tlet w0 $ \w ->+     atan2H z w + z * w++testFooLet :: Assertion+testFooLet = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627)+    (grad (kfromR @_ @Double . fooLetOld) (rscalar 1.1, rscalar 2.2, rscalar 3.3))++testFooLetPP :: Assertion+testFooLetPP = do+  resetVarCounter+  let fooLetT = fooLetOld @(AstTensor AstMethodLet FullSpan) @Double+      fooLet3 x = fooLetT (x, x, x)+      (var3, ast3) = funToAst (FTKR ZSR FTKScalar) fooLet3+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\x1 -> rfromK (tlet (kfromR x1 * sin (kfromR x1)) (\\x2 -> atan2H (kfromR x1) x2 + kfromR x1 * x2))"+  resetVarCounter+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fooLetT (FTKProduct (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar)) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct (STKProduct STKScalar STKScalar) STKScalar) (let x3 = sin (kfromR (tproject2 (tproject1 x1))) ; x4 = kfromR (tproject1 (tproject1 x1)) * x3 ; x5 = recip (kfromR (tproject2 x1) * kfromR (tproject2 x1) + x4 * x4) in tpair (let x7 = negate (kfromR (tproject2 x1)) * (x5 * kfromR dret) + kfromR (tproject2 x1) * kfromR dret in tpair (x3 * x7) (cos (kfromR (tproject2 (tproject1 x1))) * (kfromR (tproject1 (tproject1 x1)) * x7))) (x4 * (x5 * kfromR dret) + x4 * kfromR dret))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (let x4 = kfromR (tproject1 (tproject1 x1)) * sin (kfromR (tproject2 (tproject1 x1))) in atan2H (kfromR (tproject2 x1)) x4 + kfromR (tproject2 x1) * x4)"++shapedListProd :: forall k target r. (BaseTensor target, NumScalar r)+               => IxR k (target (TKS '[] r)) -> target (TKS '[] r)+shapedListProd = foldr1 (*)++testListProdPP :: Assertion+testListProdPP = do+  resetVarCounter+  let fT :: IxR 4 (AstTensor AstMethodLet FullSpan (TKS '[] Double)) -> AstTensor AstMethodLet FullSpan (TKS '[] Double)+      fT = shapedListProd+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) ftkUnit))))+  printArtifactSimple artifactRev+    @?= "\\dret x1 -> tlet (kfromS (tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1))))) (\\x2 -> tpair (tlet (kfromS (tproject1 (tproject2 x1) * sfromK x2)) (\\x3 -> sfromK x3 * dret)) (tlet (kfromS (tproject1 x1 * dret)) (\\x5 -> tpair (sfromK x2 * sfromK x5) (tlet (kfromS (tproject1 (tproject2 x1) * sfromK x5)) (\\x6 -> tpair (tproject1 (tproject2 (tproject2 (tproject2 x1))) * sfromK x6) (tpair (tproject1 (tproject2 (tproject2 x1)) * sfromK x6) (tfromPlain (STKScalar) Z1)))))))"+  printArtifactPrimalSimple artifactRev+    @?= "\\x1 -> tlet (kfromS (tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1))))) (\\x2 -> tlet (kfromS (tproject1 (tproject2 x1) * sfromK x2)) (\\x3 -> tproject1 x1 * sfromK x3))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp ConvXS (Conv0X STKScalar)) (ConvT2 (ConvCmp ConvXS (Conv0X STKScalar)) (ConvT2 (ConvCmp ConvXS (Conv0X STKScalar)) (ConvT2 (ConvCmp ConvXS (Conv0X STKScalar)) ConvId)))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar)))) (let x2 = kfromS (tproject1 (tproject2 (tproject2 x1))) * kfromS (tproject1 (tproject2 (tproject2 (tproject2 x1)))) in tpair (kfromS (tproject1 (tproject2 x1)) * (x2 * kfromS dret)) (let x5 = kfromS (tproject1 x1) * kfromS dret in tpair (x2 * x5) (let x6 = kfromS (tproject1 (tproject2 x1)) * x5 in tpair (kfromS (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x6) (tpair (kfromS (tproject1 (tproject2 (tproject2 x1))) * x6) Z1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> sfromK (kfromS (tproject1 x1) * (kfromS (tproject1 (tproject2 x1)) * (kfromS (tproject1 (tproject2 (tproject2 x1))) * kfromS (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++rankedListProdr :: forall k target r. (BaseTensor target, NumScalar r)+                => IxR k (target (TKR 0 r)) -> target (TKR 0 r)+rankedListProdr = foldr1 (*)++testListProdrPP :: Assertion+testListProdrPP = do+  resetVarCounter+  let fT :: IxR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListProdr+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId)))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar)))) (let x2 = kfromR (tproject1 (tproject2 (tproject2 x1))) * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) in tpair (kfromR (tproject1 (tproject2 x1)) * (x2 * kfromR dret)) (let x5 = kfromR (tproject1 x1) * kfromR dret in tpair (x2 * x5) (let x6 = kfromR (tproject1 (tproject2 x1)) * x5 in tpair (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x6) (tpair (kfromR (tproject1 (tproject2 (tproject2 x1))) * x6) Z1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (kfromR (tproject1 x1) * (kfromR (tproject1 (tproject2 x1)) * (kfromR (tproject1 (tproject2 (tproject2 x1))) * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++testListProdrLongPP :: Assertion+testListProdrLongPP = do+  resetVarCounter+  let fT :: IxR 13 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListProdr+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit)))))))))))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId))))))))))))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar))))))))))))) (let x2 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))))) ; x3 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x2 ; x4 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x3 ; x5 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x4 ; x6 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x5 ; x7 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x6 ; x8 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x7 ; x9 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x8 ; x10 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x9 ; x11 = kfromR (tproject1 (tproject2 (tproject2 x1))) * x10 in tpair (kfromR (tproject1 (tproject2 x1)) * (x11 * kfromR dret)) (let x14 = kfromR (tproject1 x1) * kfromR dret in tpair (x11 * x14) (let x15 = kfromR (tproject1 (tproject2 x1)) * x14 in tpair (x10 * x15) (let x16 = kfromR (tproject1 (tproject2 (tproject2 x1))) * x15 in tpair (x9 * x16) (let x17 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x16 in tpair (x8 * x17) (let x18 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x17 in tpair (x7 * x18) (let x19 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x18 in tpair (x6 * x19) (let x20 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x19 in tpair (x5 * x20) (let x21 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x20 in tpair (x4 * x21) (let x22 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x21 in tpair (x3 * x22) (let x23 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x22 in tpair (x2 * x23) (let x24 = kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x23 in tpair (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))))) * x24) (tpair (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * x24) Z1)))))))))))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (kfromR (tproject1 x1) * (kfromR (tproject1 (tproject2 x1)) * (kfromR (tproject1 (tproject2 (tproject2 x1))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * (kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * kfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))))))))))))))))"++testListProd :: Assertion+testListProd = do+  assertEqualUpToEpsilon 1e-10+    [srepl 24, srepl 12, srepl 8, srepl 6]+    (grad (kfromS @_ @Double . shapedListProd @4) [srepl 1, srepl 2, srepl 3, srepl 4])++testListProdr :: Assertion+testListProdr = do+  assertEqualUpToEpsilon 1e-10+    [rscalar 24, rscalar 12, rscalar 8, rscalar 6]+    (grad (kfromR @_ @Double . rankedListProdr @4) [rscalar 1, rscalar 2, rscalar 3, rscalar 4])++rankedListSumr :: (BaseTensor target, NumScalar r)+               => IxR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSumr = foldr1 (+)++testListSumrPP :: Assertion+testListSumrPP = do+  resetVarCounter >> resetIdCounter+  let fT :: IxR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSumr+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tpair dret (tpair dret (tpair dret (tpair dret Z1)))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (kfromR (tproject1 x1) + (kfromR (tproject1 (tproject2 x1)) + (kfromR (tproject1 (tproject2 (tproject2 x1))) + kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (DeltaShare 100000003 (DeltaAdd (DeltaConvert (ConvCmp ConvX0 ConvRX) (DeltaInput (InputId 0))) (DeltaShare 100000002 (DeltaAdd (DeltaConvert (ConvCmp ConvX0 ConvRX) (DeltaInput (InputId 1))) (DeltaShare 100000001 (DeltaAdd (DeltaConvert (ConvCmp ConvX0 ConvRX) (DeltaInput (InputId 2))) (DeltaConvert (ConvCmp ConvX0 ConvRX) (DeltaInput (InputId 3)))))))))"++-- Note that the function is not associative, so foldr vs foldl matters.+rankedListSum2r :: (BaseTensor target, NumScalar r)+                => IxR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum2r = foldr1 (\x y -> x + rscalar 2 * y)++testListSum2rPP :: Assertion+testListSum2rPP = do+  resetVarCounter+  let fT ::  IxR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum2r+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId)))) (STKProduct (STKR (SNat @0) STKScalar) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar)))) (tpair dret (tpair (2.0 * kfromR dret) (tpair (4.0 * kfromR dret) (tpair (8.0 * kfromR dret) Z1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (kfromR (tproject1 x1) + (2.0 * kfromR (tproject1 (tproject2 x1)) + (4.0 * kfromR (tproject1 (tproject2 (tproject2 x1))) + 8.0 * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++rankedListSum22r :: (BaseTensor target, NumScalar r)+                 => IxR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum22r = foldr1 (\x y -> rscalar 2 * x + rscalar 2 * y)++testListSum22rPP :: Assertion+testListSum22rPP = do+  resetVarCounter+  let fT ::  IxR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum22r+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId)))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar)))) (tpair (2.0 * kfromR dret) (tpair (4.0 * kfromR dret) (tpair (8.0 * kfromR dret) (tpair (8.0 * kfromR dret) Z1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (2.0 * kfromR (tproject1 x1) + (4.0 * kfromR (tproject1 (tproject2 x1)) + (8.0 * kfromR (tproject1 (tproject2 (tproject2 x1))) + 8.0 * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++-- Note how this tlet did not change anything, in particular the sharing.+rankedListSumk22r :: ( BaseTensor target, LetTensor target+                     , NumScalar r )+                 =>  IxR k (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSumk22r = foldr1 (\x y -> tlet (rscalar 2) (\k -> k * x + k * y))++testListSumk22rPP :: Assertion+testListSumk22rPP = do+  resetVarCounter+  let fT ::  IxR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSumk22r+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId)))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar)))) (tpair (2.0 * kfromR dret) (tpair (4.0 * kfromR dret) (tpair (8.0 * kfromR dret) (tpair (8.0 * kfromR dret) Z1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (2.0 * kfromR (tproject1 x1) + (4.0 * kfromR (tproject1 (tproject2 x1)) + (8.0 * kfromR (tproject1 (tproject2 (tproject2 x1))) + 8.0 * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++rankedListSum2xpyr :: (BaseTensor target, NumScalar r)+                   =>  IxR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum2xpyr = foldr1 (\x y -> rscalar 2 * (x + y))++testListSum2xpyrPP :: Assertion+testListSum2xpyrPP = do+  resetVarCounter+  let fT :: IxR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum2xpyr+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId)))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar)))) (tpair (2.0 * kfromR dret) (tpair (4.0 * kfromR dret) (tpair (8.0 * kfromR dret) (tpair (8.0 * kfromR dret) Z1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (2.0 * kfromR (tproject1 x1) + (4.0 * kfromR (tproject1 (tproject2 x1)) + (8.0 * kfromR (tproject1 (tproject2 (tproject2 x1))) + 8.0 * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++rankedListSum2xyr :: (BaseTensor target, NumScalar r)+                  => IxR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum2xyr = foldr1 (\x y -> rscalar 2 * (x * y))++testListSum2xyrPP :: Assertion+testListSum2xyrPP = do+  resetVarCounter+  let fT :: IxR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum2xyr+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId)))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar)))) (let x2 = kfromR (tproject1 (tproject2 (tproject2 x1))) * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) ; x6 = 8.0 * kfromR dret in tpair (kfromR (tproject1 (tproject2 x1)) * (x2 * x6)) (let x7 = kfromR (tproject1 x1) * x6 in tpair (x2 * x7) (let x8 = kfromR (tproject1 (tproject2 x1)) * x7 in tpair (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x8) (tpair (kfromR (tproject1 (tproject2 (tproject2 x1))) * x8) Z1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (8.0 * (kfromR (tproject1 x1) * (kfromR (tproject1 (tproject2 x1)) * (kfromR (tproject1 (tproject2 (tproject2 x1))) * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1))))))))"++ranked2xy :: (BaseTensor target, NumScalar r)+          => (target (TKR 0 r), target (TKR 0 r)) -> target (TKR 0 r)+ranked2xy = \(x, y) -> rscalar 2 * x * y++test2xyPP :: Assertion+test2xyPP = do+  resetVarCounter+  let fT :: (AstTensor AstMethodLet FullSpan (TKR 0 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = ranked2xy+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (let x4 = 2.0 * kfromR dret in tpair (kfromR (tproject2 x1) * x4) (kfromR (tproject1 x1) * x4))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (2.0 * (kfromR (tproject1 x1) * kfromR (tproject2 x1)))"++-- Note that the function is not associative, so foldr vs foldl matters.+rankedListSum23r :: forall k target r. (BaseTensor target, NumScalar r)+                 => IxR k (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum23r = foldr1 (\x y -> rscalar 2 * x + rscalar 3 * y)++testListSum23rPP :: Assertion+testListSum23rPP = do+  resetVarCounter+  let fT :: IxR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum23r+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId)))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar)))) (tpair (2.0 * kfromR dret) (tpair (6.0 * kfromR dret) (tpair (18.0 * kfromR dret) (tpair (27.0 * kfromR dret) Z1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (2.0 * kfromR (tproject1 x1) + (6.0 * kfromR (tproject1 (tproject2 x1)) + (18.0 * kfromR (tproject1 (tproject2 (tproject2 x1))) + 27.0 * kfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++ranked23 :: (BaseTensor target, NumScalar r)+         => (target (TKR 0 r), target (TKR 0 r)) -> target (TKR 0 r)+ranked23 = \(x, y) -> rscalar 2 * x + rscalar 3 * y++test23PP :: Assertion+test23PP = do+  resetVarCounter+  let fT :: (AstTensor AstMethodLet FullSpan (TKR 0 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = ranked23+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (tpair (2.0 * kfromR dret) (3.0 * kfromR dret))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\x1 -> rfromK (2.0 * kfromR (tproject1 x1) + 3.0 * kfromR (tproject2 x1))"++reluPrimal+  :: forall target n r.+     (ADReady target, NumScalar r, KnownNat n, Differentiable r)+  => target (TKR n r) -> target (TKR n r)+reluPrimal v =+  let oneIfGtZero = rmap0N (\x -> ifH (x <=. 0) 0.0 1.0)+                           (rprimalPart v)+  in scale2 oneIfGtZero v++scale2 :: forall target r n.+          (ADReady target, NumScalar r, KnownNat n)+       => PrimalOf target (TKR n r) -> target (TKR n r) -> target (TKR n r)+scale2 a d = rfromPrimal @target a * d++testReluPP :: Assertion+testReluPP = do+  resetVarCounter >> resetIdCounter+  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT = reluPrimal+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) reluT+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (tfromPlain (STKS [3,4] STKScalar) (sgather @[3,4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i4, i3] -> [ifH (0.0 <=. negate (tplainPart (sfromR m1) `sindex0` [i4, i3])) 0 1])) * sfromR m1)"+  resetVarCounter+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> rfromS (sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (0.0 <=. negate (splainPart (sfromR m1) `sindex0` [i5, i6])) 0 1])) * sfromR dret)"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (0.0 <=. negate (splainPart (sfromR m1) `sindex0` [i5, i6])) 0 1])) * sfromR m1)"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000001 (DeltaScale <primal> (DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0)))))"++testReluPP2 :: Assertion+testReluPP2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      reluT2 (t, r) = reluPrimal (t * rreplicate 5 r)+      (var3, ast3) = funToAst (FTKR [5] FTKScalar) (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> rfromS (tfromPlain (STKS [5] STKScalar) (sgather1 @5 (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\i2 -> [ifH (0.0 <=. (-7.0) * tplainPart (sfromR v1) `sindex0` [i2]) 0 1])) * (sfromR v1 * tfromPlain (STKS [5] STKScalar) (sreplicate0N @[5] 7.0)))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct (STKS [5] STKScalar) STKScalar) (let v7 = sfromPlain (sgather1 @5 (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\i3 -> [ifH (0.0 <=. negate (splainPart (sfromR (tproject1 v1)) `sindex0` [i3]) * kplainPart (kfromR (tproject2 v1))) 0 1])) * sfromR dret in tpair (sreplicate0N @[5] (kfromR (tproject2 v1)) * v7) (sdot0 (sfromR (tproject1 v1)) v7))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromS (sfromPlain (sgather1 @5 (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\i3 -> [ifH (0.0 <=. negate (splainPart (sfromR (tproject1 v1)) `sindex0` [i3]) * kplainPart (kfromR (tproject2 v1))) 0 1])) * (sfromR (tproject1 v1) * sreplicate0N @[5] (kfromR (tproject2 v1))))"++testReluSimpler :: Assertion+testReluSimpler = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3, 4] [1.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,0.0,0.0,1.0,1.0])+    (rev' @Double @2 relu (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12]))++testReluSimplerPP :: Assertion+testReluSimplerPP = do+  resetVarCounter >> resetIdCounter+  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT = relu+      (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> rfromS (sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (0.0 <=. negate (splainPart (sfromR m1) `sindex0` [i5, i6])) 0 1])) * sfromR dret)"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (0.0 <=. negate (splainPart (sfromR m1) `sindex0` [i5, i6])) 0 1])) * sfromR m1)"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000001 (DeltaScale <primal> (DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0)))))"++testReluSimplerPP2 :: Assertion+testReluSimplerPP2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      reluT2 (t, r) = relu (t * rreplicate 5 r)+      (var3, ast3) = funToAst (FTKR [5] FTKScalar) (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> rfromS (tlet (sfromR v1 * tfromPlain (STKS [5] STKScalar) (sreplicate0N @[5] 7.0)) (\\v2 -> tfromPlain (STKS [5] STKScalar) (sgather1 @5 (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\i3 -> [ifH (0.0 <=. negate (tplainPart v2 `sindex0` [i3])) 0 1])) * v2))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct (STKS [5] STKScalar) STKScalar) (let v4 = sfromR (tproject1 v1) * sreplicate0N @[5] (kfromR (tproject2 v1)) ; v8 = sfromPlain (sgather1 @5 (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\i5 -> [ifH (0.0 <=. negate (splainPart v4 `sindex0` [i5])) 0 1])) * sfromR dret in tpair (sreplicate0N @[5] (kfromR (tproject2 v1)) * v8) (sdot0 (sfromR (tproject1 v1)) v8))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromS (let v4 = sfromR (tproject1 v1) * sreplicate0N @[5] (kfromR (tproject2 v1)) in sfromPlain (sgather1 @5 (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\i5 -> [ifH (0.0 <=. negate (splainPart v4 `sindex0` [i5])) 0 1])) * v4)"++testReluSimplerPP3 :: Assertion+testReluSimplerPP3 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = relu (t * rreplicate 3 (rreplicate 4 r))+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (tlet (tfromPlain (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0)) * sfromR m1) (\\m2 -> tfromPlain (STKS [3,4] STKScalar) (sgather @[3,4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (0.0 <=. negate (tplainPart m2 `sindex0` [i5, i4])) 0 1])) * m2))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [3, 4] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct (STKS [3,4] STKScalar) STKScalar) (let m6 = sfromR (tproject1 m1) * sreplicate0N @[3, 4] (kfromR (tproject2 m1)) ; m11 = sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (0.0 <=. negate (splainPart m6 `sindex0` [i7, i8])) 0 1])) * sfromR dret in tpair (sreplicate0N @[3, 4] (kfromR (tproject2 m1)) * m11) (sdot0 (sfromR (tproject1 m1)) m11))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (let m6 = sfromR (tproject1 m1) * sreplicate0N @[3, 4] (kfromR (tproject2 m1)) in sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (0.0 <=. negate (splainPart m6 `sindex0` [i7, i8])) 0 1])) * m6)"++testReluSimpler3 :: Assertion+testReluSimpler3 = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = relu (t * rreplicate 3 (rreplicate 4 r))+  assertEqualUpToEpsilon 1e-10+    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , rscalar 57.1 )+    (grad (rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))++testReluSimplerPP4 :: Assertion+testReluSimplerPP4 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = relu (t * rreplicate0N [3, 4] (kfromR r))+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (tlet (tfromPlain (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0)) * sfromR m1) (\\m2 -> tfromPlain (STKS [3,4] STKScalar) (sgather @[3,4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (0.0 <=. negate (tplainPart m2 `sindex0` [i5, i4])) 0 1])) * m2))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [3, 4] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct (STKS [3,4] STKScalar) STKScalar) (let m6 = sfromR (tproject1 m1) * sreplicate0N @[3, 4] (kfromR (tproject2 m1)) ; m11 = sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (0.0 <=. negate (splainPart m6 `sindex0` [i7, i8])) 0 1])) * sfromR dret in tpair (sreplicate0N @[3, 4] (kfromR (tproject2 m1)) * m11) (sdot0 (sfromR (tproject1 m1)) m11))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (let m6 = sfromR (tproject1 m1) * sreplicate0N @[3, 4] (kfromR (tproject2 m1)) in sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (0.0 <=. negate (splainPart m6 `sindex0` [i7, i8])) 0 1])) * m6)"++testReluSimpler4 :: Assertion+testReluSimpler4 = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = relu (t * rreplicate0N [3, 4] (kfromR r))+  assertEqualUpToEpsilon 1e-10+    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , rscalar 57.1 )+    (grad (rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))++testReluSimplerPP4s :: Assertion+testReluSimplerPP4s = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Float), AstTensor AstMethodLet FullSpan (TKS '[] Float))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Float)+      reluT2 (t, r) = reluS (t * sreplicate0N (kfromS r))+      (var3, ast3) = funToAst (FTKS knownShS FTKScalar) (\t -> reluT2 (t, srepl 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> tlet (tfromPlain (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0)) * m1) (\\m2 -> tfromPlain (STKS [3,4] STKScalar) (sgather @[3,4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (0.0 <=. negate (tplainPart m2 `sindex0` [i5, i4])) 0 1])) * m2)"++testReluSimplerPP4s2 :: Assertion+testReluSimplerPP4s2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)+      -- This is tweaked compared to above to avoid test artifacts coming+      -- from counter resets, which are inherently unsafe (cse, etc.).+      reluT2 (t, r) = reluS (t * sreplicate0N (kfromS r))+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKS [3, 4] FTKScalar) (FTKS ZSS FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m6 = tproject1 m1 * sreplicate0N @[3, 4] (kfromS (tproject2 m1)) ; m9 = sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (0.0 <=. negate (splainPart m6 `sindex0` [i7, i8])) 0 1]) ; m11 = sfromPlain m9 * dret in tpair (sreplicate0N @[3, 4] (kfromS (tproject2 m1)) * m11) (sfromK (ssum0 @[3, 4] (tproject1 m1 * m11)))"+  printArtifactPrimalPretty artifactRev+    @?=  "\\m1 -> let m6 = tproject1 m1 * sreplicate0N @[3, 4] (kfromS (tproject2 m1)) ; m9 = sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (0.0 <=. negate (splainPart m6 `sindex0` [i7, i8])) 0 1]) in sfromPlain m9 * m6"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 ConvId (ConvCmp ConvXS (Conv0X STKScalar))) (STKProduct (STKS [3,4] STKScalar) STKScalar) (let m6 = tproject1 m1 * sreplicate0N @[3, 4] (kfromS (tproject2 m1)) ; m11 = sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (0.0 <=. negate (splainPart m6 `sindex0` [i7, i8])) 0 1])) * dret in tpair (sreplicate0N @[3, 4] (kfromS (tproject2 m1)) * m11) (sdot0 (tproject1 m1) m11))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> let m6 = tproject1 m1 * sreplicate0N @[3, 4] (kfromS (tproject2 m1)) in sfromPlain (sgather @[3, 4] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (0.0 <=. negate (splainPart m6 `sindex0` [i7, i8])) 0 1])) * m6"++testReluSimpler4s :: Assertion+testReluSimpler4s = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)+      reluT2 (t, r) = reluS (t * sreplicate0N (kfromS r))+  assertEqualUpToEpsilon 1e-10+    ( sconcrete+      $ Nested.sfromListPrimLinear @'[3, 4] knownShS [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , srepl 57.1 )+    (grad (ssum0 . reluT2) (sconcrete $ Nested.sfromListPrimLinear @'[3, 4] knownShS [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], srepl 7))++testReluSimplerPP7s2 :: Assertion+testReluSimplerPP7s2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)+      reluT2 (t, r) = smap0N (flip tlet $ \x -> ifH (x <=. 0) 0.0 x) (t * sreplicate0N (kfromS r))+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKS [3, 4] FTKScalar) (FTKS ZSS FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m9 = tproject1 m1 * sreplicate0N @[3, 4] (kfromS (tproject2 m1)) ; t15 = stranspose @[2, 0, 1] (sscatter @[3, 4] dret (\\[i13, i14] -> [i13, i14, ifH (0.0 <=. negate (splainPart m9 `sindex0` [i13, i14])) 0 1])) ; m16 = t15 !$ [1] in tpair (sreplicate0N @[3, 4] (kfromS (tproject2 m1)) * m16) (sfromK (ssum0 @[3, 4] (tproject1 m1 * m16)))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> let m9 = tproject1 m1 * sreplicate0N @[3, 4] (kfromS (tproject2 m1)) in sgather @[3, 4] (stranspose @[1, 2, 0] (sfromVector @2 (fromList [sconcrete (sreplicate [3,4] 0.0), m9]))) (\\[i10, i11] -> [i10, i11, ifH (0.0 <=. negate (splainPart m9 `sindex0` [i10, i11])) 0 1])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 ConvId (ConvCmp ConvXS (Conv0X STKScalar))) (STKProduct (STKS [3,4] STKScalar) STKScalar) (let m9 = tproject1 m1 * sreplicate0N @[3, 4] (kfromS (tproject2 m1)) ; m16 = sscatter @[3, 4] dret (\\[i13, i14] -> [ifH (0.0 <=. negate (splainPart m9 `sindex0` [i13, i14])) 0 1, i13, i14]) !$ [1] in tpair (sreplicate0N @[3, 4] (kfromS (tproject2 m1)) * m16) (sdot0 (tproject1 m1) m16))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> let m9 = tproject1 m1 * sreplicate0N @[3, 4] (kfromS (tproject2 m1)) in sgather @[3, 4] (sfromVector @2 (fromList [sconcrete (sreplicate [3,4] 0.0), m9])) (\\[i10, i11] -> [ifH (0.0 <=. negate (splainPart m9 `sindex0` [i10, i11])) 0 1, i10, i11])"++testReluSimpler7s :: Assertion+testReluSimpler7s = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)+      reluT2 (t, r) = smap0N (flip tlet $ \x -> ifH (x <=. 0) 0.0 x) (t * sreplicate0N (kfromS r))+  assertEqualUpToEpsilon 1e-10+    ( sconcrete+      $ Nested.sfromListPrimLinear @'[3, 4] knownShS [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , srepl 57.1 )+    (grad (ssum0 . reluT2) (sconcrete $ Nested.sfromListPrimLinear @'[3, 4] knownShS [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], srepl 7))++reluMax :: forall target n r. (ADReady target, NumScalar r, KnownNat n)+        => target (TKR n r) -> target (TKR n r)+reluMax = rmap0N (maxH 0)++testReluMax :: Assertion+testReluMax = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3, 4] [1.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,0.0,0.0,1.0,1.0])+    (rev' @Double @2 reluMax (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12]))++testReluMaxPP :: Assertion+testReluMaxPP = do+  resetVarCounter >> resetIdCounter+  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT = reluMax+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) reluT+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (sgather @[3,4] (ttranspose (makePerm @[1, 2, 0]) (sfromVector @2 (fromList [tfromPlain (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 0.0)), sfromR m1]))) (\\[i8, i6] -> [i8, i6, ifH (0.0 <=. negate (tplainPart (sfromR m1) `sindex0` [i8, i6])) 0 1]))"+  resetVarCounter+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> rfromS (sscatter @[3, 4] (sfromR dret) (\\[i12, i13] -> [ifH (0.0 <=. negate (splainPart (sfromR m1) `sindex0` [i12, i13])) 0 1, i12, i13]) !$ [1])"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (sgather @[3, 4] (sfromVector @2 (fromList [sconcrete (sreplicate [3,4] 0.0), sfromR m1])) (\\[i9, i10] -> [ifH (0.0 <=. negate (splainPart (sfromR m1) `sindex0` [i9, i10])) 0 1, i9, i10]))"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000003 (DeltaGatherS [3,4] [] [3,4,2] (DeltaShare 100000002 (DeltaTransposeS (PCons (SNat @1) (PCons (SNat @2) (PCons (SNat @0) PNil))) (DeltaShare 100000001 (DeltaFromVectorS [2] [DeltaZero (FTKS [3,4] FTKScalar),DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0))])))) <function>))"++testReluMaxPP2 :: Assertion+testReluMaxPP2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      reluT2 (t, r) = reluMax (t * rreplicate 5 r)+      (var3, ast3) = funToAst (FTKR [5] FTKScalar) (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> rfromS (tlet (tfromPlain (STKS [5] STKScalar) (sconcrete (sreplicate [5] 7.0)) * sfromR v1) (\\v3 -> sgather1 @5 (str (sfromVector @2 (fromList [tfromPlain (STKS [5] STKScalar) (sconcrete (sreplicate [5] 0.0)), v3]))) (\\i5 -> [i5, ifH (0.0 <=. negate (tplainPart v3 `sindex0` [i5])) 0 1])))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct (STKS [5] STKScalar) STKScalar) (let v6 = sfromR (tproject1 v1) * sreplicate0N @[5] (kfromR (tproject2 v1)) ; v11 = sscatter1 @5 (sfromR dret) (\\i9 -> [ifH (0.0 <=. negate (splainPart v6 `sindex0` [i9])) 0 1, i9]) !$ [1] in tpair (sreplicate0N @[5] (kfromR (tproject2 v1)) * v11) (sdot0 (sfromR (tproject1 v1)) v11))"+  printArtifactPretty artifactRev+    @?= "\\dret v1 -> let v6 = sfromR (tproject1 v1) * sreplicate0N @[5] (kfromR (tproject2 v1)) ; m10 = str (sscatter1 @5 (sfromR dret) (\\i9 -> [i9, ifH (0.0 <=. negate (splainPart v6 `sindex0` [i9])) 0 1])) ; v11 = m10 !$ [1] in tpair (rfromS (sreplicate0N @[5] (kfromR (tproject2 v1)) * v11)) (rfromK (ssum0 @[5] (sfromR (tproject1 v1) * v11)))"+  printArtifactPrimalPretty artifactRev+    @?= "\\v1 -> rfromS (let v6 = sfromR (tproject1 v1) * sreplicate0N @[5] (kfromR (tproject2 v1)) in sgather1 @5 (str (sfromVector @2 (fromList [sconcrete (sreplicate [5] 0.0), v6]))) (\\i7 -> [i7, ifH (0.0 <=. negate (splainPart v6 `sindex0` [i7])) 0 1]))"++testReluMax3 :: Assertion+testReluMax3 = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = reluMax (t * rreplicate 3 (rreplicate 4 r))+  assertEqualUpToEpsilon 1e-10+    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , rscalar 57.1 )+    (grad (rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))++testDot1PP :: Assertion+testDot1PP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent (uncurry (rdot0 @1 @Double))+                 (FTKProduct (FTKR [3] FTKScalar) (FTKR [3] FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret v1 -> tpair (rfromS (sfromR (tproject2 v1) * sreplicate0N @[3] dret)) (rfromS (sfromR (tproject1 v1) * sreplicate0N @[3] dret))"+  printArtifactPrimalPretty artifactRev+    @?= "\\v1 -> ssum0 @[3] (sfromR (tproject1 v1) * sfromR (tproject2 v1))"++testDot2PP :: Assertion+testDot2PP = do+  resetVarCounter+  let (artifactRev, _deltas) =+        revArtifactDelta UseIncomingCotangent (uncurry (rdot0 @2 @Double))+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [2, 3] FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> tpair (rfromS (sfromR (tproject2 m1) * sreplicate0N @[2, 3] dret)) (rfromS (sfromR (tproject1 m1) * sreplicate0N @[2, 3] dret))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> ssum0 @[2, 3] (sfromR (tproject1 m1) * sfromR (tproject2 m1))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,3] STKScalar) (STKS [2,3] STKScalar)) (tpair (sfromR (tproject2 m1) * sreplicate0N @[2, 3] dret) (sfromR (tproject1 m1) * sreplicate0N @[2, 3] dret))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> sdot0 (sfromR (tproject1 m1)) (sfromR (tproject2 m1))"++testMatvecmulPP :: Assertion+testMatvecmulPP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry rmatvecmul)+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3] FTKScalar))+  printArtifactPretty @_ @(TKR 1 Double) artifactRev+    @?= "\\dret m1 -> tpair (rfromS (str (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret)))) (rfromS (ssum @2 (str (str (sfromR (tproject1 m1)) * sreplicate @3 (sfromR dret)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @3 (str (sreplicate @2 (sfromR (tproject2 m1))) * str (sfromR (tproject1 m1))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,3] STKScalar) (STKS [3] STKScalar)) (tpair (sreplicate @2 (sfromR (tproject2 m1)) * str (sreplicate @3 (sfromR dret))) (sdot1In (str (sfromR (tproject1 m1))) (sreplicate @3 (sfromR dret))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (sdot1In (sreplicate @2 (sfromR (tproject2 m1))) (sfromR (tproject1 m1)))"++testMatmul2PP :: Assertion+testMatmul2PP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry rmatmul2)+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))+  printArtifactPretty @_ @(TKR 2 Double) artifactRev+    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2, 1, 0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2, 1, 0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2, 1, 0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,3] STKScalar) (STKS [3,4] STKScalar)) (tpair (smatmul2 (sfromR dret) (str (sfromR (tproject2 m1)))) (smatmul2 (str (sfromR (tproject1 m1))) (sfromR dret)))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (smatmul2 (sfromR (tproject1 m1)) (sfromR (tproject2 m1)))"++testMatmul2FromMatvecmulPP :: Assertion+testMatmul2FromMatvecmulPP = do+  resetVarCounter+  let rmatmul2F :: (BaseTensor target, ConvertTensor target, NumScalar r)+                => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)+      rmatmul2F m1 m2 =+        rbuild1 (rwidth m1) (\i -> rmatvecmul (rtr m2) (m1 ! [i]))+      (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry rmatmul2F)+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))+  printArtifactPretty @_ @(TKR 2 Double) artifactRev+    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2, 1, 0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2, 1, 0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2, 1, 0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"++testMatmul2PaperPP :: Assertion+testMatmul2PaperPP = do+  resetVarCounter+  let rmatmul2P :: (BaseTensor target, NumScalar r)+                => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)+      rmatmul2P a b =+        let k :$: m :$: _ = rshape a+            _ :$: n :$: _ = rshape b+        in rbuild1 k (\i ->+             rbuild1 n (\j ->+               rsum (rbuild1 m (\p -> a ! [i, p] * b ! [p, j]))))+      (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry rmatmul2P)+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))+  printArtifactPretty @_ @(TKR 2 Double) artifactRev+    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2, 1, 0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2, 1, 0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2, 1, 0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"++testMatmul2PPS :: Assertion+testMatmul2PPS = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry smatmul2)+                 (FTKProduct (FTKS (SNat @2 :$$ SNat @3 :$$ ZSS) (FTKScalar @Float)) (FTKS (SNat @3 :$$ SNat @4 :$$ ZSS) (FTKScalar @Float)))+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> tpair (ssum @4 (stranspose @[2, 1, 0] (str (sreplicate @2 (tproject2 m1)) * sreplicate @3 dret))) (ssum @2 (str (stranspose @[2, 1, 0] (sreplicate @4 (tproject1 m1)) * sreplicate @3 dret)))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> ssum @3 (stranspose @[2, 1, 0] (sreplicate @4 (tproject1 m1)) * str (sreplicate @2 (tproject2 m1)))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tpair (smatmul2 dret (str (tproject2 m1))) (smatmul2 (str (tproject1 m1)) dret)"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> smatmul2 (tproject1 m1) (tproject2 m1)"++testAddSpeedBig :: Assertion+testAddSpeedBig =+  let m1 :: Concrete (TKR 2 Double)+      m1 = ringestData [1000,1000] [1 .. 1000000]+      m2 = m1 + m1+  in assertEqualUpToEpsilon 1e-10 m2 m2++testMatmul2SpeedBig :: Assertion+testMatmul2SpeedBig =+  let m1 :: Concrete (TKR 2 Double)+      m1 = ringestData [1000,1000] [1 .. 1000000]+      m2 = rmatmul2 m1 m1+  in assertEqualUpToEpsilon 1e-10 m2 m2++bar :: forall a. RealFloatH a => (a, a) -> a+bar (x, y) =+  let w = foo2 (x, y, x) * sin y+  in atan2H x w + y * w++barF :: forall a. RealFloatH a => (a, a) -> a+barF (x, y) =+  let w = foo2 (x, y, x) * sin y+  in atan2H x w + y * w++testBar :: Assertion+testBar =+  assertEqualUpToEpsilon 1e-9+    (rscalar 3.1435239435581166,rscalar (-1.1053869545195814))+    (cgrad (kfromR . bar @(ADVal Concrete (TKR 0 Double))) (rscalar 1.1, rscalar 2.2))++testBarS :: Assertion+testBarS =+  assertEqualUpToEpsilon 1e-9+    (srepl 3.1435239435581166, srepl (-1.1053869545195814))+    (cgrad (kfromS . barF @(ADVal Concrete (TKS '[] Double))) (srepl 1.1, srepl 2.2))++testBar2S :: Assertion+testBar2S =+  assertEqualUpToEpsilon 1e-9+    (srepl 3.1435239435581166, srepl (-1.1053869545195814))+    (grad (ssum0 . barF @(AstTensor AstMethodLet FullSpan (TKS '[52, 2, 2, 1, 1, 3] Double))) (srepl 1.1, srepl 2.2))++testBarCFwd :: Assertion+testBarCFwd =+  assertEqualUpToEpsilon 1e-9+    (rscalar 9.327500345189534e-2)+    (cjvp (bar @(ADVal Concrete (TKR 0 Double))) (rscalar 1.1, rscalar 2.2) (rscalar 0.1, rscalar 0.2))++testBarFwd :: Assertion+testBarFwd =+  assertEqualUpToEpsilon 1e-9+    (rscalar 9.327500345189534e-2)+    (jvp @_ @(TKR 0 Double)+         bar (rscalar 1.1, rscalar 2.2) (rscalar 0.1, rscalar 0.2))++barADVal2 :: forall a. RealFloatH a+          => (a, a, a) -> a+barADVal2 (x,y,z) =+  let w = foo2 (x,y,z) * sin y+  in atan2H z w + z * w++-- A check if gradient computation is re-entrant.+-- This test fails (not on first run, but on repetition) if old terms,+-- from before current iteration of gradient descent, are not soundly+-- handled in new computations (due to naive impurity, TH, plugins,+-- or just following the papers that assume a single isolated run).+-- This example requires monomorphic types and is contrived,+-- but GHC does optimize and factor out some accidentally constant+-- subterms in real examples (presumably after it monomorphizes them)+-- causing exactly the same danger.+-- This example also tests unused parameters (x), another common cause+-- of crashes in naive gradient computing code.+baz :: ( ADVal Concrete (TKR 0 Double)+       , ADVal Concrete (TKR 0 Double)+       , ADVal Concrete (TKR 0 Double) )+    -> ADVal Concrete (TKR 0 Double)+baz (_x,y,z) =+  let w = fooFromPrimal * barADVal2 (y,y,z) * sin y+  in atan2H z w + z * w++-- An "old term", computed once, stored at top level.+fooFromPrimal :: ADVal Concrete (TKR 0 Double)+fooFromPrimal = foo2 (rscalar 7, rscalar 8, rscalar 9)++testBaz :: Assertion+testBaz =+  assertEqualUpToEpsilon 1e-9+    (rscalar 0, rscalar (-5219.20995030263), rscalar 2782.276274462047)+    (cgrad (kfromR . baz) (rscalar 1.1, rscalar 2.2, rscalar 3.3))++-- If terms are numbered and @z@ is, wrongly, decorated with number 0,+-- here @fooFromPrimal@ is likely to clash with @z@, since it was numbered+-- starting from 0, too.+-- The test fails if term numbering is reset to 0 between gradient computation+-- runs (verified) as well as when delta-expression evaluation assumes+-- it can only encounter terms numbered in the current run and,+-- e.g., allocates an array with only that much space (not verified).+-- Actually, with term counter zeroed just before @cgrad@ application,+-- even a single @testBaz@ execution produces wrong results, but this test+-- is likely to fail in less naive implementations, as well.+testBazRenumbered :: Assertion+testBazRenumbered =+  assertEqualUpToEpsilon 1e-9+    (rscalar 0, rscalar (-5219.20995030263), rscalar 2783.276274462047)+    (cgrad (kfromR . (\(x,y,z) -> z + baz (x,y,z))) (rscalar 1.1, rscalar 2.2, rscalar 3.3))++-- A dual-number and list-based version of a function that goes+-- from `R^3` to `R`.+fooD :: forall r. r ~ Double+     => IxR 3 (ADVal Concrete (TKR 0 r)) -> ADVal Concrete (TKR 0 r)+fooD (x :.: y :.: z :.: ZIR) =+  let w = x * sin y+  in atan2H z w + z * w++testFooD :: Assertion+testFooD =+  assertEqualUpToEpsilon 1e-10+    (fromList [rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627])+    (cgrad (kfromR . fooD) (fromList [rscalar 1.1, rscalar 2.2, rscalar 3.3]))++fooBuild1 :: (ADReady target, NumScalar r, Differentiable r)+          => target (TKR 1 r) -> target (TKR 1 r)+fooBuild1 v =+  let r = rfromK $ rsum0 v+      v' = rfromK $ rminimum v+  in rbuild1 3 $ \ix ->+       r * foo2 ( rscalar 3+               , rscalar 5 * r+               , r * rfromK (rminimum v) * v')+       + bar (r, rindex v [ix + 1])++testFooBuildDt :: Assertion+testFooBuildDt =+  assertEqualUpToEpsilon 1e-5+    (rconcrete $ Nested.rfromListPrimLinear [4] [-189890.46351219364,-233886.08744601303,-222532.22669716467,-206108.68889329425])+    (vjp @_ @(TKR 1 Double)+           fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [3] 42))++testFooBuildCFwd :: Assertion+testFooBuildCFwd =+  assertEqualUpToEpsilon 1e-5+    (rconcrete $ Nested.rfromListPrimLinear [3] [-296584.8166864211,-290062.472288043,-265770.1775742018])+    (cjvp @_ @(TKR 1 Double) @_ @Concrete+          fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [4] 42))++testFooBuildFwd :: Assertion+testFooBuildFwd =+  assertEqualUpToEpsilon 1e-5+    (rconcrete $ Nested.rfromListPrimLinear [3] [-296584.8166864211,-290062.472288043,-265770.1775742018])+    (jvp @_ @(TKR 1 Double)+         fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [4] 42))++testFooBuild :: Assertion+testFooBuild =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [4] [-4521.201512195087,-5568.7163677622175,-5298.386349932494,-4907.349735554627])+    (rev' @Double @1 fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]))++fooNoGo :: forall target r. (ADReady target, NumScalar r, Differentiable r)+        => target (TKR 1 r) -> target (TKR 1 r)+fooNoGo v =+  let r = rsum0 v+  in rbuild1 3 (\ix' -> let ix :: PlainOf target (TKS '[] Int)+                            ix = sfromR $ rfromK ix' in+       bar (rscalar 3.14, bar (rscalar 3.14, rindex v [kfromS $ (ix + (tplainPart . sfloor . sfromK) r) `minH` sscalar 2 `maxH` sscalar 0]))+       + ifH ((&&*) (rindex @1 v [ix' * 2] <=. rscalar 0)+                    (sscalar 6 >. abs ix))+               (rfromK r) (rscalar 5 * rfromK r))+     / rslice 1 3 (rmap0N (\x -> ifH (x >. r) r x) v)+     * rbuild1 3 (const (rscalar 1))++testFooNoGo0 :: Assertion+testFooNoGo0 =+  assertEqualUpToEpsilon' 1e-6+   (ringestData [5] [5.037878787878788,-14.394255484765257,43.23648655081373,-0.8403418295960368,5.037878787878788])+   (rev' @Double @1 fooNoGo+         (ringestData [5] [1.1 :: Double, 2.2, 3.3, 4, 5]))++nestedBuildMap :: forall target r.+                  (ADReady target, NumScalar r, Differentiable r)+               => target (TKR 0 r) -> target (TKR 1 r)+nestedBuildMap r =+  let w = rreplicate0N [4] . kfromR+      v0' = rreplicate0N [177] (kfromR r) :: target (TKR 1 r)+  in tlet v0' $ \v' ->+    let nestedMap x0 = tlet x0 $ \x -> rmap0N (kfromR x /) (w x)+        variableLengthBuild iy = rbuild1 7 (\ix -> rindex v' (ix + iy :.: ZIR))+        doublyBuild = rbuild1 5 (rfromK . rminimum . variableLengthBuild)+        fooMap1 :: target (TKR 0 r) -> target (TKR 1 r)+        fooMap1 r2 =+          let v = fooBuild1 $ rreplicate0N [130] (kfromR r2)+          in rmap0N (\x -> x * kfromR r2 + 5) v+    in rmap0N (\x0 -> tlet x0 $ \x -> x * (rsum0+                           (rbuild1 3 (\ix -> bar (rfromK x, rindex v' [ix]))+                            + (tlet (nestedMap $ rfromK x) $ \x3 -> fooBuild1 x3)+                            / (tlet (rfromK x) $ \x4 -> fooMap1 x4)))+              ) doublyBuild++testNestedBuildMap1 :: Assertion+testNestedBuildMap1 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 107.25984443006627)+    (rev' @Double @1 nestedBuildMap (rscalar 1.1))++nestedSumBuildm1 :: ADReady target+                  => target (TKR 0 CInt) -> target (TKScalar CInt)+nestedSumBuildm1 _ =+ rsum0+ $ tlet (rbuild1 2 (rfromIntegral . rfromK . tfromPlain STKScalar))+        (\x -> x)++testNestedSumBuildm1 :: Assertion+testNestedSumBuildm1 =+  assertEqualUpToEpsilon' 1e-8+    (rscalar 0)+    (rev' (rfromK . nestedSumBuildm1) (rscalar 1))++nestedSumBuild0 :: ADReady target+                => target (TKR 0 CInt) -> target (TKScalar CInt)+nestedSumBuild0 _ =+ rsum0+ $ rreplicate 2 $ rsum+ $ tlet (rbuild1 2 (rfromIntegral . rfromK . tfromPlain STKScalar))+        (\x -> x)++testNestedSumBuild0 :: Assertion+testNestedSumBuild0 =+  assertEqualUpToEpsilon' 1e-8+    (rscalar 0)+    (rev' (rfromK . nestedSumBuild0) (rscalar 1))++nestedSumBuild :: (ADReady target, NumScalar r, Differentiable r)+               => target (TKR 1 r) -> target (TKR 1 r)+nestedSumBuild v0 = tlet v0 $ \v ->+ tlet (rsum (rbuild1 9 rfromIndex0)) (\tbtf ->+  (rbuild1 13 (\ix ->+    rsum (rbuild1 4 (\ix2 ->+      flip rindex [ix2]+        (tlet (rsum v) $ \tsumv -> rbuild1 5 (const tsumv)+         * rfromList+             [ rfromIndex0 ix+             , tbtf+             , rsum (rbuild1 6 (\_ -> rsum v))+             , rindex v [ix2]+             , rsum (rbuild1 3 (\ix7 ->+                 rsum (rreplicate 5 (rfromIndex0 ix7))))+             ]))))))+ + tlet (nestedBuildMap (rfromK $ rsum0 v)) (\nbmt -> (rbuild1 13 (\ix ->+     nbmt `rindex` [minH ix 4])))++testNestedSumBuild :: Assertion+testNestedSumBuild =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [5] [-14084.715065313612,-14084.715065313612,-14084.715065313612,-14014.775065313623,-14084.715065313612])+    (rev' @Double @1 nestedSumBuild (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))++nestedBuildIndex :: forall target r. (ADReady target, GoodScalar r)+                 => target (TKR 1 r) -> target (TKR 1 r)+nestedBuildIndex v =+  rbuild1 2 $ \ix2 -> rindex (rbuild1 3 $ \ix3 -> rindex (rbuild1 4 $ \ix4 -> rindex v (ix4 :.: ZIR)) [ix3]) (ix2 :.: ZIR)++testNestedBuildIndex :: Assertion+testNestedBuildIndex =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [5] [1,1,0,0,0])+    (rev' @Double @1 nestedBuildIndex (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))++barRelu+  :: ( ADReady target, NumScalar r, KnownNat n, Differentiable r )+  => target (TKR n r) -> target (TKR n r)+barRelu x = relu $ bar (x, relu x)++testBarReluDt :: Assertion+testBarReluDt =+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [191.20462646925841])+    (vjp @_ @(TKR 0 Double)+           barRelu (rscalar 1.1) (rscalar 42.2))++testBarRelu :: Assertion+testBarRelu =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [] [4.5309153191767395])+    (rev' @Double @0 barRelu (ringestData [] [1.1]))++testBarRelu3 :: Assertion+testBarRelu3 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 1, 2] [4.5309153191767395,4.5302138998556,-9.39547533946234,95.29759282497125])+    (rev' @Double @3 barRelu (ringestData [2, 1, 2] [1.1, 2, 3, 4.2]))++barReluMax0+  :: ( ADReady target, NumScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+  => target (TKR n r) -> target (TKR n r)+barReluMax0 x = reluMax $ bar (x, x)++barReluMax+  :: ( ADReady target, NumScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+  => target (TKR n r) -> target (TKR n r)+barReluMax x = reluMax $ bar (x, reluMax x)++testBarReluMaxDt :: Assertion+testBarReluMaxDt =+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [191.20462646925841])+    (vjp @_ @(TKR 0 Double)+           barReluMax (ringestData [] [1.1]) (rscalar 42.2))++testBarReluMax :: Assertion+testBarReluMax =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [] [4.5309153191767395])+    (rev' @Double @0 barReluMax (ringestData [] [1.1]))++testBarReluMax30 :: Assertion+testBarReluMax30 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [4.5309153191767395])+    (rev' @Double @1 barReluMax0 (ringestData [1] [1.1]))++testBarReluMax31 :: Assertion+testBarReluMax31 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 1, 2] [4.5309153191767395,4.5302138998556,-9.39547533946234,95.29759282497125])+    (rev' @Double @3 barReluMax (ringestData [2, 1, 2] [1.1, 2, 3, 4.2]))++testBarReluMax3CFwd :: Assertion+testBarReluMax3CFwd =+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2, 1, 2] [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])+    (cjvp @_ @(TKR 3 Double) @_ @Concrete+          barReluMax+                     (rconcrete $ Nested.rfromListPrimLinear (fromList [2, 1, 2]) [1.1, 2, 3, 4.2])+                     (ringestData [2, 1, 2] [0.1, 0.2, 0.3, 0.42]))++reluMaxS :: forall target sh r. (ADReady target, NumScalar r, KnownShS sh)+         => target (TKS sh r) -> target (TKS sh r)+reluMaxS = smap0N (maxH 0)++barReluMaxS+  :: ( ADReady target, NumScalar r, KnownShS sh+     , RealFloatH (target (TKS sh r)) )+  => target (TKS sh r) -> target (TKS sh r)+barReluMaxS x = reluMaxS $ barF (x, reluMaxS x)++-- Previously the shape of FromIxR[DeltaZero] couldn't be determined+-- in buildDerivative, so this was needed. See below that it now works fine.+testBarReluMax3FwdS :: Assertion+testBarReluMax3FwdS =+  assertEqualUpToEpsilon 1e-10+    (sconcrete $ Nested.sfromListPrimLinear @'[2, 1, 2] knownShS [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])+    (jvp @_ @(TKS '[2, 1, 2] Double)+         barReluMaxS+         (sconcrete $ Nested.sfromListPrimLinear @'[2, 1, 2] knownShS [1.1, 2, 3, 4.2])+         (sconcrete $ Nested.sfromListPrimLinear @'[2, 1, 2] knownShS [0.1, 0.2, 0.3, 0.42]))++testBarReluMax3FwdFrom :: Assertion+testBarReluMax3FwdFrom =+  assertEqualUpToEpsilon 1e-10+    (sconcrete $ Nested.sfromListPrimLinear @'[2, 1, 2] knownShS [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])+    (jvp @_ @(TKS '[2, 1, 2] Double)+         (sfromR . barReluMax . rfromS)+         (sconcrete $ Nested.sfromListPrimLinear @'[2, 1, 2] knownShS [1.1, 2, 3, 4.2])+         (sconcrete $ Nested.sfromListPrimLinear @'[2, 1, 2] knownShS [0.1, 0.2, 0.3, 0.42]))++testBarReluMax3FwdR :: Assertion+testBarReluMax3FwdR =+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2, 1, 2] [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])+    (jvp @_ @(TKR 3 Double)+         barReluMax+         (ringestData [2, 1, 2] [1.1, 2, 3, 4.2])+         (ringestData [2, 1, 2] [0.1, 0.2, 0.3, 0.42]))++f1 :: (ADReady target, NumScalar r) => target (TKR 0 r) -> target (TKR 0 r)+f1 = \arg -> rfromK $ rsum0 (rbuild1 10 (\i -> arg * rfromIndex0 i))++testF1 :: Assertion+testF1 =+  assertEqualUpToEpsilon 1e-10+    (rscalar 45.0)+    (grad (kfromR @_ @Double . f1) (rscalar 1.1))++testF11 :: Assertion+testF11 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 45.0)+    (rev' @Double @0 f1 (rscalar 1.1))++f2 :: forall target r. (ADReady target, NumScalar r)+   => target (TKR 0 r) -> target (TKR 0 r)+f2 = \arg ->+  let fun1 i = arg * rfromIndex0 i+      v1a = rfromK $ rsum0 (rbuild1 10 fun1)+      v1b = rfromK $ rsum0 (rbuild1 20 fun1)+      fun2 y i = y * rfromIndex0 @r i+      v2a = rfromK $ rsum0 (rbuild1 10 (fun2 arg))+      v2b = rfromK $ rsum0 (rbuild1 20 (fun2 (arg + rscalar 1)))+  in v1a + v1b + v2a + v2b++testF2 :: Assertion+testF2 =+  assertEqualUpToEpsilon 1e-10+    (rscalar 470)+    (grad (kfromR @_ @Double . f2) (rscalar 1.1))++testF21 :: Assertion+testF21 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 470)+    (rev' @Double @0 f2 (rscalar 1.1))++testF2CFwd :: Assertion+testF2CFwd =+  assertEqualUpToEpsilon 1e-10+    (rscalar 47)+    (cjvp @_ @(TKR 0 Double) @_ @Concrete+          f2 (rscalar 1.1) (rscalar 0.1))++testF2Fwd :: Assertion+testF2Fwd =+  assertEqualUpToEpsilon 1e-10+    (rscalar 47)+    (jvp @_ @(TKR 0 Double)+         f2 (rscalar 1.1) (rscalar 0.1))++braidedBuilds :: forall target r.+                 (ADReady target, NumScalar r, Differentiable r)+              => target (TKR 0 r) -> target (TKR 2 r)+braidedBuilds r =+  rbuild1 3 (\ix1 ->+    rbuild1 4 (\ix2 -> rindex (rfromList+      [rfromIndex0 ix2, rscalar 7, r, rscalar (-0.2)]) (ix1 :.: ZIR)))++testBraidedBuilds1 :: Assertion+testBraidedBuilds1 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 4.0)+    (rev' @Double @2 braidedBuilds (rscalar 3.4))++recycled :: (ADReady target, NumScalar r, Differentiable r)+         => target (TKR 0 r) -> target (TKR 5 r)+recycled r =+  tlet (nestedSumBuild (rreplicate0N [7] (kfromR r))) $ \nsb ->+    rbuild1 2 $ \_ -> rbuild1 4 $ \_ -> rbuild1 2 $ \_ -> rbuild1 3 $ const nsb++testRecycled1 :: Assertion+testRecycled1 =+  assertEqualUpToEpsilon 1e-6+    (rscalar 348356.9278600814)+    (grad (rsum0 @_ @Double . recycled) (rscalar 0.0000001))++concatBuild :: (ADReady target, NumScalar r) => target (TKR 0 r) -> target (TKR 2 r)+concatBuild r =+  rbuild1 7 (\i ->+    rappend (rappend (rbuild1 5 (const r))+                     (rreplicate 1 (rfromIndex0 i)))+            (rbuild1 13 (const r)))++testConcatBuild1 :: Assertion+testConcatBuild1 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 126.0)+    (rev' @Double @2 concatBuild (rscalar 3.4))++concatBuild2 :: (ADReady target, NumScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild2 r =+  tlet (rfromList [r, rscalar 1, rscalar 2, rscalar 3, rscalar 4]) $ \a ->+    rbuild1 10 (\i -> ifH (i <. 5) (rindex a [i]) (rindex a [i - 5]))++testConcatBuild2 :: Assertion+testConcatBuild2 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.0)+    (rev' @Double @1 concatBuild2 (rscalar 3.4))++concatBuild3 :: (ADReady target, NumScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild3 r =+  tlet (rfromList [r, rscalar 1, rscalar 2, rscalar 3, rscalar 4]) $ \a ->+    rbuild1 10 (\i -> ifH (i <. 5) (rindex a [i]) (rindex a [i - 5 + (1 `quotH` maxH 1 (i - 5))]))++testConcatBuild3 :: Assertion+testConcatBuild3 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1)+    (rev' @Double @1 concatBuild3 (rscalar 3.4))++concatBuild4 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild4 r =+  tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (4 + i)) :.: ZIR)) $ \a ->+    rappend a a++testConcatBuild4 :: Assertion+testConcatBuild4 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 10)+    (rev' @Double @1 concatBuild4 (rscalar 3.4))++concatBuild5 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild5 r =+  tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (5 + i)) :.: ZIR)) $ \a ->+    (rappend a (rgather1 5 (rreplicate 1 r)+                           (\i -> (1 `quotH` (5 + i)) :.: ZIR)))++testConcatBuild5 :: Assertion+testConcatBuild5 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 10)+    (rev' @Double @1 concatBuild5 (rscalar 3.4))++concatBuild6 :: (ADReady target, NumScalar r) => target (TKR 0 r) -> target (TKR 2 r)+concatBuild6 r =+  rbuild1 7 (\j ->+    rappend (rappend+             (tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (4 + i)) :.: ZIR)) $ \a ->+    (rappend (rgather1 5 (rreplicate 1 r)+                         (\i -> (1 `quotH` (100 * maxH 1 (i - j))) :.: ZIR)) a))+                     (rreplicate 1 (rfromIndex0 j)))+            (rbuild1 13 (const r)))++testConcatBuild6 :: Assertion+testConcatBuild6 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 161)+    (rev' @Double @2 concatBuild6 (rscalar 3.4))++concatBuild7 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild7 r =+  rbuild1 10 $ \j ->+    (rappend (rreplicate 5 r) (rgather1 5 (rreplicate 1 r)+                                 (\i -> (1 `quotH` maxH 1 (j - i)) :.: ZIR)))+     ! (j :.: ZIR)++testConcatBuild7 :: Assertion+testConcatBuild7 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 10)+    (rev' @Double @1 concatBuild7 (rscalar 3.4))++-- These tests show that term rewriting changes the value+-- of out-of-bounds indexing, e.g., of gather(replicate)+-- that reduces to a non-gather. Vectorization is not needed for that.+concatBuild8 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild8 r =+  tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (5 - i)) :.: ZIR)) $ \a ->+    (rappend a (rgather1 5 (rreplicate 1 r)+                           (\i -> (1 `quotH` (5 - i)) :.: ZIR)))++testConcatBuild8 :: Assertion+testConcatBuild8 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 8)+    (rev' @Double @1 concatBuild8 (rscalar 3.4))++concatBuild9 :: (ADReady target, NumScalar r) => target (TKR 0 r) -> target (TKR 2 r)+concatBuild9 r =+  rbuild1 7 (\j ->+    rappend (rappend+             (tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (4 - i)) :.: ZIR)) $ \a ->+    (rappend (rgather1 5 (rreplicate 1 r)+                         (\i -> (1 `quotH` (100 * maxH 0 (i - j))) :.: ZIR)) a))+                     (rreplicate 1 (rfromIndex0 j)))+            (rbuild1 13 (const r)))++testConcatBuild9 :: Assertion+testConcatBuild9 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 122)+    (rev' @Double @2 concatBuild9 (rscalar 3.4))++concatBuild10 :: (ADReady target, NumScalar r)+              => target (TKR 0 r) -> target (TKR 2 r)+concatBuild10 r =+  rbuild1 7 (\j ->+    rappend (rappend+               (tlet (rgather1 5 (rreplicate 1 r)+                       (\_i -> 10000 :.: ZIR)) $ \a ->+      (rappend (rgather1 5 (rreplicate 1 r)+                 (\_i -> (-1) :.: ZIR)) a))+                   (rreplicate 1 (rfromIndex0 j)))+            (rbuild1 13 (const r)))++testConcatBuild10 :: Assertion+testConcatBuild10 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 91)+    (rev' @Double @2 concatBuild10 (rscalar 3.4))++concatBuild11 :: (ADReady target, GoodScalar r)+              => target (TKR 0 r) -> target (TKR 1 r)+concatBuild11 r =+  rgather1 5 (rreplicate 1 r) (\_i -> (-1) :.: ZIR)++testConcatBuild11 :: Assertion+testConcatBuild11 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @1 concatBuild11 (rscalar 3.4))++concatBuild12 :: (ADReady target, GoodScalar r)+              => target (TKR 0 r) -> target (TKR 0 r)+concatBuild12 r =+  rindex (rreplicate 1 r) ((-1) :.: ZIR)++testConcatBuild12 :: Assertion+testConcatBuild12 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @0 concatBuild12 (rscalar 3.4))++concatBuild13 :: (ADReady target, NumScalar r)+              => target (TKR 0 r) -> target (TKR 1 r)+concatBuild13 r =+  tlet (rfromList [r, rscalar 1, rscalar 2, rscalar 3, rscalar 4]) $ \a ->+    rbuild1 10 (\i -> ifH (notB $ 5 <=. i) (rindex a [i]) (rindex a [i - 5]))++testConcatBuild13 :: Assertion+testConcatBuild13 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.0)+    (rev' @Double @1 concatBuild13 (rscalar 100))++concatBuild14 :: (ADReady target, NumScalar r)+              => target (TKR 0 r) -> target (TKR 1 r)+concatBuild14 r =+  tlet (rfromList [r, rscalar 1, rscalar 2, rscalar 3, rscalar 4]) $ \a ->+    rbuild1 10 (\i -> ifH (i <. 5) (rindex a [i - 5]) (rindex a [i]))++testConcatBuild14 :: Assertion+testConcatBuild14 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Int @1 concatBuild14 (rscalar 100))++testConcatBuild15 :: Assertion+testConcatBuild15 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @1 concatBuild14 (rscalar 100))++testConcatBuild16 :: Assertion+testConcatBuild16 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @1 (\r -> rslice 0 0 (rfromVector (V.fromList [r, r]))) (rscalar 100))++testConcatBuild17 :: Assertion+testConcatBuild17 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @1 (\r -> ifH (rslice 0 0 (rfromVector (V.fromList [r, r])) ==. rslice 0 0 (rfromVector (V.fromList [r, rscalar 0]))) (rslice 0 0 (rfromVector (V.fromList [r]))) (rslice 0 0 (rfromVector (V.fromList [rscalar 0])))) (rscalar 100))++-- TODO: copy-paste a variant of emptyArgs with r not NumScalar+-- or maybe generalize emptyArgs and then in half of the tests do TKScalar+emptyArgs :: forall target r. (ADReady target, NumScalar r, Differentiable r)+          => target (TKR 1 r) -> target (TKR 1 r)+emptyArgs t =+  emptyTensor+  - ringestData (rshape @target @_ @(TKScalar r) emptyTensor) []+  + rconcrete Nested.remptyArray+  - rsum (ringestData (0 :$: rshape @target @_ @(TKScalar r) emptyTensor) [])+  * rsum (rconcrete $ Nested.rreshape (0 :$: 0 :$: ZSR) Nested.remptyArray)+  * rconcrete (Nested.rsumOuter1Prim $ Nested.rfromListOuter+               $ NonEmpty.fromList [Nested.remptyArray])+  * rsum (rconcrete $ Nested.rfromListOuter+          $ NonEmpty.fromList [Nested.remptyArray])+  * rsum (rfromList [rconcrete Nested.remptyArray])+  * rsum (rfromList [emptyTensor])+  * rsum (rfromList [rsum (ringestData+                             (0 :$: rshape @target @_ @(TKScalar r)+                                      emptyTensor) [])])+  - rindex (ringestData (0 :$: 0 :$: ZSR) []) (42 :.: ZIR)+  - rindex (ringestData (0 :$: rshape @target @_ @(TKScalar r)+                                     emptyTensor) []) (42 :.: ZIR)+  - rreshape @1 [0] emptyTensor+  - rsum (rreshape @1 [0, 0] emptyTensor)+  * rslice 0 0 t+  * rslice 0 0 (rfromVector (V.fromList [rscalar 1, t ! (0 :.: ZIR)]))+  + rbuild1 0 (\i -> t ! (i :.: ZIR))+  + rbuild1 0 (\i -> t ! [fromIntegral (rlength t) `quotH` i] / rfromIndex0 i)+  + rbuild @1 (0 :$: ZSR) (const $ rscalar 73)+  - rsum (rbuild @0 ZSR+                 (const (rreplicate 1 emptyTensor)))+  + rfromS+      (srepl @'[0] 0+       + sconcrete (Nested.semptyArray ZSS)+       - ssum @0 (srepl @'[0, 0] 0)+       * ssum @0 (sconcrete $ Nested.semptyArray (SNat @0 :$$ ZSS))+       * sconcrete (Nested.ssumOuter1Prim $ Nested.sfromListOuter (SNat @1)+                    $ NonEmpty.fromList [Nested.semptyArray ZSS])+       * ssum @1 (sconcrete $ Nested.sfromListOuter SNat+                        $ NonEmpty.fromList [Nested.semptyArray ZSS])+       * ssum @1 (sfromList [sconcrete $ Nested.semptyArray ZSS])+       * ssum @1 (sfromList [sfromR @_ @'[0] emptyTensor])+       * ssum @1 (sfromList [ssum @0 (srepl @'[0, 0] 0)])+       - sindex @'[0] (srepl @'[0, 0] 0) (42 :.$ ZIS)+       - sreshape @_ @'[0] (sfromR @_ @'[0] emptyTensor)+       - ssum (sreshape @_ @'[0, 0] (sfromR @_ @'[0] emptyTensor))+       * sslice (SNat @0) (SNat @0) (SNat @1) (sfromR $ rslice 0 1 (rappend t (rreplicate 1 (rscalar 1))))+       * sslice (SNat @1) (SNat @0) (SNat @1)+                (sfromVectorLinear+                   [2] (V.fromList [1, t `rindex0` (0 :.: ZIR)]))+       * sbuild1 @0 (\i -> sfromR @_ @'[0] (rslice 0 0 t) !$ (i :.$ ZIS))+       + sbuild1 @0 (\i -> sfromR @_ @'[0] (rslice 0 0 t)+                              !$ (fromIntegral (rlength t) `quotH` i :.$ ZIS)+                              / tfromPlain knownSTK (sfromIntegral (sfromK i)))+       + sbuild @'[0] (const $ sscalar 73)+       - ssum (sbuild @'[] (const (sreplicate @1 (sfromR emptyTensor)))))+  + rfromX+      (xingestData (SKnown (SNat @0) :$% ZSX) []+       + xconcrete (Nested.memptyArray ZSX)+       - xsum @0 (xingestData+                          (SKnown (SNat @0) :$% SKnown (SNat @0) :$% ZSX) [])+       * xsum @0 (xconcrete+                        $ Nested.memptyArray (SKnown (SNat @0) :$% ZSX))+       * xconcrete (Nested.msumOuter1Prim $ Nested.mfromListOuter+                    $ NonEmpty.fromList [Nested.memptyArray ZSX])+       * xsum @1 (xconcrete+                        $ Nested.mcast+                            (SKnown (SNat @1) :!% SKnown (SNat @0) :!% ZKX)+                        $ Nested.mfromListOuter+                        $ NonEmpty.fromList [Nested.memptyArray ZSX])+       * xsum @1 (xfromList [xconcrete $ Nested.memptyArray ZSX])+       * xsum @1 (xfromList [xfromR @_ @'[Just 0] emptyTensor])+       * xsum @1 (xfromList [xsum @0 (xingestData+                                        (SKnown (SNat @0)+                                         :$% xshape @target @_ @(TKScalar r)+                                               (xfromR @_ @'[Just 0]+                                                  emptyTensor)) [])])+       - xindex @_ @'[Just 0] (xingestData+                                    (SKnown (SNat @0)+                                     :$% SKnown (SNat @0)+                                     :$% ZSX) []) (42 :.% ZIX)+       - xindex @_ @'[Just 0] (xingestData+                                    (SKnown (SNat @0)+                                     :$% xshape @target @_ @(TKScalar r)+                                             (xfromR @_ @'[Just 0]+                                                emptyTensor)) []) (42 :.% ZIX)+       - xreshape (SKnown (SNat @0) :$% ZSX) (xfromR @_ @'[Just 0] emptyTensor)+       - xsum (xreshape (SKnown (SNat @0) :$% SKnown (SNat @0) :$% ZSX) (xfromR @_ @'[Just 0] emptyTensor))+       * xslice (SKnown $ SNat @0)+                (SKnown $ SNat @0)+                (SKnown $ SNat @1)+                (xfromR @_ @'[Just 1] $ rslice 0 1 (rappend t (rreplicate 1 (rscalar 1))))+       * xbuild1 @0 (\i -> xfromR @_ @'[Nothing] (rslice 0 0 t)+                            `xindex` (i :.% ZIX))+       + xbuild1 @0 (\i -> xslice (SKnown $ SNat @0)+                                  (SKnown $ SNat @0)+                                  (SUnknown 1)+                                  (xfromR @_ @'[Nothing] $ rslice 0 1 (rappend t (rreplicate 1 (rscalar 1))))+                              `xindex`+                              (fromIntegral (rlength t) `quotH` i :.% ZIX)+                              / xfromIndex0 i)+       + xbuild @'[Just 0] (SKnown (SNat @0) :$% ZSX)+                (const $ xscalar 73)+       - xsum @0 (xbuild @'[] ZSX+                         (const (xreplicate (SNat @0) (xfromR emptyTensor)))))+--  - rgather1 0 emptyTensor (:.: ZIR)+--  - rsum (rgather1 0 emptyTensor (const ZIR))+--  - rsum (rgather @target @(TKScalar r) @2 (0 :$: 0 :$: ZSR) emptyTensor (const (0 :.: ZIR)))+--  - rsum (rgather @target @(TKScalar r) @2 @0 @1 [0, 0] emptyTensor (const [0]))+--  * rflatten (rtr (rgather1 0 t (const ZIR)))+--  - rsum (rbuild @2 (0 :$: 0 :$: ZSR) (const 73))+--  - rsum (rbuild @1 (0 :$: 0 :$: ZSR) (const emptyTensor))+       -- these fail and rightly so; TODO: make them fail earlier+ where+  emptyTensor :: target (TKR 1 r)+  emptyTensor = rconcrete $ Nested.rfromListPrimLinear (fromList [0]) []++testEmptyArgs0 :: Assertion+testEmptyArgs0 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [0] [])+    (rev' @Double @1 emptyArgs (ringestData [0] []))++testEmptyArgs1 :: Assertion+testEmptyArgs1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [0])+    (rev' @Float @1 emptyArgs (ringestData [1] [0.24]))++testEmptyArgs2 :: Assertion+testEmptyArgs2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [0])+    (rev' @Double @1+          (\t -> emptyArgs $ rfromVector (V.fromList [ t ! (0 :.: ZIR)+                                                     , rscalar 0 ]))+          (ringestData [1] [0.24]))++testEmptyArgs3 :: Assertion+testEmptyArgs3 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [0])+    (rev' @Double @1+          (\t -> emptyArgs $ rfromVectorLinear+                               [2] (V.fromList [0, t `rindex0` (0 :.: ZIR)]))+          (ringestData [1] [0.24]))++testEmptyArgs4 :: Assertion+testEmptyArgs4 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [0])+    (rev' @Double @1+          (\t -> rbuild [] (const $ emptyArgs+                            $ rfromVector (V.fromList [ rscalar 7+                                                      , t ! (0 :.: ZIR) ])))+          (ringestData [1] [0.24]))++testEmptyArgs5 :: Assertion+testEmptyArgs5 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [0])+    (rev' @Double @1+          (\t -> rbuild [] (const $ emptyArgs+                            $ rfromVectorLinear+                                [2] (V.fromList [t `rindex0` (0 :.: ZIR), 1])))+          (ringestData [1] [0.24]))++testEmptyArgs6 :: Assertion+testEmptyArgs6 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [0])+    (rev' @Double @1+          (\t -> rbuild [] (const $ emptyArgs t))+          (ringestData [1] [0.24]))++filterPositiveFail :: ADReady target+                   => target (TKR 1 Double) -> target (TKR 1 Double)+filterPositiveFail v =+  let l = runravelToList v+      -- l2 = filter (\x -> x >= 0) l+      -- Could not deduce ‘Ord (target Double 0)’+      -- l2 = filter (\x -> x >=. 0) l+      -- Could not deduce ‘BoolOf target ~ Bool’+      l2 = take 12 l  -- the most I can do+  in rfromList $ NonEmpty.fromList l2++testFilterPositiveFail0 :: Assertion+testFilterPositiveFail0 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2] [1.0,1.0])+    (rev' @Double @1+          filterPositiveFail+          (ringestData [2] [0.24, 52]))++testFilterPositiveFail0PP :: Assertion+testFilterPositiveFail0PP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent filterPositiveFail+                         (FTKR (2 :$: ZSR) (FTKScalar @Double))+  printArtifactPrimalPretty artifactRev+    @?= "\\v1 -> rfromS (sfromVector0N @[2] (fromList [sfromR v1 `sindex0` [0], sfromR v1 `sindex0` [1]]))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromS (sfromVector0N @[2] (fromList [sfromR v1 `sindex0` [0], sfromR v1 `sindex0` [1]]))"+  printArtifactPretty artifactRev+    @?= "\\dret v1 -> rfromS (soneHot (sfromK (sfromR dret `sindex0` [0])) [0] + soneHot (sfromK (sfromR dret `sindex0` [1])) [1])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> rfromS (sfromVector0N @[2] (fromList [sfromR dret `sindex0` [0], sfromR dret `sindex0` [1]]))"++testFilterPositiveFail1 :: Assertion+testFilterPositiveFail1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3] [1.0,1.0,1.0])+    (rev' @Double @1+          filterPositiveFail+          (ringestData [3] [0.24, 52, -0.5]))++testFilterPositiveFail1PP :: Assertion+testFilterPositiveFail1PP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent filterPositiveFail+                         (FTKR (3 :$: ZSR) (FTKScalar @Double))+  printArtifactPrimalPretty artifactRev+    @?= "\\v1 -> rfromS (sfromVector0N @[3] (fromList [sfromR v1 `sindex0` [0], sfromR v1 `sindex0` [1], sfromR v1 `sindex0` [2]]))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromS (sfromVector0N @[3] (fromList [sfromR v1 `sindex0` [0], sfromR v1 `sindex0` [1], sfromR v1 `sindex0` [2]]))"+  printArtifactPretty artifactRev+    @?= "\\dret v1 -> rfromS (soneHot (sfromK (sfromR dret `sindex0` [0])) [0] + (soneHot (sfromK (sfromR dret `sindex0` [1])) [1] + soneHot (sfromK (sfromR dret `sindex0` [2])) [2]))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> rfromS (sfromVector0N @[3] (fromList [sfromR dret `sindex0` [0], sfromR dret `sindex0` [1], sfromR dret `sindex0` [2]]))"++testFilterPositiveFail :: Assertion+testFilterPositiveFail =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [15] [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,0.0,0.0,0.0])+    (rev' @Double @1+          filterPositiveFail+          (ringestData [15] [0.24, 52, -0.5, 0.33, 0.1, 0.24, 52, -0.5, 0.33, 0.1, 0.24, 52, -0.5, 0.33, 0.1]))++testFilterPositiveFailPP :: Assertion+testFilterPositiveFailPP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent filterPositiveFail+                         (FTKR (15 :$: ZSR) (FTKScalar @Double))+  printArtifactPrimalPretty artifactRev+    @?= "\\v1 -> rfromS (sfromVector0N @[12] (fromList [sfromR v1 `sindex0` [0], sfromR v1 `sindex0` [1], sfromR v1 `sindex0` [2], sfromR v1 `sindex0` [3], sfromR v1 `sindex0` [4], sfromR v1 `sindex0` [5], sfromR v1 `sindex0` [6], sfromR v1 `sindex0` [7], sfromR v1 `sindex0` [8], sfromR v1 `sindex0` [9], sfromR v1 `sindex0` [10], sfromR v1 `sindex0` [11]]))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromS (sfromVector0N @[12] (fromList [sfromR v1 `sindex0` [0], sfromR v1 `sindex0` [1], sfromR v1 `sindex0` [2], sfromR v1 `sindex0` [3], sfromR v1 `sindex0` [4], sfromR v1 `sindex0` [5], sfromR v1 `sindex0` [6], sfromR v1 `sindex0` [7], sfromR v1 `sindex0` [8], sfromR v1 `sindex0` [9], sfromR v1 `sindex0` [10], sfromR v1 `sindex0` [11]]))"+  printArtifactPretty artifactRev+    @?= "\\dret v1 -> rfromS (soneHot (sfromK (sfromR dret `sindex0` [0])) [0] + (soneHot (sfromK (sfromR dret `sindex0` [1])) [1] + (soneHot (sfromK (sfromR dret `sindex0` [2])) [2] + (soneHot (sfromK (sfromR dret `sindex0` [3])) [3] + (soneHot (sfromK (sfromR dret `sindex0` [4])) [4] + (soneHot (sfromK (sfromR dret `sindex0` [5])) [5] + (soneHot (sfromK (sfromR dret `sindex0` [6])) [6] + (soneHot (sfromK (sfromR dret `sindex0` [7])) [7] + (soneHot (sfromK (sfromR dret `sindex0` [8])) [8] + (soneHot (sfromK (sfromR dret `sindex0` [9])) [9] + (soneHot (sfromK (sfromR dret `sindex0` [10])) [10] + soneHot (sfromK (sfromR dret `sindex0` [11])) [11])))))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> rfromS (sappend (sfromVector0N @[12] (fromList [sfromR dret `sindex0` [0], sfromR dret `sindex0` [1], sfromR dret `sindex0` [2], sfromR dret `sindex0` [3], sfromR dret `sindex0` [4], sfromR dret `sindex0` [5], sfromR dret `sindex0` [6], sfromR dret `sindex0` [7], sfromR dret `sindex0` [8], sfromR dret `sindex0` [9], sfromR dret `sindex0` [10], sfromR dret `sindex0` [11]])) (sconcrete (sreplicate [3] 0.0)))"++-- Catastrophic loss of sharing prevented.+fblowup :: forall target r. (ADReady target, NumScalar r, Differentiable r)+        => Int -> target (TKR 1 r) -> target (TKR 0 r)+fblowup k inputs =+  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)+      blowup 0 y = y - rfromIndex0 0+      blowup n y =+        let ysum = y + y - rfromIndex0 0+            yscaled = rscalar 0.499999985 * ysum+              -- without the scaling we'd get NaN at once+        in blowup (pred n) yscaled+      y0 = (inputs ! [0]) / (inputs ! [1])+  in blowup k y0++fblowupLet :: forall target r. (ADReady target, NumScalar r, Differentiable r)+           => IntOf target -> Int -> target (TKR 1 r) -> target (TKR 0 r)+fblowupLet i k inputs =+  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)+      blowup 0 y = y - rfromIndex0 i+      blowup n y1 = tlet y1 $ \y ->+        let ysum = y + y - rfromIndex0 i+            yscaled = rscalar 0.499999985 * ysum+              -- without the scaling we'd get NaN at once+        in blowup (pred n) yscaled+      y0 = (inputs ! [0]) / (inputs ! [1])+  in blowup k y0++-- Catastrophic loss of sharing prevented also with non-trivial multiplication.+fblowupMult :: forall target r. (ADReady target, NumScalar r, Differentiable r)+            => Int -> target (TKR 1 r) -> target (TKR 0 r)+fblowupMult k inputs =+  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)+      blowup 0 y = y+      blowup n y =+        let ysum = y + y * y / (y - rscalar 0.000000001)+            yscaled = sqrt $ rscalar 0.499999985 * rscalar 0.499999985 * ysum * ysum+              -- without the scaling we'd get NaN at once+        in blowup (pred n) yscaled - rfromIndex0 0+      y0 = (inputs ! [0 `remH` 2]) * (inputs ! [1])+  in blowup k y0++fblowupMultLet :: forall target r.+                  (ADReady target, NumScalar r, Differentiable r)+               => IntOf target -> Int -> target (TKR 1 r) -> target (TKR 0 r)+fblowupMultLet i k inputs =+  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)+      blowup 0 y = y+      blowup n y1 = tlet y1 $ \y ->+        let ysum0 = y + y * y / (y - rscalar 0.000001)+            yscaled = tlet ysum0 $ \ysum ->+                        sqrt $ rscalar 0.499999985 * rscalar 0.499999985 * ysum * ysum+              -- without the scaling we'd get NaN at once+        in blowup (pred n) yscaled - rfromIndex0 i+      y0 = (inputs ! [i `remH` 2]) * (inputs ! [1])+  in blowup k y0++fblowupPP :: Assertion+fblowupPP = do+  resetVarCounter+  let fblowupT = fblowup @(AstTensor AstMethodLet FullSpan) @Double 1+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupT (FTKR [4] FTKScalar)+  printArtifactSimple (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> rfromS (tlet (sfromR v1 `sindex0` [1]) (\\x4 -> tlet (sfromR v1 `sindex0` [1]) (\\x7 -> sappend (sfromVector0N @[2] (fromList [(tfromPlain (STKScalar) 0.499999985 * kfromR dret) / x4, (tfromPlain (STKScalar) (-0.499999985) * (sfromR v1 `sindex0` [0] * kfromR dret)) / (x4 * x4)])) (tfromPlain (STKS [2] STKScalar) (sconcrete (sreplicate [2] 0.0))) + sappend (sfromVector0N @[2] (fromList [(tfromPlain (STKScalar) 0.499999985 * kfromR dret) / x7, (tfromPlain (STKScalar) (-0.499999985) * (sfromR v1 `sindex0` [0] * kfromR dret)) / (x7 * x7)])) (tfromPlain (STKS [2] STKScalar) (sconcrete (sreplicate [2] 0.0))))))"+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> rfromS (tlet (sfromR v1 `sindex0` [0]) (\\x2 -> tlet (tfromPlain (STKScalar) 0.499999985 * x2) (\\x3 -> tlet (sfromR v1 `sindex0` [1]) (\\x4 -> tlet (sfromR v1 `sindex0` [0]) (\\x5 -> tlet (tfromPlain (STKScalar) 0.499999985 * x5) (\\x6 -> tlet (sfromR v1 `sindex0` [1]) (\\x7 -> soneHot (sfromK (tfromPlain (STKScalar) 0.499999985 * (recip x4 * kfromR dret))) [0] + (soneHot (sfromK ((negate x3 / (x4 * x4)) * kfromR dret)) [1] + (soneHot (sfromK (tfromPlain (STKScalar) 0.499999985 * (recip x7 * kfromR dret))) [0] + soneHot (sfromK ((negate x6 / (x7 * x7)) * kfromR dret)) [1])))))))))"+  printArtifactPrimalSimple (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromK ((tfromPlain (STKScalar) 0.499999985 * sfromR v1 `sindex0` [0]) / sfromR v1 `sindex0` [1] + (tfromPlain (STKScalar) 0.499999985 * sfromR v1 `sindex0` [0]) / sfromR v1 `sindex0` [1])"++fblowupLetPP :: Assertion+fblowupLetPP = do+  resetVarCounter+  let fblowupLetT = fblowupLet @(AstTensor AstMethodLet FullSpan) @Double 0 1+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupLetT (FTKR [4] FTKScalar)+  printArtifactSimple (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> rfromS (tlet (sfromR v1 `sindex0` [1]) (\\x4 -> tlet (tfromPlain (STKScalar) 0.99999997 * kfromR dret) (\\x7 -> sappend (sfromVector0N @[2] (fromList [x7 / x4, (negate (sfromR v1 `sindex0` [0]) * x7) / (x4 * x4)])) (tfromPlain (STKS [2] STKScalar) (sconcrete (sreplicate [2] 0.0))))))"+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> rfromS (tlet (kfromS (sfromR v1 !$ [0])) (\\x3 -> tlet (kfromS (sfromR v1 !$ [1])) (\\x4 -> tlet (tfromPlain (STKScalar) 0.99999997 * kfromR dret) (\\x7 -> soneHot (recip (sfromK x4) * sfromK x7) [0] + soneHot ((negate (sfromK x3) / (sfromK x4 * sfromK x4)) * sfromK x7) [1]))))"++fblowupLetPP23 :: Assertion+fblowupLetPP23 = do+  resetVarCounter+  let fblowupLetT = fblowupLet @(AstTensor AstMethodLet FullSpan) @Double 4 6+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupLetT (FTKR [4] FTKScalar)+  printArtifactSimple (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> rfromS (tlet (sfromR v1 `sindex0` [1]) (\\x9 -> tlet (tfromPlain (STKScalar) 0.9999998200000132 * kfromR dret) (\\x17 -> sappend (sfromVector0N @[2] (fromList [x17 / x9, (negate (sfromR v1 `sindex0` [0]) * x17) / (x9 * x9)])) (tfromPlain (STKS [2] STKScalar) (sconcrete (sreplicate [2] 0.0))))))"+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> rfromS (tlet (kfromS (sfromR v1 !$ [0])) (\\x8 -> tlet (kfromS (sfromR v1 !$ [1])) (\\x9 -> tlet (tfromPlain (STKScalar) 0.99999997 * (tfromPlain (STKScalar) 0.99999997 * (tfromPlain (STKScalar) 0.99999997 * (tfromPlain (STKScalar) 0.99999997 * (tfromPlain (STKScalar) 0.99999997 * (tfromPlain (STKScalar) 0.99999997 * kfromS (sfromK (kfromR dret)))))))) (\\x17 -> soneHot (recip (sfromK x9) * sfromK x17) [0] + soneHot ((negate (sfromK x8) / (sfromK x9 * sfromK x9)) * sfromK x17) [1]))))"+  printArtifactPrimalSimple (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromK (tfromPlain (STKScalar) (-15.999998740000063) + (tfromPlain (STKScalar) 0.9999998200000132 * sfromR v1 `sindex0` [0]) / sfromR v1 `sindex0` [1])"++-- TODO: should do 1000000 in a few seconds+blowupTests :: TestTree+blowupTests = inOrderTestGroup "Catastrophic blowup avoidance tests"+  [ testCase "blowup prim 7" $ do+      assertEqualUpToEpsilon' 1e-5+        (ringestData [2] [0.3333332333333467,-0.22222215555556446])+        (rev' @Double @0 (fblowup 7) (ringestData [2] [2, 3]))+  , testCase "blowupLet 2000" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [0.3333133339329949,-0.22220888928866325])+        (grad (rsum0 @_ @Double . (\intputs -> fblowupLet 1 2000 intputs))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet prim 200" $ do+      assertEqualUpToEpsilon' 1e-10+        (ringestData [2] [0.33333133333930065,-0.2222208888928671])+        (rev' @Double @0 (fblowupLet 1 200) (ringestData [2] [2, 3]))+  , testCase "blowupLet 7000" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [0.3332633406816766,-0.22217556045445108])+        (grad (rsum0 @_ @Double . fblowupLet 0 7000) (ringestData [2] [2, 3]))+  , testCase "blowupLet tbuild0" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [333.2633406816765,-222.175560454451])+        (grad (rsum0 @_ @Double . (\intputs -> rbuild1 1000 (\_ -> fblowupLet 0 7000 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet tbuild2" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [333.2633406816765,-222.175560454451])+        (grad (rsum0 @_ @Double . (\intputs -> rbuild1 1000 (\_ -> fblowupLet 2 7000 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet tbuildi" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [333.2983351701977,-222.19889011346513])+        (grad (rsum0 @_ @Double+               . (\intputs -> rbuild1 1000 (\i -> fblowupLet i 3500 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet tbuildc" $ do+      assertEqualUpToEpsilon 1e-7+        (ringestData [2] [333.326333406717,-222.21755560448116])+        (cgrad @_ @_ @_ @Concrete+               (rsum0 @_ @Double+                . (\intputs -> rbuild1 1000 (\i -> fblowupLet i 700 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet prim tbuild" $ do+      assertEqualUpToEpsilon 1e-7+        (ringestData [2] [33.33263334067178,-22.221755560447928])+        (grad (rsum0 @_ @Double+               . (\intputs -> rbuild1 100 (\i -> fblowupLet i 700 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupMult 3" $ do+      assertEqualUpToEpsilon' 1e-5+        (ringestData [2] [2.999999730000007,1.9999998200000046])+        (rev' @Double @0 (fblowupMult 3) (ringestData [2] [2, 3]))+  , testCase "blowupMultLet 5" $ do+      assertEqualUpToEpsilon' 1e-10+        (ringestData [2] [2.9999995500000267,1.9999997000000178])+        (rev' @Double @0 (fblowupMultLet 0 5)+                                   (ringestData [2] [2, 3]))+  , testCase "blowupMultLet 50" $ do+      assertEqualUpToEpsilon' 1e-10+        (ringestData [2] [2.999995500001215,1.99999700000081])+        (rev' @Double @0 (fblowupMultLet 0 50)+                                   (ringestData [2] [2, 3]))+  , testCase "blowupMultLet tbuild1" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [14.9999773958889,39.9999398380561])+        (grad (rsum0 @_ @Double+               . (\intputs -> rbuild1 100 (\i -> fblowupMultLet i 50 intputs)))+              (ringestData [2] [0.2, 0.3]))+  ]++concatBuild33 :: (ADReady target, NumScalar r)+             => target (TKR 1 r) -> target (TKR 2 r)+concatBuild33 _r =+  rbuild1 5 (\i ->+    rbuild1 2 (\j -> rfromIndex0 (maxH j (i `quotH` (j + 1)))))++testConcatBuild3PP :: Assertion+testConcatBuild3PP = do+  resetVarCounter+  let t = concatBuild33 @(AstTensor AstMethodLet FullSpan) @Float+      (var3, ast3) = funToAst (FTKR [3] FTKScalar) t+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> tfromPlain (STKR (SNat @2) STKScalar) (rfromS (tlet (quotH (str (sreplicate @2 (siota (SNat @5)))) (sreplicate @5 (sconcrete (sreplicate [2] 1) + siota (SNat @2)))) (\\m4 -> sgather @[5,2] (ttranspose (makePerm @[1, 2, 0]) (sfromIntegral (sfromVector @2 (fromList [sreplicate @5 (siota (SNat @2)), m4])))) (\\[i8, i6] -> [i8, i6, ifH (m4 `sindex0` [i8, i6] <=. i6) 0 1]))))"++testConcatBuild3PP2 :: Assertion+testConcatBuild3PP2 = do+  resetVarCounter+  let t = concatBuild33 @(AstTensor AstMethodLet FullSpan) @Double+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent t (FTKR [3] FTKScalar)+  printArtifactPrimalSimple (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> tfromPlain (STKR (SNat @2) STKScalar) (rfromS (sgather @[5,2] (sconcrete (sfromListLinear [2,5,2] [0.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0,0.0,0.0,1.0,0.0,2.0,1.0,3.0,1.0,4.0,2.0])) (\\[i10, i11] -> [ifH (sconcrete (sfromListLinear [5,2] [0,0,1,0,2,1,3,1,4,2]) `sindex0` [i10, i11] <=. i11) 0 1, i10, i11])))"+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> tfromPlain (STKR (SNat @1) STKScalar) (rfromS (sconcrete (sreplicate [3] 0.0)))"+  printArtifactPrimalSimple artifactRev+    @?= "\\v1 -> tfromPlain (STKR (SNat @2) STKScalar) (rfromS (tlet (quotH (str (sreplicate @2 (siota (SNat @5)))) (sreplicate @5 (sconcrete (sreplicate [2] 1) + siota (SNat @2)))) (\\m9 -> sgather @[5,2] (ttranspose (makePerm @[1, 2, 0]) (sfromIntegral (sfromVector @2 (fromList [sreplicate @5 (siota (SNat @2)), m9])))) (\\[i10, i11] -> [i10, i11, ifH (m9 `sindex0` [i10, i11] <=. i11) 0 1]))))"
+ test/simplified/TestConvQuickCheck.hs view
@@ -0,0 +1,1286 @@+{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | QuickCheck tests of convolution AD derivatives vs handwritten derivatives.+module TestConvQuickCheck (testTrees, conv2dSame_dKrn) where++import Prelude++import Data.Bifunctor (first, second)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:))+import GHC.TypeLits (KnownNat, type (+), type (-), type (<=), type (<=?))+import System.Random+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)+import Test.Tasty.QuickCheck hiding (label, shuffle)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstInterpret++import EqEpsilon+import Shared++import MnistData+import MnistFcnnRanked2 (XParams2)+import MnistFcnnRanked2 qualified++testTrees :: [TestTree]+testTrees =+  [ tensorADOnceMnistTests2+  , testCase "conv2dSameVjp dInp" test_conv2dSameVjp_dInp+  , testCase "conv2dSameVjp dKrn" test_conv2dSameVjp_dKrn+  , testProperty "conv2dSameVjp Quickcheck Double"+                 (quickcheck_conv2dSameVjp @Double)+  , testProperty "conv2dSameVjp Quickcheck Float"+                 (quickcheck_conv2dSameVjp @Float)+  , testCase "conv2dShrinkingVjp dInp" test_conv2dShrinkingVjp_dInp+  , testCase "conv2dShrinkingVjp dKrn" test_conv2dShrinkingVjp_dKrn+  , testProperty "conv2dShrinkingVjp Quickcheck Double"+                 (quickcheck_conv2dShrinkingVjp @Double)+  , testProperty "conv2dShrinkingVjp Quickcheck Float"+                 (quickcheck_conv2dShrinkingVjp @Float)+  , testCase "conv2dPaddedVjp dInp" test_conv2dPaddedVjp_dInp+  , testCase "conv2dPaddedVjp dKrn" test_conv2dPaddedVjp_dKrn+  , testProperty "conv2dPaddedVjp Quickcheck Double"+                 (quickcheck_conv2dPaddedVjp @Double)+  , testProperty "conv2dPaddedVjp Quickcheck Float"+                 (quickcheck_conv2dPaddedVjp @Float)+  , testProperty "conv2dSameJvp Quickcheck Double"+                 (quickcheck_conv2dSameJvp @Double)+  , testProperty "conv2dSameJvp Quickcheck Float"+                 (quickcheck_conv2dSameJvp @Float)+  , testProperty "conv2dShrinkingJvp Quickcheck Double"+                 (quickcheck_conv2dShrinkingJvp @Double)+  , testProperty "conv2dShrinkingJvp Quickcheck Float"+                 (quickcheck_conv2dShrinkingJvp @Float)+  , testProperty "conv2dPaddedJvp Quickcheck Double"+                 (quickcheck_conv2dPaddedJvp @Double)+  , testProperty "conv2dPaddedJvp Quickcheck Float"+                 (quickcheck_conv2dPaddedJvp @Float)+  , testProperty "conv2dSameVjp Bench dKrn Handwritten warmup"+                 (quickcheck_conv2dSameVjpKrnHandwritten @Double)+  , testProperty "conv2dSameVjp Bench dKrn Handwritten"+                 (quickcheck_conv2dSameVjpKrnHandwritten @Double)+  , testProperty "conv2dSameVjp Bench dKrn HandwrittenVectorized"+                 (quickcheck_conv2dSameVjpKrnHandwrittenVectorized @Double)+  , testProperty "conv2dSameVjp Bench dKrn Symbolic"+                 (quickcheck_conv2dSameVjpKrnSymbolic @Double)+  , testProperty "conv2dSameVjp Bench dKrn Concrete"+                 (quickcheck_conv2dSameVjpKrnConcrete @Double)+  , testProperty "conv2dSameVjp Bench dInp Handwritten"+                 (quickcheck_conv2dSameVjpInpHandwritten @Double)+  , testProperty "conv2dSameVjp Bench dInp HandwrittenVectorized"+                 (quickcheck_conv2dSameVjpInpHandwrittenVectorized @Double)+  , testProperty "conv2dSameVjp Bench dInp Symbolic"+                 (quickcheck_conv2dSameVjpInpSymbolic @Double)+  , testProperty "conv2dSameVjp Bench dInp Concrete"+                 (quickcheck_conv2dSameVjpInpConcrete @Double)+  ]++-- This one is not convolution-related, but it's also QuickCheck.+tensorADOnceMnistTests2 :: TestTree+tensorADOnceMnistTests2 = inOrderTestGroup "Ranked2 Once MNIST tests"+  [ testProperty "VTO2 grad vs fwd" $+    \seed0 ->+    forAllShrink (chooseInt (0, 600)) shrinkIntegral $ \width1Hidden ->+    forAllShrink (chooseInt (0, 200)) shrinkIntegral $ \width1Hidden2 ->+    forAllShrink (chooseInt (0, 5)) shrinkIntegral $ \simp ->+    forAll (choose (0.01, 1)) $ \range ->+    forAll (choose (0.01, 1)) $ \range2 ->+    forAll (choose (0.5, 1.5)) $ \dt ->+    forAll (choose (0, 1e-7)) $ \(perturbation :: Double) ->+    withSNat (1 + width1Hidden) $ \(SNat @widthHidden) ->+    withSNat (1 + width1Hidden2) $ \(SNat @widthHidden2) ->+    let (glyph0, seed2) = randomValue @(Concrete (TKS '[SizeMnistGlyph] Double))+                                      0.5 (mkStdGen seed0)+        (label0, seed3) = randomValue @(Concrete (TKS '[SizeMnistLabel] Double))+                                      5 seed2+        (glyph, label) = ( rmap1 (rscalar 0.5 +) $ forgetShape glyph0+                         , rmap1 (rscalar 5 + ) $ forgetShape label0 )+        ds :: Concrete (XParams2 Double Double)+        (ds, seed4) = first forgetShape $+          randomValue+            @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped+                             Concrete widthHidden widthHidden2 Double Double)))+            range seed3+        (targetInit, artRaw) =+          MnistFcnnRanked2.mnistTrainBench2VTOGradient+            @Double (Proxy @Double) UseIncomingCotangent+            range2 seed4 (1 + width1Hidden) (1 + width1Hidden2)+        art = iterate simplifyArtifactRev artRaw !! simp+        stk = knownSTK @(XParams2 Double Double)+        ftk = tftk @Concrete stk targetInit+        parametersAndInput = tpair targetInit (tpair glyph label)+        (value0, _gradient0) = second tproject1 $+          revInterpretArtifact art parametersAndInput Nothing+        (value1, gradient1) = second tproject1 $+          revInterpretArtifact art parametersAndInput (Just $ kconcrete dt)+        f :: ADVal Concrete (XParams2 Double Double)+          -> ADVal Concrete (TKScalar Double)+        f adinputs =+          MnistFcnnRanked2.afcnnMnistLoss2+            (rfromPrimal glyph, rfromPrimal label) (fromTarget adinputs)+        (value2, derivative2) = cjvp2 f targetInit ds+--        goodDt :: forall r. GoodScalar r => r+--        goodDt = ifDifferentiable @r (realToFrac dt) 0+--        targetDt :: Concrete (XParams2 Double Double)+--        targetDt = replTarget goodDt ftk+        goodPerturbation :: forall r. NumScalar r => r+        goodPerturbation = ifDifferentiable @r (realToFrac perturbation) 0+        targetPerturbed :: Concrete (XParams2 Double Double)+        targetPerturbed = treplTarget goodPerturbation ftk+        targetInitPerturbed :: Concrete (XParams2 Double Double)+        targetInitPerturbed = taddTarget stk targetInit targetPerturbed+        (value3, derivative3) = cjvp2 f targetInit targetPerturbed+        value4 :: Concrete (TKScalar Double)+        value4 = MnistFcnnRanked2.afcnnMnistLoss2+                   (rfromPrimal glyph, rfromPrimal label)+                   (fromTarget targetInitPerturbed)+    in+      conjoin+        [ counterexample+            ("Objective function value from grad and jvp matches: "+             ++ show (value1, value2, value1 - value2))+            (abs (value1 - value2) < 1e-10)+        , counterexample+            ("Gradient and derivative agrees: "+             ++ show ( dt, derivative2, tdot0Target ftk gradient1 ds+                     , tdot0Target FTKScalar (kconcrete dt) derivative2+                       - tdot0Target ftk gradient1 ds ))+            (abs (tdot0Target FTKScalar (kconcrete dt) derivative2+                  - tdot0Target ftk gradient1 ds) < 1e-10)+--        , counterexample  -- this is implied by the other clauses+--            "Gradient is a linear function"+--            (gradient1 === tmultTarget stk targetDt gradient0)+        , counterexample+            "Objective function value unaffected by incoming cotangent"+            (value0 === value1)+        , counterexample+            "Objective function value unaffected by derivative perturbation"+            (value2 === value3)+        , counterexample+            ("Derivative approximates the perturbation of value: "+             ++ show ( value2, derivative3, value4+                     , (value3 + derivative3) - value4) )+            (abs ((value3 + derivative3) - value4) < 1e-6)+        ]+  ]++allClose :: (GoodScalar r, Fractional r, Linearizable a r, Linearizable b r)+         => a -> b -> Rational -> Bool+allClose expected actual eps =+  all (\(a, b) -> abs (a - b) <= fromRational eps)+  $ zip (linearize expected) (linearize actual)++flip42 :: (ADReady target, GoodScalar r)+       => target (TKS '[nCout, nCinp, nKh, nKw] r)+       -> target (TKS '[nCout, nCinp, nKh, nKw] r)+flip42 arr =+  stranspose @'[1, 2, 0]+  $ sreverse+  $ stranspose @'[3, 1, 2, 0]+  $ sreverse+  $ stranspose @'[3, 0, 1, 2] arr++-- | Hand-written reverse derivative of full convolution with respect+-- to the input image.+-- Example code that horde-ad generates for the same is in testSameCNNOPP0bW.+conv2dSame_dInp+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB shB1+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh, KnownNat nAw, KnownNat nKh, KnownNat nKw+     , ADReady target, NumScalar r+     , shK  ~ '[nCout, nCinp, nKh, nKw]+     , shA  ~ '[nImgs, nCinp, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh, nAw]+     , shB1 ~ '[1,     nCout,  nKh, nKw] )+  => target (TKS shK r)+  -> target (TKS shB r)+  -> target (TKS shA r)+conv2dSame_dInp arrK arrB =+  let arrKFlipped = flip42 arrK+      nKh = valueOf @nKh+      nKw = valueOf @nKw+  in sbuild @shA $ \case+    [iImg, iCinp, iAh, iAw] ->+      let arrBt = slicezS @shB1 arrB+                          [iImg,  0, iAh - nKh + 1, iAw - nKw + 1]+          arrKt = slicezS (stranspose @'[1, 0] arrKFlipped)+                          [iCinp, 0 , 0, 0]+      in sfromK $ sdot0 arrBt arrKt+    _ -> error "conv2dSame_dInp: impossible pattern needlessly required"+-- Note that+-- > ... in conv2dSameS (stranspose @'[1, 0] arrKFlipped) arrB+-- type-checks above, but test fails due to the lack of @- nKh + 1@.++-- | Hand-written reverse derivative of full convolution with respect+-- to the kernels.+-- This code vectorized is pretty-printed in test testSameCNNOPPKrnHandwritten.+-- Example code that horde-ad generates for the same is in testSameCNNOPP0cW.+conv2dSame_dKrn+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB shB1+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh, KnownNat nAw, KnownNat nKh, KnownNat nKw+     , ADReady target, NumScalar r+     , shK  ~ '[nCout, nCinp, nKh, nKw]+     , shA  ~ '[nImgs, nCinp, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh, nAw]+     , shB1 ~ '[nImgs, 1,     nAh, nAw] )+  => target (TKS shA r)+  -> target (TKS shB r)+  -> target (TKS shK r)+conv2dSame_dKrn arrA arrB =+  sbuild @shK $ \case+    [iCout, iCinp, iKh, iKw] ->+      let arrBt = slicezS @shB1 arrB+                          [0, iCout, 0, 0]+          arrAt = slicezS arrA+                          [0, iCinp, iKh, iKw]+      in sfromK $ sdot0 arrBt arrAt+    _ -> error "conv2dSame_dKrn: impossible pattern needlessly required"++test_conv2dSameVjp_dInp :: Assertion+test_conv2dSameVjp_dInp =+  let -- Input of shape: batch x chas x height x width+      arrA :: Nested.Shaped '[5, 2, 4, 8] Double+      arrA = Nested.sreplicatePrim knownShS 1.1+      -- Filters of shape: num_filters x chas x kernel_height x kernel_width+      arrK :: Nested.Shaped '[7, 2, 1, 3] Double+      arrK = Nested.sreplicatePrim knownShS (-2.2)+      -- Output gradient of shape: batch x chas x output_height x output_width+      arrB :: Nested.Shaped '[5, 7, 4, 8] Double+      arrB = Nested.sreplicatePrim knownShS (-3.3)+      -- Compare the AD version against the manual derivative.+      dInp :: Concrete (TKS '[5, 2, 4, 8] Double)+      dInp = conv2dSame_dInp (sconcrete arrK) (sconcrete arrB)+      vjpInp = cvjp @_ @_ @_ @Concrete+                    (conv2dSameS (sconcrete arrK))+                    (sconcrete arrA) (sconcrete arrB)+  in assertEqualUpToEpsilon 1e-7 dInp vjpInp++test_conv2dSameVjp_dKrn :: Assertion+test_conv2dSameVjp_dKrn =+  let -- Input of shape: batch x chas x height x width+      arrA :: Nested.Shaped '[5, 2, 4, 8] Double+      arrA = Nested.sreplicatePrim knownShS 1.1+      -- Filters of shape: num_filters x chas x kernel_height x kernel_width+      arrK :: Nested.Shaped '[7, 2, 1, 3] Double+      arrK = Nested.sreplicatePrim knownShS 2.2+      -- Output gradient of shape: batch x chas x output_height x output_width+      arrB :: Nested.Shaped '[5, 7, 4, 8] Double+      arrB = Nested.sreplicatePrim knownShS (-3.3)+      -- Compare the AD version against the manual derivative.+      dKrn :: Concrete (TKS '[7, 2, 1, 3] Double)+      dKrn = conv2dSame_dKrn (sconcrete arrA) (sconcrete arrB)+      vjpKrn = cvjp @_ @_ @_ @Concrete+                    (`conv2dSameS` sconcrete arrA)+                    (sconcrete arrK) (sconcrete arrB)+  in assertEqualUpToEpsilon 1e-7 dKrn vjpKrn++static_conv2dSameVjp+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+       -- ^ Filters of shape: num_filters x chas x kernel_height x kernel_width+  -> Nested.Shaped shA r+       -- ^ Input of shape: batch x chas x height x width+  -> Nested.Shaped shB r+       -- ^ Output gradient of shape:+       --     batch x chas x output_height x output_width+  -> Bool+static_conv2dSameVjp SNat SNat SNat SNat SNat SNat SNat arrK arrA arrB =+  let -- Compare the AD version against the manual derivative.+      -- Note that manual versions don't take one of the arguments (the point+      -- at which the gradient is taken), because maths (something about+      -- convolution being linear and so gradient the same everywhere).+      -- First, the gradient wrt the input image taken at point @arrA@.+      dInp :: Concrete (TKS shA r)+      dInp = conv2dSame_dInp+               (sconcrete arrK) (sconcrete arrB)  -- handwritten+      vjpInp = vjp (conv2dSameS (sconcrete arrK))+                   (sconcrete arrA) (sconcrete arrB)  -- symbolic+      cvjpInp = cvjp @_ @_ @_ @Concrete+                     (conv2dSameS (sconcrete arrK))+                     (sconcrete arrA) (sconcrete arrB)  -- concrete+      -- Second, the gradient wrt the kernels taken at point @arrK@.+      dKrn :: Concrete (TKS shK r)+      dKrn = conv2dSame_dKrn+               (sconcrete arrA) (sconcrete arrB)  -- handwritten+      vjpKrn = vjp (`conv2dSameS` sconcrete arrA)+                   (sconcrete arrK) (sconcrete arrB)  -- symbolic+      cvjpKrn = cvjp @_ @_ @_ @Concrete+                     (`conv2dSameS` sconcrete arrA)+                     (sconcrete arrK) (sconcrete arrB)  -- concrete+  in allClose vjpInp dInp 1e-5  -- 1e-7 is too much for Float+     && allClose cvjpInp dInp 1e-5+     && allClose vjpKrn dKrn 1e-5+     && allClose cvjpKrn dKrn 1e-5++quickcheck_conv2dSameVjp+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjp =+  forAll (choose (0, 5)) $ \nImgs' ->+  forAll (choose (0, 5)) $ \nCinp' ->+  forAll (choose (0, 5)) $ \nCout' ->+  forAll (choose (0, 5)) $ \nAh' ->+  forAll (choose (0, 5)) $ \nAw' ->+  forAll (choose (0, 5)) $ \nKh' ->+  forAll (choose (0, 5)) $ \nKw' ->+    -- The glue below is needed to bridge the dependently-typed+    -- vs normal world.+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjp+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)++-- | Hand-written reverse derivative of full shrinking convolution with respect+-- to the input image.+-- The @nKh1@ type variable reads \"nKh - 1\", while @nAh_nKh1@+-- reads \" nAh - nKh1\". We don't use explicit subtraction, because+-- it requires extra constraints to ensure type leven nats are not negative.+-- Example code that horde-ad generates for the same is+-- in testShrinkingCNNOPP0bW.+conv2dShrinking_dInp+  :: forall nImgs nCinp nCout nAh_nKh1 nAw_nKw1 nKh1 nKw1 shK shA shB+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh_nKh1, KnownNat nAw_nKw1, KnownNat nKh1, KnownNat nKw1+     , ADReady target, NumScalar r+     , shK  ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA  ~ '[nImgs, nCinp, nAh_nKh1 + nKh1, nAw_nKw1 + nKw1]+     , shB  ~ '[nImgs, nCout, nAh_nKh1, nAw_nKw1] )+  => target (TKS shK r)+  -> target (TKS shB r)+  -> target (TKS shA r)+conv2dShrinking_dInp arrK arrB =+  let arrKFlipped = flip42 arrK+  in conv2dPaddedS (stranspose @'[1, 0] arrKFlipped) arrB++-- | Hand-written reverse derivative of full shrinking convolution with respect+-- to the kernels.+-- The @nKh1@ type variable reads \"nKh - 1\", while @nAh_nKh1@+-- reads \" nAh - nKh1\". We don't use explicit subtraction, because+-- it requires extra constraints to ensure type leven nats are not negative.+-- Example code that horde-ad generates for the same is+-- in testShrinkingCNNOPP0cW.+conv2dShrinking_dKrn+  :: forall nImgs nCinp nCout nAh_nKh1 nAw_nKw1 nKh1 nKw1 shK shA shB+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh_nKh1, KnownNat nAw_nKw1, KnownNat nKh1, KnownNat nKw1+     , ADReady target, NumScalar r+     , 1 <= nAh_nKh1, 1 <= nAw_nKw1+     , shK  ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA  ~ '[nImgs, nCinp, nAh_nKh1 + nKh1, nAw_nKw1 + nKw1]+     , shB  ~ '[nImgs, nCout, nAh_nKh1, nAw_nKw1] )+  => target (TKS shA r)+  -> target (TKS shB r)+  -> target (TKS shK r)+conv2dShrinking_dKrn arrA arrB =+  stranspose @'[1, 0]+             (conv2dShrinkingS @_ @_ @_ @_ @_ @_ @(nAh_nKh1 - 1) @(nAw_nKw1 - 1)+                               (stranspose @'[1, 0] arrB)+                               (stranspose @'[1, 0] arrA))++test_conv2dShrinkingVjp_dInp :: Assertion+test_conv2dShrinkingVjp_dInp =+  let -- Input of shape: batch x chas x height x width+      arrA :: Nested.Shaped '[5, 2, 4, 8] Double+      arrA = Nested.sreplicatePrim knownShS 1.1+      -- Filters of shape: num_filters x chas x kernel_height x kernel_width+      arrK :: Nested.Shaped '[7, 2, 1, 3] Double+      arrK = Nested.sreplicatePrim knownShS (-2.2)+      -- Output gradient of shape: batch x chas x output_height x output_width+      arrB :: Nested.Shaped '[5, 7, 4, 6] Double+      arrB = Nested.sreplicatePrim knownShS 3.3+      -- Compare the AD version against the manual derivative.+      dInp :: Concrete (TKS '[5, 2, 4, 8] Double)+      dInp = conv2dShrinking_dInp (sconcrete arrK) (sconcrete arrB)+      vjpInp = cvjp @_ @_ @_ @Concrete+                    (conv2dShrinkingS (sconcrete arrK))+                    (sconcrete arrA) (sconcrete arrB)+  in assertEqualUpToEpsilon 1e-7 dInp vjpInp++test_conv2dShrinkingVjp_dKrn :: Assertion+test_conv2dShrinkingVjp_dKrn =+  let -- Input of shape: batch x chas x height x width+      arrA :: Nested.Shaped '[5, 2, 4, 8] Double+      arrA = Nested.sreplicatePrim knownShS (-1.1)+      -- Filters of shape: num_filters x chas x kernel_height x kernel_width+      arrK :: Nested.Shaped '[7, 2, 1, 3] Double+      arrK = Nested.sreplicatePrim knownShS 2.2+      -- Output gradient of shape: batch x chas x output_height x output_width+      arrB :: Nested.Shaped '[5, 7, 4, 6] Double+      arrB = Nested.sreplicatePrim knownShS 3.3+      -- Compare the AD version against the manual derivative.+      dKrn :: Concrete (TKS '[7, 2, 1, 3] Double)+      dKrn = conv2dShrinking_dKrn (sconcrete arrA) (sconcrete arrB)+      vjpKrn = cvjp @_ @_ @_ @Concrete+                    (`conv2dShrinkingS` sconcrete arrA)+                    (sconcrete arrK) (sconcrete arrB)+  in assertEqualUpToEpsilon 1e-7 dKrn vjpKrn++static_conv2dShrinkingVjp+  :: forall nImgs nCinp nCout nAh_nKh1 nAw_nKw1 nKh1 nKw1 shK shA shB r.+     ( NumScalar r, Differentiable r+     , 1 <= nAh_nKh1, 1 <= nAw_nKw1+     , shK ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA ~ '[nImgs, nCinp, nAh_nKh1 + nKh1, nAw_nKw1 + nKw1]+     , shB ~ '[nImgs, nCout, nAh_nKh1, nAw_nKw1] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh_nKh1 -> SNat nAw_nKw1 -> SNat nKh1 -> SNat nKw1+  -> Nested.Shaped shK r+       -- ^ Filters of shape: num_filters x chas x kernel_height x kernel_width+  -> Nested.Shaped shA r+       -- ^ Input of shape: batch x chas x height x width+  -> Nested.Shaped shB r+       -- ^ Output gradient of shape:+       --     batch x chas x output_height x output_width+  -> Bool+static_conv2dShrinkingVjp SNat SNat SNat SNat SNat SNat SNat arrK arrA arrB =+  let -- Compare the AD version against the manual derivative.+      -- Note that manual versions don't take one of the arguments (the point+      -- at which the gradient is taken), because maths (something about+      -- convolution being linear and so gradient the same everywhere).+      -- First, the gradient wrt the input image taken at point @arrA@.+      dInp :: Concrete (TKS shA r)+      dInp = conv2dShrinking_dInp+               (sconcrete arrK) (sconcrete arrB)  -- handwritten+      vjpInp = vjp (conv2dShrinkingS (sconcrete arrK))+                   (sconcrete arrA) (sconcrete arrB)  -- symbolic+      cvjpInp = cvjp @_ @_ @_ @Concrete+                     (conv2dShrinkingS (sconcrete arrK))+                     (sconcrete arrA) (sconcrete arrB)  -- concrete+      -- Second, the gradient wrt the kernels taken at point @arrK@.+      dKrn :: Concrete (TKS shK r)+      dKrn = conv2dShrinking_dKrn+               (sconcrete arrA) (sconcrete arrB) -- handwritten+      vjpKrn = vjp (`conv2dShrinkingS` sconcrete arrA)+                   (sconcrete arrK) (sconcrete arrB)  -- symbolic+      cvjpKrn = cvjp @_ @_ @_ @Concrete+                     (`conv2dShrinkingS` sconcrete arrA)+                     (sconcrete arrK) (sconcrete arrB)  -- concrete+  in allClose vjpInp dInp 1e-5  -- 1e-7 is too much for Float+     && allClose cvjpInp dInp 1e-5+     && allClose vjpKrn dKrn 1e-5+     && allClose cvjpKrn dKrn 1e-5++quickcheck_conv2dShrinkingVjp+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dShrinkingVjp =+  forAll chooseAny $ \(seed0 :: Int) ->+  forAll (choose (0, 5)) $ \nImgs' ->+  forAll (choose (0, 5)) $ \nCinp' ->+  forAll (choose (0, 5)) $ \nCout' ->+  forAll (choose (1, 5)) $ \nAh_nKh1' ->+  forAll (choose (1, 5)) $ \nAw_nKw1' ->+  forAll (choose (0, 5)) $ \nKh1' ->+  forAll (choose (0, 5)) $ \nKw1' ->+    -- The glue below is needed to bridge the dependently-typed+    -- vs normal world.+    -- The @b@ is needed for GHC 9.10 to type-check this code.+    let b :: Bool+        b =+          withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+          withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+          withSNat nCout' $ \(nCout :: SNat nCout) ->+          withSNat nAh_nKh1' $ \(nAh_nKh1 :: SNat nAh_nKh1) ->+          withSNat nAw_nKw1' $ \(nAw_nKw1 :: SNat nAw_nKw1) ->+          withSNat nKh1' $ \(nKh1 :: SNat nKh1) ->+          withSNat nKw1' $ \(nKw1 :: SNat nKw1) ->+            gcastWith (unsafeCoerceRefl :: (1 <=? nAh_nKh1) :~: True) $+            gcastWith (unsafeCoerceRefl :: (1 <=? nAw_nKw1) :~: True) $+              let arrK :: Concrete (TKS '[nCout, nCinp, nKh1 + 1, nKw1 + 1] r)+                  (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+                  arrA :: Concrete (TKS '[ nImgs, nCinp+                                         , nAh_nKh1 + nKh1, nAw_nKw1 + nKw1 ] r)+                  (arrA, seed3) = randomValue 0.5 seed2+                  arrB :: Concrete (TKS '[nImgs, nCout, nAh_nKh1, nAw_nKw1] r)+                  (arrB, _) = randomValue 0.5 seed3+              in static_conv2dShrinkingVjp+                   nImgs nCinp nCout nAh_nKh1 nAw_nKw1 nKh1 nKw1+                   (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)+    in b++-- | Hand-written reverse derivative of full padded convolution with respect+-- to the input image.+-- The @nKh1@ type variable reads \"nKh - 1\".+-- We don't use explicit subtraction, because+-- it requires extra constraints to ensure type leven nats are not negative.+-- Example code that horde-ad generates for the same is+-- in testPaddedCNNOPP0bW.+conv2dPadded_dInp+  :: forall nImgs nCinp nCout nAh nAw nKh1 nKw1 shK shA shB+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh, KnownNat nAw, KnownNat nKh1, KnownNat nKw1+     , ADReady target, NumScalar r+     , shK  ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA  ~ '[nImgs, nCinp, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh + nKh1, nAw + nKw1] )+  => target (TKS shK r)+  -> target (TKS shB r)+  -> target (TKS shA r)+conv2dPadded_dInp arrK arrB =+  let arrKFlipped = flip42 arrK+  in conv2dShrinkingS (stranspose @'[1, 0] arrKFlipped) arrB++-- | Hand-written reverse derivative of full padded convolution with respect+-- to kernels.+-- The @nKh1@ type variable reads \"nKh - 1\".+-- We don't use explicit subtraction, because+-- it requires extra constraints to ensure type leven nats are not negative.+-- Example code that horde-ad generates for the same is+-- in testPaddedCNNOPP0cW.+conv2dPadded_dKrn+  :: forall nImgs nCinp nCout nAh nAw nKh1 nKw1 shK shA shB+            target r.+     ( KnownNat nImgs, KnownNat nCinp, KnownNat nCout+     , KnownNat nAh, KnownNat nAw, KnownNat nKh1, KnownNat nKw1+     , ADReady target, NumScalar r+     , 1 <= nAh, 1 <= nAw+     , shK  ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA  ~ '[nImgs, nCinp, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh + nKh1, nAw + nKw1] )+  => target (TKS shA r)+  -> target (TKS shB r)+  -> target (TKS shK r)+conv2dPadded_dKrn arrA arrB =+  flip42+  $ conv2dShrinkingS @_ @_ @_ @_ @_ @_ @(nAh - 1) @(nAw - 1)+                     (stranspose @'[1, 0] arrA)+                     (stranspose @'[1, 0] arrB)++test_conv2dPaddedVjp_dInp :: Assertion+test_conv2dPaddedVjp_dInp =+  let -- Input of shape: batch x chas x height x width+      arrA :: Nested.Shaped '[5, 2, 4, 8] Double+      arrA = Nested.sreplicatePrim knownShS 1.1+      -- Filters of shape: num_filters x chas x kernel_height x kernel_width+      arrK :: Nested.Shaped '[7, 2, 1, 3] Double+      arrK = Nested.sreplicatePrim knownShS (-2.2)+      -- Output gradient of shape: batch x chas x output_height x output_width+      arrB :: Nested.Shaped '[5, 7, 4, 10] Double+      arrB = Nested.sreplicatePrim knownShS 3.3+      -- Compare the AD version against the manual derivative.+      dInp :: Concrete (TKS '[5, 2, 4, 8] Double)+      dInp = conv2dPadded_dInp (sconcrete arrK) (sconcrete arrB)+      vjpInp = cvjp @_ @_ @_ @Concrete+                    (conv2dPaddedS (sconcrete arrK))+                    (sconcrete arrA) (sconcrete arrB)+  in assertEqualUpToEpsilon 1e-7 dInp vjpInp++test_conv2dPaddedVjp_dKrn :: Assertion+test_conv2dPaddedVjp_dKrn =+  let -- Input of shape: batch x chas x height x width+      arrA :: Nested.Shaped '[5, 2, 4, 8] Double+      arrA = Nested.sreplicatePrim knownShS (-1.1)+      -- Filters of shape: num_filters x chas x kernel_height x kernel_width+      arrK :: Nested.Shaped '[7, 2, 1, 3] Double+      arrK = Nested.sreplicatePrim knownShS 2.2+      -- Output gradient of shape: batch x chas x output_height x output_width+      arrB :: Nested.Shaped '[5, 7, 4, 10] Double+      arrB = Nested.sreplicatePrim knownShS 3.3+      -- Compare the AD version against the manual derivative.+      dKrn :: Concrete (TKS '[7, 2, 1, 3] Double)+      dKrn = conv2dPadded_dKrn (sconcrete arrA) (sconcrete arrB)+      vjpKrn = cvjp @_ @_ @_ @Concrete+                    (`conv2dPaddedS` sconcrete arrA)+                    (sconcrete arrK) (sconcrete arrB)+  in assertEqualUpToEpsilon 1e-7 dKrn vjpKrn++static_conv2dPaddedVjp+  :: forall nImgs nCinp nCout nAh nAw nKh1 nKw1 shK shA shB r.+     ( NumScalar r, Differentiable r+     , 1 <= nAh, 1 <= nAw+     , shK  ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA  ~ '[nImgs, nCinp, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh + nKh1, nAw + nKw1] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh1 -> SNat nKw1+  -> Nested.Shaped shK r+       -- ^ Filters of shape: num_filters x chas x kernel_height x kernel_width+  -> Nested.Shaped shA r+       -- ^ Input of shape: batch x chas x height x width+  -> Nested.Shaped shB r+       -- ^ Output gradient of shape:+       --     batch x chas x output_height x output_width+  -> Bool+static_conv2dPaddedVjp SNat SNat SNat SNat SNat SNat SNat arrK arrA arrB =+  let -- Compare the AD version against the manual derivative.+      -- Note that manual versions don't take one of the arguments (the point+      -- at which the gradient is taken), because maths (something about+      -- convolution being linear and so gradient the same everywhere).+      -- First, the gradient wrt the input image taken at point @arrA@.+      dInp :: Concrete (TKS shA r)+      dInp = conv2dPadded_dInp+               (sconcrete arrK) (sconcrete arrB)  -- handwritten+      vjpInp = vjp (conv2dPaddedS (sconcrete arrK))+                   (sconcrete arrA) (sconcrete arrB)  -- symbolic+      cvjpInp = cvjp @_ @_ @_ @Concrete+                     (conv2dPaddedS (sconcrete arrK))+                     (sconcrete arrA) (sconcrete arrB)  -- concrete+      -- Second, the gradient wrt the kernels taken at point @arrK@.+      dKrn :: Concrete (TKS shK r)+      dKrn = conv2dPadded_dKrn+               (sconcrete arrA) (sconcrete arrB) -- handwritten+      vjpKrn = vjp (`conv2dPaddedS` sconcrete arrA)+                   (sconcrete arrK) (sconcrete arrB)  -- symbolic+      cvjpKrn = cvjp @_ @_ @_ @Concrete+                     (`conv2dPaddedS` sconcrete arrA)+                     (sconcrete arrK) (sconcrete arrB)  -- concrete+  in allClose vjpInp dInp 1e-5  -- 1e-7 is too much for Float+     && allClose cvjpInp dInp 1e-5+     && allClose vjpKrn dKrn 1e-5+     && allClose cvjpKrn dKrn 1e-5++quickcheck_conv2dPaddedVjp+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dPaddedVjp =+  forAll chooseAny $ \(seed0 :: Int) ->+  forAll (choose (0, 5)) $ \nImgs' ->+  forAll (choose (0, 5)) $ \nCinp' ->+  forAll (choose (0, 5)) $ \nCout' ->+  forAll (choose (1, 5)) $ \nAh' ->+  forAll (choose (1, 5)) $ \nAw' ->+  forAll (choose (0, 5)) $ \nKh1' ->+  forAll (choose (0, 5)) $ \nKw1' ->+    -- The glue below is needed to bridge the dependently-typed+    -- vs normal world.+    -- The @b@ is needed for GHC 9.10 to type-check this code.+    let b :: Bool+        b =+          withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+          withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+          withSNat nCout' $ \(nCout :: SNat nCout) ->+          withSNat nAh' $ \(nAh :: SNat nAh) ->+          withSNat nAw' $ \(nAw :: SNat nAw) ->+          withSNat nKh1' $ \(nKh1 :: SNat nKh1) ->+          withSNat nKw1' $ \(nKw1 :: SNat nKw1) ->+            gcastWith (unsafeCoerceRefl :: (1 <=? nAh) :~: True) $+            gcastWith (unsafeCoerceRefl :: (1 <=? nAw) :~: True) $+              let arrK :: Concrete (TKS '[nCout, nCinp, nKh1 + 1, nKw1 + 1] r)+                  (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+                  arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+                  (arrA, seed3) = randomValue 0.5 seed2+                  arrB+                    :: Concrete (TKS '[nImgs, nCout, nAh + nKh1, nAw + nKw1] r)+                  (arrB, _) = randomValue 0.5 seed3+              in static_conv2dPaddedVjp+                   nImgs nCinp nCout nAh nAw nKh1 nKw1+                   (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)+    in b++-- Forward derivative.+static_conv2dSameJvp+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r -> Nested.Shaped shK r+  -> Nested.Shaped shA r -> Nested.Shaped shA r+  -> Bool+static_conv2dSameJvp SNat SNat SNat SNat SNat SNat SNat+                     arrK arrK2 arrA arrA2 =+  let dInp :: Concrete (TKS shB r)+      dInp = conv2dSameS (sconcrete arrK) (sconcrete arrA2)+      jvpInp = jvp (conv2dSameS (sconcrete arrK))+                   (sconcrete arrA) (sconcrete arrA2)+      cjvpInp = cjvp @_ @_ @_ @Concrete+                     (conv2dSameS (sconcrete arrK))+                     (sconcrete arrA) (sconcrete arrA2)+      dKrn :: Concrete (TKS shB r)+      dKrn = conv2dSameS (sconcrete arrK2) (sconcrete arrA)+      jvpKrn = jvp (`conv2dSameS` sconcrete arrA)+                   (sconcrete arrK) (sconcrete arrK2)+      cjvpKrn = cjvp @_ @_ @_ @Concrete+                     (`conv2dSameS` sconcrete arrA)+                     (sconcrete arrK) (sconcrete arrK2)+  in allClose jvpInp dInp 1e-5  -- 1e-7 is too much for Float+     && allClose cjvpInp dInp 1e-5+     && allClose jvpKrn dKrn 1e-5+     && allClose cjvpKrn dKrn 1e-5++quickcheck_conv2dSameJvp+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameJvp =+  forAll (choose (0, 5)) $ \nImgs' ->+  forAll (choose (0, 5)) $ \nCinp' ->+  forAll (choose (0, 5)) $ \nCout' ->+  forAll (choose (0, 5)) $ \nAh' ->+  forAll (choose (0, 5)) $ \nAw' ->+  forAll (choose (0, 5)) $ \nKh' ->+  forAll (choose (0, 5)) $ \nKw' ->+    -- The glue below is needed to bridge the dependently-typed+    -- vs normal world.+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK, arrK2 :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            (arrK2, seed3) = randomValue 0.5 seed2+            arrA, arrA2 :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed4) = randomValue 0.5 seed3+            (arrA2, _) = randomValue 0.5 seed4+        in static_conv2dSameJvp+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrK2)+             (unConcrete arrA) (unConcrete arrA2)++-- Forward derivative.+static_conv2dShrinkingJvp+  :: forall nImgs nCinp nCout nAh_nKh1 nAw_nKw1 nKh1 nKw1 shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA ~ '[nImgs, nCinp, nAh_nKh1 + nKh1, nAw_nKw1 + nKw1]+     , shB  ~ '[nImgs, nCout, nAh_nKh1, nAw_nKw1])+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh_nKh1 -> SNat nAw_nKw1 -> SNat nKh1 -> SNat nKw1+  -> Nested.Shaped shK r -> Nested.Shaped shK r+  -> Nested.Shaped shA r -> Nested.Shaped shA r+  -> Bool+static_conv2dShrinkingJvp SNat SNat SNat SNat SNat SNat SNat+                          arrK arrK2 arrA arrA2 =+  let dInp :: Concrete (TKS shB r)+      dInp = conv2dShrinkingS (sconcrete arrK) (sconcrete arrA2)+      jvpInp = jvp (conv2dShrinkingS (sconcrete arrK))+                   (sconcrete arrA) (sconcrete arrA2)+      cjvpInp = cjvp @_ @_ @_ @Concrete+                     (conv2dShrinkingS (sconcrete arrK))+                     (sconcrete arrA) (sconcrete arrA2)+      dKrn :: Concrete (TKS shB r)+      dKrn = conv2dShrinkingS (sconcrete arrK2) (sconcrete arrA)+      jvpKrn = jvp (`conv2dShrinkingS` sconcrete arrA)+                   (sconcrete arrK) (sconcrete arrK2)+      cjvpKrn = cjvp @_ @_ @_ @Concrete+                     (`conv2dShrinkingS` sconcrete arrA)+                     (sconcrete arrK) (sconcrete arrK2)+  in allClose jvpInp dInp 1e-5  -- 1e-7 is too much for Float+     && allClose cjvpInp dInp 1e-5+     && allClose jvpKrn dKrn 1e-5+     && allClose cjvpKrn dKrn 1e-5++quickcheck_conv2dShrinkingJvp+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dShrinkingJvp =+  forAll (choose (0, 5)) $ \nImgs' ->+  forAll (choose (0, 5)) $ \nCinp' ->+  forAll (choose (0, 5)) $ \nCout' ->+  forAll (choose (0, 5)) $ \nAh_nKh1' ->+  forAll (choose (0, 5)) $ \nAw_nKw1' ->+  forAll (choose (0, 5)) $ \nKh1' ->+  forAll (choose (0, 5)) $ \nKw1' ->+    -- The glue below is needed to bridge the dependently-typed+    -- vs normal world.+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh_nKh1' $ \(nAh_nKh1 :: SNat nAh_nKh1) ->+    withSNat nAw_nKw1' $ \(nAw_nKw1 :: SNat nAw_nKw1) ->+    withSNat nKh1' $ \(nKh1 :: SNat nKh1) ->+    withSNat nKw1' $ \(nKw1 :: SNat nKw1) ->+      property $ \seed0 ->+        let arrK, arrK2 :: Concrete (TKS '[nCout, nCinp, nKh1 + 1, nKw1 + 1] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            (arrK2, seed3) = randomValue 0.5 seed2+            arrA, arrA2+              :: Concrete+                   (TKS '[nImgs, nCinp, nAh_nKh1 + nKh1, nAw_nKw1 + nKw1] r)+            (arrA, seed4) = randomValue 0.5 seed3+            (arrA2, _) = randomValue 0.5 seed4+        in static_conv2dShrinkingJvp+             nImgs nCinp nCout nAh_nKh1 nAw_nKw1 nKh1 nKw1+             (unConcrete arrK) (unConcrete arrK2)+             (unConcrete arrA) (unConcrete arrA2)++-- Forward derivative.+static_conv2dPaddedJvp+  :: forall nImgs nCinp nCout nAh nAw nKh1 nKw1 shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh1 + 1, nKw1 + 1]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB  ~ '[nImgs, nCout, nAh + nKh1, nAw + nKw1] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh1 -> SNat nKw1+  -> Nested.Shaped shK r -> Nested.Shaped shK r+  -> Nested.Shaped shA r -> Nested.Shaped shA r+  -> Bool+static_conv2dPaddedJvp SNat SNat SNat SNat SNat SNat SNat+                       arrK arrK2 arrA arrA2 =+  let dInp :: Concrete (TKS shB r)+      dInp = conv2dPaddedS (sconcrete arrK) (sconcrete arrA2)+      jvpInp = jvp (conv2dPaddedS (sconcrete arrK))+                   (sconcrete arrA) (sconcrete arrA2)+      cjvpInp = cjvp @_ @_ @_ @Concrete+                     (conv2dPaddedS (sconcrete arrK))+                     (sconcrete arrA) (sconcrete arrA2)+      dKrn :: Concrete (TKS shB r)+      dKrn = conv2dPaddedS (sconcrete arrK2) (sconcrete arrA)+      jvpKrn = jvp (`conv2dPaddedS` sconcrete arrA)+                   (sconcrete arrK) (sconcrete arrK2)+      cjvpKrn = cjvp @_ @_ @_ @Concrete+                     (`conv2dPaddedS` sconcrete arrA)+                     (sconcrete arrK) (sconcrete arrK2)+  in allClose jvpInp dInp 1e-5  -- 1e-7 is too much for Float+     && allClose cjvpInp dInp 1e-5+     && allClose jvpKrn dKrn 1e-5+     && allClose cjvpKrn dKrn 1e-5++quickcheck_conv2dPaddedJvp+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dPaddedJvp =+  forAll (choose (0, 5)) $ \nImgs' ->+  forAll (choose (0, 5)) $ \nCinp' ->+  forAll (choose (0, 5)) $ \nCout' ->+  forAll (choose (0, 5)) $ \nAh' ->+  forAll (choose (0, 5)) $ \nAw' ->+  forAll (choose (0, 5)) $ \nKh1' ->+  forAll (choose (0, 5)) $ \nKw1' ->+    -- The glue below is needed to bridge the dependently-typed+    -- vs normal world.+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh1' $ \(nKh1 :: SNat nKh1) ->+    withSNat nKw1' $ \(nKw1 :: SNat nKw1) ->+      property $ \seed0 ->+        let arrK, arrK2 :: Concrete (TKS '[nCout, nCinp, nKh1 + 1, nKw1 + 1] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            (arrK2, seed3) = randomValue 0.5 seed2+            arrA, arrA2 :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed4) = randomValue 0.5 seed3+            (arrA2, _) = randomValue 0.5 seed4+        in static_conv2dPaddedJvp+             nImgs nCinp nCout nAh nAw nKh1 nKw1+             (unConcrete arrK) (unConcrete arrK2)+             (unConcrete arrA) (unConcrete arrA2)+++-- * Tests as poor man's benchmarks of handwritten vs generated gradients++static_conv2dSameVjpKrnHandwritten+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+  -> Nested.Shaped shA r+  -> Nested.Shaped shB r+  -> Bool+static_conv2dSameVjpKrnHandwritten SNat SNat SNat SNat SNat SNat SNat+                                   !_arrK arrA arrB =+  let dKrn :: Concrete (TKS shK r)+      dKrn = conv2dSame_dKrn (sconcrete arrA) (sconcrete arrB)+  in allClose dKrn dKrn 1e-5++quickcheck_conv2dSameVjpKrnHandwritten+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjpKrnHandwritten =+  forAll (choose (3, 3)) $ \nImgs' ->+  forAll (choose (3, 3)) $ \nCinp' ->+  forAll (choose (3, 3)) $ \nCout' ->+  forAll (choose (6, 6)) $ \nAh' ->+  forAll (choose (6, 6)) $ \nAw' ->+  forAll (choose (3, 3)) $ \nKh' ->+  forAll (choose (3, 3)) $ \nKw' ->+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjpKrnHandwritten+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)++static_conv2dSameVjpKrnHandwrittenVectorized+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+  -> Nested.Shaped shA r+  -> Nested.Shaped shB r+  -> Bool+static_conv2dSameVjpKrnHandwrittenVectorized SNat SNat SNat SNat SNat SNat SNat+                                             !_arrK arrA arrB =+  let dKrn :: Concrete (TKS shK r)+      dKrn = interpretAstFull emptyEnv+             $ conv2dSame_dKrn (sconcrete arrA) (sconcrete arrB)+  in allClose dKrn dKrn 1e-5++quickcheck_conv2dSameVjpKrnHandwrittenVectorized+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjpKrnHandwrittenVectorized =+  forAll (choose (3, 3)) $ \nImgs' ->+  forAll (choose (3, 3)) $ \nCinp' ->+  forAll (choose (3, 3)) $ \nCout' ->+  forAll (choose (6, 6)) $ \nAh' ->+  forAll (choose (6, 6)) $ \nAw' ->+  forAll (choose (3, 3)) $ \nKh' ->+  forAll (choose (3, 3)) $ \nKw' ->+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjpKrnHandwrittenVectorized+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)++static_conv2dSameVjpKrnSymbolic+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+  -> Nested.Shaped shA r+  -> Nested.Shaped shB r+  -> Bool+static_conv2dSameVjpKrnSymbolic SNat SNat SNat SNat SNat SNat SNat+                                arrK arrA arrB =+  let vjpKrn :: Concrete (TKS shK r)+      vjpKrn = vjp (`conv2dSameS` sconcrete arrA)+                   (sconcrete arrK) (sconcrete arrB)+  in allClose vjpKrn vjpKrn 1e-5++quickcheck_conv2dSameVjpKrnSymbolic+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjpKrnSymbolic =+  forAll (choose (3, 3)) $ \nImgs' ->+  forAll (choose (3, 3)) $ \nCinp' ->+  forAll (choose (3, 3)) $ \nCout' ->+  forAll (choose (6, 6)) $ \nAh' ->+  forAll (choose (6, 6)) $ \nAw' ->+  forAll (choose (3, 3)) $ \nKh' ->+  forAll (choose (3, 3)) $ \nKw' ->+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjpKrnSymbolic+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)++static_conv2dSameVjpKrnConcrete+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+  -> Nested.Shaped shA r+  -> Nested.Shaped shB r+  -> Bool+static_conv2dSameVjpKrnConcrete SNat SNat SNat SNat SNat SNat SNat+                                arrK arrA arrB =+  let cvjpKrn :: Concrete (TKS shK r)+      cvjpKrn = cvjp @_ @_ @_ @Concrete+                     (`conv2dSameS` sconcrete arrA)+                     (sconcrete arrK) (sconcrete arrB)+  in allClose cvjpKrn cvjpKrn 1e-5++quickcheck_conv2dSameVjpKrnConcrete+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjpKrnConcrete =+  forAll (choose (3, 3)) $ \nImgs' ->+  forAll (choose (3, 3)) $ \nCinp' ->+  forAll (choose (3, 3)) $ \nCout' ->+  forAll (choose (6, 6)) $ \nAh' ->+  forAll (choose (6, 6)) $ \nAw' ->+  forAll (choose (3, 3)) $ \nKh' ->+  forAll (choose (3, 3)) $ \nKw' ->+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjpKrnConcrete+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)++static_conv2dSameVjpInpHandwritten+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+  -> Nested.Shaped shA r+  -> Nested.Shaped shB r+  -> Bool+static_conv2dSameVjpInpHandwritten SNat SNat SNat SNat SNat SNat SNat+                                   arrK !_arrA arrB =+  let dInp :: Concrete (TKS shA r)+      dInp = conv2dSame_dInp (sconcrete arrK) (sconcrete arrB)+  in allClose dInp dInp 1e-5++quickcheck_conv2dSameVjpInpHandwritten+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjpInpHandwritten =+  forAll (choose (3, 3)) $ \nImgs' ->+  forAll (choose (3, 3)) $ \nCinp' ->+  forAll (choose (3, 3)) $ \nCout' ->+  forAll (choose (6, 6)) $ \nAh' ->+  forAll (choose (6, 6)) $ \nAw' ->+  forAll (choose (3, 3)) $ \nKh' ->+  forAll (choose (3, 3)) $ \nKw' ->+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjpInpHandwritten+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)++static_conv2dSameVjpInpHandwrittenVectorized+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+  -> Nested.Shaped shA r+  -> Nested.Shaped shB r+  -> Bool+static_conv2dSameVjpInpHandwrittenVectorized SNat SNat SNat SNat SNat SNat SNat+                                             arrK !_arrA arrB =+  let dInp :: Concrete (TKS shA r)+      dInp = interpretAstFull emptyEnv+             $ conv2dSame_dInp (sconcrete arrK) (sconcrete arrB)+  in allClose dInp dInp 1e-5++quickcheck_conv2dSameVjpInpHandwrittenVectorized+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjpInpHandwrittenVectorized =+  forAll (choose (3, 3)) $ \nImgs' ->+  forAll (choose (3, 3)) $ \nCinp' ->+  forAll (choose (3, 3)) $ \nCout' ->+  forAll (choose (6, 6)) $ \nAh' ->+  forAll (choose (6, 6)) $ \nAw' ->+  forAll (choose (3, 3)) $ \nKh' ->+  forAll (choose (3, 3)) $ \nKw' ->+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjpInpHandwrittenVectorized+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)++static_conv2dSameVjpInpSymbolic+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+  -> Nested.Shaped shA r+  -> Nested.Shaped shB r+  -> Bool+static_conv2dSameVjpInpSymbolic SNat SNat SNat SNat SNat SNat SNat+                                arrK arrA arrB =+  let vjpInp :: Concrete (TKS shA r)+      vjpInp = vjp (conv2dSameS (sconcrete arrK))+                   (sconcrete arrA) (sconcrete arrB)+  in allClose vjpInp vjpInp 1e-5++quickcheck_conv2dSameVjpInpSymbolic+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjpInpSymbolic =+  forAll (choose (3, 3)) $ \nImgs' ->+  forAll (choose (3, 3)) $ \nCinp' ->+  forAll (choose (3, 3)) $ \nCout' ->+  forAll (choose (6, 6)) $ \nAh' ->+  forAll (choose (6, 6)) $ \nAw' ->+  forAll (choose (3, 3)) $ \nKh' ->+  forAll (choose (3, 3)) $ \nKw' ->+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjpInpSymbolic+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)++static_conv2dSameVjpInpConcrete+  :: forall nImgs nCinp nCout nAh nAw nKh nKw shK shA shB r.+     ( NumScalar r, Differentiable r+     , shK ~ '[nCout, nCinp, nKh, nKw]+     , shA ~ '[nImgs, nCinp, nAh, nAw]+     , shB ~ '[nImgs, nCout, nAh, nAw] )+  => SNat nImgs -> SNat nCinp -> SNat nCout+  -> SNat nAh -> SNat nAw -> SNat nKh -> SNat nKw+  -> Nested.Shaped shK r+  -> Nested.Shaped shA r+  -> Nested.Shaped shB r+  -> Bool+static_conv2dSameVjpInpConcrete SNat SNat SNat SNat SNat SNat SNat+                                arrK arrA arrB =+  let cvjpInp :: Concrete (TKS shA r)+      cvjpInp = cvjp @_ @_ @_ @Concrete+                     (conv2dSameS (sconcrete arrK))+                     (sconcrete arrA) (sconcrete arrB)+  in allClose cvjpInp cvjpInp 1e-5++quickcheck_conv2dSameVjpInpConcrete+  :: forall r. (NumScalar r, Differentiable r)+  => Property+quickcheck_conv2dSameVjpInpConcrete =+  forAll (choose (3, 3)) $ \nImgs' ->+  forAll (choose (3, 3)) $ \nCinp' ->+  forAll (choose (3, 3)) $ \nCout' ->+  forAll (choose (6, 6)) $ \nAh' ->+  forAll (choose (6, 6)) $ \nAw' ->+  forAll (choose (3, 3)) $ \nKh' ->+  forAll (choose (3, 3)) $ \nKw' ->+    withSNat nImgs' $ \(nImgs :: SNat nImgs) ->+    withSNat nCinp' $ \(nCinp :: SNat nCinp) ->+    withSNat nCout' $ \(nCout :: SNat nCout) ->+    withSNat nAh' $ \(nAh :: SNat nAh) ->+    withSNat nAw' $ \(nAw :: SNat nAw) ->+    withSNat nKh' $ \(nKh :: SNat nKh) ->+    withSNat nKw' $ \(nKw :: SNat nKw) ->+      property $ \seed0 ->+        let arrK :: Concrete (TKS '[nCout, nCinp, nKh, nKw] r)+            (arrK, seed2) = randomValue 0.5 (mkStdGen seed0)+            arrA :: Concrete (TKS '[nImgs, nCinp, nAh, nAw] r)+            (arrA, seed3) = randomValue 0.5 seed2+            arrB :: Concrete (TKS '[nImgs, nCout, nAh, nAw] r)+            (arrB, _) = randomValue 0.5 seed3+        in static_conv2dSameVjpInpConcrete+             nImgs nCinp nCout nAh nAw nKh nKw+             (unConcrete arrK) (unConcrete arrA) (unConcrete arrB)
test/simplified/TestConvSimplified.hs view
@@ -15,1979 +15,2094 @@  import Data.Array.Nested qualified as Nested import Data.Array.Nested.Ranked.Shape--import HordeAd-import HordeAd.Core.AstEnv-import HordeAd.Core.AstFreshId (resetVarCounter)-import HordeAd.Core.AstInterpret-import HordeAd.Core.AstTools-import HordeAd.Core.CarriersAst-import HordeAd.Core.Delta-import HordeAd.Core.Ops-import HordeAd.Core.OpsAst--import CrossTesting-import EqEpsilon--testTrees :: [TestTree]-testTrees =-  [ testCase "KonstG0Rev" testKonstG0Rev-  , testCase "KonstG0Tiny1" testKonstG0Tiny1-  , testCase "KonstG0TinyS" testKonstG0TinyS-  , testCase "KonstG0TinyA" testKonstG0TinyA-  , testCase "KonstG0LittleA" testKonstG0LittleA-  , testCase "Replicate0Rev" testReplicate0Rev-  , testCase "Replicate0Tiny1" testReplicate0Tiny1-  , testCase "Replicate0TinyS" testReplicate0TinyS-  , testCase "Replicate0TinyA" testReplicate0TinyA-  , testCase "Replicate0LittleA" testReplicate0LittleA-  , testCase "Konst5LittleB" testKonst5LittleB-  , testCase "Konst5LittleC" testKonst5LittleC-  , testCase "Konst5BigB" testKonst5BigB-  , testCase "KonstNotBigB" testKonstNotBigB-  , testCase "Konst5BigC" testKonst5BigC-  , testCase "KonstNotBigC" testKonstNotBigC-  , testCase "Konst5LittleB128b" testKonst5LittleB128b-  , testCase "Konst5LittleC128b" testKonst5LittleC128b-  , testCase "Konst5BigB128b" testKonst5BigB128b-  , testCase "KonstNotBigB128b" testKonstNotBigB128b-  , testCase "Konst5BigC128b" testKonst5BigC128b-  , testCase "KonstNotBigC128b" testKonstNotBigC128b-  , testCase "Konst5LittleB128c" testKonst5LittleB128c-  , testCase "Konst5LittleC128c" testKonst5LittleC128c-  , testCase "Konst5BigB128c" testKonst5BigB128c-  , testCase "KonstNotBigB128c" testKonstNotBigB128c-  , testCase "Konst5BigC128c" testKonst5BigC128c-  , testCase "KonstNotBigC128c" testKonstNotBigC128c---  , testCase "Konst5LittleB128bc" testKonst5LittleB128bc---  , testCase "Konst5LittleC128bc" testKonst5LittleC128bc---  , testCase "Konst5BigB128bc" testKonst5BigB128bc---  , testCase "KonstNotBigB128cb" testKonstNotBigB128cb---  , testCase "Konst5BigC128cb" testKonst5BigC128cb---  , testCase "KonstNotBigC128cb" testKonstNotBigC128cb-  , testCase "Replicate0RevLaborious" testReplicate0RevLaborious-  , testCase "Replicate0Tiny1Laborious" testReplicate0Tiny1Laborious-  , testCase "Replicate0TinySLaborious" testReplicate0TinySLaborious-  , testCase "Replicate0TinyALaborious" testReplicate0TinyALaborious-  , testCase "Replicate0LittleALaborious" testReplicate0LittleALaborious-  , testCase "Konst5LittleBLaborious" testKonst5LittleBLaborious-  , testCase "Konst5LittleCLaborious" testKonst5LittleCLaborious-  , testCase "Konst5BigBLaborious" testKonst5BigBLaborious-  , testCase "KonstNotBigBLaborious" testKonstNotBigBLaborious-  , testCase "Konst5BigCLaborious" testKonst5BigCLaborious-  , testCase "KonstNotBigCLaborious" testKonstNotBigCLaborious-  , testCase "Konst5LittleBLaborious128b" testKonst5LittleBLaborious128b-  , testCase "Konst5LittleCLaborious128b" testKonst5LittleCLaborious128b---  , testCase "Konst5BigBLaborious128b" testKonst5BigBLaborious128b---  , testCase "KonstNotBigBLaborious128b" testKonstNotBigBLaborious128b---  , testCase "Konst5BigCLaborious128b" testKonst5BigCLaborious128b---  , testCase "KonstNotBigCLaborious128b" testKonstNotBigCLaborious128b-  , testCase "Konst5LittleBLaborious128c" testKonst5LittleBLaborious128c-  , testCase "Konst5LittleCLaborious128c" testKonst5LittleCLaborious128c---  , testCase "Konst5BigBLaborious128c" testKonst5BigBLaborious128c---  , testCase "KonstNotBigBLaborious128c" testKonstNotBigBLaborious128c---  , testCase "Konst5BigCLaborious128c" testKonst5BigCLaborious128c---  , testCase "KonstNotBigCLaborious128c" testKonstNotBigCLaborious128c---  , testCase "Konst5LittleBLaborious128bc" testKonst5LittleBLaborious128bc---  , testCase "Konst5LittleCLaborious128bc" testKonst5LittleCLaborious128bc---  , testCase "Konst5BigBLaborious128bc" testKonst5BigBLaborious128bc---  , testCase "KonstNotBigBLaborious128cb" testKonstNotBigBLaborious128cb---  , testCase "Konst5BigCLaborious128cb" testKonst5BigCLaborious128cb---  , testCase "KonstNotBigCLaborious128cb" testKonstNotBigCLaborious128cb---  , testCase "Replicate0RevPadded" testReplicate0RevPadded-  , testCase "Replicate0Tiny1Padded" testReplicate0Tiny1Padded-  , testCase "Replicate0TinySPadded" testReplicate0TinySPadded-  , testCase "Replicate0TinyAPadded" testReplicate0TinyAPadded-  , testCase "Replicate0LittleAPadded" testReplicate0LittleAPadded---  , testCase "Konst5LittleBPadded" testKonst5LittleBPadded---  , testCase "Konst5LittleCPadded" testKonst5LittleCPadded---  , testCase "Konst5BigBPadded" testKonst5BigBPadded---  , testCase "KonstNotBigBPadded" testKonstNotBigBPadded---  , testCase "Konst5BigCPadded" testKonst5BigCPadded---  , testCase "KonstNotBigCPadded" testKonstNotBigCPadded---  , testCase "Konst5LittleBPadded128b" testKonst5LittleBPadded128b---  , testCase "Konst5LittleCPadded128b" testKonst5LittleCPadded128b---  , testCase "Konst5BigBPadded128b" testKonst5BigBPadded128b---  , testCase "KonstNotBigBPadded128b" testKonstNotBigBPadded128b---  , testCase "Konst5BigCPadded128b" testKonst5BigCPadded128b---  , testCase "KonstNotBigCPadded128b" testKonstNotBigCPadded128b---  , testCase "Konst5LittleBPadded128c" testKonst5LittleBPadded128c---  , testCase "Konst5LittleCPadded128c" testKonst5LittleCPadded128c---  , testCase "Konst5BigBPadded128c" testKonst5BigBPadded128c---  , testCase "KonstNotBigBPadded128c" testKonstNotBigBPadded128c---  , testCase "Konst5BigCPadded128c" testKonst5BigCPadded128c---  , testCase "KonstNotBigCPadded128c" testKonstNotBigCPadded128c---  , testCase "Konst5LittleBPadded128bc" testKonst5LittleBPadded128bc---  , testCase "Konst5LittleCPadded128bc" testKonst5LittleCPadded128bc---  , testCase "Konst5BigBPadded128bc" testKonst5BigBPadded128bc---  , testCase "KonstNotBigBPadded128cb" testKonstNotBigBPadded128cb---  , testCase "Konst5BigCPadded128cb" testKonst5BigCPadded128cb---  , testCase "KonstNotBigCPadded128cb" testKonstNotBigCPadded128cb-  , testCase "disparityKonst" test_disparityKonst-  , testCase "disparityKonst2" test_disparityKonst2-  , testCase "disparitySmall" test_disparitySmall-  , testCase "ConvTomsSliceRev" testTomsSliceRev-  , testCase "ConvTomsSlice" testTomsSlice-  , testCase "ConvTomsSlicePP" testTomsSlicePP-  , testCase "minimizedCNNOPP0c" testCNNOPP0c-  , testCase "minimizedCNNOPP0b" testCNNOPP0b-  , testCase "minimizedCNNOPP1e" testCNNOPP1e-  , testCase "minimizedCNNOPP2" testCNNOPP2-  , testCase "minimizedCNNOPP2b" testCNNOPP2b---  , testCase "minimizedCNNOPP3" testCNNOPP3-  , testCase "minimizedCNNOPP3b" testCNNOPP3b-  , testCase "minimizedCNNOPP4" testCNNOPP4-  , testCase "minimizedCNNOPP4b" testCNNOPP4b-  , testCase "minimizedCNNOPP5" testCNNOPP5-  , testCase "minimizedCNNOPP5b" testCNNOPP5b-  , testCase "minimizedCNNOPP6" testCNNOPP6-  , testCase "minimizedCNNOPP6b" testCNNOPP6b-  , testCase "minimizedCNNOPP7" testCNNOPP7-  , testCase "minimizedCNNOPP7b" testCNNOPP7b---  , testCase "minimizedPaddedCNNOPP0c" testPaddedCNNOPP0c---  , testCase "minimizedPaddedCNNOPP0b" testPaddedCNNOPP0b---  , testCase "minimizedPaddedCNNOPP1e" testPaddedCNNOPP1e-  , testCase "minimizedPaddedCNNOPP1b" testPaddedCNNOPP1b-  , testCase "minimizedPaddedCNNOPPLet" testPaddedCNNOPPLet-  , testCase "minimizedPaddedCNNOPPLet2" testPaddedCNNOPPLet2---  , testCase "minimizedPaddedCNNOPP2" testPaddedCNNOPP2-  , testCase "minimizedCNNOPP0cW" testCNNOPP0cW-  , testCase "minimizedCNNOPP0bW" testCNNOPP0bW-  , testCase "minimizedCNNOPP1bW" testCNNOPP1bW-  , testCase "minimizedCNNOPP4bW" testCNNOPP4bW-  , testCase "minimizedCNNOPP4bD" testCNNOPP4bD-  , testCase "minimizedCNNOPP5aW" testCNNOPP5aW-  , testCase "minimizedCNNOPP5bW" testCNNOPP5bW-  , testCase "minimizedCNNOPP5cW" testCNNOPP5cW-  , testCase "minimizedCNNOPP5dW" testCNNOPP5dW-  ]---- The examples reproduced and transformed in this file are borrowed--- from https://github.com/benl23x5/adops.--- Here they are defined using ranked tensors.---- * A non-laborious version (depends on indexing OOB giving 0 consistently)--conv2d1-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2d1 = conv2dUnpadded $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]--conv2dA-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dA = conv2dUnpadded $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]--conv2dB-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dB = conv2dUnpadded (rconcrete $ unConcrete t16b)--testKonstG0Rev :: Assertion-testKonstG0Rev =-  assertEqualUpToEpsilon 1e-4-    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dB) (rrepl [2, 2, 2, 2] 0))--testKonstG0Tiny1 :: Assertion-testKonstG0Tiny1 =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 1, 1, 1] [-0.2])-    (rev' @Double @4 conv2d1 (rrepl [1, 1, 1, 1] 0))--testKonstG0TinyS :: Assertion-testKonstG0TinyS =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 1, 1, 1] [582665.99432])-    (rev' @Double @4-          (conv2dUnpadded $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))-          (ringestData [1, 1, 1, 1] [0]))--testKonstG0TinyA :: Assertion-testKonstG0TinyA =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 2, 1, 1] [-0.2,25.0003])-    (rev' @Double @4 conv2dA (rrepl [1, 2, 1, 1] 0))--testKonstG0LittleA :: Assertion-testKonstG0LittleA =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [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])-    (rev' @Double @4 conv2dA (rrepl [2, 2, 2, 2] 0))--conv2dC-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dC = flip conv2dUnpadded (rconcrete $ unConcrete t16b)--conv2dB128b-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dB128b = conv2dUnpadded (rconcrete $ unConcrete t128b)--conv2dC128b-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dC128b = flip conv2dUnpadded (rconcrete $ unConcrete t128b)--conv2dB128c-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dB128c = conv2dUnpadded (rconcrete $ unConcrete t128c)--conv2dC128c-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dC128c = flip conv2dUnpadded (rconcrete $ unConcrete t128c)--testReplicate0Rev :: Assertion-testReplicate0Rev =-  assertEqualUpToEpsilon 1e-4-    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dB) (rrepl [2, 2, 2, 2] 0))--testReplicate0Tiny1 :: Assertion-testReplicate0Tiny1 =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 1, 1, 1] [-0.2])-    (rev' @Double @4 conv2d1 (rrepl [1, 1, 1, 1] 0))--testReplicate0TinyS :: Assertion-testReplicate0TinyS =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 1, 1, 1] [582665.99432])-    (rev' @Double @4-          (conv2dUnpadded $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))-          (ringestData [1, 1, 1, 1] [0]))--testReplicate0TinyA :: Assertion-testReplicate0TinyA =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 2, 1, 1] [-0.2,25.0003])-    (rev' @Double @4 conv2dA (rrepl [1, 2, 1, 1] 0))--testReplicate0LittleA :: Assertion-testReplicate0LittleA =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [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])-    (rev' @Double @4 conv2dA (rrepl [2, 2, 2, 2] 0))---- with data t16--testKonst5LittleB :: Assertion-testKonst5LittleB =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])-    (rev' @Double @4 conv2dB (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5LittleC :: Assertion-testKonst5LittleC =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])-    (rev' @Double @4 conv2dC (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5BigB :: Assertion-testKonst5BigB =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])-    (rev' @Double @4 conv2dB (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigB :: Assertion-testKonstNotBigB =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])-    (rev' @Double @4 conv2dB-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--testKonst5BigC :: Assertion-testKonst5BigC =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])-    (rev' @Double @4 conv2dC (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigC :: Assertion-testKonstNotBigC =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])-    (rev' @Double @4 conv2dC-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))---- with data t128b--testKonst5LittleB128b :: Assertion-testKonst5LittleB128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004])-    (rev' @Double @4 conv2dB128b (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5LittleC128b :: Assertion-testKonst5LittleC128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987])-    (rev' @Double @4 conv2dC128b (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5BigB128b :: Assertion-testKonst5BigB128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])-    (rev' @Double @4 conv2dB128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigB128b :: Assertion-testKonstNotBigB128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])-    (rev' @Double @4 conv2dB128b-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--testKonst5BigC128b :: Assertion-testKonst5BigC128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])-    (rev' @Double @4 conv2dC128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigC128b :: Assertion-testKonstNotBigC128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])-    (rev' @Double @4 conv2dC128b-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))---- with data t128c--testKonst5LittleB128c :: Assertion-testKonst5LittleB128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4])-    (rev' @Double @4 conv2dB128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5LittleC128c :: Assertion-testKonst5LittleC128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992])-    (rev' @Double @4 conv2dC128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5BigB128c :: Assertion-testKonst5BigB128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])-    (rev' @Double @4 conv2dB128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigB128c :: Assertion-testKonstNotBigB128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])-    (rev' @Double @4 conv2dB128c-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--testKonst5BigC128c :: Assertion-testKonst5BigC128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])-    (rev' @Double @4 conv2dC128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigC128c :: Assertion-testKonstNotBigC128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])-    (rev' @Double @4 conv2dC128c-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))---- with data t128b and t128c-{--testKonst5LittleB128bc :: Assertion-testKonst5LittleB128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])-    (rev' @Double @4 conv2dB128b t128c)--testKonst5LittleC128bc :: Assertion-testKonst5LittleC128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])-    (rev' @Double @4 conv2dC128b t128c)--testKonst5BigB128bc :: Assertion-testKonst5BigB128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])-    (rev' @Double @4 conv2dB128b t128c)---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigB128cb :: Assertion-testKonstNotBigB128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002])-    (rev' @Double @4 conv2dB128c t128b)--testKonst5BigC128cb :: Assertion-testKonst5BigC128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])-    (rev' @Double @4 conv2dC128c t128b)---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigC128cb :: Assertion-testKonstNotBigC128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])-    (rev' @Double @4 conv2dC128c t128b)--}----- * A laborious version (meaning, out of bounds indexing is handled explicitly)---- | Unpadded full convolution,---   where the output size is the same as the input size.------ It guards the out of bounds indexing behind a conditional--- to prevent changed values after vectorization,--- but the guarding is no longer needed, so this is only for testing.------ BTW, the indexing lower bounds in the code are spurious,--- so they get simplified away in the resulting AST program.-conv2dUnpaddedL-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpaddedL arrK arrA =-  let [nImgs, nCinpA, nAh, nAw] = rshape arrA-      [nCoutK, nCinpK, nKh, nKw] = rshape arrK-      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA-      shB = [nImgs, nCoutK, nAh, nAw]-      shK1 = [1, nCinp, nKh, nKw]-  in rbuild shB $ \case-    [iImg, iCout, iBh, iBw] ->-      let arrAt = slicezL shK1 arrA [iImg, 0, iBh, iBw]-          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]-      in rdot0 arrAt arrKt-    _ -> error "conv2dUnpaddedL: impossible pattern needlessly required"---- | Slice a section out of a tensor,---   given a base offset and shape of the section.------   If the slice extends out side the source array then the corresponding---   elements are set to zero.-slicezL-  :: (ADReady target, GoodScalar r, KnownNat n)-  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)-slicezL shOut d ixBase =-  rbuild shOut $ \ixResult -> indexz0L d (ixrZipWith (+) ixBase ixResult)---- | Retrieve the element at the given index,---   returning zero for out of range indices.------ Warning: this uses ix twice and within0 again uses it twice,--- so this variant without tlet should be used only when it's known--- that ix is of small constant size (e.g., if it contains conditionals--- that compare big tensors or their minimal elements, it likely is not,--- unless the tensors are under tlet and only variables representing them--- are used).-indexz0L-  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)-  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)-indexz0L d ix = ifH (within0 @target (rshape @target d) ix) (d ! ix) (rscalar 0)---- | Given an index and shape, check if the index is fully within the shape.--- Note that @ix@ is used twice, so should be shared outside.-within0-  :: forall target n. (ADReady target, KnownNat n)-  => IShR n -> IxROf target n -> BoolOf target-within0 sh ix =-  let within :: IntOf target -> IntOf target -> BoolOf target-      within i dim = 0 <=. i &&* dim >. i-  in foldr (&&*) true-     $ zipWith within (toList ix) (map fromIntegral $ toList sh)--conv2d1Laborious-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2d1Laborious = conv2dUnpaddedL $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]--conv2dALaborious-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dALaborious =-  conv2dUnpaddedL $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]--conv2dBLaborious-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dBLaborious = conv2dUnpaddedL (rconcrete $ unConcrete t16b)--conv2dCLaborious-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dCLaborious = flip conv2dUnpaddedL (rconcrete $ unConcrete t16b)--conv2dBLaborious128b-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dBLaborious128b = conv2dUnpaddedL (rconcrete $ unConcrete t128b)--conv2dCLaborious128b-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dCLaborious128b = flip conv2dUnpaddedL (rconcrete $ unConcrete t128b)--conv2dBLaborious128c-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dBLaborious128c = conv2dUnpaddedL (rconcrete $ unConcrete t128c)--conv2dCLaborious128c-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dCLaborious128c = flip conv2dUnpaddedL (rconcrete $ unConcrete t128c)--testReplicate0RevLaborious :: Assertion-testReplicate0RevLaborious =-  assertEqualUpToEpsilon 1e-4-    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBLaborious) (rrepl [2, 2, 2, 2] 0))--testReplicate0Tiny1Laborious :: Assertion-testReplicate0Tiny1Laborious =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 1, 1, 1] [-0.2])-    (rev' @Double @4 conv2d1Laborious (rrepl [1, 1, 1, 1] 0))--testReplicate0TinySLaborious :: Assertion-testReplicate0TinySLaborious =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 1, 1, 1] [582665.99432])-    (rev' @Double @4-          (conv2dUnpaddedL $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))-          (ringestData [1, 1, 1, 1] [0]))--testReplicate0TinyALaborious :: Assertion-testReplicate0TinyALaborious =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [1, 2, 1, 1] [-0.2,25.0003])-    (rev' @Double @4 conv2dALaborious (rrepl [1, 2, 1, 1] 0))--testReplicate0LittleALaborious :: Assertion-testReplicate0LittleALaborious =-  assertEqualUpToEpsilon' 1e-10-    (ringestData [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])-    (rev' @Double @4 conv2dALaborious (rrepl [2, 2, 2, 2] 0))---- with data t16--testKonst5LittleBLaborious :: Assertion-testKonst5LittleBLaborious =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])-    (rev' @Double @4 conv2dBLaborious (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5LittleCLaborious :: Assertion-testKonst5LittleCLaborious =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])-    (rev' @Double @4 conv2dCLaborious (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5BigBLaborious :: Assertion-testKonst5BigBLaborious =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])-    (rev' @Double @4 conv2dBLaborious (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigBLaborious :: Assertion-testKonstNotBigBLaborious =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])-    (rev' @Double @4 conv2dBLaborious-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--testKonst5BigCLaborious :: Assertion-testKonst5BigCLaborious =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])-    (rev' @Double @4 conv2dCLaborious (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigCLaborious :: Assertion-testKonstNotBigCLaborious =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])-    (rev' @Double @4 conv2dCLaborious-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))---- with data t128b--testKonst5LittleBLaborious128b :: Assertion-testKonst5LittleBLaborious128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004])-    (rev' @Double @4 conv2dBLaborious128b (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5LittleCLaborious128b :: Assertion-testKonst5LittleCLaborious128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987])-    (rev' @Double @4 conv2dCLaborious128b (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--{--testKonst5BigBLaborious128b :: Assertion-testKonst5BigBLaborious128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])-    (rev' @Double @4 conv2dBLaborious128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigBLaborious128b :: Assertion-testKonstNotBigBLaborious128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])-    (rev' @Double @4 conv2dBLaborious128b-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--testKonst5BigCLaborious128b :: Assertion-testKonst5BigCLaborious128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])-    (rev' @Double @4 conv2dCLaborious128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigCLaborious128b :: Assertion-testKonstNotBigCLaborious128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])-    (rev' @Double @4 conv2dCLaborious128b-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--}---- with data t128c--testKonst5LittleBLaborious128c :: Assertion-testKonst5LittleBLaborious128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4])-    (rev' @Double @4 conv2dBLaborious128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5LittleCLaborious128c :: Assertion-testKonst5LittleCLaborious128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992])-    (rev' @Double @4 conv2dCLaborious128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--{--testKonst5BigBLaborious128c :: Assertion-testKonst5BigBLaborious128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])-    (rev' @Double @4 conv2dBLaborious128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigBLaborious128c :: Assertion-testKonstNotBigBLaborious128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])-    (rev' @Double @4 conv2dBLaborious128c-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--testKonst5BigCLaborious128c :: Assertion-testKonst5BigCLaborious128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])-    (rev' @Double @4 conv2dCLaborious128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigCLaborious128c :: Assertion-testKonstNotBigCLaborious128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])-    (rev' @Double @4 conv2dCLaborious128c-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--}---- with data t128b and t128c-{--testKonst5LittleBLaborious128bc :: Assertion-testKonst5LittleBLaborious128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])-    (rev' @Double @4 conv2dBLaborious128b t128c)--testKonst5LittleCLaborious128bc :: Assertion-testKonst5LittleCLaborious128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])-    (rev' @Double @4 conv2dCLaborious128b t128c)--testKonst5BigBLaborious128bc :: Assertion-testKonst5BigBLaborious128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])-    (rev' @Double @4 conv2dBLaborious128b t128c)---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigBLaborious128cb :: Assertion-testKonstNotBigBLaborious128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002])-    (rev' @Double @4 conv2dBLaborious128c t128b)--testKonst5BigCLaborious128cb :: Assertion-testKonst5BigCLaborious128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])-    (rev' @Double @4 conv2dCLaborious128c t128b)---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigCLaborious128cb :: Assertion-testKonstNotBigCLaborious128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])-    (rev' @Double @4 conv2dCLaborious128c t128b)--}----- * A padded version (out of bounds indexing is not possible)---- | Full convolution with just enough extra external zero padding---   to ensure that the output size is the same as the input size---   and all input points are read the same number of times.------   The same result could be accomplished by tweaking indexes slightly---   in conv2dUnpadded, but here additionally all bounds checks in the code---   are spurious and will be simplified away in the resulting AST program.-conv2dPadded-  :: forall target r. (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)-conv2dPadded arrK arrA =-  let [nImgs, nCinpA, nAh, nAw] = rshape arrA-      [nCoutK, nCinpK, nKh, nKw] = rshape arrK-      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]-      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case-        [iImg, iCinp, iPh, iPw] ->-          ifH (iPh <. fromIntegral (nKh `div` 2)-               ||* iPw <. fromIntegral (nKw `div` 2)-               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)-               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))-              (rscalar 0)-              (arrA ! [ iImg-                      , iCinp-                      , iPh - fromIntegral (nKh `div` 2)-                      , iPw - fromIntegral (nKw `div` 2) ])-      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA-      shB = [nImgs, nCoutK, nAh, nAw]-      shK1 = [1, nCinp, nKh, nKw]-  in rbuild shB $ \case-    [iImg, iCout, iBh, iBw] ->-      let arrAt = slicezL shK1 arrAPadded [iImg, 0, iBh, iBw]-          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]-      in rdot0 arrAt arrKt-    _ -> error "conv2dPadded: impossible pattern needlessly required"--conv2d1Padded-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2d1Padded = conv2dPadded $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]--conv2dAPadded-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dAPadded =-  conv2dPadded $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]--conv2dBPadded-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dBPadded = conv2dPadded (rconcrete $ unConcrete t16b)--conv2dCPadded-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dCPadded = flip conv2dPadded (rconcrete $ unConcrete t16b)--conv2dBPadded128b-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dBPadded128b = conv2dPadded (rconcrete $ unConcrete t128b)--conv2dCPadded128b-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dCPadded128b = flip conv2dPadded (rconcrete $ unConcrete t128b)--_conv2dBPadded128c-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-_conv2dBPadded128c = conv2dPadded (rconcrete $ unConcrete t128c)--_conv2dCPadded128c-  :: (ADReady target, GoodScalar r, Differentiable r)-  => target (TKR 4 r) -> target (TKR 4 r)-_conv2dCPadded128c = flip conv2dPadded (rconcrete $ unConcrete t128c)---- TODO: OOMs-_testReplicate0RevPadded :: Assertion-_testReplicate0RevPadded =-  assertEqualUpToEpsilon 1e-4-    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded) (rrepl [2, 2, 2, 2] 0))--testReplicate0Tiny1Padded :: Assertion-testReplicate0Tiny1Padded =-  assertEqualUpToEpsilon 1e-10-    (ringestData [1, 1, 1, 1] [-0.2])-    (cgrad (kfromR . rsum0 @4 @(TKScalar Double) . conv2d1Padded) (rrepl [1, 1, 1, 1] 0))--testReplicate0TinySPadded :: Assertion-testReplicate0TinySPadded =-  assertEqualUpToEpsilon 1e-10-    (ringestData [1, 1, 1, 1] [582665.99432])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) .-          (conv2dPadded $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b))))-          (ringestData [1, 1, 1, 1] [0]))--testReplicate0TinyAPadded :: Assertion-testReplicate0TinyAPadded =-  assertEqualUpToEpsilon 1e-10-    (ringestData [1, 2, 1, 1] [-0.2,25.0003])-    (cgrad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dAPadded) (rrepl [1, 2, 1, 1] 0))--testReplicate0LittleAPadded :: Assertion-testReplicate0LittleAPadded =-  assertEqualUpToEpsilon 1e-10-    (ringestData [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])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dAPadded) (rrepl [2, 2, 2, 2] 0))---- with data t16---- TODO: OOMs-_testKonst5LittleBPadded :: Assertion-_testKonst5LittleBPadded =-  assertEqualUpToEpsilon 1e-8-    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded) (rreplicate0N [2, 2, 2, 2] (rscalar 5)))---- TODO: OOMs-_testKonst5LittleCPadded :: Assertion-_testKonst5LittleCPadded =-  assertEqualUpToEpsilon 1e-8-    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dCPadded) (rreplicate0N [2, 2, 2, 2] (rscalar 5)))---- TODO: OOMs-_testKonst5BigBPadded :: Assertion-_testKonst5BigBPadded =-  assertEqualUpToEpsilon 1e-8-    (ringestData [3, 2, 4, 2] [40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded) (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- TODO: OOMs--- The gradient is the same as above, because one argument is the same--- and convolution is linear.-_testKonstNotBigBPadded :: Assertion-_testKonstNotBigBPadded =-  assertEqualUpToEpsilon 1e-8-    (ringestData [3, 2, 4, 2] [40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded)-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))---- TODO: OOMs-_testKonst5BigCPadded :: Assertion-_testKonst5BigCPadded =-  assertEqualUpToEpsilon 1e-8-    (ringestData [3, 2, 4, 2] [0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dCPadded) (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- TODO: OOMs--- The gradient is the same as above, because one argument is the same--- and convolution is linear.-_testKonstNotBigCPadded :: Assertion-_testKonstNotBigCPadded =-  assertEqualUpToEpsilon 1e-8-    (ringestData [3, 2, 4, 2] [0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dCPadded)-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))---- with data t128b---- TODO: OOMs-_testKonst5LittleBPadded128b :: Assertion-_testKonst5LittleBPadded128b =-  assertEqualUpToEpsilon 1e-8-    (ringestData [2, 2, 2, 2] [578.1829600001,558.1716000002,608.0772800002001,577.7659200003001,729.1778800002002,701.1835600003001,833.9722000003002,803.9778800004001,578.1829600001,558.1716000002,608.0772800002001,577.7659200003001,729.1778800002002,701.1835600003001,833.9722000003002,803.9778800004001])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded128b) (rreplicate0N [2, 2, 2, 2] (rscalar 5)))---- TODO: OOMs-_testKonst5LittleCPadded128b :: Assertion-_testKonst5LittleCPadded128b =-  assertEqualUpToEpsilon 1e-8-    (ringestData [2, 2, 2, 2] [1113.1722000001,1412.1551500001997,1234.1494800003002,1627.8210700004993,1500.2781800001994,1870.0614400004986,2156.3671200003987,2725.0393200008984,1113.1722000001,1412.1551500001997,1234.1494800003002,1627.8210700004993,1500.2781800001994,1870.0614400004986,2156.3671200003987,2725.0393200008984])-    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dCPadded128b) (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--{--testKonst5BigBPadded128b :: Assertion-testKonst5BigBPadded128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])-    (rev' @Double @4 conv2dBPadded128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigBPadded128b :: Assertion-testKonstNotBigBPadded128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])-    (rev' @Double @4 conv2dBPadded128b-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--testKonst5BigCPadded128b :: Assertion-testKonst5BigCPadded128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])-    (rev' @Double @4 conv2dCPadded128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigCPadded128b :: Assertion-testKonstNotBigCPadded128b =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])-    (rev' @Double @4 conv2dCPadded128b-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))---- with data t128c--testKonst5LittleBPadded128c :: Assertion-testKonst5LittleBPadded128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [186.7886400001,121.7829600001,269.09432000009997,261.3943200001,210.9943200001,231.79432000010002,160.00030000000004,194.00060000000005,186.7886400001,121.7829600001,269.09432000009997,261.3943200001,210.9943200001,231.79432000010002,160.00030000000004,194.00060000000005])-    (rev' @Double @4 conv2dBPadded128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5LittleCPadded128c :: Assertion-testKonst5LittleCPadded128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2, 2, 2, 2] [1772.649480000399,2138.4267600005987,2157.0438000004983,2640.8154000007976,961.7781800001002,1359.4557500003987,1233.4728000001987,1712.044990000598,1772.649480000399,2138.4267600005987,2157.0438000004983,2640.8154000007976,961.7781800001002,1359.4557500003987,1233.4728000001987,1712.044990000598])-    (rev' @Double @4 conv2dCPadded128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))--testKonst5BigBPadded128c :: Assertion-testKonst5BigBPadded128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])-    (rev' @Double @4 conv2dBPadded128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigBPadded128c :: Assertion-testKonstNotBigBPadded128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])-    (rev' @Double @4 conv2dBPadded128c-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--testKonst5BigCPadded128c :: Assertion-testKonst5BigCPadded128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])-    (rev' @Double @4 conv2dCPadded128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigCPadded128c :: Assertion-testKonstNotBigCPadded128c =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])-    (rev' @Double @4 conv2dCPadded128c-          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))--}---- with data t128b and t128c-{--testKonst5LittleBPadded128bc :: Assertion-testKonst5LittleBPadded128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001,1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001])-    (rev' @Double @4 conv2dBPadded128b t128c)--testKonst5LittleCPadded128bc :: Assertion-testKonst5LittleCPadded128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [0.0,0.0,0.0,0.0,251.5006,417.7949200000001,494.8949100000001,382.59461000000005,482.69492000000014,778.9778800001002,952.0721900001002,742.5775700001002,610.5892400000002,1113.1722000001,1412.1551500001997,1182.6605300002002,580.6778800001001,1234.1494800003002,1627.8210700004993,1571.2321300004994,329.17728000010004,816.3545600003002,1132.9261600005002,1188.6375200005,97.98296000010004,455.17160000020016,675.7488800003999,828.6545600004001,-29.9113599999,120.97728000019995,215.6659200003,388.5716000003,0.0,0.0,0.0,0.0,65.90000000000003,106.90000000000003,173.90000000000006,170.90000000000003,365.89432000010004,593.1946200001,821.2892400002003,657.1892400002001,1060.8778800002,1500.2781800001994,1870.0614400004986,1202.8611400005,1465.6665200003995,2156.3671200003987,2725.0393200008984,1831.7390200008983,1399.7665200003996,2049.4671200003986,2551.139320000898,1660.8390200008987,1099.7722000003,1563.1725000002994,1903.750080000699,1174.5497800006997,404.7886400002001,656.0889400002,854.9778800004001,628.8778800004001,0.0,0.0,0.0,0.0,251.5006,417.7949200000001,494.8949100000001,382.59461000000005,482.69492000000014,778.9778800001002,952.0721900001002,742.5775700001002,610.5892400000002,1113.1722000001,1412.1551500001997,1182.6605300002002,580.6778800001001,1234.1494800003002,1627.8210700004993,1571.2321300004994,329.17728000010004,816.3545600003002,1132.9261600005002,1188.6375200005,97.98296000010004,455.17160000020016,675.7488800003999,828.6545600004001,-29.9113599999,120.97728000019995,215.6659200003,388.5716000003,0.0,0.0,0.0,0.0,65.90000000000003,106.90000000000003,173.90000000000006,170.90000000000003,365.89432000010004,593.1946200001,821.2892400002003,657.1892400002001,1060.8778800002,1500.2781800001994,1870.0614400004986,1202.8611400005,1465.6665200003995,2156.3671200003987,2725.0393200008984,1831.7390200008983,1399.7665200003996,2049.4671200003986,2551.139320000898,1660.8390200008987,1099.7722000003,1563.1725000002994,1903.750080000699,1174.5497800006997,404.7886400002001,656.0889400002,854.9778800004001,628.8778800004001])-    (rev' @Double @4 conv2dCPadded128b t128c)--testKonst5BigBPadded128bc :: Assertion-testKonst5BigBPadded128bc =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [2,2,8,4] [1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001,1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001])-    (rev' @Double @4 conv2dBPadded128b t128c)---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigBPadded128cb :: Assertion-testKonstNotBigBPadded128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002])-    (rev' @Double @4 conv2dBPadded128c t128b)--testKonst5BigCPadded128cb :: Assertion-testKonst5BigCPadded128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992])-    (rev' @Double @4 conv2dCPadded128c t128b)---- The gradient is the same as above, because one argument is the same--- and convolution is linear.-testKonstNotBigCPadded128cb :: Assertion-testKonstNotBigCPadded128cb =-  assertEqualUpToEpsilon' 1e-8-    (ringestData [4,2,4,4] [720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992])-    (rev' @Double @4 conv2dCPadded128c t128b)--}----- * Disparity and misc---- | Disparity cost volume.------   Take two arrays of multi channel 2d images, where the first contains---   left views of the scene and the second contains right views.------   For each pair of images, slice the right image over the left image,---   and for each offset produce the L1 distance indicating how well---   correponding---   multi-channel image elements in the right image match those in the left.------   Described in:---    Anytime Stereo Image Depth Estimation on Mobile Devices---    Wang, Lai et al, ICRA 2019---    https://arxiv.org/abs/1810.11408---    Section III b).----costVolume-  :: forall r target. (ADReady target, GoodScalar r)-  => Int -> Int -> target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)-costVolume iStart nCount arrL arrR =-  let [nImgs, nChas, nRows, nCols] = rshape arrL-      shO = [nImgs, nCount, nRows, nCols]-  in rbuild shO $ \[iImg, iDisp, iRow, iCol] ->-       let arrVecL = rbuild (nChas :$: ZSR) $ \[iCha] ->-                       rindex0 arrL [iImg, iCha, iRow, iCol]-           iSrc = iCol - fromIntegral iStart - iDisp-           arrVecR = rbuild [nChas] $ \[iCha] ->-                       rindex0 arrR [iImg, iCha, iRow, iSrc]-       in rsum0 $ rzipWith1 (\xL xR -> abs (xL - xR)) arrVecL arrVecR--test_disparityKonst :: Assertion-test_disparityKonst = do-  let arrL :: ADReady target => target (TKR 4 Double)-      arrL = rreplicate0N [1, 2, 4, 6] (rscalar (-0.2))-      arrR :: ADReady target => target (TKR 4 Double)-      arrR = rreplicate0N [1, 2, 4, 6] (rscalar 0.3)-      arrO = costVolume @Double 0 4 arrL arrR-      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO-      arrDR = vjp (\aR -> costVolume 0 4 (rfromPrimal arrL) aR) arrR arrO-  assertEqualUpToEpsilon 1e-7-    (rconcrete $ Nested.rfromListPrimLinear [1,4,4,6] [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0])-    arrO-  assertEqualUpToEpsilon 1e-7-    (rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0])-    arrDL-  assertEqualUpToEpsilon 1e-7-    (rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0])-   arrDR-  assertEqualUpToEpsilon' 1e-7-    (ringestData [1,2,4,6] [4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0])-    (rev' @Double @4 (costVolume 0 4 arrL) arrR)-  assertEqualUpToEpsilon' 1e-7-    (ringestData [1,2,4,6] [-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0])-    (rev' @Double @4 (\aL -> costVolume 0 2 aL arrR) arrL)-  assertEqualUpToEpsilon' 1e-7-    (ringestData [1,2,4,6] [2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0])-    (rev' @Double @4 (costVolume 0 2 arrL) arrR)--test_disparityKonst2 :: Assertion-test_disparityKonst2 = do-  let arrL :: (BaseTensor target, GoodScalar r, Differentiable r) => target (TKR 4 r)-      arrL = ringestData [1, 2, 4, 6] [0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0, 1.7041241452319316,1.21999,0.21355339059327375,0.7867666666666666,0.7331698975466578,0.6964466094067263,1.1,1.1041141452319316,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,1.986339795093316,0.6964466094067263]-      arrR :: (BaseTensor target, GoodScalar r, Differentiable r) => target (TKR 4 r)-      arrR = ringestData [1, 2, 4, 6] [0.2, 0.5, -0.2, 0.0001, 0.44, 0.9, -0.9, 0.00001, -0.22, -0.28, -0.34, -0.40, -0.40,-0.22,-0.28,-0.34, 0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5, -0.35355339059327373,0.16666666666666666,0.17677669529663687,-0.25, -2.808238290463863,-1.21999,-0.5672067811865474,-0.7867666666666666,-1.986339795093316,-0.6964466094067263,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,0.6964466094067263,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,0.50001,0.42000000000000004,0.2801,0.78,1.1,0.50001,0.42000000000000004,0.2801,0.78]-      arrO = rreplicate0N [1, 4, 4, 6] (rscalar (1 :: Double))-      res1 = rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [4.0,2.0,2.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,2.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,2.0,0.0,0.0,-2.0,0.0,4.0,4.0,2.0,0.0,-4.0,1.0,4.0,4.0,4.0,-4.0,2.0,4.0,2.0]-      res2 = rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [-4.0,0.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,-4.0,-2.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,4.0,4.0,-4.0,1.0,-2.0,-1.0,-2.0,3.0,2.0,-1.0,-2.0,-1.0,-2.0,0.0,-2.0,-3.0,-2.0,1.0]-      arrDL :: Concrete (TKR 4 Double)-      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO-      arrDR :: Concrete (TKR 4 Double)-      arrDR = vjp (costVolume 0 4 (rfromPrimal arrL)) arrR arrO-  assertEqualUpToEpsilon 1e-7-    res1-    arrDL-  assertEqualUpToEpsilon 1e-7-    res2-    arrDR-  assertEqualUpToEpsilon' 1e-7-    res1-    (rev' @Double @4 (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL)-  assertEqualUpToEpsilon' 1e-7-    res2-    (rev' @Double @4 (costVolume 0 4 (rfromPrimal arrL)) arrR)--test_disparitySmall :: Assertion-test_disparitySmall = do-  let arrL :: ADReady target => target (TKR 4 Double)-      arrL = ringestData [1, 2, 3, 2] [0.2, 0.5, -0.2, 0.0001, 0.44, 0.9, -0.9, 0.00001, -0.22, -0.28, -0.34, -0.40]-      arrR :: ADReady target => target (TKR 4 Double)-      arrR = ringestData [1, 2, 3, 2] [-0.40,-0.22,-0.28,-0.34, 0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5, -0.35355339059327373,0.16666666666666666,0.17677669529663687,-0.25]-      arrO = costVolume @Double 0 4 arrL arrR-      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO-      arrDR = vjp (\aR -> costVolume 0 4 (rfromPrimal arrL) aR) arrR arrO-  assertEqualUpToEpsilon 1e-7-    (rconcrete $ Nested.rfromListPrimLinear [1,4,3,2] [1.7041241452319316,1.21999,0.21355339059327375,0.7867666666666666,0.7331698975466578,0.6964466094067263,1.1,1.1041141452319316,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3])-    arrO-  assertEqualUpToEpsilon' 1e-7-    (ringestData [1,2,3,2] [-2.0,-1.0,-2.0,-1.0,-2.0,-1.0,2.0,1.0,-2.0,1.0,2.0,1.0])-    (rev' @Double @4 (costVolume 0 4 arrL) arrR)-  assertEqualUpToEpsilon 1e-7-    (rconcrete $ Nested.rfromListPrimLinear [1,2,3,2] [5.004124145231932,3.3241241452319317,-1.0464466094067264,1.7006200572599404,3.0731698975466575,4.5496165069533845,-5.004124145231932,-1.3240841452319316,-1.0464466094067264,-0.9933132760733929,-3.0731698975466575,-4.5496165069533845])-    arrDL-  assertEqualUpToEpsilon 1e-7-    (rconcrete $ Nested.rfromListPrimLinear [1,2,3,2] [-2.808238290463863,-1.21999,-0.5672067811865474,-0.7867666666666666,-1.986339795093316,-0.6964466094067263,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,1.986339795093316,0.6964466094067263])-   arrDR-  assertEqualUpToEpsilon' 1e-7-    (ringestData [1,2,3,2] [-1.0,0.0,-1.0,0.0,-1.0,0.0,1.0,0.0,-1.0,0.0,1.0,0.0])-    (rev' @Double @4 (costVolume 1 4 arrL) arrR)-  assertEqualUpToEpsilon' 1e-7-    (ringestData [1,2,3,2] [2.0,2.0,-2.0,2.0,2.0,2.0,-2.0,2.0,-2.0,-2.0,-2.0,-2.0])-    (rev' @Double @4 (\aL -> costVolume 2 2 aL arrR) arrL)-  assertEqualUpToEpsilon' 1e-7-    (ringestData [1,2,3,2] [-1.0,0.0,-1.0,0.0,-1.0,0.0,1.0,0.0,-1.0,0.0,1.0,0.0])-    (rev' @Double @4 (costVolume 1 2 arrL) arrR)--codeTomsSlice :: ADReady target-              => target (TKR 2 Double) -> target (TKR 0 Double)-codeTomsSlice a =-  let (n, m) = case rshape a of-        [n', m'] -> (n', m')-        _ -> error "codeTomsSlice"-      a1 = rbuild @2 @0 [n,m-1] (\[i',j'] -> rindex0 a [i',j'])-      a2 = rbuild [n,m-1] (\[i',j'] -> rindex0 a [i',j' + 1])-  in rsum0 @2 $ rbuild [n,m] $ \[i, _j] ->-       rfromIndex0 i * rsum0 (a1 * a2)--testTomsSliceRev :: Assertion-testTomsSliceRev = do-  assertEqualUpToEpsilon 1e-5-    (ringestData [32,4] [63686.39999999999,137292.80000000002,121222.4,79558.40000000002,192646.40000000005,223971.0617601984,228556.80000000005,116846.33088019838,63686.39999999999,137292.80000000002,127174.4,79558.40000000002,192646.40000000005,158499.06176019844,202566.40000000005,51374.330880198424,11904.0,5952.0,7936.0,1984.0,116846.33088019838,385292.8000000001,227740.66176039676,192646.40000000005,116846.33088019838,228556.80000000005,174580.73088019836,35910.399999999994,79558.40000000002,127372.79999999997,143244.80000000002,63686.39999999999,105152.0,186683.13088000007,105151.98016,107124.73088000003,-396.79999999999995,26188.8,17459.2,25990.399999999998,-7936.0,73408.0,-1995.2691200000017,57536.0,51584.0,-660672.0,55552.0,3968.0,3968.0,3571.2,3571.2,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,49203.79519999998,49203.79519999998,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,129158.9952,65472.59519999998,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,1984.0000000000146,67059.20000000001,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,-21823.99999999993,108921.6,16070.400000000005,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,771974.4,218019.0617601984,192646.40000000005,170414.3308801984,385292.8000000001,340828.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,236294.40000000005,271587.0617601984,192646.40000000005,45422.33088019842,385292.8000000001,162268.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,369222.4,220003.0617601984,192646.40000000005,104942.33088019838,385292.8000000001,215836.66176039676,192646.40000000005])-    (grad (kfromR . codeTomsSlice) (rreshape [32, 4] t128))--testTomsSlice :: Assertion-testTomsSlice = do-  assertEqualUpToEpsilon' 1e-5-    (ringestData [32,4] [63686.39999999999,137292.80000000002,121222.4,79558.40000000002,192646.40000000005,223971.0617601984,228556.80000000005,116846.33088019838,63686.39999999999,137292.80000000002,127174.4,79558.40000000002,192646.40000000005,158499.06176019844,202566.40000000005,51374.330880198424,11904.0,5952.0,7936.0,1984.0,116846.33088019838,385292.8000000001,227740.66176039676,192646.40000000005,116846.33088019838,228556.80000000005,174580.73088019836,35910.399999999994,79558.40000000002,127372.79999999997,143244.80000000002,63686.39999999999,105152.0,186683.13088000007,105151.98016,107124.73088000003,-396.79999999999995,26188.8,17459.2,25990.399999999998,-7936.0,73408.0,-1995.2691200000017,57536.0,51584.0,-660672.0,55552.0,3968.0,3968.0,3571.2,3571.2,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,49203.79519999998,49203.79519999998,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,129158.9952,65472.59519999998,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,1984.0000000000146,67059.20000000001,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,-21823.99999999993,108921.6,16070.400000000005,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,771974.4,218019.0617601984,192646.40000000005,170414.3308801984,385292.8000000001,340828.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,236294.40000000005,271587.0617601984,192646.40000000005,45422.33088019842,385292.8000000001,162268.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,369222.4,220003.0617601984,192646.40000000005,104942.33088019838,385292.8000000001,215836.66176039676,192646.40000000005])-    (rev' codeTomsSlice (rreshape [32, 4] t128))----- * PP Tests--testTomsSlicePP :: Assertion-testTomsSlicePP = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent codeTomsSlice (FTKR [32, 4] FTKScalar)-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (sscalar 4.0 * sdot0 (sconcrete (sfromListLinear [32] [0.0,1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0,15.0,16.0,17.0,18.0,19.0,20.0,21.0,22.0,23.0,24.0,25.0,26.0,27.0,28.0,29.0,30.0,31.0])) (sreplicate @32 (sdot0 (sslice (SNat @0) (SNat @3) (str (sfromR m1))) (sslice (SNat @1) (SNat @3) (str (sfromR m1))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> let v8 = sreplicate @32 (ssum @96 (sreshape @[96] (str (sslice (SNat @0) (SNat @3) (str (sfromR m1))) * str (sslice (SNat @1) (SNat @3) (str (sfromR m1)))))) in rfromS (ssum @128 (sreshape @[128] (str (sreplicate @4 (siota (SNat @32) * v8)))))"-  printArtifactPretty artifactRev-    @?= "\\dret m1 -> let m10 = sreshape @[32,3] (sreplicate @96 (ssum @32 (siota (SNat @32) * ssum @4 (str (sreshape @[32,4] (sreplicate @128 (sfromR dret))))))) in rfromS (str (sappend (sconcrete (sfromListLinear [0,32] [])) (sappend (str (str (sslice (SNat @1) (SNat @3) (str (sfromR m1))) * m10)) (sconcrete (sreplicate [1,32] 0.0)))) + str (sappend (sconcrete (sreplicate [1,32] 0.0)) (sappend (str (str (sslice (SNat @0) (SNat @3) (str (sfromR m1))) * m10)) (sconcrete (sfromListLinear [0,32] [])))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> rfromS (let x10 = sdot0 (sconcrete (sfromListLinear [32] [0.0,4.0,8.0,12.0,16.0,20.0,24.0,28.0,32.0,36.0,40.0,44.0,48.0,52.0,56.0,60.0,64.0,68.0,72.0,76.0,80.0,84.0,88.0,92.0,96.0,100.0,104.0,108.0,112.0,116.0,120.0,124.0])) (sreplicate @32 (sfromR dret)) in str (sappend (sslice (SNat @1) (SNat @3) (str (sfromR m1)) * sreplicate @3 (sreplicate @32 x10)) (sconcrete (sreplicate [1,32] 0.0))) + str (sappend (sconcrete (sreplicate [1,32] 0.0)) (sslice (SNat @0) (SNat @3) (str (sfromR m1)) * sreplicate @3 (sreplicate @32 x10))))"--testCNNOPP0c :: Assertion-testCNNOPP0c = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dCLaborious (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i86, i88] -> [i86 + i88]))) (\\[i40, i41] -> [i40 + i41])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i42, i43, i44, i45] -> [ifH (notB (2 <=. i42 + i44) &&* notB (2 <=. i43 + i45)) 0 1, i42, i43, i44, i45])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w46 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i38, i39] -> [i38 + i39]))) (\\[i40, i41] -> [i40 + i41])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i42, i43, i44, i45] -> [ifH (notB (2 <=. i42 + i44) &&* notB (2 <=. i43 + i45)) 0 1, i42, i43, i44, i45])))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w46 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w46 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i38, i39] -> [i38 + i39]))) (\\[i40, i41] -> [i40 + i41])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i42, i43, i44, i45] -> [ifH (notB (2 <=. i42 + i44) &&* notB (2 <=. i43 + i45)) 0 1, i42, i43, i44, i45])))))) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w46 * sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (ssum @2 (ssum @2 (sdot1In (stranspose @[0,1,2,5,3,4] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,4,2] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i112, i114] -> [i112 + i114]))) (\\[i40, i41] -> [i40 + i41])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i96, i97, i98, i103, i104] -> [ifH (notB (2 <=. i96 + i103) &&* notB (2 <=. i97 + i104)) 0 1, i96, i97, i103, i104]))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))) !$ [0]))))"--testCNNOPP0b :: Assertion-testCNNOPP0b = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dBLaborious (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i108, i110] -> [i108 + i110]))) (\\[i47, i48] -> [i47 + i48])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i49, i50, i51, i52] -> [ifH (notB (2 <=. i49 + i51) &&* notB (2 <=. i50 + i52)) 0 1, i49, i50, i51, i52]))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w53 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i45, i46] -> [i45 + i46]))) (\\[i47, i48] -> [i47 + i48])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i49, i50, i51, i52] -> [ifH (notB (2 <=. i49 + i51) &&* notB (2 <=. i50 + i52)) 0 1, i49, i50, i51, i52])))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * w53))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w59 = sscatter (stranspose @[1,2,4,5,0,3] (ssum @1 (stranspose @[3,0,1,2] (ssum @2 (str (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))))))))) (\\[i55, i56, i57, i58] -> [ifH (notB (2 <=. i55 + i57) &&* notB (2 <=. i56 + i58)) 0 1, i55, i56, i57, i58]) in rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[1,3,4,0,5,2] (w59 !$ [0])) (\\[i60, i61] -> [i60 + i61]))) (\\[i62, i63] -> [i62 + i63])))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[0,2,4,5,1,6,3] (sscatter (sdot1In (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0])) (stranspose @[4,2,3,6,7,0,5,1] (sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))) !$ [0])) (\\[i55, i56, i57, i58] -> [ifH (notB (2 <=. i55 + i57) &&* notB (2 <=. i56 + i58)) 0 1, i55, i56, i57, i58])) !$ [0]) (\\[i60, i61] -> [i60 + i61]))) (\\[i62, i63] -> [i62 + i63])))"--testCNNOPP1e :: Assertion-testCNNOPP1e = do-  resetVarCounter-  let f :: AstTensor AstMethodLet FullSpan-                     (TKProduct (TKR 4 Double) (TKR 4 Double))-        -> AstTensor AstMethodLet FullSpan-                     (TKR 4 Double)-      f v = conv2dUnpaddedL (tproject1 v) (tproject2 v)-      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i93, i95] -> [i93 + i95]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i32, i33, i34, i35] -> [ifH (notB (2 <=. i32 + i34) &&* notB (2 <=. i33 + i35)) 0 1, i32, i33, i34, i35])))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w36 = str (sreplicate @2 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i32, i33, i34, i35] -> [ifH (notB (2 <=. i32 + i34) &&* notB (2 <=. i33 + i35)) 0 1, i32, i33, i34, i35])))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w36 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w38 = sreshape @[2,2,2,2,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i124, i126] -> [i124 + i126]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i32, i33, i34, i35] -> [ifH (notB (2 <=. i32 + i34) &&* notB (2 <=. i33 + i35)) 0 1, i32, i33, i34, i35]))))) (stranspose @[2,3,1,4,5,6,0] w38)))) (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[0,2,4,5,1,6,3] (sscatter (sdot1In (stranspose @[2,3,5,6,0,4,1] (sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1)))))))) (stranspose @[2,3,5,6,0,4,1] w38)) (\\[i39, i40, i41, i42] -> [ifH (notB (2 <=. i39 + i41) &&* notB (2 <=. i40 + i42)) 0 1, i39, i40, i41, i42])) !$ [0]) (\\[i44, i45] -> [i44 + i45]))) (\\[i46, i47] -> [i46 + i47]))))"--testCNNOPP2 :: Assertion-testCNNOPP2 = do-  resetVarCounter-  let t = maxPool2dUnpadded2-            (rconcrete $ Nested.rreplicateScal (1 :$: 1 :$: 2 :$: 2 :$: ZSR) 1)-  printAstPretty (simplifyInlineContract t)-    @?= "rfromS (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sappend (sreplicate @1 (sgather (sreplicate @1 (stranspose @[2,0,1] (sgather (sconcrete (sfromListLinear [2,2] [1.0,1.0,1.0,1.0])) (\\[i68, i69] -> [i69 + i68])))) (\\[i44, i35, i8] -> [i8, i8, i8, 2 * i44 + i35]))) (sconcrete (sreplicate [1,2,2,2] 0.0))) !$ [0, 0])))"-  printAstPretty t-    @?= "rfromS (sreplicate @2 (sreplicate @2 (let u36 = let u41 = sgather (sgather (sreplicate @1 (let w32 = sgather (stranspose @[3,2,0,1] (sgather (sconcrete (sfromListLinear [2,3,2] [1.0,1.0,0.0,0.0,0.0,0.0,1.0,1.0,0.0,0.0,0.0,0.0])) (\\[i26, i15] -> [i26 + i15]))) (\\[i22, i16] -> [i22 + i16]) in stranspose @[1,2,3,0] (sappend (sreplicate @1 (stranspose @[2,0,4,1,3] w32 !$ [0])) (sconcrete (sreplicate [2,2,2,2,2] 0.0))))) (\\[i20] -> [i20, i20, i20, 0])) (\\[i44, i39, i35, i8] -> [2 * i39 + i8, i39, 2 * i44 + i35]) in str (sappend (sreplicate @1 (str u41 !$ [0])) (sconcrete (sreplicate [1,2,2,2] 0.0))) in stranspose @[2,3,0,1] u36 !$ [0, 0])))"--testCNNOPP2b :: Assertion-testCNNOPP2b = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded2 (FTKR [1, 1, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sappend (sreplicate @1 (sgather (sreplicate @1 (stranspose @[2,0,1] (sgather (sfromR u1 !$ [0, 0]) (\\[i92, i93] -> [i93 + i92])))) (\\[i94, i95, i96] -> [i96, i96, i96, 2 * i94 + i95]))) (sconcrete (sreplicate [1,2,2,2] 0.0))) !$ [0, 0])))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sappend (sreplicate @1 (sgather (sreplicate @1 (stranspose @[2,0,1] (sgather (sfromR u1 !$ [0, 0]) (\\[i92, i93] -> [i93 + i92])))) (\\[i94, i95, i96] -> [i96, i96, i96, 2 * i94 + i95]))) (sconcrete (sreplicate [1,2,2,2] 0.0))) !$ [0, 0])))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let u98 = stranspose @[3,2,0,1] (soneHot (ssum @2 (ssum @2 (sfromR dret))) [0, 0]) in rfromS (soneHot (sscatter (stranspose @[1,2,0] (ssum @1 (sscatter (ssum @1 (sslice (SNat @0) (SNat @1) u98)) (\\[i99, i100, i101] -> [i101, i101, i101, 2 * i99 + i100])))) (\\[i102, i103] -> [i103 + i102])) [0, 0])"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (sreplicate @1 (sreplicate @1 (sscatter (sscatter (stranspose @[3,2,0,1] (soneHot (ssum @2 (ssum @2 (sfromR dret))) [0, 0]) !$ [0]) (\\[i99, i100, i101] -> [i101, i101, 2 * i99 + i100, i101]) !$ [0]) (\\[i102, i103] -> [i103 + i102]))))"--maxPool2dUnpadded2-  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)-  => target (TKR 4 r) -> target (TKR 4 r)-maxPool2dUnpadded2 a =-  rbuild [2, 2, 2, 2] $ \case-    [_, _, iBh, iBw] ->-      let arrt = slicez2 (conv2dUnpadded2 a) [iBw, 1, 2 * iBh, 2 * iBw]-      in rmaximum2 arrt-    _ -> error "maxPool2dUnpadded2: impossible pattern needlessly required"--conv2dUnpadded2-  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded2 a =-  rbuild [3, 3, 2, 2] $ \case-    [iImg, _, iBh, iBw] ->-      let arrAt = slicez2 a [iImg, 0, iBh, iBw]-      in rindex0 arrAt [0, iBw, iBw, 0]-    _ -> error "conv2dUnpadded2: impossible pattern needlessly required"--slicez2-  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)-  => target (TKR n r) -> IxROf target n -> target (TKR n r)-slicez2 d ixBase =-  rbuild [1, 1, 2, 2] $ \ixResult -> indexz02 d (ixrZipWith (+) ixBase ixResult)--indexz02-  :: forall target r n.-     (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)-  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)-indexz02 d ix = ifH (1 >. (toList ix !! 0)) (d ! ix) (rscalar 0)--rmaximum2 :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)-         => target (TKR 4 r) -> target (TKR 0 r)-rmaximum2 t0 = tlet t0 $ \t -> rindex0 t [0, 0, 0, 0]--{- TODO: divergent result; bring back when GHC 9.10 dropped:-testCNNOPP3 :: Assertion-testCNNOPP3 = do-  resetVarCounter-  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                       (rconcrete $ Nested.rscalar 7-                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))-      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = maxPool2dUnpadded33 $ conv2dUnpadded3 blackGlyph-  printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (sreplicate @2 (sgather (stranspose @[2,1,0,4,3] (sappend (sreplicate @1 (sgather (sconcrete (sfromListLinear [2] [7.0,0.0])) (\\[i18, i22, i17, i15] -> [ifH (notB (2 <=. remH i22 4 + i18) &&* (notB (2 <=. i22 + i17) &&* notB (2 <=. i22 + i15))) 0 1]))) (sconcrete (sfromListLinear [1,2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])))) (\\[i52, i51] -> [remH i51 4, i52, i52, remH i51 4])))"-  printAstPretty afcnn2T-    @?= "rfromS (let w30 = sgather (sfromVector (fromList [stranspose @[4,0,1,2,5,3] (sgather (stranspose @[1,2,4,5,0,3] (sgather (sappend (sreplicate @1 (sgather (sconcrete (sfromListLinear [2] [7.0,0.0])) (\\[i18, i22, i17, i15] -> [ifH (notB (2 <=. remH i22 4 + i18) &&* (notB (2 <=. i22 + i17) &&* notB (2 <=. i22 + i15))) 0 1]))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sscalar 0.0))))))) (\\[i43, i38, i29, i7] -> [i43 + i7, i43 + i7, remH i38 4 + i29]))) (\\[i37, i33, i28, i8] -> [i37, i28, i33 + i8, remH i37 4 + i28])), sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sscalar 0.0))))))])) (\\[i46, i41, i36, i32, i27, i26, i24, i23] -> [ifH (notB (2 <=. remH i36 4 + i27) &&* (notB (2 <=. i46 + i26) &&* (notB (2 <=. i41 + i24) &&* notB (2 <=. i32 + i23)))) 0 1, i41, i36, i32, i27, i24, i23]) in stranspose @[4,5,6,7,0,1,2,3] w30 !$ [0, 0, 0, 0])"--}--testCNNOPP3b :: Assertion-testCNNOPP3b = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded33 . conv2dUnpadded3) (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (stranspose @[2,1,0] (sgather (sgather (sappend (sreplicate @1 (stranspose @[0,4,5,1,2,3] (sgather (sfromVector (fromList [stranspose @[5,2,3,4,0,1] (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sgather (stranspose @[3,0,2,1] (sgather (stranspose @[3,0,1,2] (sgather (stranspose @[3,0,2,1] (sfromR u1) !$ [1]) (\\[i191, i193] -> [remH i191 4 + i193]))) (\\[i195, i197] -> [i195 + i197, i195]))) (\\[i128, i129] -> [i128 + i129, i128]))))), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i130, i131, i132, i133] -> [ifH (notB (2 <=. remH i130 4 + i131) &&* (notB (2 <=. i130 + i132) &&* notB (2 <=. i130 + i133))) 0 1, i130, i131, i132, i133])))) (sconcrete (sreplicate [1,2,2,2,2,2,2] 0.0))) (\\[i134, i135, i136, i137] -> [i135, i134, i135, i137, i135, i137, i134])) (\\[i138] -> [remH i138 4])))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (stranspose @[2,1,0] (sgather (sgather (sappend (sreplicate @1 (stranspose @[0,4,5,1,2,3] (sgather (sfromVector (fromList [stranspose @[5,2,3,4,0,1] (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sgather (stranspose @[3,0,2,1] (sgather (stranspose @[3,0,1,2] (sgather (stranspose @[3,0,2,1] (sfromR u1) !$ [1]) (\\[i124, i125] -> [remH i124 4 + i125]))) (\\[i126, i127] -> [i126 + i127, i126]))) (\\[i128, i129] -> [i128 + i129, i128]))))), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i130, i131, i132, i133] -> [ifH (notB (2 <=. remH i130 4 + i131) &&* (notB (2 <=. i130 + i132) &&* notB (2 <=. i130 + i133))) 0 1, i130, i131, i132, i133])))) (sconcrete (sreplicate [1,2,2,2,2,2,2] 0.0))) (\\[i134, i135, i136, i137] -> [i135, i134, i135, i137, i135, i137, i134])) (\\[i138] -> [remH i138 4])))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w145 = sscatter (sscatter (stranspose @[2,1,0] (sfromR dret)) (\\[i140] -> [remH i140 4])) (\\[i141, i142, i143, i144] -> [i142, i141, i142, i144, i142, i144, i141]) ; w150 = sscatter (stranspose @[0,3,4,5,1,2] (ssum @1 (sslice (SNat @0) (SNat @1) w145))) (\\[i146, i147, i148, i149] -> [ifH (notB (2 <=. remH i146 4 + i147) &&* (notB (2 <=. i146 + i148) &&* notB (2 <=. i146 + i149))) 0 1, i146, i147, i148, i149]) in rfromS (stranspose @[1,3,2,0] (soneHot (sscatter (stranspose @[1,2,3,0] (sscatter (stranspose @[1,3,2,0] (sscatter (stranspose @[3,2,0,1] (ssum @2 (ssum @2 (stranspose @[4,5,1,2,3,0] (w150 !$ [0]))))) (\\[i151, i152] -> [i151 + i152, i151]))) (\\[i153, i154] -> [i153 + i154, i153]))) (\\[i155, i156] -> [remH i155 4 + i156])) [1]))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (stranspose @[1,3,2,0] (soneHot (sscatter (stranspose @[1,2,3,0] (sscatter (stranspose @[1,3,2,0] (sscatter (ssum @2 (ssum @2 (stranspose @[0,5,6,1,4,2,3] (sscatter (sscatter (sscatter (stranspose @[2,1,0] (sfromR dret)) (\\[i140] -> [remH i140 4])) (\\[i141, i142, i143, i144] -> [i142, i141, i142, i144, i141, i142, i144]) !$ [0]) (\\[i146, i147, i148, i149] -> [ifH (notB (2 <=. remH i146 4 + i147) &&* (notB (2 <=. i146 + i148) &&* notB (2 <=. i146 + i149))) 0 1, i146, i147, i148, i149])) !$ [0]))) (\\[i151, i152] -> [i151 + i152, i151]))) (\\[i153, i154] -> [i153 + i154, i153]))) (\\[i155, i156] -> [remH i155 4 + i156])) [1]))"--maxPool2dUnpadded3-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r)-maxPool2dUnpadded3 arr =-  rbuild [2, 2, 2, 2] $ \case-    [aa, bb, iBh, iBw] ->-      let arrt = slicez3 [2, 2, 2, 2] arr [iBh `quotH` 4, aa, bb, iBw]-      in rmaximum3 arrt-    _ -> error "maxPool2dUnpadded3: impossible pattern needlessly required"--maxPool2dUnpadded33-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r)-maxPool2dUnpadded33 arr =-  rbuild [2, 2, 2, 2] $ \case-    [aa, bb, iBh, iBw] ->-      let arrt = slicez33 [2, 2, 2, 2] arr [iBh `remH` 4, aa, bb, iBw]-      in rmaximum3 arrt-    _ -> error "maxPool2dUnpadded33: impossible pattern needlessly required"--conv2dUnpadded3-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded3 arrA =-  let shB = [2, 2, 2, 2]-  in rbuild shB $ \case-    [iImg, _, iBh, iBw] ->-      let arrAt = slicez33 shB arrA [iImg `remH` 4, iImg, iImg, 1]-      in rindex0 arrAt [iBh, iBw, iImg, iBh]-    _ -> error "conv2dUnpadded3: impossible pattern needlessly required"--slicez3-  :: (ADReady target, GoodScalar r, KnownNat n)-  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)-slicez3 shOut d ixBase =-  rbuild shOut $ \_ -> indexz03 d (ixrZipWith (+) ixBase ixBase)--slicez33-  :: (ADReady target, GoodScalar r, KnownNat n)-  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)-slicez33 shOut d ixBase =-  rbuild shOut $ \ixResult -> indexz03 d (ixrZipWith (+) ixBase ixResult)--indexz03-  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)-  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)-indexz03 d ix = ifH (within0 @target (rshape @target d) ix) (d ! ix) (rscalar 0)--rmaximum3 :: (BaseTensor target, LetTensor target, KnownNat n, GoodScalar r)-         => target (TKR n r) -> target (TKR 0 r)-rmaximum3 t0 = tlet t0 $ \t -> rindex0 t [0, 0, 0, 0]--testCNNOPP4 :: Assertion-testCNNOPP4 = do-  resetVarCounter-  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK-                       (rconcrete $ Nested.rscalar 7-                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))-      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = maxPool2dUnpadded4 blackGlyph-  printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (str (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sreplicate @1 (stranspose @[2,1,0] (sgather (stranspose @[3,4,5,2,6,1,0] (sgather (stranspose @[4,1,3,0,2] (sgather (stranspose @[3,0,4,1,2] (sgather (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i61, i64] -> [i61 + i64]))) (\\[i67, i69] -> [3 + (negate i69 + i67), i69]))) (\\[i71, i73, i76] -> [i71 * i73 + i76])) !$ [1, 0, 0, 0]) (\\[i84] -> [2 * i84]))))))"-      -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,7.0,7.0,7.0,7.0,0.0,0.0,0.0,0.0,7.0,7.0,7.0,7.0]))"-  printAstPretty afcnn2T-    @?= "rfromS (let w19 = sgather (sfromVector (fromList [stranspose @[3,0,5,6,1,2,4] (sgather (stranspose @[6,0,3,1,4,5,2] (sgather (stranspose @[3,0,2,1] (sgather (stranspose @[0,2,1] (sgather (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i32, i5] -> [i32 + i5]))) (\\[i31, i6] -> [i31, 3 + (negate i31 + i6)]))) (\\[i36, i26, i7] -> [i36 * i26 + i7]))) (\\[i22, i8] -> [2 * i22 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) (\\[i28, i21, i15, i12, i9] -> [ifH (notB (2 <=. i28 + i15) &&* (notB (0 <=. negate i28 + i12) &&* notB (3 <=. 2 * i21 + i9))) 0 1, i28, i21, i15, i12, i9]) in stranspose @[2,3,4,7,5,0,6,1] w19 !$ [0, 0, 0, 0])"--testCNNOPP4b :: Assertion-testCNNOPP4b = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded4 (FTKR [3, 3, 3, 3] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (str (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sreplicate @1 (stranspose @[2,1,0] (sgather (stranspose @[2,3,0,1] (sgather (stranspose @[1,0,3,2] (sreplicate @2 (stranspose @[2,3,0,1] (sreplicate @2 (stranspose @[2,1,0] (sreplicate @2 (sfromR u1 !$ [2, 2]))))))) (\\[i194, i195] -> [i195 * i194, i194, i195]))) (\\[i125] -> [i125, 2 * i125]))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (str (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sreplicate @1 (stranspose @[2,1,0] (sgather (stranspose @[3,5,6,2,4,7,1,0] (sgather (stranspose @[3,4,7,1,5,6,0,2] (sgather (stranspose @[6,0,7,4,3,2,1,5] (sgather (sslice (SNat @1) (SNat @2) (stranspose @[0,2,1] (sfromR u1))) (\\[i115, i116, i117, i118, i119] -> [i115 + i116]))) (\\[i120, i121] -> [3 + (negate i121 + i120), i121]))) (\\[i122, i123, i124] -> [i122, i123, i122 * i123 + i124])) !$ [1, 0, 0, 0]) (\\[i125] -> [i125, 2 * i125]))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> rfromS (stranspose @[0,2,1] (sappend (sconcrete (sreplicate [1,3,3,3] 0.0)) (sappend (sscatter (stranspose @[1,6,5,4,3,7,0,2] (sscatter (stranspose @[6,3,7,0,1,4,5,2] (sscatter (stranspose @[7,6,3,0,4,1,2,5] (soneHot (sscatter (stranspose @[2,1,0] (ssum @1 (sslice (SNat @1) (SNat @1) (str (sfromR dret))))) (\\[i127] -> [i127, 2 * i127])) [1, 0, 0, 0])) (\\[i128, i129, i130] -> [i128, i129, i128 * i129 + i130]))) (\\[i131, i132] -> [3 + (negate i132 + i131), i132]))) (\\[i133, i134, i135, i136, i137] -> [i133 + i134])) (sconcrete (sfromListLinear [0,3,3,3] [])))))"-      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sscatter (ssum @1 (sslice (SNat @1) (SNat @1) (str (sfromR dret)))) (\\[i86, i87, i88] -> [i86 * i87, 2 * i88])) [2, 2])"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (sappend (sconcrete (sreplicate [1,3,3,3] 0.0)) (stranspose @[0,2,1] (sscatter (stranspose @[1,6,5,4,3,7,0,2] (sscatter (stranspose @[6,3,7,0,1,4,5,2] (sscatter (stranspose @[7,6,3,0,4,1,2,5] (soneHot (sscatter (stranspose @[1,3,2,0] (sfromR dret) !$ [1]) (\\[i127] -> [i127, 2 * i127])) [1, 0, 0, 0])) (\\[i128, i129, i130] -> [i128, i129, i128 * i129 + i130]))) (\\[i131, i132] -> [3 + (negate i132 + i131), i132]))) (\\[i133, i134, i135, i136, i137] -> [i133 + i134]))))"-      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sscatter (str (sfromR dret) !$ [1]) (\\[i86, i87, i88] -> [i86 * i87, 2 * i88])) [2, 2])"--testCNNOPP5 :: Assertion-testCNNOPP5 = do-  resetVarCounter-  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @6) knownSTK-                   $ AstReplicate (SNat @6) knownSTK-                   $ AstReplicate (SNat @6) knownSTK-                   $ AstReplicate (SNat @6) knownSTK-                       (rconcrete $ Nested.rscalar 7-                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))-      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = conv2dUnpadded4 blackGlyph-  printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (sconcrete (sreplicate [1,1,2,2] 7.0))"-  printAstPretty afcnn2T-    @?= "rfromS (sconcrete (sreplicate [1,1,2,2] 7.0))"--testCNNOPP5b :: Assertion-testCNNOPP5b = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dUnpadded4 (FTKR [5, 5, 5, 5] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (sfromR u1 !$ [0, 0])))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (sfromR u1 !$ [0, 0])))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> rfromS (soneHot (sappend (sconcrete (sfromListLinear [0,5] [])) (sappend (str (sappend (sconcrete (sfromListLinear [0,2] [])) (sappend (str (ssum @1 (ssum @1 (sfromR dret)))) (sconcrete (sreplicate [3,2] 0.0))))) (sconcrete (sreplicate [3,5] 0.0)))) [0, 0])"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (soneHot (sappend (str (sappend (stranspose @[0,1,3,2] (sfromR dret) !$ [0, 0]) (sconcrete (sreplicate [3,2] 0.0)))) (sconcrete (sreplicate [3,5] 0.0))) [0, 0])"--maxPool2dUnpadded4-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r)-maxPool2dUnpadded4 arr =-  rbuild [2, 2, 2, 2] $ \case-    [aa, bb, iBh, iBw] ->-      let arrt = slicez4 [2, 2, 2, 2] arr [bb + 1, 3 - bb, aa * iBh, 2 * iBw]-      in rmaximum3 arrt-    _ -> error "maxPool2dUnpadded4: impossible pattern needlessly required"--conv2dUnpadded4-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded4 arrA =-  let shB = [1, 1, 2, 2]-  in rbuild shB $ \case-    [iImg, _, iBh, iBw] ->-      let arrAt = slicez4 shB arrA [iImg, 0, iBh, iBw]-      in rindex0 arrAt [0, 0, 0, 0]-    _ -> error "conv2dUnpadded4: impossible pattern needlessly required"--slicez4-  :: (ADReady target, GoodScalar r, KnownNat n)-  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)-slicez4 shOut d ixBase =-  rbuild shOut $ \ixResult -> indexz03 d (ixrZipWith (+) ixBase ixResult)--testCNNOPP6 :: Assertion-testCNNOPP6 = do-  resetVarCounter-  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                       (rconcrete $ Nested.rscalar 7-                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))-      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = maxPool2dUnpadded3 $ conv2dUnpadded3z blackGlyph-  printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"-  printAstPretty afcnn2T-    @?= "rfromS (stranspose @[1,2,0] (sreplicate @2 (let t30 = sgather (stranspose @[2,1,0] (sgather (str (sgather (sreplicate @2 (str (sreplicate @2 (let m21 = sgather (str (sgather (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i9] -> [2 * i9, 2 * i9, 2 * i9]))) (\\[i12] -> [2 * i12]) in sappend (sreplicate @1 (sappend (sreplicate @1 (m21 !$ [0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))))) (\\[i1] -> [2 * i1, 0]))) (\\[i2] -> [2 * i2]))) (\\[i4] -> [2 * i4]) in sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (t30 !$ [0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"---testCNNOPP6b :: Assertion-testCNNOPP6b = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dUnpadded3z) (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let t34 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) in rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t34)))))) [0, 0, 0, 0])"---  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (soneHot (ssum0 (stranspose @[0,1,3,2] (sfromR dret) !$ [0, 0, 0])) [0, 0, 0, 0])"--conv2dUnpadded3z-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded3z arrA =-  let shB = [2, 2, 2, 2]-  in rbuild shB $ \case-    [iImg, _, iBh, iBw] ->-      let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBw]-      in rindex0 arrAt [iBh, iBw, iImg, iBh]-    _ -> error "conv2dUnpadded3z: impossible pattern needlessly required"--testCNNOPP7 :: Assertion-testCNNOPP7 = do-  resetVarCounter-  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                       (rconcrete $ Nested.rscalar 7-                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))-      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = maxPool2dUnpadded3y $ conv2dUnpadded3y blackGlyph-  printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"-  printAstPretty afcnn2T-    @?= "rfromS (let u27 = sgather (stranspose @[3,2,0,1] (sgather (stranspose @[1,2,0] (sgather (sreplicate @2 (stranspose @[1,2,0] (sreplicate @2 (let m21 = sgather (str (sgather (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i9] -> [2 * i9, 2 * i9, 2 * i9]))) (\\[i11] -> [2 * i11]) in sappend (sreplicate @1 (sappend (sreplicate @1 (m21 !$ [0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))))) (\\[i1] -> [2 * i1]))) (\\[i31, i3] -> [2 * i3, 2 * i31]))) (\\[i4] -> [2 * i4]) in stranspose @[1,2,0] (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (u27 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"--testCNNOPP7b :: Assertion-testCNNOPP7b = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3y . conv2dUnpadded3y) (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (stranspose @[2,0,1] (sfromR dret)))))))))) [0, 0, 0, 0])"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (soneHot (sfromR dret !$ [0, 0, 0, 0]) [0, 0, 0, 0])"--maxPool2dUnpadded3y-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r)-maxPool2dUnpadded3y arr =-  rbuild [2, 2, 2, 2] $ \case-    [aa, bb, iBh, iBw] ->-      let arrt = slicez3 [2, 2, 2, 2] arr [iBh, aa, bb, iBw]-      in rmaximum3 arrt-    _ -> error "maxPool2dUnpadded3y: impossible pattern needlessly required"--conv2dUnpadded3y-  :: (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded3y arrA =-  let shB = [2, 2, 2, 2]-  in rbuild shB $ \case-    [iImg, _, iBh, iBw] ->-      let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBh]-      in rindex0 arrAt [iBh, iBw, iImg, iBh]-    _ -> error "conv2dUnpadded3y: impossible pattern needlessly required"---- TODO: OOMs-_testPaddedCNNOPP0c :: Assertion-_testPaddedCNNOPP0c = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dCPadded (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[5,0,1,4,2,3] (sgather (sconcrete (sfromListLinear [4,4,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i37, i38, i39, i40] -> [i37 + i39, i38 + i40])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w41 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[5,0,1,4,2,3] (sgather (sconcrete (sfromListLinear [4,4,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i37, i38, i39, i40] -> [i37 + i39, i38 + i40])))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w41 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w41 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[5,0,1,4,2,3] (sgather (sconcrete (sfromListLinear [4,4,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i37, i38, i39, i40] -> [i37 + i39, i38 + i40])))))) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w41 * sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (ssum @2 (ssum @2 (sdot1In (stranspose @[0,1,2,6,3,4,5] (sgather (sconcrete (sfromListLinear [4,4,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,-2.0,13.1,582934.0,2.0,0.0,9.0,2.99432,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,6.0,0.1,8.0,-335.0,1.0,-0.2,-4.0,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i66, i67, i68, i73, i74] -> [i66 + i73, i67 + i74]))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))) !$ [0]))))"---- TODO: OOMs-_testPaddedCNNOPP0b :: Assertion-_testPaddedCNNOPP0b = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dBPadded (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR u1)) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))))))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i50, i51, i52, i53, i54, i55] -> [i51 + i54, i51, i54, i50, i52, i53, i55, i50, i53, i52 + i55])))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w56 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR u1)) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))))))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i50, i51, i52, i53, i54, i55] -> [i51 + i54, i51, i54, i50, i52, i53, i55, i50, i53, i52 + i55]))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * w56))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w64 = sscatter (ssum @1 (stranspose @[3,0,1,2] (ssum @2 (str (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))))))) (\\[i58, i59, i60, i61, i62, i63] -> [i59 + i62, i59, i62, i58, i60, i61, i63, i58, i61, i60 + i63]) ; u65 = ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (stranspose @[6,1,2,3,4,5,9,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @3) (stranspose @[9,3,1,4,5,2,6,7,8,0] (sslice (SNat @1) (SNat @3) w64))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)))))))))) in rfromS (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (sslice (SNat @1) (SNat @3) u65))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (stranspose @[1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (stranspose @[6,1,2,3,4,5,9,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @3) (stranspose @[9,3,1,4,5,2,6,7,8,0] (sslice (SNat @1) (SNat @3) (sscatter (sdot1In (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0])) (stranspose @[4,0,2,3,5,6,7,1] (sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))) !$ [0])) (\\[i58, i59, i60, i61, i62, i63] -> [i59 + i62, i59, i62, i58, i60, i61, i63, i58, i61, i60 + i63]))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)))))))))))))))"---- TODO: OOMs-_testPaddedCNNOPP1e :: Assertion-_testPaddedCNNOPP1e = do-  resetVarCounter-  let f :: AstTensor AstMethodLet FullSpan-                     (TKProduct (TKR 4 Double) (TKR 4 Double))-        -> AstTensor AstMethodLet FullSpan-                     (TKR 4 Double)-      f v = conv2dPadded (tproject1 v) (tproject2 v)-      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))))))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i44, i75, i45, i46, i76, i47] -> [i75 + i76, i75, i76, i44, i45, i46, i47, i44, i46, i45 + i47]))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let u41 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; w48 = str (sreplicate @2 (stranspose @[0,4,1,2,5,3] (sgather (stranspose @[2,3,4,5,6,7,8,0,1] (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 u41)))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i42, i43] -> [i42 + i43, i42, i43]))) (\\[i44, i45, i46, i47] -> [i44, i45, i46, i47, i44, i46, i45 + i47])))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w48 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w50 = sreshape @[2,2,2,2,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))))))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i44, i84, i45, i46, i85, i47] -> [i84 + i85, i84, i85, i44, i45, i46, i47, i44, i46, i45 + i47])))) (stranspose @[2,3,1,4,5,6,0] w50)))) (stranspose @[1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (stranspose @[6,1,2,3,4,5,9,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @3) (stranspose @[9,3,1,4,5,2,6,7,8,0] (sslice (SNat @1) (SNat @3) (sscatter (stranspose @[7,8,0,1,2,3,4,5,6] (sscatter (sdot1In (sreplicate @2 (stranspose @[2,3,5,0,4,1] (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))) (stranspose @[0,3,4,6,2,5,1] w50)) (\\[i51, i52, i53, i54] -> [i51, i52, i53, i54, i51, i53, i52 + i54]))) (\\[i55, i56] -> [i55 + i56, i55, i56]))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0))))))))))))))))"---- This is fragile due to indexing out of bounds, see above.-testPaddedCNNOPP1b :: Assertion-testPaddedCNNOPP1b = do-  resetVarCounter-  let f :: AstTensor AstMethodLet FullSpan-                     (TKProduct (TKR 4 Double) (TKR 4 Double))-        -> AstTensor AstMethodLet FullSpan-                     (TKR 4 Double)-      f v = conv2dShrinking (tproject1 v) (tproject2 v)-      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-                       (FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) FTKScalar)-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[6,2,4,4,8] (str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i64, i66] -> [i64 + i66]))) (\\[i24, i25] -> [i24 + i25])))) * sreplicate @6 (str (sreplicate @4 (str (sreplicate @4 (sfromR (tproject1 u1))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w26 = str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[6,2,4,4,8] (w26 * sreplicate @6 (str (sreplicate @4 (str (sreplicate @4 (sfromR (tproject1 u1))))))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w26 = str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) ; w28 = sreshape @[6,2,4,4,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (rfromS (ssum @4 (str (ssum @4 (str (ssum @6 (w26 * w28))))))) (rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[2,5,0,1,3,4] (ssum @2 (str (sreplicate @6 (str (sreplicate @4 (str (sreplicate @4 (sfromR (tproject1 u1)))))) * w28)))) (\\[i29, i30] -> [i29 + i30]))) (\\[i31, i32] -> [i31 + i32]))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [6,2,6,6] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [6,2,6,6] STKScalar)) (let w28 = sreshape @[6,2,4,4,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @4 (ssum @4 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i81, i83] -> [i81 + i83]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2,3,1,4,5,6,0] w28)))) (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (sdot1In (stranspose @[3,6,0,2,4,5,1] (sreplicate @6 (str (sreplicate @4 (str (sreplicate @4 (sfromR (tproject1 u1)))))))) (stranspose @[3,6,0,2,4,5,1] w28)) (\\[i29, i30] -> [i29 + i30]))) (\\[i31, i32] -> [i31 + i32]))))"--testPaddedCNNOPPLet :: Assertion-testPaddedCNNOPPLet = do-  resetVarCounter-  let f :: AstTensor AstMethodLet FullSpan-                     (TKProduct (TKR 4 Double) (TKR 4 Double))-        -> AstTensor AstMethodLet FullSpan-                     (TKR 4 Double)-      f v = conv2dPaddedLet (tproject1 v) (tproject2 v)-      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))), sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i36, i37] -> [i36, i37, ifH (notB (2 <=. i37) &&* notB (2 <=. i36)) 0 1])))))) (\\[i81, i83] -> [i81 + i83]))) (\\[i41, i42] -> [i41 + i42])))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let u35 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; u38 = stranspose @[1,2,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] u35, sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i36, i37] -> [i36, i37, ifH (notB (2 <=. i37) &&* notB (2 <=. i36)) 0 1])))))) ; w43 = str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] u38) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w43 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w45 = sreshape @[2,2,2,2,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))), sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i36, i37] -> [i36, i37, ifH (notB (2 <=. i37) &&* notB (2 <=. i36)) 0 1])))))) (\\[i100, i102] -> [i100 + i102]))) (\\[i41, i42] -> [i41 + i42]))))) (stranspose @[2,3,1,4,5,6,0] w45)))) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,2,1,4,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,2,1,0] (sappend (sconcrete (sreplicate [1,2,2,4,2] 0.0)) (stranspose @[1,3,4,0,2] (sappend (sconcrete (sreplicate [1,3,2,2,2] 0.0)) (sscatter (sslice (SNat @1) (SNat @3) (stranspose @[3,0,1,2] (sslice (SNat @1) (SNat @3) (sscatter (stranspose @[2,4,1,3,0] (sscatter (sdot1In (stranspose @[3,6,0,2,4,5,1] (sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1)))))))) (stranspose @[3,6,0,2,4,5,1] w45)) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49]))))) (\\[i51, i52] -> [i51, i52, ifH (notB (2 <=. i52) &&* notB (2 <=. i51)) 0 1]))))))))) !$ [0]))"--conv2dPaddedLet-  :: forall target r. (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)-conv2dPaddedLet arrK arrA =-  let [nImgs, nCinpA, nAh, nAw] = rshape arrA-      [nCoutK, nCinpK, nKh, nKw] = rshape arrK-      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]-      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case-        [iImg, iCinp, iPh, iPw] ->-          ifH (iPh <. fromIntegral (nKh `div` 2)-               ||* iPw <. fromIntegral (nKw `div` 2)-               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)-               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))-              (rscalar 0)-              (arrA ! [ iImg-                      , iCinp-                      , iPh - fromIntegral (nKh `div` 2)-                      , iPw - fromIntegral (nKw `div` 2) ])-      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA-      shB = [nImgs, nCoutK, nAh, nAw]-      shK1 = [1, nCinp, nKh, nKw]-  in tlet arrAPadded $ \arrAPadded2 -> rbuild shB $ \case-    [iImg, iCout, iBh, iBw] ->-      let arrAt = slicezL shK1 arrAPadded2 [iImg, 0, iBh, iBw]-          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]-      in rdot0 arrAt arrKt-    _ -> error "conv2dPaddedLet: impossible pattern needlessly required"--testPaddedCNNOPPLet2 :: Assertion-testPaddedCNNOPPLet2 = do-  resetVarCounter-  let f :: AstTensor AstMethodLet FullSpan-                     (TKProduct (TKR 4 Double) (TKR 4 Double))-        -> AstTensor AstMethodLet FullSpan-                     (TKR 4 Double)-      f v = conv2dPaddedLet2 (tproject1 v) (tproject2 v)-      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[3,5,2,0,4,1] (sgather (stranspose @[1,2,0] (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))), sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i53, i54] -> [i53, i54, ifH (notB (2 <=. i54) &&* notB (2 <=. i53)) 0 1])))))))) (\\[i180, i182] -> [i180 + i182]))) (\\[i62, i63] -> [i62, i62 + i63])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR (tproject1 u1))))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let u52 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; u55 = stranspose @[1,2,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] u52, sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i53, i54] -> [i53, i54, ifH (notB (2 <=. i54) &&* notB (2 <=. i53)) 0 1])))))) ; w64 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[3,5,2,0,4,1] (sgather (stranspose @[1,2,0] (sgather (stranspose @[1,5,0,2,3,4] (sreplicate @2 (stranspose @[3,0,1,2] (sgather (stranspose @[2,3,0,1] (sreplicate @2 u55)) (\\[i56, i57, i58] -> [i58, i56]))))) (\\[i59] -> [i59, i59]))) (\\[i60, i61] -> [i60, i60 + i61]))) (\\[i62, i63] -> [i62, i62 + i63])))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w64 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR (tproject1 u1))))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w66 = sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[5,0,1,4,3,2] (sgather (stranspose @[2,4,0,3,1] (sgather (str (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,4,2] 0.0)) (stranspose @[3,2,0,1] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))), sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i53, i54] -> [i53, i54, ifH (notB (2 <=. i54) &&* notB (2 <=. i53)) 0 1])))))))) (\\[i253, i255] -> [i255 + i253]))) (\\[i222, i223, i229] -> [i222, i222 + i229]))) (stranspose @[4,2,3,1,5,6,7,0] w66 !$ [0])))) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,2,1,4,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,2,1,0] (sappend (sconcrete (sreplicate [1,2,2,4,2] 0.0)) (stranspose @[1,3,4,0,2] (sappend (sconcrete (sreplicate [1,3,2,2,2] 0.0)) (sscatter (sslice (SNat @1) (SNat @3) (stranspose @[3,0,1,2] (sslice (SNat @1) (SNat @3) (ssum @2 (stranspose @[2,1,3,0] (sscatter (ssum @2 (stranspose @[2,3,4,5,0,1] (sscatter (stranspose @[2,0,1] (sscatter (stranspose @[3,5,2,0,4,1] (sscatter (sdot1In (sreplicate @2 (stranspose @[2,0,1,4,3] (sreplicate @2 (sreplicate @2 (stranspose @[3,2,1,0] (sfromR (tproject1 u1))))))) (stranspose @[4,3,7,0,2,5,6,1] w66 !$ [0])) (\\[i67, i68] -> [i67, i67 + i68]))) (\\[i69, i70] -> [i69, i69 + i70]))) (\\[i71] -> [i71, i71])))) (\\[i72, i73, i74] -> [i74, i72]))))))) (\\[i76, i77] -> [i76, i77, ifH (notB (2 <=. i77) &&* notB (2 <=. i76)) 0 1]))))))))) !$ [0]))"--conv2dPaddedLet2-  :: forall target r. (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)-conv2dPaddedLet2 arrK arrA =-  let [nImgs, nCinpA, nAh, nAw] = rshape arrA-      [nCoutK, nCinpK, nKh, nKw] = rshape arrK-      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]-      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case-        [iImg, iCinp, iPh, iPw] ->-          ifH (iPh <. fromIntegral (nKh `div` 2)-               ||* iPw <. fromIntegral (nKw `div` 2)-               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)-               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))-              (rscalar 0)-              (arrA ! [ iImg-                      , iCinp-                      , iPh - fromIntegral (nKh `div` 2)-                      , iPw - fromIntegral (nKw `div` 2) ])-      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA-      shB = [nImgs, nCoutK, nAh, nAw]-      shK1 = [1, nCinp, nKh, nKw]-  in rbuild shB $ \case-    [iImg, iCout, iBh, iBw] ->-      let arrAt = tlet arrAPadded $ \arrAPadded2 ->-                    slicezL shK1 arrAPadded2 [iImg, 0, iBh, iBw]-          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]-      in rdot0 arrAt arrKt-    _ -> error "conv2dPaddedLet2: impossible pattern needlessly required"---- TODO: OOMs-_testPaddedCNNOPP2 :: Assertion-_testPaddedCNNOPP2 = do-  resetVarCounter-  let f :: AstTensor AstMethodLet FullSpan-                     (TKProduct (TKR 4 Double) (TKR 4 Double))-        -> AstTensor AstMethodLet FullSpan-                     (TKR 4 Double)-      f v = conv2dPadded2 (tproject1 v) (tproject2 v)-      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[5,0,1,4,2,3] (sgather (sappend (sconcrete (sreplicate [1,4,2,2] 0.0)) (sappend (stranspose @[3,0,2,1] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,4,2,2] 0.0)))) (\\[i24, i25, i27, i28] -> [i25 + i28, i24 + i27])))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w29 = str (sreplicate @2 (sgather (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))) (\\[i23, i24, i25, i26, i27, i28] -> [i25 + i28, i23, i26, i24 + i27]))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w29 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w31 = sreshape @[2,2,2,2,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (stranspose @[5,0,1,4,2,3] (sgather (sappend (sconcrete (sreplicate [1,4,2,2] 0.0)) (sappend (stranspose @[3,0,2,1] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,4,2,2] 0.0)))) (\\[i24, i25, i27, i28] -> [i25 + i28, i24 + i27]))))) (stranspose @[2,3,1,4,5,6,0] w31)))) (stranspose @[1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (sscatter (sdot1In (sreplicate @2 (sreplicate @2 (sreplicate @2 (stranspose @[1,2,3,0] (sfromR (tproject1 u1)))))) (stranspose @[0,2,3,4,5,6,1] w31)) (\\[i32, i33, i34, i35, i36, i37] -> [i34 + i37, i32, i35, i33 + i36])))))))"--conv2dPadded2-  :: forall target r. (ADReady target, GoodScalar r)-  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)-conv2dPadded2 arrK arrA =-  let [nImgs, nCinpA, nAh, nAw] = rshape arrA-      [nCoutK, nCinpK, nKh, nKw] = rshape arrK-      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]-      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case-        [iImg, iCinp, iPh, iPw] ->-               arrA ! [ iImg-                      , iCinp-                      , iPh - fromIntegral (nKh `div` 2)-                      , iPw - fromIntegral (nKw `div` 2) ]-      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA-      shB = [nImgs, nCoutK, nAh, nAw]-      shK1 = [1, nCinp, nKh, nKw]-  in rbuild shB $ \case-    [iImg, iCout, iBh, iBw] ->-      let arrAt = slicezL shK1 arrAPadded [iImg, 0, iBh, iBw]-          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]-      in rdot0 arrAt arrKt-    _ -> error "conv2dPadded2: impossible pattern needlessly required"----- * Non-laborious CNN PP tests---- Convolution differentiated wrt the kernel.-testCNNOPP0cW :: Assertion-testCNNOPP0cW = do-  resetVarCounter-  let ftk = FTKR (7 :$: 5 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000-      var = AstVar varName-      ftk2 = FTKR (5 :$: 5 :$: 5 :$: 5 :$: ZSR) (FTKScalar @Double)-      f = simplifyInline . flip conv2dUnpadded var-      env =-        extendEnv varName (dDnotShared (AstRaw var) (DeltaZero ftk)) emptyEnv-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f env) ftk2-  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u0 -> \\u1 -> rfromS (ssum @125 (stranspose @[4,0,1,2,3] (sreshape @[7,5,7,7,125] (str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u0)) (\\[i81, i83] -> [i81 + i83]))) (\\[i41, i42] -> [i41 + i42])))) * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u1)))))))))"-  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev-    @?= "\\u0 -> \\u1 -> let w43 = str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u0)) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) in rfromS (ssum @125 (stranspose @[4,0,1,2,3] (sreshape @[7,5,7,7,125] (w43 * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u1)))))))))"-  "\\u0 -> " ++ printArtifactPretty artifactRev-    @?= "\\u0 -> \\dret u1 -> let w43 = str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u0)) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) ; w45 = sreshape @[7,5,7,7,5,5,5] (stranspose @[1,2,3,4,0] (sreplicate @125 (sfromR dret))) in rfromS (ssum @7 (str (ssum @7 (str (ssum @7 (w43 * w45))))))"-  "\\u0 -> " ++ printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\u0 -> \\dret u1 -> rfromS (ssum @7 (ssum @7 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u0)) (\\[i98, i100] -> [i98 + i100]))) (\\[i41, i42] -> [i41 + i42]))))) (stranspose @[2,3,1,4,5,6,0] (sreshape @[7,5,7,7,5,5,5] (stranspose @[1,2,3,4,0] (sreplicate @125 (sfromR dret))))))))"---- Convolution differentiated wrt the data.-testCNNOPP0bW :: Assertion-testCNNOPP0bW = do-  resetVarCounter-  let ftk = FTKR (5 :$: 5 :$: 5 :$: 5 :$: ZSR) (FTKScalar @Double)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000-      var = AstVar varName-      ftk2 = FTKR (7 :$: 5 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)-      f = simplifyInline . conv2dUnpadded var-      env =-        extendEnv varName (dDnotShared (AstRaw var) (DeltaZero ftk)) emptyEnv-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f env) ftk2-  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u0 -> \\u1 -> rfromS (ssum @125 (stranspose @[4,0,1,2,3] (sreshape @[7,5,7,7,125] (str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i64, i66] -> [i64 + i66]))) (\\[i41, i42] -> [i41 + i42])))) * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u0)))))))))"-  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev-    @?= "\\u0 -> \\u1 -> let w43 = str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) in rfromS (ssum @125 (stranspose @[4,0,1,2,3] (sreshape @[7,5,7,7,125] (w43 * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u0)))))))))"-  "\\u0 -> " ++ printArtifactPretty artifactRev-    @?= "\\u0 -> \\dret u1 -> let w45 = sreshape @[7,5,7,7,5,5,5] (stranspose @[1,2,3,4,0] (sreplicate @125 (sfromR dret))) in rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[2,5,0,1,3,4] (ssum @5 (str (sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u0))))) * w45)))) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49])))"-  "\\u0 -> " ++ printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\u0 -> \\dret u1 -> rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (sdot1In (stranspose @[3,6,0,2,4,5,1] (sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u0))))))) (stranspose @[3,6,0,2,4,5,1] (sreshape @[7,5,7,7,5,5,5] (stranspose @[1,2,3,4,0] (sreplicate @125 (sfromR dret)))))) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49])))"--testCNNOPP1bW :: Assertion-testCNNOPP1bW = do-  resetVarCounter-  let f :: AstTensor AstMethodLet FullSpan-                     (TKProduct (TKR 4 Double) (TKR 4 Double))-        -> AstTensor AstMethodLet FullSpan-                     (TKR 4 Double)-      f v = simplifyInline $ conv2dUnpadded (tproject1 v) (tproject2 v)-      ftk = FTKProduct (FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) FTKScalar)-                       (FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) FTKScalar)-      (artifactRev, _) =-        revArtifactFromForwardPass-          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (ssum @343 (stranspose @[4,0,1,2,3] (sreshape @[7,7,7,7,343] (str (sreplicate @7 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i81, i83] -> [i81 + i83]))) (\\[i41, i42] -> [i41 + i42])))) * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR (tproject1 u1))))))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w43 = str (sreplicate @7 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) in rfromS (ssum @343 (stranspose @[4,0,1,2,3] (sreshape @[7,7,7,7,343] (w43 * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR (tproject1 u1))))))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w43 = str (sreplicate @7 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) ; w45 = sreshape @[7,7,7,7,7,7,7] (stranspose @[1,2,3,4,0] (sreplicate @343 (sfromR dret))) in tpair (rfromS (ssum @7 (str (ssum @7 (str (ssum @7 (w43 * w45))))))) (rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[2,5,0,1,3,4] (ssum @7 (str (sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR (tproject1 u1)))))) * w45)))) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49]))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [7,7,7,7] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [7,7,7,7] FTKScalar)) ConvSX))) (STKProduct (STKS [7,7,7,7] STKScalar) (STKS [7,7,7,7] STKScalar)) (let w45 = sreshape @[7,7,7,7,7,7,7] (stranspose @[1,2,3,4,0] (sreplicate @343 (sfromR dret))) in tpair (ssum @7 (ssum @7 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @7 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i98, i100] -> [i98 + i100]))) (\\[i41, i42] -> [i41 + i42]))))) (stranspose @[2,3,1,4,5,6,0] w45)))) (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (sdot1In (stranspose @[3,6,0,2,4,5,1] (sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR (tproject1 u1)))))))) (stranspose @[3,6,0,2,4,5,1] w45)) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49]))))"--testCNNOPP4bW :: Assertion-testCNNOPP4bW = do-  resetVarCounter-  let !artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2) (FTKR [7, 7, 7, 7] (FTKScalar @Double))-      !artSimp = simplifyArtifact artifactRev-  let ftk1 = FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)-      ftkDt = FTKR (7 :$: 7 :$: 3 :$: 3 :$: ZSR) (FTKScalar @Double)-      env = extendEnv (artVarDtRev artSimp)-                      (tconcrete ftkDt (treplTarget 7 ftkDt))-            $ extendEnv (artVarDomainRev artSimp)-                        (tconcrete ftk1 (treplTarget 42 ftk1)) emptyEnv-  interpretAstPrimal @Concrete env (artPrimalRev artifactRev)-    @?= interpretAstPrimal @Concrete env (artPrimalRev artSimp)-  interpretAstPrimal @Concrete env (artDerivativeRev artifactRev)-    @?= interpretAstPrimal @Concrete env (artDerivativeRev artSimp)-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i137, i138] -> [2 * i137 + i138]))) (\\[i45, i46] -> [2 * i45 + i46]))) in sgather w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kfromS (smaxIndex (w47 !$ [i48, i49, i50, i51]))]))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) in rfromS (sgather w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kfromS (smaxIndex (w47 !$ [i48, i49, i50, i51]))]))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) in rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (w47 !$ [i53, i54, i55, i56]))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"-  -- The remH comes from the indexing of reshape rule and it looks terrible,-  -- but w42 looks even worse, depending on available primitives,-  -- so the rule is probably fine.-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (sgather (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i159, i160] -> [2 * i159 + i160]))) (\\[i45, i46] -> [2 * i45 + i46])) (\\[i151] -> [remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 16) 3, remH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 4, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 1008) 7, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 144) 7, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 48) 3, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 4) 4])))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"-  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)-    @?= "rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR u52) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) !$ [i53, i54, i55, i56]))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"--testCNNOPP4bD :: Assertion-testCNNOPP4bD = do-  resetVarCounter-  setTotalSharing True-  let !artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2) (FTKR [7, 7, 7, 7] (FTKScalar @Double))-      !artSimp = simplifyArtifact artifactRev-  setTotalSharing False-  let ftk1 = FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)-      ftkDt = FTKR (7 :$: 7 :$: 3 :$: 3 :$: ZSR) (FTKScalar @Double)-      env = extendEnv (artVarDtRev artSimp)-                      (tconcrete ftkDt (treplTarget 7 ftkDt))-            $ extendEnv (artVarDomainRev artSimp)-                        (tconcrete ftk1 (treplTarget 42 ftk1)) emptyEnv-  interpretAstPrimal @Concrete env (artPrimalRev artifactRev)-    @?= interpretAstPrimal @Concrete env (artPrimalRev artSimp)-  interpretAstPrimal @Concrete env (artDerivativeRev artifactRev)-    @?= interpretAstPrimal @Concrete env (artDerivativeRev artSimp)-  printArtifactPrimalPretty artSimp-    @?= "\\u1 -> rfromS (let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i98, i99] -> [2 * i98 + i99]))) (\\[i45, i46] -> [2 * i45 + i46]))) in sgather w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kfromS (smaxIndex (w47 !$ [i48, i49, i50, i51]))]))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) in rfromS (sgather w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kfromS (smaxIndex (w47 !$ [i48, i49, i50, i51]))]))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) in rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (w47 !$ [i53, i54, i55, i56]))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"-  printArtifactPretty artSimp-    @?= "\\dret u1 -> rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (sgather (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i81, i82] -> [2 * i81 + i82]))) (\\[i45, i46] -> [2 * i45 + i46])) (\\[i73] -> [remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 16) 3, remH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 4, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 1008) 7, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 144) 7, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 48) 3, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 4) 4])))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"-  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)-    @?= "rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR u52) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) !$ [i53, i54, i55, i56]))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"--testCNNOPP5aW :: Assertion-testCNNOPP5aW = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let t49 = sreshape @[2,7,16] (ssum @98 (stranspose @[4,1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i109, i111] -> [i109 + i111]))) (\\[i46, i47] -> [i46 + i47])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))))) in stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (sgather t49 (\\[i50, i51] -> [i50, i51, kfromS (smaxIndex (t49 !$ [i50, i51]))]))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w48 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i44, i45] -> [i44 + i45]))) (\\[i46, i47] -> [i46 + i47])))))) ; t49 = sreshape @[2,7,16] (ssum @98 (stranspose @[4,1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (w48 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))))) in rfromS (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (sgather t49 (\\[i50, i51] -> [i50, i51, kfromS (smaxIndex (t49 !$ [i50, i51]))]))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w48 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i44, i45] -> [i44 + i45]))) (\\[i46, i47] -> [i46 + i47])))))) ; t49 = sreshape @[2,7,16] (ssum @98 (stranspose @[4,1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (w48 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))))) ; w55 = stranspose @[4,1,2,3,0] (sreplicate @98 (sreshape @[2,7,4,4] (sscatter (ssum @1 (stranspose @[2,0,1] (ssum @1 (stranspose @[2,0,1] (sfromR dret))))) (\\[i53, i54] -> [i53, i54, kfromS (smaxIndex (t49 !$ [i53, i54]))])))) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w48 * sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) w55))))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (let w48 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i155, i157] -> [i155 + i157]))) (\\[i46, i47] -> [i46 + i47]) in ssum @2 (ssum @2 (sdot1In (stranspose @[1,6,0,4,5,3,2] (sreplicate @7 (stranspose @[3,2,1,4,5,0] w48))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[4,1,2,3,0] (sreplicate @98 (sreshape @[2,7,4,4] (sscatter (stranspose @[2,3,0,1] (sfromR dret) !$ [0, 0]) (\\[i53, i54] -> [i53, i54, kfromS (smaxIndex (ssum @98 (str (sgather (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w48)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))) (\\[i129] -> [remH ((112 * i53 + 16 * i54) + i129) 4, remH (quotH ((112 * i53 + 16 * i54) + i129) 112) 2, remH (quotH ((112 * i53 + 16 * i54) + i129) 16) 7, remH (quotH ((112 * i53 + 16 * i54) + i129) 4) 4])))))])))))))))) !$ [0]))))"-  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)-    @?= "rfromS (let w48 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i44, i45] -> [i44 + i45]))) (\\[i46, i47] -> [i46 + i47]) in ssum @2 (ssum @2 (sdot1In (stranspose @[4,2,3,0,5,6,7,1] (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w48)))) !$ [0]) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[4,1,2,3,0] (sreplicate @98 (sreshape @[2,7,4,4] (sscatter (stranspose @[2,3,0,1] (sfromR u52) !$ [0, 0]) (\\[i53, i54] -> [i53, i54, kfromS (smaxIndex (sreshape @[2,7,16] (ssum @98 (stranspose @[4,1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w48)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))))) !$ [i53, i54]))])))))))))) !$ [0]))))"--testCNNOPP5bW :: Assertion-testCNNOPP5bW = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . relu) (FTKR [7, 2, 7, 7] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let w75 = sreshape @[7,2,3,3,16] (stranspose @[4,5,0,1,2,3] (sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i65, i67]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i66, i68])])) * stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i69, i71]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i70, i72])]))) in sgather w75 (\\[i76, i77, i78, i79] -> [i76, i77, i78, i79, kfromS (smaxIndex (w75 !$ [i76, i77, i78, i79]))]))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let m59 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; m60 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; w73 = stranspose @[4,5,0,1,2,3] (sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (m59 !$ [i65, i67]), kfromS (m60 !$ [i66, i68])])) ; w74 = stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (m59 !$ [i69, i71]), kfromS (m60 !$ [i70, i72])])) ; w75 = sreshape @[7,2,3,3,16] (w73 * w74) in rfromS (sgather w75 (\\[i76, i77, i78, i79] -> [i76, i77, i78, i79, kfromS (smaxIndex (w75 !$ [i76, i77, i78, i79]))]))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let m59 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; m60 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; w73 = stranspose @[4,5,0,1,2,3] (sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (m59 !$ [i65, i67]), kfromS (m60 !$ [i66, i68])])) ; w74 = stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (m59 !$ [i69, i71]), kfromS (m60 !$ [i70, i72])])) ; w75 = sreshape @[7,2,3,3,16] (w73 * w74) in rfromS (stranspose @[2,3,0,1] (sscatter (stranspose @[2,3,4,5,0,1] (w73 * sreshape @[7,2,3,3,4,4] (sscatter (sfromR dret) (\\[i81, i82, i83, i84] -> [i81, i82, i83, i84, kfromS (smaxIndex (w75 !$ [i81, i82, i83, i84]))])))) (\\[i85, i86, i87, i88] -> [kfromS (m59 !$ [i85, i87]), kfromS (m60 !$ [i86, i88])])))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (let w73 = sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i65, i67]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i66, i68])]) in stranspose @[2,3,0,1] (sscatter (w73 * stranspose @[2,3,4,5,0,1] (sreshape @[7,2,3,3,4,4] (sscatter (sfromR dret) (\\[i81, i82, i83, i84] -> [i81, i82, i83, i84, kfromS (smaxIndex (sgather (stranspose @[4,5,0,1,2,3] w73 * stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i69, i71]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i70, i72])]))) (\\[i118] -> [remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 288) 7, remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 144) 2, remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 48) 3, remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 16) 3, remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 4) 4, remH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 4])))])))) (\\[i85, i86, i87, i88] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i85, i87]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i86, i88])])))"-  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)-    @?= "rfromS (let w73 = sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i65, i67]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i66, i68])]) in stranspose @[2,3,0,1] (sscatter (w73 * stranspose @[2,3,4,5,0,1] (sreshape @[7,2,3,3,4,4] (sscatter (sfromR u80) (\\[i81, i82, i83, i84] -> [i81, i82, i83, i84, kfromS (smaxIndex (sreshape @[7,2,3,3,16] (stranspose @[4,5,0,1,2,3] w73 * stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i69, i71]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i70, i72])]))) !$ [i81, i82, i83, i84]))])))) (\\[i85, i86, i87, i88] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i85, i87]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i86, i88])])))"--testCNNOPP5cW :: Assertion-testCNNOPP5cW = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (relu . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let u45 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i114, i116] -> [i114 + i116]))) (\\[i42, i43] -> [i42 + i43])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (u45 !$ [i46, i47, i48, i49])) 0 1]) * u45)"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w44 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i40, i41] -> [i40 + i41]))) (\\[i42, i43] -> [i42 + i43])))))) ; u45 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (w44 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u50 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (u45 !$ [i46, i47, i48, i49])) 0 1]) in rfromS (u50 * u45)"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w44 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i40, i41] -> [i40 + i41]))) (\\[i42, i43] -> [i42 + i43])))))) ; u45 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (w44 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u50 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (u45 !$ [i46, i47, i48, i49])) 0 1]) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w44 * sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (u50 * sfromR dret)))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (let w44 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i169, i171] -> [i169 + i171]))) (\\[i42, i43] -> [i42 + i43]) in ssum @2 (ssum @2 (sdot1In (stranspose @[1,6,0,4,5,3,2] (sreplicate @7 (stranspose @[3,2,1,4,5,0] w44))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (ssum0 (sgather (stranspose @[6,1,0,2,4,5,3] (sreplicate @7 (stranspose @[3,2,1,4,5,0] w44)) * sreplicate @2 (sreplicate @2 (str (sreplicate @2 (sfromR u1))))) (\\[i132] -> [remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 98) 2, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 196) 2, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 392) 7, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 2744) 2, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 49) 2, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 7) 7, remH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 7])))) 0 1]) * sfromR dret)))) !$ [0]))))"-  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)-    @?= "rfromS (let w44 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i40, i41] -> [i40 + i41]))) (\\[i42, i43] -> [i42 + i43]) in ssum @2 (ssum @2 (sdot1In (stranspose @[4,2,3,0,5,6,7,1] (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w44)))) !$ [0]) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (ssum0 (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w44)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))) !$ [i46, i47, i48, i49]))) 0 1]) * sfromR u51)))) !$ [0]))))"--testCNNOPP5dW :: Assertion-testCNNOPP5dW = do-  resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . relu . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i130, i132] -> [i130 + i132]))) (\\[i64, i65] -> [i64 + i65])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; t74 = sreshape @[2,7,16] (stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) * stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))) in stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (sgather t74 (\\[i75, i76] -> [i75, i76, kfromS (smaxIndex (t74 !$ [i75, i76]))]))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w66 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i62, i63] -> [i62 + i63]))) (\\[i64, i65] -> [i64 + i65])))))) ; u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (w66 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u72 = stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; u73 = stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; t74 = sreshape @[2,7,16] (u72 * u73) in rfromS (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (sgather t74 (\\[i75, i76] -> [i75, i76, kfromS (smaxIndex (t74 !$ [i75, i76]))]))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w66 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i62, i63] -> [i62 + i63]))) (\\[i64, i65] -> [i64 + i65])))))) ; u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (w66 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u72 = stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; u73 = stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; t74 = sreshape @[2,7,16] (u72 * u73) ; u80 = stranspose @[3,0,1,2] (u72 * sreshape @[2,7,4,4] (sscatter (ssum @1 (stranspose @[2,0,1] (ssum @1 (stranspose @[2,0,1] (sfromR dret))))) (\\[i78, i79] -> [i78, i79, kfromS (smaxIndex (t74 !$ [i78, i79]))]))) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w66 * sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) u80))))))))))))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (let w66 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i172, i174] -> [i172 + i174]))) (\\[i64, i65] -> [i64 + i65]) ; u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w66)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u72 = sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)) in ssum @2 (ssum @2 (sdot1In (stranspose @[1,6,0,4,5,3,2] (sreplicate @7 (stranspose @[3,2,1,4,5,0] w66))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) u72))) * stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,0,1,2] (sreshape @[2,7,4,4] (sscatter (stranspose @[2,3,0,1] (sfromR dret) !$ [0, 0]) (\\[i78, i79] -> [i78, i79, kfromS (smaxIndex (sgather (stranspose @[1,2,3,0] u72 * stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))) (\\[i146] -> [remH (quotH ((112 * i78 + 16 * i79) + i146) 112) 2, remH (quotH ((112 * i78 + 16 * i79) + i146) 16) 7, remH (quotH ((112 * i78 + 16 * i79) + i146) 4) 4, remH ((112 * i78 + 16 * i79) + i146) 4])))]))))))))))) !$ [0]))))"-  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)-    @?= "rfromS (let w66 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i62, i63] -> [i62 + i63]))) (\\[i64, i65] -> [i64 + i65]) ; u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w66)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u72 = sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)) in ssum @2 (ssum @2 (sdot1In (stranspose @[4,2,3,0,5,6,7,1] (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w66)))) !$ [0]) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) u72))) * stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,0,1,2] (sreshape @[2,7,4,4] (sscatter (stranspose @[2,3,0,1] (sfromR u77) !$ [0, 0]) (\\[i78, i79] -> [i78, i79, kfromS (smaxIndex (sreshape @[2,7,16] (stranspose @[1,2,3,0] u72 * stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))) !$ [i78, i79]))]))))))))))) !$ [0]))))"+import Data.Array.Nested.Shaped.Shape++import HordeAd+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId (resetVarCounter)+import HordeAd.Core.AstInterpret+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersAst+import HordeAd.Core.Delta+import HordeAd.Core.Ops+import HordeAd.Core.OpsAst++import TestConvQuickCheck (conv2dSame_dKrn)++import CrossTesting+import EqEpsilon++testTrees :: [TestTree]+testTrees =+  [ testCase "KonstG0Rev" testKonstG0Rev+  , testCase "KonstG0Tiny1" testKonstG0Tiny1+  , testCase "KonstG0TinyS" testKonstG0TinyS+  , testCase "KonstG0TinyA" testKonstG0TinyA+  , testCase "KonstG0LittleA" testKonstG0LittleA+  , testCase "Replicate0Rev" testReplicate0Rev+  , testCase "Replicate0Tiny1" testReplicate0Tiny1+  , testCase "Replicate0TinyS" testReplicate0TinyS+  , testCase "Replicate0TinyA" testReplicate0TinyA+  , testCase "Replicate0LittleA" testReplicate0LittleA+  , testCase "Konst5LittleB" testKonst5LittleB+  , testCase "Konst5LittleC" testKonst5LittleC+  , testCase "Konst5BigB" testKonst5BigB+  , testCase "KonstNotBigB" testKonstNotBigB+  , testCase "Konst5BigC" testKonst5BigC+  , testCase "KonstNotBigC" testKonstNotBigC+  , testCase "Konst5LittleB128b" testKonst5LittleB128b+  , testCase "Konst5LittleC128b" testKonst5LittleC128b+  , testCase "Konst5BigB128b" testKonst5BigB128b+  , testCase "KonstNotBigB128b" testKonstNotBigB128b+  , testCase "Konst5BigC128b" testKonst5BigC128b+  , testCase "KonstNotBigC128b" testKonstNotBigC128b+  , testCase "Konst5LittleB128c" testKonst5LittleB128c+  , testCase "Konst5LittleC128c" testKonst5LittleC128c+  , testCase "Konst5BigB128c" testKonst5BigB128c+  , testCase "KonstNotBigB128c" testKonstNotBigB128c+  , testCase "Konst5BigC128c" testKonst5BigC128c+  , testCase "KonstNotBigC128c" testKonstNotBigC128c+--  , testCase "Konst5LittleB128bc" testKonst5LittleB128bc+--  , testCase "Konst5LittleC128bc" testKonst5LittleC128bc+--  , testCase "Konst5BigB128bc" testKonst5BigB128bc+--  , testCase "KonstNotBigB128cb" testKonstNotBigB128cb+--  , testCase "Konst5BigC128cb" testKonst5BigC128cb+--  , testCase "KonstNotBigC128cb" testKonstNotBigC128cb+  , testCase "Replicate0RevLaborious" testReplicate0RevLaborious+  , testCase "Replicate0Tiny1Laborious" testReplicate0Tiny1Laborious+  , testCase "Replicate0TinySLaborious" testReplicate0TinySLaborious+  , testCase "Replicate0TinyALaborious" testReplicate0TinyALaborious+  , testCase "Replicate0LittleALaborious" testReplicate0LittleALaborious+  , testCase "Konst5LittleBLaborious" testKonst5LittleBLaborious+  , testCase "Konst5LittleCLaborious" testKonst5LittleCLaborious+  , testCase "Konst5BigBLaborious" testKonst5BigBLaborious+  , testCase "KonstNotBigBLaborious" testKonstNotBigBLaborious+  , testCase "Konst5BigCLaborious" testKonst5BigCLaborious+  , testCase "KonstNotBigCLaborious" testKonstNotBigCLaborious+  , testCase "Konst5LittleBLaborious128b" testKonst5LittleBLaborious128b+  , testCase "Konst5LittleCLaborious128b" testKonst5LittleCLaborious128b+--  , testCase "Konst5BigBLaborious128b" testKonst5BigBLaborious128b+--  , testCase "KonstNotBigBLaborious128b" testKonstNotBigBLaborious128b+--  , testCase "Konst5BigCLaborious128b" testKonst5BigCLaborious128b+--  , testCase "KonstNotBigCLaborious128b" testKonstNotBigCLaborious128b+  , testCase "Konst5LittleBLaborious128c" testKonst5LittleBLaborious128c+  , testCase "Konst5LittleCLaborious128c" testKonst5LittleCLaborious128c+--  , testCase "Konst5BigBLaborious128c" testKonst5BigBLaborious128c+--  , testCase "KonstNotBigBLaborious128c" testKonstNotBigBLaborious128c+--  , testCase "Konst5BigCLaborious128c" testKonst5BigCLaborious128c+--  , testCase "KonstNotBigCLaborious128c" testKonstNotBigCLaborious128c+--  , testCase "Konst5LittleBLaborious128bc" testKonst5LittleBLaborious128bc+--  , testCase "Konst5LittleCLaborious128bc" testKonst5LittleCLaborious128bc+--  , testCase "Konst5BigBLaborious128bc" testKonst5BigBLaborious128bc+--  , testCase "KonstNotBigBLaborious128cb" testKonstNotBigBLaborious128cb+--  , testCase "Konst5BigCLaborious128cb" testKonst5BigCLaborious128cb+--  , testCase "KonstNotBigCLaborious128cb" testKonstNotBigCLaborious128cb+  , testCase "Replicate0RevPadded" testReplicate0RevPadded+  , testCase "Replicate0Tiny1Padded" testReplicate0Tiny1Padded+  , testCase "Replicate0TinySPadded" testReplicate0TinySPadded+  , testCase "Replicate0TinyAPadded" testReplicate0TinyAPadded+  , testCase "Replicate0LittleAPadded" testReplicate0LittleAPadded+  , testCase "Konst5LittleBPadded" testKonst5LittleBPadded+  , testCase "Konst5LittleCPadded" testKonst5LittleCPadded+  , testCase "Konst5BigBPadded" testKonst5BigBPadded+  , testCase "KonstNotBigBPadded" testKonstNotBigBPadded+  , testCase "Konst5BigCPadded" testKonst5BigCPadded+  , testCase "KonstNotBigCPadded" testKonstNotBigCPadded+  , testCase "Konst5LittleBPadded128b" testKonst5LittleBPadded128b+  , testCase "Konst5LittleCPadded128b" testKonst5LittleCPadded128b+--  , testCase "Konst5BigBPadded128b" testKonst5BigBPadded128b+--  , testCase "KonstNotBigBPadded128b" testKonstNotBigBPadded128b+--  , testCase "Konst5BigCPadded128b" testKonst5BigCPadded128b+--  , testCase "KonstNotBigCPadded128b" testKonstNotBigCPadded128b+--  , testCase "Konst5LittleBPadded128c" testKonst5LittleBPadded128c+--  , testCase "Konst5LittleCPadded128c" testKonst5LittleCPadded128c+--  , testCase "Konst5BigBPadded128c" testKonst5BigBPadded128c+--  , testCase "KonstNotBigBPadded128c" testKonstNotBigBPadded128c+--  , testCase "Konst5BigCPadded128c" testKonst5BigCPadded128c+--  , testCase "KonstNotBigCPadded128c" testKonstNotBigCPadded128c+--  , testCase "Konst5LittleBPadded128bc" testKonst5LittleBPadded128bc+--  , testCase "Konst5LittleCPadded128bc" testKonst5LittleCPadded128bc+--  , testCase "Konst5BigBPadded128bc" testKonst5BigBPadded128bc+--  , testCase "KonstNotBigBPadded128cb" testKonstNotBigBPadded128cb+--  , testCase "Konst5BigCPadded128cb" testKonst5BigCPadded128cb+--  , testCase "KonstNotBigCPadded128cb" testKonstNotBigCPadded128cb+  , testCase "disparityKonst" test_disparityKonst+  , testCase "disparityKonst2" test_disparityKonst2+  , testCase "disparitySmall" test_disparitySmall+  , testCase "ConvTomsSliceRev" testTomsSliceRev+  , testCase "ConvTomsSliceCRev" testTomsSliceCRev+  , testCase "ConvTomsSlice" testTomsSlice+  , testCase "ConvTomsSlicePP" testTomsSlicePP+  , testCase "minimizedCNNOPP0c" testCNNOPP0c+  , testCase "minimizedCNNOPP0b" testCNNOPP0b+  , testCase "minimizedCNNOPP1e" testCNNOPP1e+  , testCase "minimizedCNNOPP2" testCNNOPP2+  , testCase "minimizedCNNOPP2b" testCNNOPP2b+  , testCase "minimizedCNNOPP3" testCNNOPP3+  , testCase "minimizedCNNOPP3b" testCNNOPP3b+  , testCase "minimizedCNNOPP4" testCNNOPP4+  , testCase "minimizedCNNOPP4b" testCNNOPP4b+  , testCase "minimizedCNNOPP5" testCNNOPP5+  , testCase "minimizedCNNOPP5b" testCNNOPP5b+  , testCase "minimizedCNNOPP6" testCNNOPP6+  , testCase "minimizedCNNOPP6b" testCNNOPP6b+  , testCase "minimizedCNNOPP7" testCNNOPP7+  , testCase "minimizedCNNOPP7b" testCNNOPP7b+  , testCase "minimizedPaddedCNNOPP0c" testPaddedCNNOPP0c+  , testCase "minimizedPaddedCNNOPP0b" testPaddedCNNOPP0b+  , testCase "minimizedPaddedCNNOPP1e" testPaddedCNNOPP1e+  , testCase "minimizedPaddedCNNOPPLet" testPaddedCNNOPPLet+  , testCase "minimizedPaddedCNNOPPLet2" testPaddedCNNOPPLet2+  , testCase "minimizedPaddedCNNOPP2" testPaddedCNNOPP2+  , testCase "minimizedSameCNNOPPKrnHandwritten" testSameCNNOPPKrnHandwritten+  , testCase "minimizedSameCNNOPP0cW" testSameCNNOPP0cW+  , testCase "minimizedSameCNNOPP0bW" testSameCNNOPP0bW+  , testCase "minimizedSameCNNOPP1bW" testSameCNNOPP1bW+  , testCase "minimizedShrinkingCNNOPP0cW" testShrinkingCNNOPP0cW+  , testCase "minimizedShrinkingCNNOPP0bW" testShrinkingCNNOPP0bW+  , testCase "minimizedShrinkingCNNOPP1bW" testShrinkingCNNOPP1bW+  , testCase "minimizedPaddedCNNOPP0cW" testPaddedCNNOPP0cW+  , testCase "minimizedPaddedCNNOPP0bW" testPaddedCNNOPP0bW+  , testCase "minimizedPaddedCNNOPP1bW" testPaddedCNNOPP1bW+  , testCase "minimizedCNNOPP4bW" testCNNOPP4bW+  , testCase "minimizedCNNOPP4bD" testCNNOPP4bD+  , testCase "minimizedCNNOPP5aW" testCNNOPP5aW+  , testCase "minimizedCNNOPP5bW" testCNNOPP5bW+  , testCase "minimizedCNNOPP5cW" testCNNOPP5cW+  , testCase "minimizedCNNOPP5dW" testCNNOPP5dW+  ]++-- The examples reproduced and transformed in this file are borrowed+-- from https://github.com/benl23x5/adops.+-- Here they are defined using ranked tensors.++-- * A non-laborious version (depends on indexing OOB giving 0 consistently)++conv2d1+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2d1 = conv2dSame $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]++conv2dA+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dA = conv2dSame $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]++conv2dB+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dB = conv2dSame (rconcrete $ unConcrete t16b)++testKonstG0Rev :: Assertion+testKonstG0Rev =+  assertEqualUpToEpsilon 1e-4+    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (grad (rsum0 @4 @Double . conv2dB) (rrepl [2, 2, 2, 2] 0))++testKonstG0Tiny1 :: Assertion+testKonstG0Tiny1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [-0.2])+    (rev' @Double @4 conv2d1 (rrepl [1, 1, 1, 1] 0))++testKonstG0TinyS :: Assertion+testKonstG0TinyS =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [582665.99432])+    (rev' @Double @4+          (conv2dSame $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))+          (ringestData [1, 1, 1, 1] [0]))++testKonstG0TinyA :: Assertion+testKonstG0TinyA =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 2, 1, 1] [-0.2,25.0003])+    (rev' @Double @4 conv2dA (rrepl [1, 2, 1, 1] 0))++testKonstG0LittleA :: Assertion+testKonstG0LittleA =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [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])+    (rev' @Double @4 conv2dA (rrepl [2, 2, 2, 2] 0))++conv2dC+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dC = flip conv2dSame (rconcrete $ unConcrete t16b)++conv2dB128b+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dB128b = conv2dSame (rconcrete $ unConcrete t128b)++conv2dC128b+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dC128b = flip conv2dSame (rconcrete $ unConcrete t128b)++conv2dB128c+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dB128c = conv2dSame (rconcrete $ unConcrete t128c)++conv2dC128c+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dC128c = flip conv2dSame (rconcrete $ unConcrete t128c)++testReplicate0Rev :: Assertion+testReplicate0Rev =+  assertEqualUpToEpsilon 1e-4+    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (grad (rsum0 @4 @Double . conv2dB) (rrepl [2, 2, 2, 2] 0))++testReplicate0Tiny1 :: Assertion+testReplicate0Tiny1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [-0.2])+    (rev' @Double @4 conv2d1 (rrepl [1, 1, 1, 1] 0))++testReplicate0TinyS :: Assertion+testReplicate0TinyS =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [582665.99432])+    (rev' @Double @4+          (conv2dSame $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))+          (ringestData [1, 1, 1, 1] [0]))++testReplicate0TinyA :: Assertion+testReplicate0TinyA =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 2, 1, 1] [-0.2,25.0003])+    (rev' @Double @4 conv2dA (rrepl [1, 2, 1, 1] 0))++testReplicate0LittleA :: Assertion+testReplicate0LittleA =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [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])+    (rev' @Double @4 conv2dA (rrepl [2, 2, 2, 2] 0))++-- with data t16++testKonst5LittleB :: Assertion+testKonst5LittleB =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dB (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleC :: Assertion+testKonst5LittleC =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])+    (rev' @Double @4 conv2dC (rreplicate0N [2, 2, 2, 2] 5))++testKonst5BigB :: Assertion+testKonst5BigB =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dB (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigB :: Assertion+testKonstNotBigB =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dB+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigC :: Assertion+testKonst5BigC =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dC (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigC :: Assertion+testKonstNotBigC =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dC+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++-- with data t128b++testKonst5LittleB128b :: Assertion+testKonst5LittleB128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004])+    (rev' @Double @4 conv2dB128b (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleC128b :: Assertion+testKonst5LittleC128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987])+    (rev' @Double @4 conv2dC128b (rreplicate0N [2, 2, 2, 2] 5))++testKonst5BigB128b :: Assertion+testKonst5BigB128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dB128b (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigB128b :: Assertion+testKonstNotBigB128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dB128b+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigC128b :: Assertion+testKonst5BigC128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dC128b (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigC128b :: Assertion+testKonstNotBigC128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dC128b+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++-- with data t128c++testKonst5LittleB128c :: Assertion+testKonst5LittleB128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4])+    (rev' @Double @4 conv2dB128c (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleC128c :: Assertion+testKonst5LittleC128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992])+    (rev' @Double @4 conv2dC128c (rreplicate0N [2, 2, 2, 2] 5))++testKonst5BigB128c :: Assertion+testKonst5BigB128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dB128c (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigB128c :: Assertion+testKonstNotBigB128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dB128c+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigC128c :: Assertion+testKonst5BigC128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dC128c (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigC128c :: Assertion+testKonstNotBigC128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dC128c+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++-- with data t128b and t128c+{-+testKonst5LittleB128bc :: Assertion+testKonst5LittleB128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])+    (rev' @Double @4 conv2dB128b t128c)++testKonst5LittleC128bc :: Assertion+testKonst5LittleC128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dC128b t128c)++testKonst5BigB128bc :: Assertion+testKonst5BigB128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])+    (rev' @Double @4 conv2dB128b t128c)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigB128cb :: Assertion+testKonstNotBigB128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002])+    (rev' @Double @4 conv2dB128c t128b)++testKonst5BigC128cb :: Assertion+testKonst5BigC128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])+    (rev' @Double @4 conv2dC128c t128b)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigC128cb :: Assertion+testKonstNotBigC128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])+    (rev' @Double @4 conv2dC128c t128b)+-}+++-- * A laborious version (meaning, out of bounds indexing is handled explicitly)++-- | Unpadded full convolution,+--   where the output size is the same as the input size.+--+-- It guards the out of bounds indexing behind a conditional+-- to prevent changed values after vectorization,+-- but the guarding is no longer needed, so this is only for testing.+--+-- BTW, the indexing lower bounds in the code are spurious,+-- so they get simplified away in the resulting AST program.+conv2dSameL+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dSameL arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezL shK1 arrA [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rfromK $ rdot0 arrAt arrKt+    _ -> error "conv2dSameL: impossible pattern needlessly required"++-- | Slice a section out of a tensor,+--   given a base offset and shape of the section.+--+--   If the slice extends out side the source array then the corresponding+--   elements are set to zero.+slicezL+  :: (ADReady target, NumScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicezL shOut d ixBase =+  rbuild shOut $ \ixResult -> indexz0L d (ixrZipWith (+) ixBase ixResult)++-- | Retrieve the element at the given index,+--   returning zero for out of range indices.+--+-- Warning: this uses ix twice and within0 again uses it twice,+-- so this variant without tlet should be used only when it's known+-- that ix is of small constant size (e.g., if it contains conditionals+-- that compare big tensors or their minimal elements, it likely is not,+-- unless the tensors are under tlet and only variables representing them+-- are used).+indexz0L+  :: forall target r n. (ADReady target, NumScalar r, KnownNat n)+  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)+indexz0L d ix = ifH (within0 @target (rshape @target d) ix) (d ! ix) (rscalar 0)++-- | Given an index and shape, check if the index is fully within the shape.+-- Note that @ix@ is used twice, so should be shared outside.+within0+  :: forall target n. (ADReady target, KnownNat n)+  => IShR n -> IxROf target n -> BoolOf target+within0 sh ix =+  let within :: IntOf target -> IntOf target -> BoolOf target+      within i dim = 0 <=. i &&* dim >. i+  in foldr (&&*) true+     $ zipWith within (toList ix) (map fromIntegral $ toList sh)++conv2d1Laborious+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2d1Laborious = conv2dSameL $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]++conv2dALaborious+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dALaborious =+  conv2dSameL $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]++conv2dBLaborious+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBLaborious = conv2dSameL (rconcrete $ unConcrete t16b)++conv2dCLaborious+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCLaborious = flip conv2dSameL (rconcrete $ unConcrete t16b)++conv2dBLaborious128b+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBLaborious128b = conv2dSameL (rconcrete $ unConcrete t128b)++conv2dCLaborious128b+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCLaborious128b = flip conv2dSameL (rconcrete $ unConcrete t128b)++conv2dBLaborious128c+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBLaborious128c = conv2dSameL (rconcrete $ unConcrete t128c)++conv2dCLaborious128c+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCLaborious128c = flip conv2dSameL (rconcrete $ unConcrete t128c)++testReplicate0RevLaborious :: Assertion+testReplicate0RevLaborious =+  assertEqualUpToEpsilon 1e-4+    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (grad (rsum0 @4 @Double . conv2dBLaborious) (rrepl [2, 2, 2, 2] 0))++testReplicate0Tiny1Laborious :: Assertion+testReplicate0Tiny1Laborious =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [-0.2])+    (rev' @Double @4 conv2d1Laborious (rrepl [1, 1, 1, 1] 0))++testReplicate0TinySLaborious :: Assertion+testReplicate0TinySLaborious =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [582665.99432])+    (rev' @Double @4+          (conv2dSameL $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))+          (ringestData [1, 1, 1, 1] [0]))++testReplicate0TinyALaborious :: Assertion+testReplicate0TinyALaborious =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 2, 1, 1] [-0.2,25.0003])+    (rev' @Double @4 conv2dALaborious (rrepl [1, 2, 1, 1] 0))++testReplicate0LittleALaborious :: Assertion+testReplicate0LittleALaborious =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [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])+    (rev' @Double @4 conv2dALaborious (rrepl [2, 2, 2, 2] 0))++-- with data t16++testKonst5LittleBLaborious :: Assertion+testKonst5LittleBLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dBLaborious (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleCLaborious :: Assertion+testKonst5LittleCLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])+    (rev' @Double @4 conv2dCLaborious (rreplicate0N [2, 2, 2, 2] 5))++testKonst5BigBLaborious :: Assertion+testKonst5BigBLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dBLaborious (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBLaborious :: Assertion+testKonstNotBigBLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dBLaborious+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigCLaborious :: Assertion+testKonst5BigCLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dCLaborious (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCLaborious :: Assertion+testKonstNotBigCLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dCLaborious+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++-- with data t128b++testKonst5LittleBLaborious128b :: Assertion+testKonst5LittleBLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004])+    (rev' @Double @4 conv2dBLaborious128b (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleCLaborious128b :: Assertion+testKonst5LittleCLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987])+    (rev' @Double @4 conv2dCLaborious128b (rreplicate0N [2, 2, 2, 2] 5))++{-+testKonst5BigBLaborious128b :: Assertion+testKonst5BigBLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dBLaborious128b (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBLaborious128b :: Assertion+testKonstNotBigBLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dBLaborious128b+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigCLaborious128b :: Assertion+testKonst5BigCLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCLaborious128b (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCLaborious128b :: Assertion+testKonstNotBigCLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCLaborious128b+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))+-}++-- with data t128c++testKonst5LittleBLaborious128c :: Assertion+testKonst5LittleBLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4])+    (rev' @Double @4 conv2dBLaborious128c (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleCLaborious128c :: Assertion+testKonst5LittleCLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992])+    (rev' @Double @4 conv2dCLaborious128c (rreplicate0N [2, 2, 2, 2] 5))++{-+testKonst5BigBLaborious128c :: Assertion+testKonst5BigBLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dBLaborious128c (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBLaborious128c :: Assertion+testKonstNotBigBLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dBLaborious128c+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigCLaborious128c :: Assertion+testKonst5BigCLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dCLaborious128c (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCLaborious128c :: Assertion+testKonstNotBigCLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dCLaborious128c+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))+-}++-- with data t128b and t128c+{-+testKonst5LittleBLaborious128bc :: Assertion+testKonst5LittleBLaborious128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])+    (rev' @Double @4 conv2dBLaborious128b t128c)++testKonst5LittleCLaborious128bc :: Assertion+testKonst5LittleCLaborious128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dCLaborious128b t128c)++testKonst5BigBLaborious128bc :: Assertion+testKonst5BigBLaborious128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])+    (rev' @Double @4 conv2dBLaborious128b t128c)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBLaborious128cb :: Assertion+testKonstNotBigBLaborious128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002])+    (rev' @Double @4 conv2dBLaborious128c t128b)++testKonst5BigCLaborious128cb :: Assertion+testKonst5BigCLaborious128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])+    (rev' @Double @4 conv2dCLaborious128c t128b)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCLaborious128cb :: Assertion+testKonstNotBigCLaborious128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])+    (rev' @Double @4 conv2dCLaborious128c t128b)+-}+++-- * A padded version (out of bounds indexing is not possible)++-- | Full convolution with just enough extra external zero padding+--   to ensure that the output size is the same as the input size+--   and all input points are read the same number of times.+--+--   The same result could be accomplished by tweaking indexes slightly+--   in conv2dSame, but here additionally all bounds checks in the code+--   are spurious and will be simplified away in the resulting AST program.+conv2dPaddedB+  :: forall target r. (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPaddedB arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]+      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case+        [iImg, iCinp, iPh, iPw] ->+          ifH (iPh <. fromIntegral (nKh `div` 2)+               ||* iPw <. fromIntegral (nKw `div` 2)+               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)+               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))+              (rscalar 0)+              (arrA ! [ iImg+                      , iCinp+                      , iPh - fromIntegral (nKh `div` 2)+                      , iPw - fromIntegral (nKw `div` 2) ])+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezL shK1 arrAPadded [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rfromK $ rdot0 arrAt arrKt+    _ -> error "conv2dPaddedB: impossible pattern needlessly required"++conv2d1Padded+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2d1Padded = conv2dPaddedB $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]++conv2dAPadded+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dAPadded =+  conv2dPaddedB $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]++conv2dBPadded+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBPadded = conv2dPaddedB (rconcrete $ unConcrete t16b)++conv2dCPadded+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCPadded = flip conv2dPaddedB (rconcrete $ unConcrete t16b)++conv2dBPadded128b+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBPadded128b = conv2dPaddedB (rconcrete $ unConcrete t128b)++conv2dCPadded128b+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCPadded128b = flip conv2dPaddedB (rconcrete $ unConcrete t128b)++_conv2dBPadded128c+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+_conv2dBPadded128c = conv2dPaddedB (rconcrete $ unConcrete t128c)++_conv2dCPadded128c+  :: (ADReady target, NumScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+_conv2dCPadded128c = flip conv2dPaddedB (rconcrete $ unConcrete t128c)++testReplicate0RevPadded :: Assertion+testReplicate0RevPadded =+  assertEqualUpToEpsilon 1e-4+    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])+    (grad (rsum0 @4 @Double . conv2dBPadded) (rrepl [2, 2, 2, 2] 0))++testReplicate0Tiny1Padded :: Assertion+testReplicate0Tiny1Padded =+  assertEqualUpToEpsilon 1e-10+    (ringestData [1, 1, 1, 1] [-0.2])+    (cgrad @_ @_ @_ @Concrete+           (rsum0 @4 @Double . conv2d1Padded) (rrepl [1, 1, 1, 1] 0))++testReplicate0TinySPadded :: Assertion+testReplicate0TinySPadded =+  assertEqualUpToEpsilon 1e-10+    (ringestData [1, 1, 1, 1] [582665.99432])+    (grad (rsum0 @4 @Double .+          (conv2dPaddedB $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b))))+          (ringestData [1, 1, 1, 1] [0]))++testReplicate0TinyAPadded :: Assertion+testReplicate0TinyAPadded =+  assertEqualUpToEpsilon 1e-10+    (ringestData [1, 2, 1, 1] [-0.2,25.0003])+    (cgrad @_ @_ @_ @Concrete+           (rsum0 @4 @Double . conv2dAPadded) (rrepl [1, 2, 1, 1] 0))++testReplicate0LittleAPadded :: Assertion+testReplicate0LittleAPadded =+  assertEqualUpToEpsilon 1e-10+    (ringestData [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])+    (grad (rsum0 @4 @Double . conv2dAPadded) (rrepl [2, 2, 2, 2] 0))++-- with data t16++testKonst5LittleBPadded :: Assertion+testKonst5LittleBPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])+    (grad (rsum0 @4 @Double . conv2dBPadded) (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleCPadded :: Assertion+testKonst5LittleCPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (grad (rsum0 @4 @Double . conv2dCPadded) (rreplicate0N [2, 2, 2, 2] 5))++testKonst5BigBPadded :: Assertion+testKonst5BigBPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])+    (grad (rsum0 @4 @Double . conv2dBPadded) (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBPadded :: Assertion+testKonstNotBigBPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])+    (grad (rsum0 @4 @Double . conv2dBPadded)+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigCPadded :: Assertion+testKonst5BigCPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [3, 2, 4, 2] [0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997])+    (grad (rsum0 @4 @Double . conv2dCPadded) (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCPadded :: Assertion+testKonstNotBigCPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [3, 2, 4, 2] [0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997])+    (grad (rsum0 @4 @Double . conv2dCPadded)+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++-- with data t128b++testKonst5LittleBPadded128b :: Assertion+testKonst5LittleBPadded128b =+  assertEqualUpToEpsilon 1e-8+    (ringestData [2, 2, 2, 2] [578.1829600001,558.1716000002,608.0772800002001,577.7659200003001,729.1778800002002,701.1835600003001,833.9722000003002,803.9778800004001,578.1829600001,558.1716000002,608.0772800002001,577.7659200003001,729.1778800002002,701.1835600003001,833.9722000003002,803.9778800004001])+    (grad (rsum0 @4 @Double . conv2dBPadded128b) (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleCPadded128b :: Assertion+testKonst5LittleCPadded128b =+  assertEqualUpToEpsilon 1e-8+    (ringestData [2, 2, 2, 2] [1113.1722000001,1412.1551500001997,1234.1494800003002,1627.8210700004993,1500.2781800001994,1870.0614400004986,2156.3671200003987,2725.0393200008984,1113.1722000001,1412.1551500001997,1234.1494800003002,1627.8210700004993,1500.2781800001994,1870.0614400004986,2156.3671200003987,2725.0393200008984])+    (grad (rsum0 @4 @Double . conv2dCPadded128b) (rreplicate0N [2, 2, 2, 2] 5))++{-+testKonst5BigBPadded128b :: Assertion+testKonst5BigBPadded128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dBPadded128b (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBPadded128b :: Assertion+testKonstNotBigBPadded128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dBPadded128b+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigCPadded128b :: Assertion+testKonst5BigCPadded128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCPadded128b (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCPadded128b :: Assertion+testKonstNotBigCPadded128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCPadded128b+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++-- with data t128c++testKonst5LittleBPadded128c :: Assertion+testKonst5LittleBPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [186.7886400001,121.7829600001,269.09432000009997,261.3943200001,210.9943200001,231.79432000010002,160.00030000000004,194.00060000000005,186.7886400001,121.7829600001,269.09432000009997,261.3943200001,210.9943200001,231.79432000010002,160.00030000000004,194.00060000000005])+    (rev' @Double @4 conv2dBPadded128c (rreplicate0N [2, 2, 2, 2] 5))++testKonst5LittleCPadded128c :: Assertion+testKonst5LittleCPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [1772.649480000399,2138.4267600005987,2157.0438000004983,2640.8154000007976,961.7781800001002,1359.4557500003987,1233.4728000001987,1712.044990000598,1772.649480000399,2138.4267600005987,2157.0438000004983,2640.8154000007976,961.7781800001002,1359.4557500003987,1233.4728000001987,1712.044990000598])+    (rev' @Double @4 conv2dCPadded128c (rreplicate0N [2, 2, 2, 2] 5))++testKonst5BigBPadded128c :: Assertion+testKonst5BigBPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dBPadded128c (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBPadded128c :: Assertion+testKonstNotBigBPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dBPadded128c+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))++testKonst5BigCPadded128c :: Assertion+testKonst5BigCPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dCPadded128c (rreplicate0N [3, 2, 4, 2] 5))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCPadded128c :: Assertion+testKonstNotBigCPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dCPadded128c+          (ringestData [3, 2, 4, 2] [37, 36 .. -10]))+-}++-- with data t128b and t128c+{-+testKonst5LittleBPadded128bc :: Assertion+testKonst5LittleBPadded128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001,1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001])+    (rev' @Double @4 conv2dBPadded128b t128c)++testKonst5LittleCPadded128bc :: Assertion+testKonst5LittleCPadded128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [0.0,0.0,0.0,0.0,251.5006,417.7949200000001,494.8949100000001,382.59461000000005,482.69492000000014,778.9778800001002,952.0721900001002,742.5775700001002,610.5892400000002,1113.1722000001,1412.1551500001997,1182.6605300002002,580.6778800001001,1234.1494800003002,1627.8210700004993,1571.2321300004994,329.17728000010004,816.3545600003002,1132.9261600005002,1188.6375200005,97.98296000010004,455.17160000020016,675.7488800003999,828.6545600004001,-29.9113599999,120.97728000019995,215.6659200003,388.5716000003,0.0,0.0,0.0,0.0,65.90000000000003,106.90000000000003,173.90000000000006,170.90000000000003,365.89432000010004,593.1946200001,821.2892400002003,657.1892400002001,1060.8778800002,1500.2781800001994,1870.0614400004986,1202.8611400005,1465.6665200003995,2156.3671200003987,2725.0393200008984,1831.7390200008983,1399.7665200003996,2049.4671200003986,2551.139320000898,1660.8390200008987,1099.7722000003,1563.1725000002994,1903.750080000699,1174.5497800006997,404.7886400002001,656.0889400002,854.9778800004001,628.8778800004001,0.0,0.0,0.0,0.0,251.5006,417.7949200000001,494.8949100000001,382.59461000000005,482.69492000000014,778.9778800001002,952.0721900001002,742.5775700001002,610.5892400000002,1113.1722000001,1412.1551500001997,1182.6605300002002,580.6778800001001,1234.1494800003002,1627.8210700004993,1571.2321300004994,329.17728000010004,816.3545600003002,1132.9261600005002,1188.6375200005,97.98296000010004,455.17160000020016,675.7488800003999,828.6545600004001,-29.9113599999,120.97728000019995,215.6659200003,388.5716000003,0.0,0.0,0.0,0.0,65.90000000000003,106.90000000000003,173.90000000000006,170.90000000000003,365.89432000010004,593.1946200001,821.2892400002003,657.1892400002001,1060.8778800002,1500.2781800001994,1870.0614400004986,1202.8611400005,1465.6665200003995,2156.3671200003987,2725.0393200008984,1831.7390200008983,1399.7665200003996,2049.4671200003986,2551.139320000898,1660.8390200008987,1099.7722000003,1563.1725000002994,1903.750080000699,1174.5497800006997,404.7886400002001,656.0889400002,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCPadded128b t128c)++testKonst5BigBPadded128bc :: Assertion+testKonst5BigBPadded128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001,1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001])+    (rev' @Double @4 conv2dBPadded128b t128c)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBPadded128cb :: Assertion+testKonstNotBigBPadded128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002])+    (rev' @Double @4 conv2dBPadded128c t128b)++testKonst5BigCPadded128cb :: Assertion+testKonst5BigCPadded128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992])+    (rev' @Double @4 conv2dCPadded128c t128b)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCPadded128cb :: Assertion+testKonstNotBigCPadded128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992])+    (rev' @Double @4 conv2dCPadded128c t128b)+-}+++-- * Disparity and misc++-- | Disparity cost volume.+--+--   Take two arrays of multi channel 2d images, where the first contains+--   left views of the scene and the second contains right views.+--+--   For each pair of images, slice the right image over the left image,+--   and for each offset produce the L1 distance indicating how well+--   correponding+--   multi-channel image elements in the right image match those in the left.+--+--   Described in:+--    Anytime Stereo Image Depth Estimation on Mobile Devices+--    Wang, Lai et al, ICRA 2019+--    https://arxiv.org/abs/1810.11408+--    Section III b).+--+costVolume+  :: forall r target. (ADReady target, NumScalar r)+  => Int -> Int -> target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+costVolume iStart nCount arrL arrR =+  let [nImgs, nChas, nRows, nCols] = rshape arrL+      shO = [nImgs, nCount, nRows, nCols]+  in rbuild shO $ \[iImg, iDisp, iRow, iCol] ->+       let arrVecL = rbuild (nChas :$: ZSR) $ \[iCha] ->+                       rindex @_ @0 arrL [iImg, iCha, iRow, iCol]+           iSrc = iCol - fromIntegral iStart - iDisp+           arrVecR = rbuild [nChas] $ \[iCha] ->+                       rindex @_ @0 arrR [iImg, iCha, iRow, iSrc]+       in rfromK $ rsum0 $ rzipWith1 (\xL xR -> abs (xL - xR)) arrVecL arrVecR++test_disparityKonst :: Assertion+test_disparityKonst = do+  let arrL :: ADReady target => target (TKR 4 Double)+      arrL = rreplicate0N [1, 2, 4, 6] (-0.2)+      arrR :: ADReady target => target (TKR 4 Double)+      arrR = rreplicate0N [1, 2, 4, 6] 0.3+      arrO = costVolume @Double 0 4 arrL arrR+      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO+      arrDR = vjp (\aR -> costVolume 0 4 (rfromPrimal arrL) aR) arrR arrO+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,4,4,6] [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0])+    arrO+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0])+    arrDL+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0])+   arrDR+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,4,6] [4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0])+    (rev' @Double @4 (costVolume 0 4 arrL) arrR)+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,4,6] [-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0])+    (rev' @Double @4 (\aL -> costVolume 0 2 aL arrR) arrL)+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,4,6] [2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0])+    (rev' @Double @4 (costVolume 0 2 arrL) arrR)++test_disparityKonst2 :: Assertion+test_disparityKonst2 = do+  let arrL :: (BaseTensor target, GoodScalar r, Differentiable r) => target (TKR 4 r)+      arrL = ringestData [1, 2, 4, 6] [0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0, 1.7041241452319316,1.21999,0.21355339059327375,0.7867666666666666,0.7331698975466578,0.6964466094067263,1.1,1.1041141452319316,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,1.986339795093316,0.6964466094067263]+      arrR :: (BaseTensor target, GoodScalar r, Differentiable r) => target (TKR 4 r)+      arrR = ringestData [1, 2, 4, 6] [0.2, 0.5, -0.2, 0.0001, 0.44, 0.9, -0.9, 0.00001, -0.22, -0.28, -0.34, -0.40, -0.40,-0.22,-0.28,-0.34, 0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5, -0.35355339059327373,0.16666666666666666,0.17677669529663687,-0.25, -2.808238290463863,-1.21999,-0.5672067811865474,-0.7867666666666666,-1.986339795093316,-0.6964466094067263,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,0.6964466094067263,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,0.50001,0.42000000000000004,0.2801,0.78,1.1,0.50001,0.42000000000000004,0.2801,0.78]+      arrO = rreplicate0N [1, 4, 4, 6] (1 :: Concrete (TKScalar Double))+      res1 = rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [4.0,2.0,2.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,2.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,2.0,0.0,0.0,-2.0,0.0,4.0,4.0,2.0,0.0,-4.0,1.0,4.0,4.0,4.0,-4.0,2.0,4.0,2.0]+      res2 = rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [-4.0,0.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,-4.0,-2.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,4.0,4.0,-4.0,1.0,-2.0,-1.0,-2.0,3.0,2.0,-1.0,-2.0,-1.0,-2.0,0.0,-2.0,-3.0,-2.0,1.0]+      arrDL :: Concrete (TKR 4 Double)+      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO+      arrDR :: Concrete (TKR 4 Double)+      arrDR = vjp (costVolume 0 4 (rfromPrimal arrL)) arrR arrO+  assertEqualUpToEpsilon 1e-7+    res1+    arrDL+  assertEqualUpToEpsilon 1e-7+    res2+    arrDR+  assertEqualUpToEpsilon' 1e-7+    res1+    (rev' @Double @4 (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL)+  assertEqualUpToEpsilon' 1e-7+    res2+    (rev' @Double @4 (costVolume 0 4 (rfromPrimal arrL)) arrR)++test_disparitySmall :: Assertion+test_disparitySmall = do+  let arrL :: ADReady target => target (TKR 4 Double)+      arrL = ringestData [1, 2, 3, 2] [0.2, 0.5, -0.2, 0.0001, 0.44, 0.9, -0.9, 0.00001, -0.22, -0.28, -0.34, -0.40]+      arrR :: ADReady target => target (TKR 4 Double)+      arrR = ringestData [1, 2, 3, 2] [-0.40,-0.22,-0.28,-0.34, 0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5, -0.35355339059327373,0.16666666666666666,0.17677669529663687,-0.25]+      arrO = costVolume @Double 0 4 arrL arrR+      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO+      arrDR = vjp (\aR -> costVolume 0 4 (rfromPrimal arrL) aR) arrR arrO+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,4,3,2] [1.7041241452319316,1.21999,0.21355339059327375,0.7867666666666666,0.7331698975466578,0.6964466094067263,1.1,1.1041141452319316,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3])+    arrO+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,3,2] [-2.0,-1.0,-2.0,-1.0,-2.0,-1.0,2.0,1.0,-2.0,1.0,2.0,1.0])+    (rev' @Double @4 (costVolume 0 4 arrL) arrR)+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,2,3,2] [5.004124145231932,3.3241241452319317,-1.0464466094067264,1.7006200572599404,3.0731698975466575,4.5496165069533845,-5.004124145231932,-1.3240841452319316,-1.0464466094067264,-0.9933132760733929,-3.0731698975466575,-4.5496165069533845])+    arrDL+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,2,3,2] [-2.808238290463863,-1.21999,-0.5672067811865474,-0.7867666666666666,-1.986339795093316,-0.6964466094067263,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,1.986339795093316,0.6964466094067263])+   arrDR+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,3,2] [-1.0,0.0,-1.0,0.0,-1.0,0.0,1.0,0.0,-1.0,0.0,1.0,0.0])+    (rev' @Double @4 (costVolume 1 4 arrL) arrR)+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,3,2] [2.0,2.0,-2.0,2.0,2.0,2.0,-2.0,2.0,-2.0,-2.0,-2.0,-2.0])+    (rev' @Double @4 (\aL -> costVolume 2 2 aL arrR) arrL)+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,3,2] [-1.0,0.0,-1.0,0.0,-1.0,0.0,1.0,0.0,-1.0,0.0,1.0,0.0])+    (rev' @Double @4 (costVolume 1 2 arrL) arrR)++codeTomsSlice :: ADReady target+              => target (TKR 2 Double) -> target (TKR 0 Double)+codeTomsSlice a =+  let (n, m) = case rshape a of+        [n', m'] -> (n', m')+        _ -> error "codeTomsSlice"+      a1 = rbuild @2 @0 [n,m-1] (\[i',j'] -> rindex a [i',j'])+      a2 = rbuild [n,m-1] (\[i',j'] -> rindex @_ @0 a [i',j' + 1])+  in rfromK $ rsum0 @2 $ rbuild [n,m] $ \[i, _j] ->+       tfromPlain knownSTK (rfromIntegral (rfromK i)) * rfromK (rsum0 (a1 * a2))++testTomsSliceRev :: Assertion+testTomsSliceRev = do+  assertEqualUpToEpsilon 1e-5+    (ringestData [32,4] [63686.39999999999,137292.80000000002,121222.4,79558.40000000002,192646.40000000005,223971.0617601984,228556.80000000005,116846.33088019838,63686.39999999999,137292.80000000002,127174.4,79558.40000000002,192646.40000000005,158499.06176019844,202566.40000000005,51374.330880198424,11904.0,5952.0,7936.0,1984.0,116846.33088019838,385292.8000000001,227740.66176039676,192646.40000000005,116846.33088019838,228556.80000000005,174580.73088019836,35910.399999999994,79558.40000000002,127372.79999999997,143244.80000000002,63686.39999999999,105152.0,186683.13088000007,105151.98016,107124.73088000003,-396.79999999999995,26188.8,17459.2,25990.399999999998,-7936.0,73408.0,-1995.2691200000017,57536.0,51584.0,-660672.0,55552.0,3968.0,3968.0,3571.2,3571.2,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,49203.79519999998,49203.79519999998,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,129158.9952,65472.59519999998,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,1984.0000000000146,67059.20000000001,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,-21823.99999999993,108921.6,16070.400000000005,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,771974.4,218019.0617601984,192646.40000000005,170414.3308801984,385292.8000000001,340828.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,236294.40000000005,271587.0617601984,192646.40000000005,45422.33088019842,385292.8000000001,162268.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,369222.4,220003.0617601984,192646.40000000005,104942.33088019838,385292.8000000001,215836.66176039676,192646.40000000005])+    (grad (kfromR . codeTomsSlice) (rreshape [32, 4] t128))++testTomsSliceCRev :: Assertion+testTomsSliceCRev = do+  assertEqualUpToEpsilon 1e-5+    (ringestData [32,4] [63686.39999999999,137292.80000000002,121222.4,79558.40000000002,192646.40000000005,223971.0617601984,228556.80000000005,116846.33088019838,63686.39999999999,137292.80000000002,127174.4,79558.40000000002,192646.40000000005,158499.06176019844,202566.40000000005,51374.330880198424,11904.0,5952.0,7936.0,1984.0,116846.33088019838,385292.8000000001,227740.66176039676,192646.40000000005,116846.33088019838,228556.80000000005,174580.73088019836,35910.399999999994,79558.40000000002,127372.79999999997,143244.80000000002,63686.39999999999,105152.0,186683.13088000007,105151.98016,107124.73088000003,-396.79999999999995,26188.8,17459.2,25990.399999999998,-7936.0,73408.0,-1995.2691200000017,57536.0,51584.0,-660672.0,55552.0,3968.0,3968.0,3571.2,3571.2,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,49203.79519999998,49203.79519999998,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,129158.9952,65472.59519999998,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,1984.0000000000146,67059.20000000001,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,-21823.99999999993,108921.6,16070.400000000005,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,771974.4,218019.0617601984,192646.40000000005,170414.3308801984,385292.8000000001,340828.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,236294.40000000005,271587.0617601984,192646.40000000005,45422.33088019842,385292.8000000001,162268.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,369222.4,220003.0617601984,192646.40000000005,104942.33088019838,385292.8000000001,215836.66176039676,192646.40000000005])+    (cgrad (kfromR . codeTomsSlice) (rreshape [32, 4] t128))++testTomsSlice :: Assertion+testTomsSlice = do+  assertEqualUpToEpsilon' 1e-5+    (ringestData [12,4] [7682.400000000001,13252.8,18268.8,8474.4,14254.50048,39362.399999999994,29802.6009600264,25634.399999999998,7682.400000000001,13252.8,18268.8,8474.4,14518.50048,40946.399999999994,28482.6009600264,25634.399999999998,528.0,2904.0,792.0,1584.0,242.88000000000002,-501.6,190.08,26.400000000000002,2376.0,5570.4,1320.0,2112.0,790.50048,264.0,7654.50048,-8712.0,528.0,1056.0,1056.0,528.0,-52.800000000000004,-105.60000000000001,-105.60000000000001,-52.800000000000004,-52.800000000000004,6547.2792,-105.60000000000001,-52.800000000000004,6600.0792,13200.1584,13200.1584,6600.0792])+    (rev' codeTomsSlice (rreshape [12, 4] t48))+++-- * PP Tests++testTomsSlicePP :: Assertion+testTomsSlicePP = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent codeTomsSlice (FTKR [32, 4] FTKScalar)+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromK (1984.0 * sdot0 (sslice (SNat @0) (SNat @3) (str (sfromR m1))) (sslice (SNat @1) (SNat @3) (str (sfromR m1))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromK (let v8 = sreplicate0N @[32] (ssum0 @[32, 3] (str (sslice (SNat @0) (SNat @3) (str (sfromR m1))) * str (sslice (SNat @1) (SNat @3) (str (sfromR m1))))) ; x9 = ssum0 @[32] (sfromPlain (siota (SNat @32)) * v8) in 4.0 * x9)"+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> rfromS (let m11 = sreplicate0N @[32, 3] (ssum0 @[32] (sfromPlain (siota (SNat @32)) * sreplicate0N @[32] (4.0 * kfromR dret))) in str (sappend (sconcrete (sfromListLinear [0,32] [])) (sappend (str (str (sslice (SNat @1) (SNat @3) (str (sfromR m1))) * m11)) (sconcrete (sreplicate [1,32] 0.0)))) + str (sappend (sconcrete (sreplicate [1,32] 0.0)) (sappend (str (str (sslice (SNat @0) (SNat @3) (str (sfromR m1))) * m11)) (sconcrete (sfromListLinear [0,32] [])))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> rfromS (let x11 = 1984.0 * kfromR dret in str (sappend (sslice (SNat @1) (SNat @3) (str (sfromR m1)) * sreplicate0N @[3, 32] x11) (sconcrete (sreplicate [1,32] 0.0))) + str (sappend (sconcrete (sreplicate [1,32] 0.0)) (sslice (SNat @0) (SNat @3) (str (sfromR m1)) * sreplicate0N @[3, 32] x11)))"++testCNNOPP0c :: Assertion+testCNNOPP0c = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dCLaborious (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (sfromPlain (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (sfromVector @2 (fromList [stranspose @[3, 0, 5, 1, 2, 4] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i72, i74] -> [i72 + i74]))) (\\[i49, i50] -> [i49 + i50])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i45, i46, i47, i48] -> [ifH ((-1) <=. negate i45 + negate i47 &&* (-1) <=. negate i46 + negate i48) 0 1, i45, i46, i47, i48])))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1)))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w53 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (sfromVector @2 (fromList [stranspose @[3, 0, 5, 1, 2, 4] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i51, i52] -> [i51 + i52]))) (\\[i49, i50] -> [i49 + i50])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i45, i46, i47, i48] -> [ifH ((-1) <=. negate i45 + negate i47 &&* (-1) <=. negate i46 + negate i48) 0 1, i45, i46, i47, i48])))) in sfromPlain w53 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (ssum @2 (stranspose @[4, 0, 1, 2, 3] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (ssum @2 (stranspose @[3, 0, 1, 2] (let w53 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (sfromVector @2 (fromList [stranspose @[3, 0, 5, 1, 2, 4] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i51, i52] -> [i51 + i52]))) (\\[i49, i50] -> [i49 + i50])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i45, i46, i47, i48] -> [ifH ((-1) <=. negate i45 + negate i47 &&* (-1) <=. negate i46 + negate i48) 0 1, i45, i46, i47, i48])))) in sfromPlain w53 * sreplicateN @[2, 2, 2] (sfromR dret))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (ssumN @[2, 2] (sdot1In (sfromPlain (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (sfromVector @2 (fromList [stranspose @[3, 0, 5, 1, 2, 4] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i82, i84] -> [i82 + i84]))) (\\[i49, i50] -> [i49 + i50])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i45, i46, i47, i48] -> [ifH ((-1) <=. negate i45 + negate i47 &&* (-1) <=. negate i46 + negate i48) 0 1, i45, i46, i47, i48])))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret)))))"++testCNNOPP0b :: Assertion+testCNNOPP0b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dBLaborious (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,-2.0,5.0,-2.0,5.0,-2.0,5.0,-2.0,13.1,582934.0,13.1,582934.0,13.1,582934.0,13.1,582934.0,5.0,-2.0,5.0,-2.0,5.0,-2.0,5.0,-2.0,13.1,582934.0,13.1,582934.0,13.1,582934.0,13.1,582934.0,2.0,0.0,2.0,0.0,2.0,0.0,2.0,0.0,9.0,2.99432,9.0,2.99432,9.0,2.99432,9.0,2.99432,2.0,0.0,2.0,0.0,2.0,0.0,2.0,0.0,9.0,2.99432,9.0,2.99432,9.0,2.99432,9.0,2.99432,6.0,0.1,6.0,0.1,6.0,0.1,6.0,0.1,8.0,-335.0,8.0,-335.0,8.0,-335.0,8.0,-335.0,6.0,0.1,6.0,0.1,6.0,0.1,6.0,0.1,8.0,-335.0,8.0,-335.0,8.0,-335.0,8.0,-335.0,1.0,-0.2,1.0,-0.2,1.0,-0.2,1.0,-0.2,-4.0,26.0,-4.0,26.0,-4.0,26.0,-4.0,26.0,1.0,-0.2,1.0,-0.2,1.0,-0.2,1.0,-0.2,-4.0,26.0,-4.0,26.0,-4.0,26.0,-4.0,26.0])) (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (sfromVector @2 (fromList [stranspose @[3, 0, 5, 1, 2, 4] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i99, i101] -> [i99 + i101]))) (\\[i54, i55] -> [i54 + i55])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i56, i57, i58, i59] -> [ifH ((-1) <=. negate i56 + negate i58 &&* (-1) <=. negate i57 + negate i59) 0 1, i56, i57, i58, i59])))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w60 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (sfromVector @2 (fromList [stranspose @[3, 0, 5, 1, 2, 4] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i52, i53] -> [i52 + i53]))) (\\[i54, i55] -> [i54 + i55])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i56, i57, i58, i59] -> [ifH ((-1) <=. negate i56 + negate i58 &&* (-1) <=. negate i57 + negate i59) 0 1, i56, i57, i58, i59])))) in sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,5.0,5.0,5.0,13.1,13.1,13.1,13.1,5.0,5.0,5.0,5.0,13.1,13.1,13.1,13.1,2.0,2.0,2.0,2.0,9.0,9.0,9.0,9.0,2.0,2.0,2.0,2.0,9.0,9.0,9.0,9.0,6.0,6.0,6.0,6.0,8.0,8.0,8.0,8.0,6.0,6.0,6.0,6.0,8.0,8.0,8.0,8.0,1.0,1.0,1.0,1.0,-4.0,-4.0,-4.0,-4.0,1.0,1.0,1.0,1.0,-4.0,-4.0,-4.0,-4.0,-2.0,-2.0,-2.0,-2.0,582934.0,582934.0,582934.0,582934.0,-2.0,-2.0,-2.0,-2.0,582934.0,582934.0,582934.0,582934.0,0.0,0.0,0.0,0.0,2.99432,2.99432,2.99432,2.99432,0.0,0.0,0.0,0.0,2.99432,2.99432,2.99432,2.99432,0.1,0.1,0.1,0.1,-335.0,-335.0,-335.0,-335.0,0.1,0.1,0.1,0.1,-335.0,-335.0,-335.0,-335.0,-0.2,-0.2,-0.2,-0.2,26.0,26.0,26.0,26.0,-0.2,-0.2,-0.2,-0.2,26.0,26.0,26.0,26.0]) * w60))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[2, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[2, 2] (stranspose @[1, 3, 4, 0, 5, 2] (sscatter @[2, 2, 2, 2] (stranspose @[4, 5, 2, 3, 0, 1] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,5.0,5.0,5.0,13.1,13.1,13.1,13.1,5.0,5.0,5.0,5.0,13.1,13.1,13.1,13.1,2.0,2.0,2.0,2.0,9.0,9.0,9.0,9.0,2.0,2.0,2.0,2.0,9.0,9.0,9.0,9.0,6.0,6.0,6.0,6.0,8.0,8.0,8.0,8.0,6.0,6.0,6.0,6.0,8.0,8.0,8.0,8.0,1.0,1.0,1.0,1.0,-4.0,-4.0,-4.0,-4.0,1.0,1.0,1.0,1.0,-4.0,-4.0,-4.0,-4.0,-2.0,-2.0,-2.0,-2.0,582934.0,582934.0,582934.0,582934.0,-2.0,-2.0,-2.0,-2.0,582934.0,582934.0,582934.0,582934.0,0.0,0.0,0.0,0.0,2.99432,2.99432,2.99432,2.99432,0.0,0.0,0.0,0.0,2.99432,2.99432,2.99432,2.99432,0.1,0.1,0.1,0.1,-335.0,-335.0,-335.0,-335.0,0.1,0.1,0.1,0.1,-335.0,-335.0,-335.0,-335.0,-0.2,-0.2,-0.2,-0.2,26.0,26.0,26.0,26.0,-0.2,-0.2,-0.2,-0.2,26.0,26.0,26.0,26.0]) * sreplicateN @[2, 2, 2] (sfromR dret))))) (\\[i62, i63, i64, i65] -> [ifH ((-1) <=. negate i62 + negate i64 &&* (-1) <=. negate i63 + negate i65) 0 1, i62, i63, i64, i65]) !$ [0])) (\\[i67, i68] -> [i67 + i68]))) (\\[i69, i70] -> [i69 + i70])))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[2, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[2, 2] (stranspose @[0, 2, 4, 5, 1, 6, 3] (sscatter @[2, 2, 2, 2] (sdot1In (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0])) (stranspose @[5, 6, 1, 2, 3, 0, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i62, i63, i64, i65] -> [ifH ((-1) <=. negate i62 + negate i64 &&* (-1) <=. negate i63 + negate i65) 0 1, i62, i63, i64, i65])) !$ [0]) (\\[i67, i68] -> [i67 + i68]))) (\\[i69, i70] -> [i69 + i70])))"++testCNNOPP1e :: Assertion+testCNNOPP1e = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dSameL (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (sfromVector @2 (fromList [stranspose @[3, 0, 5, 1, 2, 4] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i106, i108] -> [i106 + i108]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i32, i33, i34, i35] -> [ifH ((-1) <=. negate i32 + negate i34 &&* (-1) <=. negate i33 + negate i35) 0 1, i32, i33, i34, i35]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w36 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (stranspose @[2, 3, 5, 6, 0, 1, 4] (sfromVector @2 (fromList [stranspose @[2, 3, 0, 4, 5, 1] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)]))) (\\[i32, i33, i34, i35] -> [i32, i33, i34, i35, ifH ((-1) <=. negate i32 + negate i34 &&* (-1) <=. negate i33 + negate i35) 0 1])))) in w36 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1)))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (tpair (ssumN @[2, 2] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[4, 5, 2, 3, 0, 1] (sgather @[2, 2, 2, 2] (sfromVector @2 (fromList [stranspose @[3, 0, 5, 1, 2, 4] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i144, i146] -> [i144 + i146]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i32, i33, i34, i35] -> [ifH ((-1) <=. negate i32 + negate i34 &&* (-1) <=. negate i33 + negate i35) 0 1, i32, i33, i34, i35]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret))))) (stranspose @[1, 2, 0] (sscatter @[2, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[2, 2] (stranspose @[4, 1, 3, 5, 0, 6, 2] (sscatter @[2, 2, 2, 2] (sdot1In (stranspose @[5, 6, 2, 3, 0, 1, 4] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1)))))))) (stranspose @[5, 6, 1, 2, 3, 0, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i38, i39, i40, i41] -> [i38, i39, i40, i41, ifH ((-1) <=. negate i38 + negate i40 &&* (-1) <=. negate i39 + negate i41) 0 1])) !$ [0]) (\\[i43, i44] -> [i43 + i44]))) (\\[i45, i46] -> [i45 + i46]))))"++testCNNOPP2 :: Assertion+testCNNOPP2 = do+  resetVarCounter+  let t = maxPool2dUnpadded2+            (rconcrete $ Nested.rreplicatePrim (1 :$: 1 :$: 2 :$: 2 :$: ZSR) 1)+  printAstPretty (simplifyInlineContract t)+    @?= "rfromPlain (rfromS (sreplicateN @[2, 2] (stranspose @[2, 3, 1, 0] (sfromVector @2 (fromList [sgather @[2, 2, 2] (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (sconcrete (sreplicate [2,2] 1.0)) (\\[i62, i64] -> [i64 + i62])))) (\\[i44, i32, i8] -> [i8, i8, i8, 2 * i44 + i32]), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0, 0])))"+  printAstPretty t+    @?= "rfromPlain (rfromS (sreplicateN @[2, 2] (let t38 = sgather @[2, 2, 2] (sreplicate @1 (let u29 = sgather @[2, 2] (stranspose @[2, 0, 1] (sgather @[2, 2] (sconcrete (sreplicate [2,2] 1.0)) (\\[i26, i15] -> [i26 + i15]))) (\\[i22, i16] -> [i22 + i16]) in sappend (sreplicate @1 (stranspose @[2, 0, 3, 1] u29)) (sconcrete (sreplicate [2,2,2,2,2] 0.0)))) (\\[i44, i32, i8] -> [i8, 0, 2 * i44 + i32, i8, i8, 0]) in stranspose @[2, 3, 1, 0] (sfromVector @2 (fromList [t38, sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0, 0])))"++testCNNOPP2b :: Assertion+testCNNOPP2b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded2 (FTKR [1, 1, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (sreplicateN @[2, 2] (stranspose @[2, 3, 1, 0] (sfromVector @2 (fromList [sgather @[2, 2, 2] (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (sfromR u1 !$ [0, 0]) (\\[i69, i70] -> [i70 + i69])))) (\\[i71, i72, i73] -> [i73, i73, i73, 2 * i71 + i72]), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0, 0]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (sreplicateN @[2, 2] (stranspose @[2, 3, 1, 0] (sfromVector @2 (fromList [sgather @[2, 2, 2] (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (sfromR u1 !$ [0, 0]) (\\[i69, i70] -> [i70 + i69])))) (\\[i71, i72, i73] -> [i73, i73, i73, 2 * i71 + i72]), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0, 0]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (soneHot (sscatter @[2, 2] (stranspose @[0, 2, 1] (ssum @1 (sscatter @[2, 2, 2] (stranspose @[3, 2, 0, 1] (soneHot (ssumN @[2, 2] (sfromR dret)) [0, 0]) !$ [0]) (\\[i76, i77, i78] -> [i78, i78, i78, 2 * i76 + i77])))) (\\[i79, i80] -> [i80 + i79])) [0, 0])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (sreplicateN @[1, 1] (sscatter @[2, 2] (sscatter @[2, 2, 2] (stranspose @[3, 2, 0, 1] (soneHot (ssumN @[2, 2] (sfromR dret)) [0, 0]) !$ [0]) (\\[i76, i77, i78] -> [i78, i78, 2 * i76 + i77, i78]) !$ [0]) (\\[i79, i80] -> [i80 + i79])))"++maxPool2dUnpadded2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded2 a =+  rbuild [2, 2, 2, 2] $ \case+    [_, _, iBh, iBw] ->+      let arrt = slicez2 (conv2dSame2 a) [iBw, 1, 2 * iBh, 2 * iBw]+      in rmaximum2 arrt+    _ -> error "maxPool2dUnpadded2: impossible pattern needlessly required"++conv2dSame2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dSame2 a =+  rbuild [3, 3, 2, 2] $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez2 a [iImg, 0, iBh, iBw]+      in rindex @_ @0 arrAt [0, iBw, iBw, 0]+    _ -> error "conv2dSame2: impossible pattern needlessly required"++slicez2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)+  => target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez2 d ixBase =+  rbuild [1, 1, 2, 2] $ \ixResult -> indexz02 d (ixrZipWith (+) ixBase ixResult)++indexz02+  :: forall target r n.+     (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)+  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)+indexz02 d ix = ifH (1 >. (toList ix !! 0)) (d ! ix) (rscalar 0)++rmaximum2 :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+         => target (TKR 4 r) -> target (TKR 0 r)+rmaximum2 t0 = tlet t0 $ \t -> rindex t [0, 0, 0, 0]++testCNNOPP3 :: Assertion+testCNNOPP3 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ rreplicateN (2 :$: 2 :$: 2 :$: 2 :$: ZSR)+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded33 $ conv2dSame3 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromPlain (rfromS (sfromVector @2 (fromList [str (sgather1 @2 (sgather @[2, 2, 2] (sfromVector @2 (fromList [sgather @[2, 2] (sappend (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 1, 2] (sgather @[2, 2] (sconcrete (sreplicate [2,2,2] 7.0)) (\\[i97, i99] -> [remH i97 4 + i99]))) (\\[i101, i103] -> [i101 + i103, i101]))) (\\[i33, i15] -> [i33 + i15, i33])))) (sconcrete (sreplicate [1,2,2,2,2] 0.0))) (\\[i32, i27] -> [i27, i32, i27, i32]), sconcrete (sreplicate [2,2,2] 0.0)])) (\\[i31, i28, i26] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i31, i28, i26], i31, i28, i26])) (\\i123 -> [remH i123 4])), str (sgather1 @2 (sgather @[2, 2, 2] (sfromVector @2 (fromList [sgather @[2, 2] (sappend (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 1, 2] (sgather @[2, 2] (sconcrete (sreplicate [2,2,2] 7.0)) (\\[i127, i129] -> [remH i127 4 + i129]))) (\\[i131, i133] -> [i131 + i133, i131]))) (\\[i33, i15] -> [i33 + i15, i33])))) (sconcrete (sreplicate [1,2,2,2,2] 0.0))) (\\[i32, i27] -> [i27, i32, i27, i32]), sconcrete (sreplicate [2,2,2] 0.0)])) (\\[i31, i28, i26] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i31, i28, i26], i31, i28, i26])) (\\i153 -> [remH i153 4]))])))"+  printAstPretty afcnn2T+    @?= "rfromPlain (rfromS (let w48 = sgather @[2, 2, 2, 2, 2, 2, 2, 2] (stranspose @[1, 2, 3, 4, 5, 6, 7, 8, 0] (sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[6, 2, 0, 3, 4, 5, 1] (sgather @[2, 2] (stranspose @[4, 0, 1, 2, 3] (sgather @[2, 2, 2, 2] (sreplicate @2 (sgather @[2, 2, 2] (stranspose @[1, 2, 3, 0] (sfromVector @2 (fromList [sgather @[2, 2] (str (sappend (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (stranspose @[0, 3, 4, 2, 1] (sgather @[2, 2] (stranspose @[0, 2, 1] (sgather @[2, 2] (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i35, i13] -> [remH i35 4 + i13]))) (\\[i34, i14] -> [i34, i34 + i14]))) (\\[i33, i15] -> [i33, i33 + i15, 1])))) (sconcrete (sreplicate [1,2,2,2,2] 0.0)))) (\\[i32, i27] -> [i32, i27, i27, i32]), sconcrete (sreplicate [2,2,2] 0.0)]))) (\\[i31, i28, i26] -> [i31, i28, i26, sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i31, i28, i26]]))) (\\[i65, i56, i46, i6] -> [i65 + i6, remH i56 4 + i46]))) (\\[i60, i7] -> [i60 + i7]))) (\\[i51, i8] -> [i51 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)]))) (\\[i62, i58, i53, i50, i43, i40, i38, i37] -> [i62, i58, i53, i50, i43, i40, i38, i37, ifH ((-1) <=. remH (negate i53) 4 + negate i43 &&* ((-1) <=. negate i62 + negate i40 &&* ((-1) <=. negate i58 + negate i38 &&* (-1) <=. negate i50 + negate i37))) 0 1]) in stranspose @[4, 5, 6, 7, 0, 1, 2, 3] w48 !$ [0, 0, 0, 0]))"++testCNNOPP3b :: Assertion+testCNNOPP3b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded33 . conv2dSame3) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (sfromVector @2 (fromList [str (sgather1 @2 (sgather @[2, 2, 2] (sfromVector @2 (fromList [sgather @[2, 2] (sappend (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 1, 2] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sfromR u1) !$ [1]) (\\[i255, i257] -> [remH i255 4 + i257]))) (\\[i259, i261] -> [i259 + i261, i259]))) (\\[i159, i160] -> [i159 + i160, i159])))) (sconcrete (sreplicate [1,2,2,2,2] 0.0))) (\\[i161, i162] -> [i162, i161, i162, i161]), sconcrete (sreplicate [2,2,2] 0.0)])) (\\[i163, i164, i165] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i163, i164, i165], i163, i164, i165])) (\\i166 -> [remH i166 4])), str (sgather1 @2 (sgather @[2, 2, 2] (sfromVector @2 (fromList [sgather @[2, 2] (sappend (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 1, 2] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sfromR u1) !$ [1]) (\\[i267, i269] -> [remH i267 4 + i269]))) (\\[i271, i273] -> [i271 + i273, i271]))) (\\[i171, i172] -> [i171 + i172, i171])))) (sconcrete (sreplicate [1,2,2,2,2] 0.0))) (\\[i173, i174] -> [i174, i173, i174, i173]), sconcrete (sreplicate [2,2,2] 0.0)])) (\\[i175, i176, i177] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i175, i176, i177], i175, i176, i177])) (\\i178 -> [remH i178 4]))]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (sfromVector @2 (fromList [str (sgather1 @2 (sgather @[2, 2, 2] (sfromVector @2 (fromList [sgather @[2, 2] (sappend (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 1, 2] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sfromR u1) !$ [1]) (\\[i155, i156] -> [remH i155 4 + i156]))) (\\[i157, i158] -> [i157 + i158, i157]))) (\\[i159, i160] -> [i159 + i160, i159])))) (sconcrete (sreplicate [1,2,2,2,2] 0.0))) (\\[i161, i162] -> [i162, i161, i162, i161]), sconcrete (sreplicate [2,2,2] 0.0)])) (\\[i163, i164, i165] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i163, i164, i165], i163, i164, i165])) (\\i166 -> [remH i166 4])), str (sgather1 @2 (sgather @[2, 2, 2] (sfromVector @2 (fromList [sgather @[2, 2] (sappend (sreplicate @1 (stranspose @[0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sgather @[2, 2] (stranspose @[3, 0, 1, 2] (sgather @[2, 2] (stranspose @[3, 0, 2, 1] (sfromR u1) !$ [1]) (\\[i167, i168] -> [remH i167 4 + i168]))) (\\[i169, i170] -> [i169 + i170, i169]))) (\\[i171, i172] -> [i171 + i172, i171])))) (sconcrete (sreplicate [1,2,2,2,2] 0.0))) (\\[i173, i174] -> [i174, i173, i174, i173]), sconcrete (sreplicate [2,2,2] 0.0)])) (\\[i175, i176, i177] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i175, i176, i177], i175, i176, i177])) (\\i178 -> [remH i178 4]))]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (let u184 = sscatter @[2, 2, 2] (sscatter1 @2 (str (sfromR dret !$ [1])) (\\i180 -> [remH i180 4])) (\\[i181, i182, i183] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i181, i182, i183], i181, i182, i183]) ; w187 = sscatter @[2, 2] (u184 !$ [0]) (\\[i185, i186] -> [i186, i185, i186, i185]) ; u198 = sscatter @[2, 2, 2] (sscatter1 @2 (str (sfromR dret !$ [0])) (\\i194 -> [remH i194 4])) (\\[i195, i196, i197] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i195, i196, i197], i195, i196, i197]) ; w201 = sscatter @[2, 2] (u198 !$ [0]) (\\[i199, i200] -> [i200, i199, i200, i199]) in stranspose @[1, 3, 2, 0] (soneHot (sscatter @[2, 2] (stranspose @[1, 2, 3, 0] (sscatter @[2, 2] (stranspose @[1, 3, 2, 0] (sscatter @[2, 2] (stranspose @[0, 2, 1] (ssum @1 (sslice (SNat @0) (SNat @1) w201))) (\\[i202, i203] -> [i202 + i203, i202]))) (\\[i204, i205] -> [i204 + i205, i204]))) (\\[i206, i207] -> [remH i206 4 + i207])) [1]) + stranspose @[1, 3, 2, 0] (soneHot (sscatter @[2, 2] (stranspose @[1, 2, 3, 0] (sscatter @[2, 2] (stranspose @[1, 3, 2, 0] (sscatter @[2, 2] (stranspose @[0, 2, 1] (ssum @1 (sslice (SNat @0) (SNat @1) w187))) (\\[i188, i189] -> [i188 + i189, i188]))) (\\[i190, i191] -> [i190 + i191, i190]))) (\\[i192, i193] -> [remH i192 4 + i193])) [1]))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[1, 3, 2, 0] (sscatter @[2, 2] (stranspose @[1, 2, 3, 0] (sscatter @[2, 2] (stranspose @[1, 3, 2, 0] (sscatter @[2, 2] (sscatter @[2, 2] (sscatter @[2, 2, 2] (sscatter1 @2 (stranspose @[0, 2, 1] (sfromR dret) !$ [0]) (\\i194 -> [remH i194 4])) (\\[i195, i196, i197] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i195, i196, i197], i195, i196, i197]) !$ [0]) (\\[i199, i200] -> [i200, i199, i199, i200]) !$ [0]) (\\[i202, i203] -> [i202 + i203, i202]))) (\\[i204, i205] -> [i204 + i205, i204]))) (\\[i206, i207] -> [1, remH i206 4 + i207])) + stranspose @[1, 3, 2, 0] (sscatter @[2, 2] (stranspose @[1, 2, 3, 0] (sscatter @[2, 2] (stranspose @[1, 3, 2, 0] (sscatter @[2, 2] (sscatter @[2, 2] (sscatter @[2, 2, 2] (sscatter1 @2 (stranspose @[0, 2, 1] (sfromR dret) !$ [1]) (\\i180 -> [remH i180 4])) (\\[i181, i182, i183] -> [sconcrete (sfromListLinear [2,2,2] [0,0,1,1,1,1,1,1]) `sindex0` [i181, i182, i183], i181, i182, i183]) !$ [0]) (\\[i185, i186] -> [i186, i185, i185, i186]) !$ [0]) (\\[i188, i189] -> [i188 + i189, i188]))) (\\[i190, i191] -> [i190 + i191, i190]))) (\\[i192, i193] -> [1, remH i192 4 + i193])))"++maxPool2dUnpadded3+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded3 arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez3 [2, 2, 2, 2] arr [iBh `quotH` 4, aa, bb, iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded3: impossible pattern needlessly required"++maxPool2dUnpadded33+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded33 arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez33 [2, 2, 2, 2] arr [iBh `remH` 4, aa, bb, iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded33: impossible pattern needlessly required"++conv2dSame3+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dSame3 arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez33 shB arrA [iImg `remH` 4, iImg, iImg, 1]+      in rindex @_ @0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dSame3: impossible pattern needlessly required"++slicez3+  :: (ADReady target, NumScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez3 shOut d ixBase =+  rbuild shOut $ \_ -> indexz03 d (ixrZipWith (+) ixBase ixBase)++slicez33+  :: (ADReady target, NumScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez33 shOut d ixBase =+  rbuild shOut $ \ixResult -> indexz03 d (ixrZipWith (+) ixBase ixResult)++indexz03+  :: forall target r n. (ADReady target, NumScalar r, KnownNat n)+  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)+indexz03 d ix = ifH (within0 @target (rshape @target d) ix) (d ! ix) (rscalar 0)++rmaximum3 :: (BaseTensor target, LetTensor target, KnownNat n, GoodScalar r)+         => target (TKR n r) -> target (TKR 0 r)+rmaximum3 t0 = tlet t0 $ \t -> rindex t [0, 0, 0, 0]++testCNNOPP4 :: Assertion+testCNNOPP4 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet PlainSpan (TKR 4 Double)+      blackGlyph = rreplicateN (3 :$: 3 :$: 3 :$: 3 :$: ZSR)+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PlainSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet PlainSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded4 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (str (let w68 = stranspose @[5, 6, 8, 4, 2, 1, 3, 7, 0] (sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[6, 0, 3, 1, 4, 5, 2] (sgather @[2, 2, 2] (stranspose @[4, 1, 3, 0, 2] (sgather @[2, 2] (stranspose @[3, 0, 4, 1, 2] (sgather @[2, 2] (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i81, i84] -> [i81 + i84]))) (\\[i87, i89] -> [3 + (negate i89 + i87), i89]))) (\\[i91, i93, i96] -> [i91 * i93 + i96]))) (\\[i22, i8] -> [2 * i22 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0] in sfromVector @2 (fromList [stranspose @[1, 4, 5, 2, 3, 0] w68 !$ [0, 0, 1], stranspose @[1, 4, 5, 2, 3, 0] w68 !$ [1, 0, 0]])))"+      -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,7.0,7.0,7.0,7.0,0.0,0.0,0.0,0.0,7.0,7.0,7.0,7.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (let w19 = sgather @[2, 2, 2, 2, 2] (stranspose @[2, 4, 5, 6, 8, 0, 1, 3, 7] (sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[6, 0, 3, 1, 4, 5, 2] (sgather @[2, 2, 2] (stranspose @[3, 0, 2, 1] (sgather @[2, 2] (stranspose @[0, 2, 1] (sgather @[2, 2] (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i32, i5] -> [i32 + i5]))) (\\[i31, i6] -> [i31, 3 + (negate i31 + i6)]))) (\\[i36, i26, i7] -> [i36 * i26 + i7]))) (\\[i22, i8] -> [2 * i22 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)]))) (\\[i28, i21, i15, i12, i9] -> [i28, i21, i15, i12, i9, ifH ((-1) <=. negate i28 + negate i15 &&* (1 <=. i28 + negate i12 &&* (-2) <=. (-2) * i21 + negate i9)) 0 1]) in stranspose @[2, 3, 4, 7, 5, 0, 6, 1] w19 !$ [0, 0, 0, 0])"++testCNNOPP4b :: Assertion+testCNNOPP4b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded4 (FTKR [3, 3, 3, 3] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (str (let w158 = stranspose @[5, 6, 8, 4, 2, 1, 3, 7, 0] (sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[4, 7, 0, 3, 1, 5, 6, 2] (sgather @[2, 2, 2] (stranspose @[3, 4, 7, 1, 5, 6, 0, 2] (sgather @[2, 2] (stranspose @[6, 0, 7, 4, 3, 2, 1, 5] (sgather @[2, 2, 2, 2, 2] (stranspose @[0, 2, 1] (sfromR u1)) (\\[i267, i269, i270, i271, i272] -> [1 + sconcrete (sfromListLinear [2,2] [0,1,1,2]) `sindex0` [i267, i269]]))) (\\[i278, i281] -> [3 + (negate i281 + i278), i281]))) (\\[i284, i286, i289] -> [i284, i286, i284 * i286 + i289]))) (\\[i156, i157] -> [i156, 2 * i156 + i157])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0] in sfromVector @2 (fromList [stranspose @[1, 4, 5, 2, 3, 0] w158 !$ [0, 0, 1], stranspose @[1, 4, 5, 2, 3, 0] w158 !$ [1, 0, 0]])))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (str (let m145 = str (sreplicate @2 (siota (SNat @2))) + sreplicate @2 (siota (SNat @2)) ; w158 = stranspose @[5, 6, 8, 4, 2, 1, 3, 7, 0] (sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[4, 7, 0, 3, 1, 5, 6, 2] (sgather @[2, 2, 2] (stranspose @[3, 4, 7, 1, 5, 6, 0, 2] (sgather @[2, 2] (stranspose @[6, 0, 7, 4, 3, 2, 1, 5] (sgather @[2, 2, 2, 2, 2] (stranspose @[0, 2, 1] (sfromR u1)) (\\[i146, i147, i148, i149, i150] -> [1 + m145 `sindex0` [i146, i147]]))) (\\[i151, i152] -> [3 + (negate i152 + i151), i152]))) (\\[i153, i154, i155] -> [i153, i154, i153 * i154 + i155]))) (\\[i156, i157] -> [i156, 2 * i156 + i157])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0] in sfromVector @2 (fromList [stranspose @[1, 4, 5, 2, 3, 0] w158 !$ [0, 0, 1], stranspose @[1, 4, 5, 2, 3, 0] w158 !$ [1, 0, 0]])))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (stranspose @[0, 2, 1] (let m145 = str (sreplicate @2 (siota (SNat @2))) + sreplicate @2 (siota (SNat @2)) in sscatter @[2, 2, 2, 2, 2] (stranspose @[1, 6, 5, 4, 3, 7, 0, 2] (sscatter @[2, 2] (stranspose @[6, 3, 7, 0, 1, 4, 5, 2] (sscatter @[2, 2, 2] (stranspose @[2, 4, 7, 3, 0, 5, 6, 1] (sscatter @[2, 2] (stranspose @[3, 7, 0, 1, 2, 4, 5, 6] (stranspose @[8, 5, 4, 6, 3, 0, 1, 7, 2] (soneHot (stranspose @[5, 0, 3, 4, 1, 2] (soneHot (str (sfromR dret) !$ [0]) [0, 0, 1]) + stranspose @[5, 0, 3, 4, 1, 2] (soneHot (str (sfromR dret) !$ [1]) [1, 0, 0])) [0, 0, 0]) !$ [0])) (\\[i161, i162] -> [i161, 2 * i161 + i162]))) (\\[i163, i164, i165] -> [i163, i164, i163 * i164 + i165]))) (\\[i166, i167] -> [3 + (negate i167 + i166), i167]))) (\\[i168, i169, i170, i171, i172] -> [1 + m145 `sindex0` [i168, i169]])))"+      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sscatter (ssum @1 (sslice (SNat @1) (SNat @1) (str (sfromR dret)))) (\\[i86, i87, i88] -> [i86 * i87, 2 * i88])) [2, 2])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[0, 2, 1] (sscatter @[2, 2, 2, 2, 2] (stranspose @[1, 6, 5, 4, 3, 7, 0, 2] (sscatter @[2, 2] (stranspose @[6, 3, 7, 0, 1, 4, 5, 2] (sscatter @[2, 2, 2] (stranspose @[2, 4, 7, 3, 0, 5, 6, 1] (sscatter @[2, 2] (stranspose @[8, 3, 2, 5, 4, 6, 0, 1, 7] (soneHot (stranspose @[5, 0, 3, 4, 1, 2] (soneHot (str (sfromR dret) !$ [0]) [0, 0, 1]) + stranspose @[5, 0, 3, 4, 1, 2] (soneHot (str (sfromR dret) !$ [1]) [1, 0, 0])) [0, 0, 0]) !$ [0]) (\\[i161, i162] -> [i161, 2 * i161 + i162]))) (\\[i163, i164, i165] -> [i163, i164, i163 * i164 + i165]))) (\\[i166, i167] -> [3 + (negate i167 + i166), i167]))) (\\[i168, i169, i170, i171, i172] -> [1 + sconcrete (sfromListLinear [2,2] [0,1,1,2]) `sindex0` [i168, i169]])))"+      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sscatter (str (sfromR dret) !$ [1]) (\\[i86, i87, i88] -> [i86 * i87, 2 * i88])) [2, 2])"++testCNNOPP5 :: Assertion+testCNNOPP5 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet PrimalSpan (TKR 4 Double)+      blackGlyph = rreplicateN (6 :$: 6 :$: 6 :$: 6 :$: ZSR)+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet PrimalSpan (TKR 4 Double)+      afcnn2T = conv2dSame4 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sconcrete (sreplicate [1,1,2,2] 7.0))"+  printAstPretty afcnn2T+    @?= "rfromS (sconcrete (sreplicate [1,1,2,2] 7.0))"++testCNNOPP5b :: Assertion+testCNNOPP5b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dSame4 (FTKR [5, 5, 5, 5] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (sreplicateN @[1, 1] (str (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (sfromR u1 !$ [0, 0]))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (sreplicateN @[1, 1] (str (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (sfromR u1 !$ [0, 0]))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (soneHot (sappend (sconcrete (sfromListLinear [0,5] [])) (sappend (str (sappend (sconcrete (sfromListLinear [0,2] [])) (sappend (str (ssumN @[1, 1] (sfromR dret))) (sconcrete (sreplicate [3,2] 0.0))))) (sconcrete (sreplicate [3,5] 0.0)))) [0, 0])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (sappend (str (sappend (stranspose @[0, 1, 3, 2] (sfromR dret) !$ [0, 0]) (sconcrete (sreplicate [3,2] 0.0)))) (sconcrete (sreplicate [3,5] 0.0))) [0, 0])"++maxPool2dUnpadded4+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded4 arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez4 [2, 2, 2, 2] arr [bb + 1, 3 - bb, aa * iBh, 2 * iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded4: impossible pattern needlessly required"++conv2dSame4+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dSame4 arrA =+  let shB = [1, 1, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez4 shB arrA [iImg, 0, iBh, iBw]+      in rindex @_ @0 arrAt [0, 0, 0, 0]+    _ -> error "conv2dSame4: impossible pattern needlessly required"++slicez4+  :: (ADReady target, NumScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez4 shOut d ixBase =+  rbuild shOut $ \ixResult -> indexz03 d (ixrZipWith (+) ixBase ixResult)++testCNNOPP6 :: Assertion+testCNNOPP6 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPlain $ rreplicateN (2 :$: 2 :$: 2 :$: 2 :$: ZSR)+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PlainSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded3 $ conv2dSame3z blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromPlain (rfromS (stranspose @[1, 2, 0] (sreplicate @2 (let t31 = sgather1 @2 (stranspose @[2, 1, 0] (sgather1 @2 (stranspose @[1, 2, 0] (sgather1 @2 (sreplicateN @[2, 2] (let m21 = sgather1 @2 (str (sgather1 @2 (sreplicate0N @[2, 2, 2, 2] 7.0) (\\i9 -> [2 * i9, 2 * i9, 2 * i9]))) (\\i12 -> [2 * i12]) in sfromVector @2 (fromList [sfromVector0N @[2] (fromList [m21 `sindex0` [0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]))) (\\i1 -> [2 * i1]))) (\\i2 -> [2 * i2, 0]))) (\\i4 -> [2 * i4]) in sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [t31 `sindex0` [0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)])))))"+++testCNNOPP6b :: Assertion+testCNNOPP6b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dSame3z) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @2 (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @2 (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (soneHot (sfromK (((ssum @2 (stranspose @[2, 0, 1] (sfromR dret)) !$ [0]) !$ [0]) `sindex0` [0])) [0, 0, 0, 0])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (sfromK (ssum0 @[2] (stranspose @[0, 1, 3, 2] (sfromR dret) !$ [0, 0, 0]))) [0, 0, 0, 0])"++conv2dSame3z+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dSame3z arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBw]+      in rindex @_ @0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dSame3z: impossible pattern needlessly required"++testCNNOPP7 :: Assertion+testCNNOPP7 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ rreplicateN (2 :$: 2 :$: 2 :$: 2 :$: ZSR)+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded3y $ conv2dSame3y blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromPlain (rfromS (let u28 = sgather1 @2 (stranspose @[3, 2, 0, 1] (sgather @[2, 2] (stranspose @[1, 2, 0] (sgather1 @2 (sreplicate @2 (stranspose @[1, 2, 0] (sreplicate @2 (let m21 = sgather1 @2 (str (sgather1 @2 (sreplicate0N @[2, 2, 2, 2] 7.0) (\\i9 -> [2 * i9, 2 * i9, 2 * i9]))) (\\i11 -> [2 * i11]) in sfromVector @2 (fromList [sfromVector0N @[2] (fromList [m21 `sindex0` [0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]))))) (\\i1 -> [2 * i1]))) (\\[i32, i3] -> [2 * i3, 2 * i32]))) (\\i4 -> [2 * i4]) in stranspose @[1, 2, 0] (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [u28 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]), sconcrete (sreplicate [2,2,2] 0.0)]))))"++testCNNOPP7b :: Assertion+testCNNOPP7b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3y . conv2dSame3y) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]), sconcrete (sreplicate [2,2,2] 0.0)])))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]), sconcrete (sreplicate [2,2,2] 0.0)])))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (soneHot (sfromK ((((stranspose @[2, 0, 1] (sfromR dret) !$ [0]) !$ [0]) !$ [0]) `sindex0` [0])) [0, 0, 0, 0])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (sfromK (sfromR dret `sindex0` [0, 0, 0, 0])) [0, 0, 0, 0])"++maxPool2dUnpadded3y+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded3y arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez3 [2, 2, 2, 2] arr [iBh, aa, bb, iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded3y: impossible pattern needlessly required"++conv2dSame3y+  :: (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dSame3y arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBh]+      in rindex @_ @0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dSame3y: impossible pattern needlessly required"++testPaddedCNNOPP0c :: Assertion+testPaddedCNNOPP0c = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dCPadded (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (sfromPlain (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [4,2,2,4] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,2.0,0.0,0.0,-2.0,0.0,0.0,0.0,13.1,9.0,0.0,0.0,582934.0,2.99432,0.0,0.0,6.0,1.0,0.0,0.0,0.1,-0.2,0.0,0.0,8.0,-4.0,0.0,0.0,-335.0,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i65, i67] -> [i65 + i67]))) (\\[i42, i43] -> [i42 + i43])))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1)))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w46 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [4,2,2,4] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,2.0,0.0,0.0,-2.0,0.0,0.0,0.0,13.1,9.0,0.0,0.0,582934.0,2.99432,0.0,0.0,6.0,1.0,0.0,0.0,0.1,-0.2,0.0,0.0,8.0,-4.0,0.0,0.0,-335.0,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i44, i45] -> [i44 + i45]))) (\\[i42, i43] -> [i42 + i43])))) in sfromPlain w46 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (ssum @2 (stranspose @[4, 0, 1, 2, 3] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (ssum @2 (stranspose @[3, 0, 1, 2] (let w46 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [4,2,2,4] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,2.0,0.0,0.0,-2.0,0.0,0.0,0.0,13.1,9.0,0.0,0.0,582934.0,2.99432,0.0,0.0,6.0,1.0,0.0,0.0,0.1,-0.2,0.0,0.0,8.0,-4.0,0.0,0.0,-335.0,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i44, i45] -> [i44 + i45]))) (\\[i42, i43] -> [i42 + i43])))) in sfromPlain w46 * sreplicateN @[2, 2, 2] (sfromR dret))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (ssumN @[2, 2] (sdot1In (sfromPlain (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [4,2,2,4] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,2.0,0.0,0.0,-2.0,0.0,0.0,0.0,13.1,9.0,0.0,0.0,582934.0,2.99432,0.0,0.0,6.0,1.0,0.0,0.0,0.1,-0.2,0.0,0.0,8.0,-4.0,0.0,0.0,-335.0,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i75, i77] -> [i75 + i77]))) (\\[i42, i43] -> [i42 + i43])))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret)))))"++testPaddedCNNOPP0b :: Assertion+testPaddedCNNOPP0b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dBPadded (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,-2.0,5.0,-2.0,5.0,-2.0,5.0,-2.0,13.1,582934.0,13.1,582934.0,13.1,582934.0,13.1,582934.0,5.0,-2.0,5.0,-2.0,5.0,-2.0,5.0,-2.0,13.1,582934.0,13.1,582934.0,13.1,582934.0,13.1,582934.0,2.0,0.0,2.0,0.0,2.0,0.0,2.0,0.0,9.0,2.99432,9.0,2.99432,9.0,2.99432,9.0,2.99432,2.0,0.0,2.0,0.0,2.0,0.0,2.0,0.0,9.0,2.99432,9.0,2.99432,9.0,2.99432,9.0,2.99432,6.0,0.1,6.0,0.1,6.0,0.1,6.0,0.1,8.0,-335.0,8.0,-335.0,8.0,-335.0,8.0,-335.0,6.0,0.1,6.0,0.1,6.0,0.1,6.0,0.1,8.0,-335.0,8.0,-335.0,8.0,-335.0,8.0,-335.0,1.0,-0.2,1.0,-0.2,1.0,-0.2,1.0,-0.2,-4.0,26.0,-4.0,26.0,-4.0,26.0,-4.0,26.0,1.0,-0.2,1.0,-0.2,1.0,-0.2,1.0,-0.2,-4.0,26.0,-4.0,26.0,-4.0,26.0,-4.0,26.0])) (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (sgather @[2, 2, 2, 2, 2, 2] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @2) (stranspose @[9, 5, 4, 3, 2, 1, 0, 7, 8, 6] (sreplicateN @[2, 2, 2, 2, 2, 2] (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR u1)) (sconcrete (sreplicate [1,2,2,2] 0.0))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i37, i38, i39, i40, i41, i42] -> [i41 + i39, i37, i41, i42, i38, i39, i40, i37, i38, i42 + i40]))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w43 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (sgather @[2, 2, 2, 2, 2, 2] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @2) (stranspose @[9, 5, 4, 3, 2, 1, 0, 7, 8, 6] (sreplicateN @[2, 2, 2, 2, 2, 2] (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR u1)) (sconcrete (sreplicate [1,2,2,2] 0.0))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i37, i38, i39, i40, i41, i42] -> [i41 + i39, i37, i41, i42, i38, i39, i40, i37, i38, i42 + i40]))) in sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,5.0,5.0,5.0,13.1,13.1,13.1,13.1,5.0,5.0,5.0,5.0,13.1,13.1,13.1,13.1,2.0,2.0,2.0,2.0,9.0,9.0,9.0,9.0,2.0,2.0,2.0,2.0,9.0,9.0,9.0,9.0,6.0,6.0,6.0,6.0,8.0,8.0,8.0,8.0,6.0,6.0,6.0,6.0,8.0,8.0,8.0,8.0,1.0,1.0,1.0,1.0,-4.0,-4.0,-4.0,-4.0,1.0,1.0,1.0,1.0,-4.0,-4.0,-4.0,-4.0,-2.0,-2.0,-2.0,-2.0,582934.0,582934.0,582934.0,582934.0,-2.0,-2.0,-2.0,-2.0,582934.0,582934.0,582934.0,582934.0,0.0,0.0,0.0,0.0,2.99432,2.99432,2.99432,2.99432,0.0,0.0,0.0,0.0,2.99432,2.99432,2.99432,2.99432,0.1,0.1,0.1,0.1,-335.0,-335.0,-335.0,-335.0,0.1,0.1,0.1,0.1,-335.0,-335.0,-335.0,-335.0,-0.2,-0.2,-0.2,-0.2,26.0,26.0,26.0,26.0,-0.2,-0.2,-0.2,-0.2,26.0,26.0,26.0,26.0]) * w43))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (stranspose @[1, 2, 0] (sslice (SNat @0) (SNat @2) (sslice (SNat @1) (SNat @3) (stranspose @[3, 1, 2, 0] (ssumN @[2, 2, 2, 2, 2, 2] (stranspose @[6, 5, 4, 3, 2, 1, 9, 7, 8, 0] (let w51 = sscatter @[2, 2, 2, 2, 2, 2] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,5.0,5.0,5.0,13.1,13.1,13.1,13.1,5.0,5.0,5.0,5.0,13.1,13.1,13.1,13.1,2.0,2.0,2.0,2.0,9.0,9.0,9.0,9.0,2.0,2.0,2.0,2.0,9.0,9.0,9.0,9.0,6.0,6.0,6.0,6.0,8.0,8.0,8.0,8.0,6.0,6.0,6.0,6.0,8.0,8.0,8.0,8.0,1.0,1.0,1.0,1.0,-4.0,-4.0,-4.0,-4.0,1.0,1.0,1.0,1.0,-4.0,-4.0,-4.0,-4.0,-2.0,-2.0,-2.0,-2.0,582934.0,582934.0,582934.0,582934.0,-2.0,-2.0,-2.0,-2.0,582934.0,582934.0,582934.0,582934.0,0.0,0.0,0.0,0.0,2.99432,2.99432,2.99432,2.99432,0.0,0.0,0.0,0.0,2.99432,2.99432,2.99432,2.99432,0.1,0.1,0.1,0.1,-335.0,-335.0,-335.0,-335.0,0.1,0.1,0.1,0.1,-335.0,-335.0,-335.0,-335.0,-0.2,-0.2,-0.2,-0.2,26.0,26.0,26.0,26.0,-0.2,-0.2,-0.2,-0.2,26.0,26.0,26.0,26.0]) * sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i45, i46, i47, i48, i49, i50] -> [i49 + i47, i45, i49, i50, i46, i47, i48, i45, i46, i50 + i48]) in sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @0) (SNat @2) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @0) (SNat @2) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @3) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @3) w51)))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0))))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[1, 2, 0] (sslice (SNat @1) (SNat @2) (ssumN @[2, 2, 2, 2, 2, 2] (stranspose @[6, 5, 4, 3, 2, 1, 0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @0) (SNat @2) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @0) (SNat @2) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @3) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @3) (sscatter @[2, 2, 2, 2, 2, 2] (sdot1In (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,13.1,5.0,13.1,5.0,13.1,5.0,13.1,2.0,9.0,2.0,9.0,2.0,9.0,2.0,9.0,6.0,8.0,6.0,8.0,6.0,8.0,6.0,8.0,1.0,-4.0,1.0,-4.0,1.0,-4.0,1.0,-4.0,-2.0,582934.0,-2.0,582934.0,-2.0,582934.0,-2.0,582934.0,0.0,2.99432,0.0,2.99432,0.0,2.99432,0.0,2.99432,0.1,-335.0,0.1,-335.0,0.1,-335.0,0.1,-335.0,-0.2,26.0,-0.2,26.0,-0.2,26.0,-0.2,26.0,5.0,13.1,5.0,13.1,5.0,13.1,5.0,13.1,2.0,9.0,2.0,9.0,2.0,9.0,2.0,9.0,6.0,8.0,6.0,8.0,6.0,8.0,6.0,8.0,1.0,-4.0,1.0,-4.0,1.0,-4.0,1.0,-4.0,-2.0,582934.0,-2.0,582934.0,-2.0,582934.0,-2.0,582934.0,0.0,2.99432,0.0,2.99432,0.0,2.99432,0.0,2.99432,0.1,-335.0,0.1,-335.0,0.1,-335.0,0.1,-335.0,-0.2,26.0,-0.2,26.0,-0.2,26.0,-0.2,26.0])) (stranspose @[3, 0, 1, 2, 5, 6, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i45, i46, i47, i48, i49, i50] -> [i49 + i47, i45, i49, i50, i46, i47, i48, i45, i46, i50 + i48])))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0))))))))"++testPaddedCNNOPP1e :: Assertion+testPaddedCNNOPP1e = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPaddedB (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (sgather @[2, 2, 2, 2, 2, 2] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @2) (stranspose @[9, 5, 4, 3, 2, 1, 0, 7, 8, 6] (sreplicateN @[2, 2, 2, 2, 2, 2] (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i38, i39, i40, i41, i42, i43] -> [i42 + i40, i38, i42, i43, i39, i40, i41, i38, i39, i43 + i41])))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let u37 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; w44 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (sgather @[2, 2, 2, 2, 2, 2] (stranspose @[1, 2, 3, 4, 5, 6, 0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @2) (stranspose @[9, 5, 4, 3, 2, 1, 0, 7, 8, 6] (sreplicateN @[2, 2, 2, 2, 2, 2] (sslice (SNat @1) (SNat @2) u37))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)))))))) (\\[i38, i39, i40, i41, i42, i43] -> [i38, i42, i43, i39, i40, i41, i42 + i40, i38, i39, i43 + i41]))) in w44 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1)))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (tpair (ssumN @[2, 2] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (sgather @[2, 2, 2, 2, 2, 2] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @2) (stranspose @[9, 5, 4, 3, 2, 1, 0, 7, 8, 6] (sreplicateN @[2, 2, 2, 2, 2, 2] (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i38, i39, i40, i41, i42, i43] -> [i42 + i40, i38, i42, i43, i39, i40, i41, i38, i39, i43 + i41])))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret))))) (stranspose @[1, 2, 0] (sslice (SNat @1) (SNat @2) (ssumN @[2, 2, 2, 2, 2, 2] (stranspose @[6, 5, 4, 3, 2, 1, 0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @0) (SNat @2) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @0) (SNat @2) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @3) (stranspose @[9, 1, 2, 3, 4, 5, 6, 7, 8, 0] (sslice (SNat @1) (SNat @3) (sscatter @[2, 2, 2, 2, 2, 2] (sdot1In (sreplicate @2 (stranspose @[2, 3, 4, 0, 5, 1] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1))))))) (stranspose @[3, 0, 1, 2, 5, 6, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i46, i47, i48, i49, i50, i51] -> [i50 + i48, i46, i50, i51, i47, i48, i49, i46, i47, i51 + i49])))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))))))))"++testPaddedCNNOPPLet :: Assertion+testPaddedCNNOPPLet = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPaddedLet (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 0, 1, 2] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,3] 0.0))))))) (\\[i64, i66] -> [i64 + i66]))) (\\[i37, i38] -> [i37 + i38]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let u33 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; u34 = stranspose @[1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @2) u33)))) (sconcrete (sreplicate [1,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,3] 0.0))))))) ; w39 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (stranspose @[2, 0, 1] u34) (\\[i35, i36] -> [i35 + i36]))) (\\[i37, i38] -> [i37 + i38])))) in w39 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1)))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (tpair (ssumN @[2, 2] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 2] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 0, 1, 2] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,3] 0.0))))))) (\\[i74, i76] -> [i74 + i76]))) (\\[i37, i38] -> [i37 + i38]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret))))) (stranspose @[1, 2, 3, 0] (sslice (SNat @0) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @0) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @3) (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @3) (sscatter @[2, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[2, 2] (sdot1In (stranspose @[6, 3, 0, 5, 1, 2, 4] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1)))))))) (stranspose @[6, 2, 3, 5, 0, 1, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i41, i42] -> [i41 + i42]))) (\\[i43, i44] -> [i43 + i44])))))))))))"++conv2dPaddedLet+  :: forall target r. (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPaddedLet arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]+      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case+        [iImg, iCinp, iPh, iPw] ->+          ifH (iPh <. fromIntegral (nKh `div` 2)+               ||* iPw <. fromIntegral (nKw `div` 2)+               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)+               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))+              (rscalar 0)+              (arrA ! [ iImg+                      , iCinp+                      , iPh - fromIntegral (nKh `div` 2)+                      , iPw - fromIntegral (nKw `div` 2) ])+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in tlet arrAPadded $ \arrAPadded2 -> rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezL shK1 arrAPadded2 [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rfromK $ rdot0 arrAt arrKt+    _ -> error "conv2dPaddedLet: impossible pattern needlessly required"++testPaddedCNNOPPLet2 :: Assertion+testPaddedCNNOPPLet2 = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPaddedLet2 (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[3, 5, 2, 0, 4, 1] (sgather @[2, 2] (stranspose @[1, 2, 0] (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 0, 1, 2] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,3] 0.0))))))))) (\\[i118, i120] -> [i118 + i120]))) (\\[i58, i59] -> [i58, i58 + i59]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let u49 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; u50 = stranspose @[1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @2) u49)))) (sconcrete (sreplicate [1,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,3] 0.0))))))) ; w60 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[3, 5, 2, 0, 4, 1] (sgather @[2, 2] (sgather @[2, 4, 2, 2, 2] (stranspose @[1, 3, 4, 0, 2] (sreplicateN @[2, 2] u50)) (\\[i51, i52, i53, i54, i55] -> [i53, i55, i52, i51, i53])) (\\[i56, i57] -> [i56, i56 + i57]))) (\\[i58, i59] -> [i58, i58 + i59])))) in w60 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1)))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (tpair (ssumN @[2, 2] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 2] (stranspose @[3, 5, 2, 0, 4, 1] (sgather @[2, 2] (stranspose @[1, 2, 0] (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[3, 0, 1, 2] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,3] 0.0))))))))) (\\[i142, i144] -> [i142 + i144]))) (\\[i58, i59] -> [i58, i58 + i59]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret))))) (stranspose @[1, 2, 3, 0] (sslice (SNat @0) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @0) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @3) (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @3) (ssumN @[2, 2] (stranspose @[4, 0, 3, 1, 2] (sscatter @[4, 2, 2, 2] (stranspose @[1, 2, 3, 4, 0] (sscatter @[2, 2] (stranspose @[3, 5, 2, 0, 4, 1] (sscatter @[2, 2] (sdot1In (stranspose @[6, 3, 0, 5, 1, 2, 4] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1)))))))) (stranspose @[6, 2, 3, 5, 0, 1, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i62, i63] -> [i62, i62 + i63]))) (\\[i64, i65] -> [i64, i64 + i65]))) (\\[i67, i68, i69, i70] -> [i68, i68, i70, i67])))))))))))))"++conv2dPaddedLet2+  :: forall target r. (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPaddedLet2 arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]+      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case+        [iImg, iCinp, iPh, iPw] ->+          ifH (iPh <. fromIntegral (nKh `div` 2)+               ||* iPw <. fromIntegral (nKw `div` 2)+               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)+               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))+              (rscalar 0)+              (arrA ! [ iImg+                      , iCinp+                      , iPh - fromIntegral (nKh `div` 2)+                      , iPw - fromIntegral (nKw `div` 2) ])+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = tlet arrAPadded $ \arrAPadded2 ->+                    slicezL shK1 arrAPadded2 [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rfromK $ rdot0 arrAt arrKt+    _ -> error "conv2dPaddedLet2: impossible pattern needlessly required"++testPaddedCNNOPP2 :: Assertion+testPaddedCNNOPP2 = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPadded2 (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[5, 4, 0, 1, 2, 3] (sgather @[2, 2, 2, 2] (sappend (sconcrete (sreplicate [1,4,2,2] 0.0)) (sappend (stranspose @[3, 0, 2, 1] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,4,2,2] 0.0)))) (\\[i25, i26, i27, i28] -> [i28 + i26, i27 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w29 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (sgather @[2, 2, 2, 2, 2, 2] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3, 1, 2, 0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))) (\\[i23, i24, i25, i26, i27, i28] -> [i28 + i26, i23, i24, i27 + i25]))) in w29 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1)))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (tpair (ssumN @[2, 2] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[5, 4, 0, 1, 2, 3] (sgather @[2, 2, 2, 2] (sappend (sconcrete (sreplicate [1,4,2,2] 0.0)) (sappend (stranspose @[3, 0, 2, 1] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,4,2,2] 0.0)))) (\\[i25, i26, i27, i28] -> [i28 + i26, i27 + i25]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret))))) (stranspose @[1, 2, 0] (sslice (SNat @1) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @1) (SNat @2) (sscatter @[2, 2, 2, 2, 2, 2] (sdot1In (sreplicate @2 (stranspose @[2, 3, 4, 0, 5, 1] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR (tproject1 u1))))))) (stranspose @[3, 0, 1, 2, 5, 6, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i31, i32, i33, i34, i35, i36] -> [i36 + i34, i31, i32, i35 + i33])))))))"++conv2dPadded2+  :: forall target r. (ADReady target, NumScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPadded2 arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]+      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case+        [iImg, iCinp, iPh, iPw] ->+               arrA ! [ iImg+                      , iCinp+                      , iPh - fromIntegral (nKh `div` 2)+                      , iPw - fromIntegral (nKw `div` 2) ]+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezL shK1 arrAPadded [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rfromK $ rdot0 arrAt arrKt+    _ -> error "conv2dPadded2: impossible pattern needlessly required"+++-- * Non-laborious CNN PP tests++-- Hand-written reverse derivative of full convolution with respect+-- to the kernels.+testSameCNNOPPKrnHandwritten :: Assertion+testSameCNNOPPKrnHandwritten = do+  resetVarCounter+  let ftk = FTKS (knownShS @'[3, 3, 6, 6]) (FTKScalar @Double)+      varNameA = mkAstVarName ftk . intToAstVarId $ 100000000+      varA = AstVar varNameA+      varNameB = mkAstVarName ftk . intToAstVarId $ 100000099+      varB = AstVar varNameB+      t :: AstTensor AstMethodLet FullSpan (TKS [3, 3, 3, 3] Double)+      t = conv2dSame_dKrn varA varB+  "\\u0 -> \\u99 -> " ++ printAstPretty t+    @?= "\\u0 -> \\u99 -> ssumN @[3, 1, 6, 6] (stranspose @[2, 3, 4, 5, 1, 0] (sreplicate @3 (stranspose @[1, 2, 3, 4, 5, 0] (sreplicate @3 (stranspose @[1, 2, 3, 4, 5, 0] (sreplicate @3 (stranspose @[1, 2, 0] (sreplicate @1 (str u99)))))))) * stranspose @[1, 2, 3, 4, 0] (sreplicate @3 (stranspose @[4, 0, 5, 6, 1, 2, 3] (sreplicate @1 (stranspose @[2, 3, 0, 4, 5, 1] (sgather @[3, 6] (stranspose @[4, 2, 0, 3, 1] (sgather @[3, 6] (stranspose @[2, 1, 0] u0) (\\[i18, i11] -> [i18 + i11]))) (\\[i16, i12] -> [i16 + i12])))))))"+  "\\u0 -> \\u99 -> " ++ printAstPretty (simplifyUserCode t)+    @?= "\\u0 -> \\u99 -> ssumN @[3, 1, 6, 6] (stranspose @[2, 3, 4, 5, 1, 0] (sreplicate @3 (stranspose @[1, 2, 3, 4, 5, 0] (sreplicate @3 (stranspose @[1, 2, 3, 4, 5, 0] (sreplicate @3 (stranspose @[1, 2, 0] (sreplicate @1 (str u99)))))))) * stranspose @[1, 2, 3, 4, 0] (sreplicate @3 (stranspose @[4, 0, 5, 6, 1, 2, 3] (sreplicate @1 (stranspose @[2, 3, 0, 4, 5, 1] (sgather @[3, 6] (stranspose @[4, 2, 0, 3, 1] (sgather @[3, 6] (stranspose @[2, 1, 0] u0) (\\[i35, i37] -> [i35 + i37]))) (\\[i16, i12] -> [i16 + i12])))))))"++-- Convolution differentiated wrt the kernel.+testSameCNNOPP0cW :: Assertion+testSameCNNOPP0cW = do+  resetVarCounter+  let ftkD = FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) (FTKScalar @Double)+      varName = mkAstVarName ftkD . intToAstVarId $ 100000000+      var = AstVar varName+      ftkK = FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) (FTKScalar @Double)+      f = flip conv2dSame var+      varName2 = mkAstVarName ftkD . intToAstVarId $ 100000000+      var2 = AstVar varName2+      env =+        extendEnv varName (dDnotShared (AstRaw var2) (DeltaZero ftkD)) emptyEnv+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f env) ftkK+  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR u0)) (\\[i63, i65] -> [i63 + i65]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR u1)))))))))"+  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR u0)) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR u1))))))))"+  "\\u0 -> " ++ printArtifactPretty artifactRev+    @?= "\\u0 -> \\dret u1 -> rfromS (ssum @6 (stranspose @[4, 0, 1, 2, 3] (ssum @6 (stranspose @[4, 3, 0, 1, 2] (ssum @6 (stranspose @[3, 0, 1, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR u0)) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * sreplicateN @[2, 2, 2] (sfromR dret))))))))"+  "\\u0 -> " ++ printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\dret u1 -> rfromS (ssumN @[6, 6] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR u0)) (\\[i80, i82] -> [i80 + i82]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret)))))"++-- Convolution differentiated wrt the data.+testSameCNNOPP0bW :: Assertion+testSameCNNOPP0bW = do+  resetVarCounter+  let ftkK = FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) (FTKScalar @Double)+      varName = mkAstVarName ftkK . intToAstVarId $ 100000000+      var = AstVar varName+      ftkD = FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) (FTKScalar @Double)+      f = conv2dSame var+      varName2 = mkAstVarName ftkK . intToAstVarId $ 100000000+      var2 = AstVar varName2+      env =+        extendEnv varName (dDnotShared (AstRaw var2) (DeltaZero ftkK)) emptyEnv+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f env) ftkD+  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i46, i48] -> [i46 + i48]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR u0)))))))))"+  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR u0))))))))"+  "\\u0 -> " ++ printArtifactPretty artifactRev+    @?= "\\u0 -> \\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[6, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[6, 2] (stranspose @[5, 3, 0, 4, 1, 2] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR u0)))))) * sreplicateN @[2, 2, 2] (sfromR dret))))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31])))"+  "\\u0 -> " ++ printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[6, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[6, 2] (sdot1In (stranspose @[6, 3, 0, 5, 1, 2, 4] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR u0))))))) (stranspose @[6, 2, 3, 5, 0, 1, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31])))"++-- Convolution differentiated wrt both arguments.+testSameCNNOPP1bW :: Assertion+testSameCNNOPP1bW = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dSame (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i63, i65] -> [i63 + i65]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR (tproject1 u1)))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> tpair (rfromS (ssum @6 (stranspose @[4, 0, 1, 2, 3] (ssum @6 (stranspose @[4, 3, 0, 1, 2] (ssum @6 (stranspose @[3, 0, 1, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * sreplicateN @[2, 2, 2] (sfromR dret))))))))) (rfromS (stranspose @[1, 2, 0] (sscatter @[6, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[6, 2] (stranspose @[5, 3, 0, 4, 1, 2] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR (tproject1 u1))))))) * sreplicateN @[2, 2, 2] (sfromR dret))))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31]))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [6,2,6,6] STKScalar)) (tpair (ssumN @[6, 6] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[6, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[6, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i80, i82] -> [i80 + i82]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret))))) (stranspose @[1, 2, 0] (sscatter @[6, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[6, 2] (sdot1In (stranspose @[6, 3, 0, 5, 1, 2, 4] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @6 (stranspose @[1, 2, 3, 4, 0] (sreplicate @6 (sfromR (tproject1 u1)))))))) (stranspose @[6, 2, 3, 5, 0, 1, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31]))))"++-- Convolution differentiated wrt the kernel.+testShrinkingCNNOPP0cW :: Assertion+testShrinkingCNNOPP0cW = do+  resetVarCounter+  let ftkD = FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) (FTKScalar @Double)+      varName = mkAstVarName ftkD . intToAstVarId $ 100000000+      var = AstVar varName+      ftkK = FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) (FTKScalar @Double)+      f = flip conv2dShrinking var+      varName2 = mkAstVarName ftkD . intToAstVarId $ 100000000+      var2 = AstVar varName2+      env =+        extendEnv varName (dDnotShared (AstRaw var2) (DeltaZero ftkD)) emptyEnv+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f env) ftkK+  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR u0)) (\\[i63, i65] -> [i63 + i65]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR u1)))))))))"+  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR u0)) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR u1))))))))"+  "\\u0 -> " ++ printArtifactPretty artifactRev+    @?= "\\u0 -> \\dret u1 -> rfromS (ssum @5 (stranspose @[4, 0, 1, 2, 3] (ssum @5 (stranspose @[4, 3, 0, 1, 2] (ssum @6 (stranspose @[3, 0, 1, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR u0)) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * sreplicateN @[2, 2, 2] (sfromR dret))))))))"+  "\\u0 -> " ++ printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\dret u1 -> rfromS (ssumN @[5, 5] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR u0)) (\\[i80, i82] -> [i80 + i82]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret)))))"++-- Convolution differentiated wrt the data.+testShrinkingCNNOPP0bW :: Assertion+testShrinkingCNNOPP0bW = do+  resetVarCounter+  let ftkK = FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) (FTKScalar @Double)+      varName = mkAstVarName ftkK . intToAstVarId $ 100000000+      var = AstVar varName+      ftkD = FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) (FTKScalar @Double)+      f = conv2dShrinking var+      varName2 = mkAstVarName ftkK . intToAstVarId $ 100000000+      var2 = AstVar varName2+      env =+        extendEnv varName (dDnotShared (AstRaw var2) (DeltaZero ftkK)) emptyEnv+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f env) ftkD+  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i46, i48] -> [i46 + i48]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR u0)))))))))"+  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR u0))))))))"+  "\\u0 -> " ++ printArtifactPretty artifactRev+    @?= "\\u0 -> \\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[5, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[5, 2] (stranspose @[5, 3, 0, 4, 1, 2] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR u0)))))) * sreplicateN @[2, 2, 2] (sfromR dret))))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31])))"+  "\\u0 -> " ++ printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[5, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[5, 2] (sdot1In (stranspose @[6, 3, 0, 5, 1, 2, 4] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR u0))))))) (stranspose @[6, 2, 3, 5, 0, 1, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31])))"++-- Convolution differentiated wrt both arguments.+testShrinkingCNNOPP1bW :: Assertion+testShrinkingCNNOPP1bW = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dShrinking (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i63, i65] -> [i63 + i65]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR (tproject1 u1)))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> tpair (rfromS (ssum @5 (stranspose @[4, 0, 1, 2, 3] (ssum @5 (stranspose @[4, 3, 0, 1, 2] (ssum @6 (stranspose @[3, 0, 1, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in w26 * sreplicateN @[2, 2, 2] (sfromR dret))))))))) (rfromS (stranspose @[1, 2, 0] (sscatter @[5, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[5, 2] (stranspose @[5, 3, 0, 4, 1, 2] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR (tproject1 u1))))))) * sreplicateN @[2, 2, 2] (sfromR dret))))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31]))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [6,2,6,6] STKScalar)) (tpair (ssumN @[5, 5] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[5, 2] (stranspose @[4, 2, 0, 3, 1] (sgather @[5, 2] (stranspose @[2, 0, 1] (sfromR (tproject2 u1))) (\\[i80, i82] -> [i80 + i82]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret))))) (stranspose @[1, 2, 0] (sscatter @[5, 2] (stranspose @[2, 4, 1, 3, 0] (sscatter @[5, 2] (sdot1In (stranspose @[6, 3, 0, 5, 1, 2, 4] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @5 (stranspose @[1, 2, 3, 4, 0] (sreplicate @5 (sfromR (tproject1 u1)))))))) (stranspose @[6, 2, 3, 5, 0, 1, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31]))))"++-- Convolution differentiated wrt the kernel.+testPaddedCNNOPP0cW :: Assertion+testPaddedCNNOPP0cW = do+  resetVarCounter+  let ftkD = FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) (FTKScalar @Double)+      varName = mkAstVarName ftkD . intToAstVarId $ 100000000+      var = AstVar varName+      ftkK = FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) (FTKScalar @Double)+      f = flip conv2dPadded var+      varName2 = mkAstVarName ftkD . intToAstVarId $ 100000000+      var2 = AstVar varName2+      env =+        extendEnv varName (dDnotShared (AstRaw var2) (DeltaZero ftkD)) emptyEnv+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f env) ftkK+  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR u0))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR u1)))))))))"+  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR u0))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR u1))))))))"+  "\\u0 -> " ++ printArtifactPretty artifactRev+    @?= "\\u0 -> \\dret u1 -> rfromS (ssum @7 (stranspose @[4, 0, 1, 2, 3] (ssum @7 (stranspose @[4, 3, 0, 1, 2] (ssum @6 (stranspose @[3, 0, 1, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR u0))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25])))) in w26 * sreplicateN @[2, 2, 2] (sfromR dret))))))))"+  "\\u0 -> " ++ printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\dret u1 -> rfromS (ssumN @[7, 7] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR u0))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret)))))"++-- Convolution differentiated wrt the data.+testPaddedCNNOPP0bW :: Assertion+testPaddedCNNOPP0bW = do+  resetVarCounter+  let ftkK = FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) (FTKScalar @Double)+      varName = mkAstVarName ftkK . intToAstVarId $ 100000000+      var = AstVar varName+      ftkD = FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) (FTKScalar @Double)+      f = conv2dPadded var+      varName2 = mkAstVarName ftkK . intToAstVarId $ 100000000+      var2 = AstVar varName2+      env =+        extendEnv varName (dDnotShared (AstRaw var2) (DeltaZero ftkK)) emptyEnv+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f env) ftkD+  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR u1))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR u0)))))))))"+  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev+    @?= "\\u0 -> \\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR u1))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR u0))))))))"+  "\\u0 -> " ++ printArtifactPretty artifactRev+    @?= "\\u0 -> \\dret u1 -> rfromS (stranspose @[1, 2, 0] (sslice (SNat @1) (SNat @6) (sscatter @[7, 2] (stranspose @[2, 4, 1, 3, 0] (sslice (SNat @1) (SNat @6) (sscatter @[7, 2] (stranspose @[5, 3, 0, 4, 1, 2] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR u0)))))) * sreplicateN @[2, 2, 2] (sfromR dret))))) (\\[i28, i29] -> [i28 + i29])))) (\\[i31, i32] -> [i31 + i32]))))"+  "\\u0 -> " ++ printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\u0 -> \\dret u1 -> rfromS (stranspose @[1, 2, 0] (sslice (SNat @1) (SNat @6) (sscatter @[7, 2] (stranspose @[2, 4, 1, 3, 0] (sslice (SNat @1) (SNat @6) (sscatter @[7, 2] (sdot1In (stranspose @[6, 3, 0, 5, 1, 2, 4] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR u0))))))) (stranspose @[6, 2, 3, 5, 0, 1, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i28, i29] -> [i28 + i29])))) (\\[i31, i32] -> [i31 + i32]))))"++-- Convolution differentiated wrt both arguments.+testPaddedCNNOPP1bW :: Assertion+testPaddedCNNOPP1bW = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPadded (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (ssumN @[2, 2] (sdot1In (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR (tproject2 u1)))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (ssumN @[2, 2, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR (tproject2 u1)))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25])))) in w26 * stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR (tproject1 u1)))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> tpair (rfromS (ssum @7 (stranspose @[4, 0, 1, 2, 3] (ssum @7 (stranspose @[4, 3, 0, 1, 2] (ssum @6 (stranspose @[3, 0, 1, 2] (let w26 = stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR (tproject2 u1)))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25])))) in w26 * sreplicateN @[2, 2, 2] (sfromR dret))))))))) (rfromS (stranspose @[1, 2, 0] (sslice (SNat @1) (SNat @6) (sscatter @[7, 2] (stranspose @[2, 4, 1, 3, 0] (sslice (SNat @1) (SNat @6) (sscatter @[7, 2] (stranspose @[5, 3, 0, 4, 1, 2] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (stranspose @[1, 2, 3, 0] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR (tproject1 u1))))))) * sreplicateN @[2, 2, 2] (sfromR dret))))) (\\[i28, i29] -> [i28 + i29])))) (\\[i31, i32] -> [i31 + i32])))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [6,2,6,6] STKScalar)) (tpair (ssumN @[7, 7] (sdot1In (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @2 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,7,2,2] 0.0)) (stranspose @[4, 2, 0, 3, 1] (sgather @[7, 2] (sappend (sconcrete (sreplicate [1,6,2,6] 0.0)) (stranspose @[2, 0, 1] (sfromR (tproject2 u1)))) (\\[i22, i23] -> [i22 + i23])))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 2, 2] (sfromR dret))))) (stranspose @[1, 2, 0] (sslice (SNat @1) (SNat @6) (sscatter @[7, 2] (stranspose @[2, 4, 1, 3, 0] (sslice (SNat @1) (SNat @6) (sscatter @[7, 2] (sdot1In (stranspose @[6, 3, 0, 5, 1, 2, 4] (sreplicate @6 (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @7 (sfromR (tproject1 u1)))))))) (stranspose @[6, 2, 3, 5, 0, 1, 4] (sreplicateN @[2, 2, 2] (sfromR dret)))) (\\[i28, i29] -> [i28 + i29])))) (\\[i31, i32] -> [i31 + i32])))))"++testCNNOPP4bW :: Assertion+testCNNOPP4bW = do+  resetVarCounter+  let !artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2) (FTKR [7, 7, 7, 7] (FTKScalar @Double))+      !artSimp = simplifyArtifactRev artifactRev+  let ftk1 = FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)+      ftkDt = FTKR (7 :$: 7 :$: 3 :$: 3 :$: ZSR) (FTKScalar @Double)+      env = extendEnv (artVarDtRev artSimp)+                      (tconcrete ftkDt (treplTarget 7 ftkDt))+            $ extendEnv (artVarDomainRev artSimp)+                        (tconcrete ftk1 (treplTarget 42 ftk1)) emptyEnv+  interpretAstFull @Concrete env (artPrimalRev artifactRev)+    @?= interpretAstFull @Concrete env (artPrimalRev artSimp)+  interpretAstFull @Concrete env (artDerivativeRev artifactRev)+    @?= interpretAstFull @Concrete env (artDerivativeRev artSimp)+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (let w47 = sreshape @[7, 7, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[3, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[3, 4] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i181, i182] -> [2 * i181 + i182]))) (\\[i45, i46] -> [2 * i45 + i46]))))) in sgather @[7, 7, 3, 3] w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kargMax (splainPart w47 !$ [i48, i49, i50, i51])]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (let w47 = sreshape @[7, 7, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[3, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[3, 4] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))))) in sgather @[7, 7, 3, 3] w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kargMax (splainPart w47 !$ [i48, i49, i50, i51])]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[3, 4] (stranspose @[3, 4, 1, 2, 0] (sscatter @[3, 4] (stranspose @[3, 5, 0, 1, 2, 4] (ssumN @[1, 1] (stranspose @[4, 5, 0, 1, 2, 3] (sreshape @[7, 7, 3, 3, 1, 1, 4, 4] (let w47 = sreshape @[7, 7, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[3, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[3, 4] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))))) in sscatter @[7, 7, 3, 3] (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kargMax (splainPart w47 !$ [i53, i54, i55, i56])])))))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"+  -- The remH comes from the indexing of reshape rule and it looks terrible,+  -- but w42 looks even worse, depending on available primitives,+  -- so the rule is probably fine.+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[3, 4] (stranspose @[3, 4, 1, 2, 0] (sscatter @[3, 4] (stranspose @[4, 5, 3, 7, 0, 1, 2, 6] (sreshape @[7, 7, 3, 3, 1, 1, 4, 4] (let w226 = sreshape @[7, 7, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[3, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[3, 4] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i221, i222] -> [2 * i221 + i222]))) (\\[i45, i46] -> [2 * i45 + i46]))))) in sscatter @[7, 7, 3, 3] (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kargMax (splainPart w226 !$ [i53, i54, i55, i56])]))) !$ [0, 0]) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"++testCNNOPP4bD :: Assertion+testCNNOPP4bD = do+  resetVarCounter+  setTotalSharing True+  let !artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2) (FTKR [7, 7, 7, 7] (FTKScalar @Double))+      !artSimp = simplifyArtifactRev artifactRev+  setTotalSharing False+  let ftk1 = FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)+      ftkDt = FTKR (7 :$: 7 :$: 3 :$: 3 :$: ZSR) (FTKScalar @Double)+      env = extendEnv (artVarDtRev artSimp)+                      (tconcrete ftkDt (treplTarget 7 ftkDt))+            $ extendEnv (artVarDomainRev artSimp)+                        (tconcrete ftk1 (treplTarget 42 ftk1)) emptyEnv+  interpretAstFull @Concrete env (artPrimalRev artifactRev)+    @?= interpretAstFull @Concrete env (artPrimalRev artSimp)+  interpretAstFull @Concrete env (artDerivativeRev artifactRev)+    @?= interpretAstFull @Concrete env (artDerivativeRev artSimp)+  printArtifactPrimalPretty artSimp+    @?= "\\u1 -> rfromS (let w47 = sreshape @[7, 7, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[3, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[3, 4] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i120, i121] -> [2 * i120 + i121]))) (\\[i45, i46] -> [2 * i45 + i46]))))) in sgather @[7, 7, 3, 3] w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kargMax (splainPart w47 !$ [i48, i49, i50, i51])]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (let w47 = sreshape @[7, 7, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[3, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[3, 4] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))))) in sgather @[7, 7, 3, 3] w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kargMax (splainPart w47 !$ [i48, i49, i50, i51])]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[3, 4] (stranspose @[3, 4, 1, 2, 0] (sscatter @[3, 4] (stranspose @[3, 5, 0, 1, 2, 4] (ssumN @[1, 1] (stranspose @[4, 5, 0, 1, 2, 3] (sreshape @[7, 7, 3, 3, 1, 1, 4, 4] (let w47 = sreshape @[7, 7, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[3, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[3, 4] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))))) in sscatter @[7, 7, 3, 3] (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kargMax (splainPart w47 !$ [i53, i54, i55, i56])])))))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"+  printArtifactPretty artSimp+    @?= "\\dret u1 -> rfromS (stranspose @[1, 2, 0] (sscatter @[3, 4] (stranspose @[3, 4, 1, 2, 0] (sscatter @[3, 4] (stranspose @[4, 5, 3, 7, 0, 1, 2, 6] (sreshape @[7, 7, 3, 3, 1, 1, 4, 4] (let w104 = sreshape @[7, 7, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[3, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[3, 4] (stranspose @[2, 0, 1] (sfromR u1)) (\\[i99, i100] -> [2 * i99 + i100]))) (\\[i45, i46] -> [2 * i45 + i46]))))) in sscatter @[7, 7, 3, 3] (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kargMax (splainPart w104 !$ [i53, i54, i55, i56])]))) !$ [0, 0]) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"++testCNNOPP5aW :: Assertion+testCNNOPP5aW = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @1 (stranspose @[1, 2, 0] (sreplicate @1 (let t50 = sreshape @[2, 7, 16] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (ssumN @[7, 7] (sdot1In (sfromPlain (stranspose @[5, 6, 1, 2, 3, 0, 4] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[5, 1, 0, 6, 2, 3, 4] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i90, i92] -> [i90 + i92]))) (\\[i44, i45] -> [i44 + i45]))))) (sconcrete (sreplicate [2,2,7,2,2,7,7] 0.0)))) (sconcrete (sreplicate [2,2,7,4,2,7,7] 0.0))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[3, 4, 5, 1, 2, 0] (sappend (stranspose @[2, 1, 0] (sappend (sreplicate @2 (str (sreplicate @2 (sfromR u1)))) (sconcrete (sreplicate [2,7,2,2,7,7] 0.0)))) (sconcrete (sreplicate [2,7,4,2,7,7] 0.0)))))))))) in sgather @[2, 7] t50 (\\[i51, i52] -> [i51, i52, kargMax (splainPart t50 !$ [i51, i52])]))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @1 (stranspose @[1, 2, 0] (sreplicate @1 (let w48 = stranspose @[4, 5, 6, 1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[5, 1, 0, 6, 2, 3, 4] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i46, i47] -> [i46 + i47]))) (\\[i44, i45] -> [i44 + i45]))))) (sconcrete (sreplicate [2,2,7,2,2,7,7] 0.0)))) (sconcrete (sreplicate [2,2,7,4,2,7,7] 0.0))) ; w49 = stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[3, 4, 5, 1, 2, 0] (sappend (stranspose @[2, 1, 0] (sappend (sreplicate @2 (str (sreplicate @2 (sfromR u1)))) (sconcrete (sreplicate [2,7,2,2,7,7] 0.0)))) (sconcrete (sreplicate [2,7,4,2,7,7] 0.0))))) ; t50 = sreshape @[2, 7, 16] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (ssumN @[2, 7, 7] (sfromPlain w48 * w49)))) in sgather @[2, 7] t50 (\\[i51, i52] -> [i51, i52, kargMax (splainPart t50 !$ [i51, i52])]))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (ssum @2 (str (ssum @2 (sslice (SNat @0) (SNat @2) (stranspose @[2, 1, 0] (sslice (SNat @0) (SNat @2) (stranspose @[5, 3, 4, 0, 1, 2] (ssum @2 (stranspose @[3, 0, 1, 2] (let w48 = stranspose @[4, 5, 6, 1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[5, 1, 0, 6, 2, 3, 4] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i46, i47] -> [i46 + i47]))) (\\[i44, i45] -> [i44 + i45]))))) (sconcrete (sreplicate [2,2,7,2,2,7,7] 0.0)))) (sconcrete (sreplicate [2,2,7,4,2,7,7] 0.0))) ; w49 = stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[3, 4, 5, 1, 2, 0] (sappend (stranspose @[2, 1, 0] (sappend (sreplicate @2 (str (sreplicate @2 (sfromR u1)))) (sconcrete (sreplicate [2,7,2,2,7,7] 0.0)))) (sconcrete (sreplicate [2,7,4,2,7,7] 0.0))))) ; t50 = sreshape @[2, 7, 16] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (ssumN @[2, 7, 7] (sfromPlain w48 * w49)))) in sfromPlain w48 * sreplicateN @[2, 7, 7] (ssumN @[1, 1] (stranspose @[2, 3, 0, 1] (sreshape @[2, 7, 1, 1, 4, 4] (sscatter @[2, 7] (ssum @1 (stranspose @[2, 0, 1] (ssum @1 (stranspose @[2, 0, 1] (sfromR dret))))) (\\[i54, i55] -> [i54, i55, kargMax (splainPart t50 !$ [i54, i55])])))))))))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (ssumN @[2, 2] (sslice (SNat @0) (SNat @2) (stranspose @[2, 0, 1] (sslice (SNat @0) (SNat @2) (let w48 = sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[5, 1, 0, 6, 2, 3, 4] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i116, i118] -> [i116 + i118]))) (\\[i44, i45] -> [i44 + i45]))))) (sconcrete (sreplicate [2,2,7,2,2,7,7] 0.0)))) (sconcrete (sreplicate [2,2,7,4,2,7,7] 0.0)) ; t50 = sreshape @[2, 7, 16] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (ssumN @[7, 7] (sdot1In (sfromPlain (stranspose @[5, 6, 1, 2, 3, 0, 4] w48)) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[3, 4, 5, 1, 2, 0] (sappend (stranspose @[2, 1, 0] (sappend (sreplicate @2 (str (sreplicate @2 (sfromR u1)))) (sconcrete (sreplicate [2,7,2,2,7,7] 0.0)))) (sconcrete (sreplicate [2,7,4,2,7,7] 0.0)))))))))) in sdot1In (sfromPlain (stranspose @[0, 2, 3, 4, 5, 6, 1] w48)) (stranspose @[6, 4, 5, 0, 1, 2, 3] (sreplicateN @[2, 7, 7] (stranspose @[2, 3, 0, 1] (sreshape @[2, 7, 1, 1, 4, 4] (sscatter @[2, 7] (stranspose @[2, 3, 0, 1] (sfromR dret) !$ [0, 0]) (\\[i54, i55] -> [i54, i55, kargMax (splainPart t50 !$ [i54, i55])]))) !$ [0, 0]))))))))"++testCNNOPP5bW :: Assertion+testCNNOPP5bW = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . relu) (FTKR [7, 2, 7, 7] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (let w75 = sreshape @[7, 2, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (sfromPlain (stranspose @[4, 5, 0, 1, 2, 3] (sgather @[3, 3, 4, 4] (sgather @[7, 7, 7, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i65, i66, i67, i68] -> [ifH (0.0 <=. negate (splainPart (sfromR u1) `sindex0` [i67, i68, i65, i66])) 0 1])) (\\[i61, i62, i63, i64] -> [sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i61, i63], sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i62, i64]]))) * stranspose @[4, 5, 0, 1, 2, 3] (sgather @[3, 3, 4, 4] (stranspose @[2, 3, 0, 1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i69, i71], sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i70, i72]]))))) in sgather @[7, 2, 3, 3] w75 (\\[i76, i77, i78, i79] -> [i76, i77, i78, i79, kargMax (splainPart w75 !$ [i76, i77, i78, i79])]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (let m59 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; m60 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; w73 = stranspose @[4, 5, 0, 1, 2, 3] (sgather @[3, 3, 4, 4] (sgather @[7, 7, 7, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i65, i66, i67, i68] -> [ifH (0.0 <=. negate (splainPart (sfromR u1) `sindex0` [i67, i68, i65, i66])) 0 1])) (\\[i61, i62, i63, i64] -> [m60 `sindex0` [i61, i63], m59 `sindex0` [i62, i64]])) ; w74 = stranspose @[4, 5, 0, 1, 2, 3] (sgather @[3, 3, 4, 4] (stranspose @[2, 3, 0, 1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [m60 `sindex0` [i69, i71], m59 `sindex0` [i70, i72]])) ; w75 = sreshape @[7, 2, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (sfromPlain w73 * w74))) in sgather @[7, 2, 3, 3] w75 (\\[i76, i77, i78, i79] -> [i76, i77, i78, i79, kargMax (splainPart w75 !$ [i76, i77, i78, i79])]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (stranspose @[2, 3, 0, 1] (let m59 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; m60 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) in sscatter @[3, 3, 4, 4] (stranspose @[2, 3, 4, 5, 0, 1] (let w73 = stranspose @[4, 5, 0, 1, 2, 3] (sgather @[3, 3, 4, 4] (sgather @[7, 7, 7, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i65, i66, i67, i68] -> [ifH (0.0 <=. negate (splainPart (sfromR u1) `sindex0` [i67, i68, i65, i66])) 0 1])) (\\[i61, i62, i63, i64] -> [m60 `sindex0` [i61, i63], m59 `sindex0` [i62, i64]])) ; w74 = stranspose @[4, 5, 0, 1, 2, 3] (sgather @[3, 3, 4, 4] (stranspose @[2, 3, 0, 1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [m60 `sindex0` [i69, i71], m59 `sindex0` [i70, i72]])) ; w75 = sreshape @[7, 2, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (sfromPlain w73 * w74))) in sfromPlain w73 * ssumN @[1, 1] (stranspose @[4, 5, 0, 1, 2, 3] (sreshape @[7, 2, 3, 3, 1, 1, 4, 4] (sscatter @[7, 2, 3, 3] (sfromR dret) (\\[i81, i82, i83, i84] -> [i81, i82, i83, i84, kargMax (splainPart w75 !$ [i81, i82, i83, i84])])))))) (\\[i85, i86, i87, i88] -> [m60 `sindex0` [i85, i87], m59 `sindex0` [i86, i88]])))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[2, 3, 0, 1] (sscatter @[3, 3, 4, 4] (let w157 = sgather @[3, 3, 4, 4] (sgather @[7, 7, 7, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i65, i66, i67, i68] -> [ifH (0.0 <=. negate (splainPart (sfromR u1) `sindex0` [i67, i68, i65, i66])) 0 1])) (\\[i61, i62, i63, i64] -> [sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i61, i63], sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i62, i64]]) ; w158 = sreshape @[7, 2, 3, 3, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (sfromPlain (stranspose @[4, 5, 0, 1, 2, 3] w157) * stranspose @[4, 5, 0, 1, 2, 3] (sgather @[3, 3, 4, 4] (stranspose @[2, 3, 0, 1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i69, i71], sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i70, i72]]))))) in sfromPlain w157 * stranspose @[4, 5, 2, 3, 6, 7, 0, 1] (sreshape @[7, 2, 3, 3, 1, 1, 4, 4] (sscatter @[7, 2, 3, 3] (sfromR dret) (\\[i81, i82, i83, i84] -> [i81, i82, i83, i84, kargMax (splainPart w158 !$ [i81, i82, i83, i84])]))) !$ [0, 0]) (\\[i85, i86, i87, i88] -> [sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i85, i87], sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) `sindex0` [i86, i88]])))"++testCNNOPP5cW :: Assertion+testCNNOPP5cW = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (relu . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (let u45 = ssumN @[7, 7] (sdot1In (sfromPlain (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i69, i71] -> [i69 + i71]))) (\\[i40, i41] -> [i40 + i41])))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1)))))))) in sfromPlain (sgather @[2, 7, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (0.0 <=. negate (splainPart u45 `sindex0` [i46, i47, i48, i49])) 0 1])) * u45)"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (ssumN @[2, 2] (let w44 = sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i79, i81] -> [i79 + i81]))) (\\[i40, i41] -> [i40 + i41]) ; u45 = ssumN @[7, 7] (sdot1In (sfromPlain (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] w44)))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1)))))))) in sdot1In (sfromPlain (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] w44)))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 7, 7] (sfromPlain (sgather @[2, 7, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (0.0 <=. negate (splainPart u45 `sindex0` [i46, i47, i48, i49])) 0 1])) * sfromR dret)))))"++testCNNOPP5dW :: Assertion+testCNNOPP5dW = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . relu . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @1 (stranspose @[1, 2, 0] (sreplicate @1 (let u68 = ssumN @[7, 7] (sdot1In (sfromPlain (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i119, i121] -> [i119 + i121]))) (\\[i63, i64] -> [i63 + i64])))))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1)))))))) ; t75 = sreshape @[2, 7, 16] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (sfromPlain (stranspose @[1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (sgather @[2, 2, 7, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i69, i70, i71, i72] -> [ifH (0.0 <=. negate (splainPart u68 `sindex0` [i70, i71, i69, i72])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))) * stranspose @[1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[2, 0, 1] u68) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))))) in sgather @[2, 7] t75 (\\[i76, i77] -> [i76, i77, kargMax (splainPart t75 !$ [i76, i77])]))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @1 (stranspose @[1, 2, 0] (sreplicate @1 (let w67 = stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i65, i66] -> [i65 + i66]))) (\\[i63, i64] -> [i63 + i64])))) ; u68 = ssumN @[2, 7, 7] (sfromPlain w67 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1))))))) ; u73 = stranspose @[1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (sgather @[2, 2, 7, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i69, i70, i71, i72] -> [ifH (0.0 <=. negate (splainPart u68 `sindex0` [i70, i71, i69, i72])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; u74 = stranspose @[1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[2, 0, 1] u68) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; t75 = sreshape @[2, 7, 16] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (sfromPlain u73 * u74))) in sgather @[2, 7] t75 (\\[i76, i77] -> [i76, i77, kargMax (splainPart t75 !$ [i76, i77])]))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (ssum @2 (stranspose @[4, 0, 1, 2, 3] (ssum @2 (stranspose @[4, 3, 0, 1, 2] (ssum @2 (stranspose @[3, 0, 1, 2] (let w67 = stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] (sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i65, i66] -> [i65 + i66]))) (\\[i63, i64] -> [i63 + i64])))) ; u68 = ssumN @[2, 7, 7] (sfromPlain w67 * stranspose @[1, 2, 3, 0] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1))))))) ; u73 = stranspose @[1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (sgather @[2, 2, 7, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i69, i70, i71, i72] -> [ifH (0.0 <=. negate (splainPart u68 `sindex0` [i70, i71, i69, i72])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; u74 = stranspose @[1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[2, 0, 1] u68) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; t75 = sreshape @[2, 7, 16] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (sfromPlain u73 * u74))) ; u81 = stranspose @[3, 0, 1, 2] (sfromPlain u73 * ssumN @[1, 1] (stranspose @[2, 3, 0, 1] (sreshape @[2, 7, 1, 1, 4, 4] (sscatter @[2, 7] (ssum @1 (stranspose @[2, 0, 1] (ssum @1 (stranspose @[2, 0, 1] (sfromR dret))))) (\\[i79, i80] -> [i79, i80, kargMax (splainPart t75 !$ [i79, i80])]))))) in sfromPlain w67 * sreplicateN @[2, 7, 7] (stranspose @[1, 2, 0] (sslice (SNat @0) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @0) (SNat @2) u81)))))))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (ssumN @[2, 2] (let w67 = sgather @[2, 7] (stranspose @[4, 2, 0, 3, 1] (sgather @[2, 7] (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i149, i151] -> [i149 + i151]))) (\\[i63, i64] -> [i63 + i64]) ; u68 = ssumN @[7, 7] (sdot1In (sfromPlain (stranspose @[3, 4, 1, 0, 5, 6, 2] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] w67)))) (stranspose @[2, 3, 0, 4, 5, 6, 1] (sreplicate @2 (stranspose @[2, 3, 4, 1, 0] (sreplicate @2 (stranspose @[1, 2, 3, 4, 0] (sreplicate @2 (sfromR u1)))))))) ; u73 = sappend (stranspose @[3, 1, 2, 0] (sappend (sgather @[2, 2, 7, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i69, i70, i71, i72] -> [ifH (0.0 <=. negate (splainPart u68 `sindex0` [i70, i71, i69, i72])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)) ; t75 = sreshape @[2, 7, 16] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (sfromPlain (stranspose @[1, 2, 3, 0] u73) * stranspose @[1, 2, 3, 0] (sappend (stranspose @[3, 1, 2, 0] (sappend (stranspose @[2, 0, 1] u68) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))))) in sdot1In (sfromPlain (stranspose @[5, 6, 0, 2, 3, 4, 1] (sreplicate @7 (stranspose @[2, 4, 5, 1, 3, 0] w67)))) (stranspose @[5, 6, 4, 0, 1, 2, 3] (sreplicateN @[2, 7, 7] (sfromPlain (stranspose @[1, 2, 0] (sslice (SNat @0) (SNat @2) (stranspose @[3, 1, 2, 0] (sslice (SNat @0) (SNat @2) u73)))) * stranspose @[3, 2, 1, 0] (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (stranspose @[2, 3, 4, 5, 1, 0] (sreshape @[2, 7, 1, 1, 4, 4] (sscatter @[2, 7] (stranspose @[2, 3, 0, 1] (sfromR dret) !$ [0, 0]) (\\[i79, i80] -> [i79, i80, kargMax (splainPart t75 !$ [i79, i80])]))) !$ [0, 0])))))))))"
test/simplified/TestGatherSimplified.hs view
@@ -7,23 +7,27 @@  import Prelude -import Data.Int (Int64)+import Control.Monad.ST.Strict (ST, runST)+import Data.Int (Int64, Int8)+import Data.Type.Equality ((:~:) (Refl))+import Data.Vector.Storable qualified as VS+import Data.Vector.Storable.Mutable qualified as VSM import GHC.Exts (IsList (..)) import GHC.TypeLits (Div, KnownNat, type (<=)) import Test.Tasty import Test.Tasty.HUnit hiding (assert)  import Data.Array.Nested qualified as Nested-import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Lemmas import Data.Array.Nested.Ranked.Shape import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat')  import HordeAd import HordeAd.Core.AstEnv import HordeAd.Core.AstFreshId (resetVarCounter) import HordeAd.Core.AstInterpret import HordeAd.Core.CarriersAst-import HordeAd.Core.Ops  import CrossTesting import EqEpsilon@@ -57,6 +61,7 @@   , testCase "gatherSimp23" testGatherSimp23   , testCase "gatherTranspose33" testGatherTranspose33   , testCase "gatherTransposeBuild33" testGatherTransposeBuild33+  , testCase "gatherTransposeBuild33PP" testGatherTransposeBuild33PP   , testCase "gatherTransposeBuild331" testGatherTransposeBuild331   , testCase "gatherTransposeBuild332" testGatherTransposeBuild332   , testCase "gatherTransposeBuild333" testGatherTransposeBuild333@@ -100,6 +105,8 @@   , testCase "scatterNestedBuild12" testScatterNestedBuild12   , testCase "scatter12" testScatter12   , testCase "scatterBuild12" testScatterBuild12+  , testCase "scatterBuild12b" testScatterBuild12b+  , testCase "scatterBuild12c" testScatterBuild12c   , testCase "scatterSimpPP12" testScatterSimpPP12   , testCase "scatterSimp12" testScatterSimp12 @@ -108,8 +115,8 @@    , testCase "sminimizedCNNOPP2" testCNNOPP2   , testCase "sminimizedCNNOPP2b" testCNNOPP2b---  , testCase "sminimizedCNNOPP3" testCNNOPP3---  , testCase "sminimizedCNNOPP3b" testCNNOPP3b+  , testCase "sminimizedCNNOPP3" testCNNOPP3+  , testCase "sminimizedCNNOPP3b" testCNNOPP3b   , testCase "sminimizedCNNOPP4" testCNNOPP4   , testCase "sminimizedCNNOPP4b" testCNNOPP4b   , testCase "sminimizedCNNOPP5" testCNNOPP5@@ -119,8 +126,15 @@   , testCase "sminimizedCNNOPP7" testCNNOPP7   , testCase "sminimizedCNNOPP7b" testCNNOPP7b   , testCase "minimizedCNNOPP4bU" testCNNOPP4bU+  , testCase "detSquarePP" testDetSquarePP+  , testCase "detSquare3" testDetSquare3+  , testCase "detSquare9" testDetSquare9   ] +rfromIndex0 :: forall r target.+               (BaseTensor target, ConvertTensor target, NumScalar r)+            => IntOf target -> target (TKR 0 r)+rfromIndex0 = rfromIntegral . rfromK . tfromPlain STKScalar  -- * Gathers @@ -129,8 +143,8 @@ gatherNested1 t =   rgather @1           (2 :$: ZSR)-          (rgather @1-                   (4 :$: 2 :$: ZSR) t+          (rgather @1 @1 @1+                   (4 :$: ZSR) t                    (\(k3 :.: ZIR) -> k3 :.: ZIR))           (\(i2 :.: ZIR) -> i2 + i2 :.: i2 :.: ZIR) @@ -180,17 +194,17 @@ testGatherSimpPP1 :: Assertion testGatherSimpPP1 = do   resetVarCounter-  let !t1 = gatherNested1 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 315+  let !t1 = gatherNested1 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 317   resetVarCounter-  let !t2 = gather1 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 315+  let !t2 = gather1 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 317   length (show (simplifyInlineContract @(TKR 1 Float) t1))     @?= length (show (simplifyInlineContract @(TKR 1 Float) @PrimalSpan t2))  testGatherSimp1 :: Assertion testGatherSimp1 = do-  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]       env = extendEnv varName (ringestData [7, 2] vals) emptyEnv@@ -222,7 +236,7 @@ gatherNested02 t =   rgather @1           (1 :$: ZSR)-          (rgather @1+          (rgather @1 @0 @1                    (2 :$: ZSR) t                    (\(k3 :.: ZIR) -> k3 + k3 :.: ZIR))           (\(i1 :.: ZIR) -> i1 + i1 + i1 :.: ZIR)@@ -238,8 +252,8 @@ gatherNested2 t =   rgather @2           (2 :$: 3 :$: ZSR)-          (rgather @3-                   (2 :$: 3 :$: 4 :$: 2 :$: ZSR) t+          (rgather @3 @1 @1+                   (2 :$: 3 :$: 4 :$: ZSR) t                    (\(k1 :.: k2 :.: k3 :.: ZIR) -> k1 + k2 + k3 :.: ZIR))           (\(i1 :.: i2 :.: ZIR) -> i1 :.: i2 :.: i1 + i2 :.: i1 :.: ZIR) @@ -258,7 +272,7 @@                  [6.0,0.0,0.0,0.0,6.0,6.0,0.0,0.0,6.0,6.0,0.0,0.0,0.0,6.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherNested2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+             gatherNested2 (t * rreplicate0N [7, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 7 $ ringestData [2] [0, 1]))  gather2 :: forall target r. (ADReady target, GoodScalar r)@@ -283,23 +297,23 @@                  [6.0,0.0,0.0,0.0,6.0,6.0,0.0,0.0,6.0,6.0,0.0,0.0,0.0,6.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gather2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+             gather2 (t * rreplicate0N [7, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 7 $ ringestData [2] [0, 1]))  testGatherSimpPP2 :: Assertion testGatherSimpPP2 = do   resetVarCounter-  let !t1 = gatherNested2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 582+  let !t1 = gatherNested2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 570   resetVarCounter-  let !t2 = gather2 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 394-  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t1)) @?= 582-  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t2)) @?= 394+  let !t2 = gather2 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 390+  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t1)) @?= 570+  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t2)) @?= 390  testGatherSimp2 :: Assertion testGatherSimp2 = do-  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]       env = extendEnv varName (ringestData [7, 2] vals) emptyEnv@@ -330,8 +344,8 @@                => target (TKR 2 r) -> target (TKR 2 r) gatherNested12 t =   rgather @1-          (2 :$: 4 :$: ZSR)-          (rgather @3+          (2 :$: ZSR)+          (rgather @3 @0 @2                    (2 :$: 3 :$: 4 :$: ZSR) t                    (\(k1 :.: k2 :.: k3 :.: ZIR) -> k1 + k2 + k3 :.: k1 :.: ZIR))           (\(i1 :.: ZIR) -> i1 :.: i1 + i1 :.: ZIR)@@ -383,17 +397,17 @@ testGatherSimpPP12 :: Assertion testGatherSimpPP12 = do   resetVarCounter-  let !t1 = gatherNested12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 515+  let !t1 = gatherNested12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 337   resetVarCounter-  let !t2 = gather12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 341-  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 515-  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 341+  let !t2 = gather12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 337+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 337+  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 337  testGatherSimp12 :: Assertion testGatherSimp12 = do-  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]       env = extendEnv varName (ringestData [7, 2] vals) emptyEnv@@ -444,24 +458,24 @@                  [6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherReshape22 (t * rreplicate0N [6, 2] (rfromIndex0 i))))+             gatherReshape22 (t * rreplicate0N [6, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 6 $ ringestData [2] [0, 1]))  testGatherSimpPP22 :: Assertion testGatherSimpPP22 = do   resetVarCounter-  let !t1 = gatherReshape22 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  let !t1 = gatherReshape22 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) . intToAstVarId $ 100000000)   length (show t1) @?= 159   length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 159   resetVarCounter   let !t2 = rreshape @2 @2 [2, 6]-            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) . intToAstVarId $ 100000000)   length (show t2) @?= 159   length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t2)) @?= 159  testGatherSimp22 :: Assertion testGatherSimp22 = do-  let varName = mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [6, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0]       env = extendEnv varName (ringestData [6, 2] vals) emptyEnv@@ -493,36 +507,36 @@   resetVarCounter   let !t1 = (\t -> rbuild1 4 (\i ->               gatherReshape22 @(AstTensor AstMethodLet PrimalSpan)-                (t * rreplicate0N [6, 2] (rfromIndex0 i))))-            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 450-  length (show (simplifyInlineContract @(TKR 3 Float) t1)) @?= 456+                (t * rreplicate0N [6, 2] (kfromR $ rfromIndex0 i))))+            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 326+  length (show (simplifyInlineContract @(TKR 3 Float) t1)) @?= 443   resetVarCounter   let !t2 = (\t -> rbuild1 4 (\i ->               rreshape @2 @2 [2, 6]-                (t * rreplicate0N [6, 2] (rfromIndex0 i))))-            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 450-  length (show (simplifyInlineContract @(TKR 3 Float) @PrimalSpan t2)) @?= 456+                (t * rreplicate0N [6, 2] (kfromR $ rfromIndex0 i))))+            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 326+  length (show (simplifyInlineContract @(TKR 3 Float) @PrimalSpan t2)) @?= 443  testGatherSimp23 :: Assertion testGatherSimp23 = do-  let varName = mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [6, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0]       env = extendEnv varName (ringestData [6, 2] vals) emptyEnv   let !t1 = (\t -> rbuild1 4 (\i ->               gatherReshape22 @(AstTensor AstMethodLet PrimalSpan)-                (t * rreplicate0N [6, 2] (rfromIndex0 i)))) var+                (t * rreplicate0N [6, 2] (kfromR $ rfromIndex0 i)))) var   let !t2 = (\t -> rbuild1 4 (\i ->               rreshape @2 @2 [2, 6]-                (t * rreplicate0N [6, 2] (rfromIndex0 i)))) (ringestData [6, 2] vals)+                (t * rreplicate0N [6, 2] (kfromR $ rfromIndex0 i)))) (ringestData [6, 2] vals)   let !t1n = unAstNoSimplify $ (\t -> rbuild1 4 (\i ->               gatherReshape22-                (t * rreplicate0N [6, 2] (rfromIndex0 i)))) $ AstNoSimplify var+                (t * rreplicate0N [6, 2] (kfromR $ rfromIndex0 i)))) $ AstNoSimplify var   let !t2n = unAstNoSimplify $ (\t -> rbuild1 4 (\i ->               rreshape @2 @2 [2, 6]-                (t * rreplicate0N [6, 2] (rfromIndex0 i)))) $ AstNoSimplify var+                (t * rreplicate0N [6, 2] (kfromR $ rfromIndex0 i)))) $ AstNoSimplify var   interpretAstPrimal @Concrete env t1     @?= interpretAstPrimal @Concrete env t1n   interpretAstPrimal @Concrete env t1n@@ -544,7 +558,7 @@  -- Depending on if and how transpose it desugared, this may or may not result -- in dozens of nested gathers that should vanish after simplification.-gatherTranspose33 :: forall target r. (ADReady target, GoodScalar r, RealFloat r)+gatherTranspose33 :: forall target r. (ADReady target, NumScalar r, RealFloat r)                   => target (TKR 10 r) -> target (TKR 2 r) gatherTranspose33 t =   rmatmul2 (rreshape [6, 8] (rconcrete $ unConcrete t48))@@ -586,16 +600,31 @@     (ringestData [1,2,2,1,2,2,2,2,2,1] [487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,480.0,474.0,480.0,474.0,480.0,474.0,480.0,474.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,480.0,474.0,480.0,474.0,480.0,474.0,480.0,474.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,480.0,474.0,480.0,474.0,480.0,474.0,480.0,474.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,480.0,474.0,480.0,474.0,480.0,474.0,480.0,474.0,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherTranspose33 (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (rfromIndex0 i))))+             gatherTranspose33 (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (kfromR $ rfromIndex0 i))))           t128) +testGatherTransposeBuild33PP :: Assertion+testGatherTransposeBuild33PP = do+  resetVarCounter+  let artifactRev =+        revArtifactAdapt+          UseIncomingCotangent+          (\t -> rbuild1 4 (\i ->+             gatherTranspose33+               (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1]+                                 (kfromR $ rfromIndex0 i))))+          (FTKR @10 [1, 2, 2, 1, 2, 2, 2, 2, 2, 1]+                (FTKScalar @Double))+  printArtifactPrimalSimple (simplifyArtifactRev artifactRev)+    @?= "\\w1 -> rfromS (sdot1In (tfromPlain (STKS [4,6,16,8] STKScalar) (sconcrete (sfromListLinear [4,6,16,8] [18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,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,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,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,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,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,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,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,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003]))) (ttranspose (makePerm @[1, 0, 3, 2]) (sreplicate @6 (ttranspose (makePerm @[0, 2, 1]) (sreshape @[4, 16, 8] (ttranspose (makePerm @[0, 1, 3, 2]) (sreshape @[4, 2, 2, 8, 4] (sreplicate @4 (ttranspose (makePerm @[3, 7, 0, 1, 2, 4, 6, 5]) (sfromR w1)) * tfromPlain (STKS [4,1,2,1,2,2,2,2,2,2,1] STKScalar) (ttranspose (makePerm @[10, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) (sreplicateN @[1, 2, 1, 2, 2, 2, 2, 2, 2, 1] (sconcrete (sfromListLinear [4] [0.0,1.0,2.0,3.0]))))))))))))"+ testGatherTransposeBuild331 :: Assertion testGatherTransposeBuild331 =   assertEqualUpToEpsilon' 1e-10     (ringestData [2, 3] [1,1,1,1,1,1])     (rev' @Double @3           (\t -> rbuild1 2 (\i ->-             rtranspose [1, 0] (t * rreplicate0N [2, 3] (rfromIndex0 i))))+             rtranspose [1, 0] (t * rreplicate0N [2, 3] (kfromR $ rfromIndex0 i))))           (ringestData [2, 3] [1,2,3,4,5,6]))  testGatherTransposeBuild332 :: Assertion@@ -604,7 +633,7 @@     (ringestData [2, 3] [1,1,1,1,1,1])     (rev' @Double @3           (\t -> rbuild1 2 (\i ->-             rtranspose [1, 0] (t * rreplicate0N [2, 3] (rfromIndex0 i))))+             rtranspose [1, 0] (t * rreplicate0N [2, 3] (kfromR $ rfromIndex0 i))))           (ringestData [2, 3] [1,2,3,4,5,6]))  testGatherTransposeBuild333 :: Assertion@@ -613,7 +642,7 @@     (ringestData [2] [1,1])     (rev' @Double @2           (\t -> rbuild1 2 (\i ->-             t * rreplicate0N [2] (rfromIndex0 i)))+             t * rreplicate0N [2] (kfromR $ rfromIndex0 i)))           (ringestData [2] [0,0]))  testGatherTransposeBuild334 :: Assertion@@ -651,32 +680,34 @@ testGatherSimpPP33 = do   resetVarCounter   let !t1 = gatherTranspose33 @(AstTensor AstMethodLet PrimalSpan)-            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 1117-  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 847+            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 1046+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 862+  printAstSimple (simplifyInlineContract @(TKR 2 Float) t1)+    @?= "rfromS (smatmul2 (tfromPlain (STKS [6,8] STKScalar) (sconcrete (sfromListLinear [6,8] [18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.89432,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.89432,5.0,2.0,6.0,1.0,-2.0,0.92,0.1,-0.2,13.1,9.0,8.0,-4.0,34.0,2.99432,-33.0,26.0,2.0,2.0,2.0,2.0,-0.2,-0.2,-0.2,-0.2,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003]))) (str (sreshape @[16, 8] (ttranspose (makePerm @[0, 2, 1]) (sreshape @[2, 2, 8, 4] (ttranspose (makePerm @[3, 7, 0, 1, 2, 4, 6, 5]) (sfromR w0)))))))"   resetVarCounter   let !t2 = (\t -> rmatmul2 (rreshape [6, 8] (rconcrete $ unConcrete t48))                             (rreshape @10 [8, 16] t))-            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 796-  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t2)) @?= 526+            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 725+  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t2)) @?= 541  testGatherSimpPP34 :: Assertion testGatherSimpPP34 = do   resetVarCounter   let !t1 = (\t -> rbuild1 4 (\i ->-             gatherTranspose33 @(AstTensor AstMethodLet PrimalSpan) (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (rfromIndex0 i))))-            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 2504-  length (show (simplifyInlineContract @(TKR 3 Float) t1)) @?= 19838+             gatherTranspose33 @(AstTensor AstMethodLet PrimalSpan) (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (kfromR $ rfromIndex0 i))))+            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 1507+  length (show (simplifyInlineContract @(TKR 3 Float) t1)) @?= 17999   resetVarCounter   let !t2 = (\t -> rbuild1 4 (\i ->               (\t' -> rmatmul2 (rreshape [6, 8] (rconcrete $ unConcrete t48))                                (rreshape @10 [8, 16] t'))-                (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (rfromIndex0 i))))-            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 2145-  length (show (simplifyInlineContract @(TKR 3 Float) @PrimalSpan t2)) @?= 19479+                (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (kfromR $ rfromIndex0 i))))+            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 1148+  length (show (simplifyInlineContract @(TKR 3 Float) @PrimalSpan t2)) @?= 17640  {- TODO: re-enable the tests once we drop GHC 9.10    (they don't type-check with 9.10)@@ -701,7 +732,7 @@                  [6.0,6.0,6.0,6.0,6.0,6.0,12.0,0.0,6.0,6.0,6.0,6.0,6.0,6.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherCond (t * rreplicate0N [7, 2] (rfromIndex0 i))))+             gatherCond (t * rreplicate0N [7, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 7 $ ringestData [2] [0, 1]))  gatherCond2 :: forall target r. (ADReady target, GoodScalar r)@@ -725,12 +756,12 @@                  [6.0,6.0,6.0,6.0,6.0,6.0,12.0,0.0,6.0,6.0,6.0,6.0,6.0,6.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherCond2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+             gatherCond2 (t * rreplicate0N [7, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 7 $ ringestData [2] [0, 1]))  testGatherSimpCond :: Assertion testGatherSimpCond = do-  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]       env = extendEnv varName (ringestData [7, 2] vals) emptyEnv@@ -778,7 +809,7 @@                  [6.0,0.0,6.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherCond3 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+             gatherCond3 (t * rreplicate0N [7, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 7 $ ringestData [2] [0, 1]))  gatherCond4 :: forall target r. (ADReady target, GoodScalar r)@@ -802,12 +833,12 @@                  [6.0,0.0,6.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherCond4 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+             gatherCond4 (t * rreplicate0N [7, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 7 $ ringestData [2] [0, 1]))  testGatherSimpCond3 :: Assertion testGatherSimpCond3 = do-  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]       env = extendEnv varName (ringestData [7, 2] vals) emptyEnv@@ -854,7 +885,7 @@                  [6.0,0.0,0.0,0.0,0.0,12.0,0.0,0.0,6.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherCond5 (t * rreplicate0N [2,4,2] (rfromIndex0 i))))+             gatherCond5 (t * rreplicate0N [2,4,2] (kfromR $ rfromIndex0 i))))           (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))  gatherCond6 :: forall target r. (ADReady target, GoodScalar r)@@ -878,12 +909,12 @@                  [6.0,0.0,0.0,0.0,0.0,12.0,0.0,0.0,6.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             gatherCond6 (t * rreplicate0N [2,4,2] (rfromIndex0 i))))+             gatherCond6 (t * rreplicate0N [2,4,2] (kfromR $ rfromIndex0 i))))           (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))  testGatherSimpCond5 :: Assertion testGatherSimpCond5 = do-  let varName = mkAstVarName (FTKR [2,4,2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [2,4,2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1,0,2.0,2.13,0.2,11.0,-17.0,23.0,29.0,-35.0,41.0,1.4,-0.33,33.0,0.1,0.007]       env = extendEnv varName (ringestData [2,4,2] vals) emptyEnv@@ -914,13 +945,13 @@  -- * Scatters instead of gathers -scatterNested1 :: forall target r. (ADReady target, GoodScalar r)+scatterNested1 :: forall target r. (ADReady target, NumScalar r)                => target (TKR 2 r) -> target (TKR 1 r) scatterNested1 t =   rscatter @2           (2 :$: ZSR)           (rscatter @1-                   (7 :$: 2 :$: ZSR) t+                   (7 :$: ZSR) t                    (\(k3 :.: ZIR) -> k3 :.: ZIR))           (\(i1 :.: i2 :.: ZIR) -> i2 `quotH` (1 + i1) :.: ZIR) @@ -941,7 +972,7 @@              ifH (i >. 2) (scatterNested1 t) (t ! [i])))           (rreplicate 7 $ ringestData [2] [0, 1])) -scatter1 :: forall target r. (ADReady target, GoodScalar r)+scatter1 :: forall target r. (ADReady target, NumScalar r)          => target (TKR 2 r) -> target (TKR 1 r) scatter1 t =   rscatter @2@@ -969,17 +1000,17 @@ testScatterSimpPP1 :: Assertion testScatterSimpPP1 = do   resetVarCounter-  let !t1 = scatterNested1 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 397+  let !t1 = scatterNested1 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 312   resetVarCounter-  let !t2 = scatter1 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 492-  length (show (simplifyInlineContract @(TKR 1 Float) @PrimalSpan t1)) @?= 397-  length (show (simplifyInlineContract @(TKR 1 Float) @PrimalSpan t2)) @?= 492+  let !t2 = scatter1 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 453+  length (show (simplifyInlineContract @(TKR 1 Float) @PrimalSpan t1)) @?= 312+  length (show (simplifyInlineContract @(TKR 1 Float) @PrimalSpan t2)) @?= 453  testScatterSimp1 :: Assertion testScatterSimp1 = do-  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]       env = extendEnv varName (ringestData [7, 2] vals) emptyEnv@@ -1007,13 +1038,13 @@                      (simplifyInlineContract @(TKR 1 Float) t2n)     @?= interpretAstPrimal @Concrete env t2n -scatterNested2 :: forall target r. (ADReady target, GoodScalar r)-              => target (TKR 2 r) -> target (TKR 2 r)+scatterNested2 :: forall target r. (ADReady target, NumScalar r)+               => target (TKR 2 r) -> target (TKR 2 r) scatterNested2 t =   rscatter @4           (2 :$: 3 :$: ZSR)           (rscatter @1-                   (2 :$: 3 :$: 4 :$: 2 :$: ZSR) t+                   (2 :$: 3 :$: 4 :$: ZSR) t                    (\(k1 :.: ZIR) -> minH k1 1 :.: minH k1 2  :.: minH k1 3 :.: ZIR))           (\(i1 :.: i2 :.: _i3 :.: i4 :.: ZIR) ->             minH (i1 + i2) 1 :.: minH (i4 + i1) 2 :.: ZIR)@@ -1032,11 +1063,11 @@                  [6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             scatterNested2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+             scatterNested2 (t * rreplicate0N [7, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 7 $ ringestData [2] [0, 1])) -scatter2 :: forall target r. (ADReady target, GoodScalar r)-        => target (TKR 2 r) -> target (TKR 2 r)+scatter2 :: forall target r. (ADReady target, NumScalar r)+         => target (TKR 2 r) -> target (TKR 2 r) scatter2 t =   rscatter @2           (2 :$: 3 :$: ZSR)@@ -1057,23 +1088,23 @@                  [6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0])     (rev' @Double @3           (\t -> rbuild1 4 (\i ->-             scatter2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+             scatter2 (t * rreplicate0N [7, 2] (kfromR $ rfromIndex0 i))))           (rreplicate 7 $ ringestData [2] [0, 1]))  testScatterSimpPP2 :: Assertion testScatterSimpPP2 = do   resetVarCounter-  let !t1 = scatterNested2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 1022+  let !t1 = scatterNested2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 1090   resetVarCounter-  let !t2 = scatter2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 782-  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 1022-  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 782+  let !t2 = scatter2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 621+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 1090+  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 621  testScatterSimp2 :: Assertion testScatterSimp2 = do-  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]       env = extendEnv varName (ringestData [7, 2] vals) emptyEnv@@ -1101,11 +1132,11 @@                      (simplifyInlineContract @(TKR 2 Float) t2n)     @?= interpretAstPrimal @Concrete env t2n -scatterNested12 :: forall target r. (ADReady target, GoodScalar r)+scatterNested12 :: forall target r. (ADReady target, NumScalar r)                => target (TKR 2 r) -> target (TKR 2 r) scatterNested12 t =   rscatter @2-          (2 :$: 4 :$: ZSR)+          (2 :$: ZSR)           (rscatter @2                    (2 :$: 3 :$: 4 :$: ZSR) t                    (\(k1 :.: k2 :.: ZIR) ->@@ -1130,7 +1161,7 @@                           (rtranspose [1, 0] $ rreplicate 4 $ t ! [i]))) [1])           (rreplicate 7 $ ringestData [2] [0, 1])) -scatter12 :: forall target r. (ADReady target, GoodScalar r)+scatter12 :: forall target r. (ADReady target, NumScalar r)          => target (TKR 2 r) -> target (TKR 2 r) scatter12 t =   rscatter @2@@ -1156,20 +1187,48 @@                           (rtranspose [1, 0] $ rreplicate 4 $ t ! [i]))) [1])           (rreplicate 7 $ ringestData [2] [0, 1])) +testScatterBuild12b :: Assertion+testScatterBuild12b =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [0.0,0.0,4.0,4.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2+          (\t -> rindex (rbuild1 5 (\i ->+             tlet (tfromPlain STKScalar $ i >. 2) $ \b ->+               ifH (tplainPart b &&* tplainPart b)+                   (scatter12 t)+                   (rtranspose [1, 0] $ rreplicate 4 $ t ! [i]))+                        ) [1])+          (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterBuild12c :: Assertion+testScatterBuild12c =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [0.0,0.0,4.0,4.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2+          (\t -> rindex (rbuild1 5 (\i ->+             tletPlain (i >. 2) $ \b ->+               ifH (b &&* b)+                   (scatter12 t)+                   (rtranspose [1, 0] $ rreplicate 4 $ t ! [i]))+                       ) [1])+          (rreplicate 7 $ ringestData [2] [0, 1]))+ testScatterSimpPP12 :: Assertion testScatterSimpPP12 = do   resetVarCounter-  let !t1 = scatterNested12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 952+  let !t1 = scatterNested12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 685   resetVarCounter-  let !t2 = scatter12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 625-  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 952-  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 625+  let !t2 = scatter12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 568+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 685+  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 568  testScatterSimp12 :: Assertion testScatterSimp12 = do-  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) . intToAstVarId $ 100000000       var = AstVar varName       vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]       env = extendEnv varName (ringestData [7, 2] vals) emptyEnv@@ -1207,59 +1266,62 @@   in atan2H x w + y * w  barRelu-  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+  :: ( ADReady target, NumScalar r, KnownNat n, Differentiable r )   => target (TKR n r) -> target (TKR n r)-barRelu x = let t = rreplicate0N (rshape x) (rscalar 0.001) * x+barRelu x = let t = rreplicate0N (rshape x) 0.001 * x             in relu $ bar (t, relu t)  barRelu10xSlower-  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+  :: ( ADReady target, NumScalar r, KnownNat n, Differentiable r )   => target (TKR n r) -> target (TKR n r)-barRelu10xSlower x = let t = rmap0N (* rscalar 0.001) x+barRelu10xSlower x = let t = rmap0N (* 0.001) x                      in relu $ bar (t, relu t)  testBarReluADVal320 :: Assertion testBarReluADVal320 =   assertEqualUpToEpsilon 1e-10     (ringestData [1,2,2,1,2,2,2,2,2,1] [2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885034988157713e-4,2.885923176600045e-4,2.887454843457817e-4,2.886097295122454e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.8851415976532735e-4,2.885923176600045e-4,2.887454843457817e-4,2.8849246223035154e-4,2.884182085399516e-4,2.884075468755327e-4,2.8842176240868867e-4,2.8840399312321096e-4,0.0,2.887454843457817e-4,2.886097295122454e-4,2.887454843457817e-4,2.88599069218435e-4,2.887454843457817e-4,2.886097295122454e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.885145151321922e-4,2.8854294397024206e-4,2.8858878438222746e-4,2.885923176600045e-4,0.0,2.884007943794131e-4,0.0,2.884469945274759e-4,2.8843242392031246e-4,2.884288700806792e-4,0.0,2.885034988157713e-4,2.884110805753153e-4,0.0,2.8849283778617973e-4,2.884075468755327e-4,2.884075468755327e-4,2.884075468755327e-4,2.884075468755327e-4,0.0,0.0,0.0,0.0,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,0.0,0.0,0.0,0.0,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,2.8854294397024206e-4,2.884288700806792e-4,2.884395315486472e-4,0.0,2.8849246223035154e-4,2.8850276789489724e-4,0.0,2.8849212704517413e-4,2.8854294397024206e-4,2.884288700806792e-4,2.884395315486472e-4,0.0,2.8849246223035154e-4,2.8850276789489724e-4,0.0,2.8849212704517413e-4,2.8842922547482884e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.894378297730782e-4,2.885923176600045e-4,2.887454843457817e-4,2.88599069218435e-4,2.887454843457817e-4,2.887056688523444e-4,2.887454843457817e-4,2.887056688523444e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.884786229769816e-4,2.885923176600045e-4,2.887454843457817e-4,2.886950092188272e-4,2.887454843457817e-4,2.884818011261814e-4,2.887454843457817e-4,2.886097295122454e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.887167039107226e-4,2.885923176600045e-4,2.887454843457817e-4,2.8860262265516213e-4,2.887454843457817e-4,2.885884088500461e-4,2.887454843457817e-4,2.88599069218435e-4])-    (grad (kfromR . rsum0 @10 @(TKScalar Double) . barRelu10xSlower)-          (rmap0N (* rscalar 0.001) t128))+    (grad (rsum0 @10 @Double . barRelu10xSlower)+          (rmap0N (* 0.001) t128))  testReluSimpPP :: Assertion testReluSimpPP = do   resetVarCounter   let !t1 = barRelu10xSlower @(AstTensor AstMethodLet PrimalSpan)-            $ AstVar (mkAstVarName (FTKR [1,2,2,1,2,2,2,2,2,1] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t1) @?= 22675-  length (show (simplifyInlineContract @(TKR 10 Float) t1)) @?= 22675+            $ AstVar (mkAstVarName (FTKR [1,2,2,1,2,2,2,2,2,1] FTKScalar) . intToAstVarId $ 100000000)+  length (show t1) @?= 19762+  length (show (simplifyInlineContract @(TKR 10 Float) t1)) @?= 19810   resetVarCounter   let !t2 = barRelu @(AstTensor AstMethodLet PrimalSpan)-            $ AstVar (mkAstVarName (FTKR [1,2,2,1,2,2,2,2,2,1] FTKScalar) Nothing . intToAstVarId $ 100000000)-  length (show t2) @?= 12258-  length (show (simplifyInlineContract @(TKR 10 Float) t2)) @?= 12258+            $ AstVar (mkAstVarName (FTKR [1,2,2,1,2,2,2,2,2,1] FTKScalar) . intToAstVarId $ 100000000)+  length (show t2) @?= 11480+  length (show (simplifyInlineContract @(TKR 10 Float) t2)) @?= 11480  testCNNOPP2 :: Assertion testCNNOPP2 = do   resetVarCounter   let t = maxPool2dUnpadded2-            (rconcrete $ Nested.rreplicateScal (1 :$: 1 :$: 2 :$: 2 :$: ZSR) 1)+            (rconcrete $ Nested.rreplicatePrim (1 :$: 1 :$: 2 :$: 2 :$: ZSR) 1)   printAstPretty (simplifyInlineContract t)-    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+    @?= "rfromPlain (rfromS (sreplicateN @[2, 2] (str (sgather1 @2 (stranspose @[2, 3, 0, 1] (sfromVector @2 (fromList [sconcrete (sreplicate [2,2,2,2] 0.0), sgather @[2, 2, 2, 2] (sgather1 @2 (stranspose @[3, 0, 1, 2] (sreplicate @1 (str (sgather1 @3 (stranspose @[1, 0, 2, 4, 3] (sfromVector @2 (fromList [stranspose @[1, 2, 4, 0, 3] (sgather @[2, 2] (stranspose @[3, 2, 0, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [2,3,2] [1.0,1.0,0.0,0.0,0.0,0.0,1.0,1.0,0.0,0.0,0.0,0.0])) (\\[i82, i83] -> [i82 + i83]))) (\\[i27, i19] -> [i27 + i19])), sconcrete (sreplicate [2,3,2,2,2] 0.0)])) !$ [0]) (\\i89 -> [ifH (0 <=. sconcrete (sfromListLinear [3] [1,0,-1]) `sindex0` [i89] &&* 0 <=. negate (sconcrete (sfromListLinear [3] [1,0,-1]) `sindex0` [i89])) 0 1, i89]))))) (\\i24 -> [i24, i24, i24])) (\\[i50, i44, i40, i54] -> [2 * i50 + i40, i50, 2 * i54 + i44])])) !$ [0, 0]) (\\i79 -> [ifH (1 <=. sconcrete (sfromListLinear [2] [1,0]) `sindex0` [i79]) 0 1, i79])))))"+    -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"   printAstPretty t-    @?= "rfromS (sreplicate @2 (sreplicate @2 (let u37 = let u42 = sgather (sgather (sreplicate @1 (let w32 = sgather (stranspose @[3,2,0,1] (sgather (sconcrete (sfromListLinear [2,3,2] [1.0,1.0,0.0,0.0,0.0,0.0,1.0,1.0,0.0,0.0,0.0,0.0])) (\\[i26, i15] -> [i26 + i15]))) (\\[i22, i16] -> [i22 + i16]) in stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,2,2] 0.0)) (sappend (sreplicate @1 (stranspose @[2,0,4,1,3] w32 !$ [1])) (sconcrete (sreplicate [1,2,2,2,2] 0.0)))))) (\\[i20] -> [i20, i20, i20, 0])) (\\[i45, i40, i36, i8] -> [2 * i40 + i8, i40, 2 * i45 + i36]) in str (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sreplicate @1 (str u42 !$ [1]))) in stranspose @[2,3,0,1] u37 !$ [0, 0])))"+    @?= "rfromPlain (rfromS (sreplicateN @[2, 2] (let u46 = let v39 = sconcrete (sreplicate [2] 1) + negate (siota (SNat @2)) in str (sgather1 @2 (stranspose @[2, 0, 1] (sfromVector @2 (fromList [sconcrete (sreplicate [2,2,2,2] 0.0), sgather @[2, 2, 2, 2] (sgather1 @2 (stranspose @[3, 0, 4, 5, 1, 2] (sreplicate @1 (let v18 = sconcrete (sreplicate [3] 1) + negate (siota (SNat @3)) in sgather1 @3 (str (sfromVector @2 (fromList [stranspose @[2, 3, 0, 4, 1] (sgather @[2, 2] (stranspose @[3, 2, 0, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [2,3,2] [1.0,1.0,0.0,0.0,0.0,0.0,1.0,1.0,0.0,0.0,0.0,0.0])) (\\[i32, i21] -> [i32 + i21]))) (\\[i27, i19] -> [i27 + i19])), sconcrete (sreplicate [3,2,2,2,2] 0.0)]))) (\\i35 -> [i35, ifH (0 <=. v18 `sindex0` [i35] &&* 0 <=. negate (v18 `sindex0` [i35])) 0 1])))) (\\i24 -> [i24, i24, i24, 0])) (\\[i54, i50, i44, i40] -> [2 * i50 + i40, i50, 2 * i54 + i44])]))) (\\i49 -> [i49, ifH (1 <=. v39 `sindex0` [i49]) 0 1])) in stranspose @[2, 3, 0, 1] u46 !$ [0, 0])))"  testCNNOPP2b :: Assertion testCNNOPP2b = do   resetVarCounter   let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded2 (FTKR [1, 1, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (sreplicateN @[2, 2] (str (sgather1 @2 (stranspose @[2, 3, 0, 1] (sfromVector @2 (fromList [sconcrete (sreplicate [2,2,2,2] 0.0), sgather @[2, 2, 2, 2] (sgather1 @2 (stranspose @[3, 0, 1, 2] (sreplicate @1 (str (sgather1 @3 (stranspose @[1, 0, 2, 4, 3] (sfromVector @2 (fromList [stranspose @[1, 2, 4, 0, 3] (sgather @[2, 2] (stranspose @[3, 2, 0, 1] (sgather @[2, 2] (str (sappend (sreplicate @1 (sfromR u1 !$ [0, 0])) (sconcrete (sreplicate [2,2,2] 0.0)))) (\\[i205, i206] -> [i205 + i206]))) (\\[i107, i108] -> [i107 + i108])), sconcrete (sreplicate [2,3,2,2,2] 0.0)])) !$ [0]) (\\i212 -> [ifH (0 <=. sconcrete (sfromListLinear [3] [1,0,-1]) `sindex0` [i212] &&* 0 <=. negate (sconcrete (sfromListLinear [3] [1,0,-1]) `sindex0` [i212])) 0 1, i212]))))) (\\i110 -> [i110, i110, i110])) (\\[i111, i112, i113, i114] -> [2 * i111 + i113, i111, 2 * i114 + i112])])) !$ [0, 0]) (\\i115 -> [ifH (1 <=. sconcrete (sfromListLinear [2] [1,0]) `sindex0` [i115]) 0 1, i115]))))"+    -- TODO: was once "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"   printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+    @?= "\\u1 -> rfromS (sreplicateN @[2, 2] (str (let v103 = sconcrete (sreplicate [2] 1) + negate (siota (SNat @2)) in sgather1 @2 (stranspose @[2, 3, 0, 1] (let v104 = sconcrete (sreplicate [3] 1) + negate (siota (SNat @3)) in sfromVector @2 (fromList [sconcrete (sreplicate [2,2,2,2] 0.0), sgather @[2, 2, 2, 2] (sgather1 @2 (stranspose @[3, 0, 1, 2] (sreplicate @1 (str (sgather1 @3 (stranspose @[1, 0, 2, 4, 3] (sfromVector @2 (fromList [stranspose @[1, 2, 4, 0, 3] (sgather @[2, 2] (stranspose @[3, 2, 0, 1] (sgather @[2, 2] (str (sappend (sreplicate @1 (sfromR u1 !$ [0, 0])) (sconcrete (sreplicate [2,2,2] 0.0)))) (\\[i105, i106] -> [i105 + i106]))) (\\[i107, i108] -> [i107 + i108])), sconcrete (sreplicate [2,3,2,2,2] 0.0)])) !$ [0]) (\\i109 -> [ifH (0 <=. v104 `sindex0` [i109] &&* 0 <=. negate (v104 `sindex0` [i109])) 0 1, i109]))))) (\\i110 -> [i110, i110, i110])) (\\[i111, i112, i113, i114] -> [2 * i111 + i113, i111, 2 * i114 + i112])])) !$ [0, 0]) (\\i115 -> [ifH (1 <=. v103 `sindex0` [i115]) 0 1, i115]))))"   printArtifactPretty artifactRev-    @?= "\\dret u1 -> rfromS (sconcrete (sreplicate [1,1,2,2] 0.0))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (sconcrete (sreplicate [1,1,2,2] 0.0))"+    @?= "\\dret u1 -> rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (str (sscatter @[2, 2] (stranspose @[2, 3, 1, 0] (sscatter @[2, 2] (stranspose @[3, 0, 1, 4, 2] (stranspose @[1, 0, 2, 4, 3] (soneHot (let v103 = sconcrete (sreplicate [2] 1) + negate (siota (SNat @2)) ; v104 = sconcrete (sreplicate [3] 1) + negate (siota (SNat @3)) in sscatter1 @3 (str (ssum @1 (stranspose @[1, 2, 3, 0] (sscatter1 @2 (sscatter @[2, 2, 2, 2] (stranspose @[2, 3, 0, 1] (soneHot (sscatter1 @2 (str (ssumN @[2, 2] (sfromR dret))) (\\i117 -> [ifH (1 <=. v103 `sindex0` [i117]) 0 1, i117])) [0, 0]) !$ [1]) (\\[i119, i120, i121, i122] -> [2 * i119 + i121, i119, 2 * i122 + i120])) (\\i123 -> [i123, i123, i123]))))) (\\i124 -> [ifH (0 <=. v104 `sindex0` [i124] &&* 0 <=. negate (v104 `sindex0` [i124])) 0 1, i124])) [0]) !$ [0])) (\\[i126, i127] -> [i126 + i127]))) (\\[i128, i129] -> [i128 + i129]))))) [0, 0])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (sreplicateN @[1, 1] (str (sscatter @[2, 2] (stranspose @[2, 3, 1, 0] (sscatter @[2, 2] (stranspose @[1, 3, 0, 2, 5, 4] (sscatter1 @3 (stranspose @[1, 3, 2, 0] (sscatter1 @2 (sscatter @[2, 2, 2, 2] (sscatter1 @2 (ssumN @[2, 2] (stranspose @[0, 1, 3, 2] (sfromR dret))) (\\i117 -> [ifH (1 <=. sconcrete (sfromListLinear [2] [1,0]) `sindex0` [i117]) 0 1, i117, 0, 0]) !$ [1]) (\\[i119, i120, i121, i122] -> [2 * i119 + i121, i119, 2 * i122 + i120])) (\\i123 -> [i123, i123, i123])) !$ [0]) (\\i124 -> [0, ifH (0 <=. sconcrete (sfromListLinear [3] [1,0,-1]) `sindex0` [i124] &&* 0 <=. negate (sconcrete (sfromListLinear [3] [1,0,-1]) `sindex0` [i124])) 0 1, i124])) !$ [0]) (\\[i126, i127] -> [i126 + i127]))) (\\[i128, i129] -> [i128 + i129])) !$ [0]))"+    -- TODO: was once "\\dret u1 -> rfromS (sconcrete (sreplicate [1,1,2,2] 0.0))"  maxPool2dUnpadded2   :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)@@ -1267,19 +1329,19 @@ maxPool2dUnpadded2 a =   rbuild [2, 2, 2, 2] $ \case     [_, _, iBh, iBw] ->-      let arrt = slicez2 (conv2dUnpadded2 a) [iBw, 1, 2 * iBh, 2 * iBw]+      let arrt = slicez2 (conv2dSame2 a) [iBw, 1, 2 * iBh, 2 * iBw]       in rmaximum2 arrt     _ -> error "maxPool2dUnpadded2: impossible pattern needlessly required" -conv2dUnpadded2+conv2dSame2   :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)   => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded2 a =+conv2dSame2 a =   rbuild [3, 3, 2, 2] $ \case     [iImg, _, iBh, iBw] ->       let arrAt = slicez2 a [iImg, 0, iBh, iBw]-      in rindex0 arrAt [0, iBw, iBw, 0]-    _ -> error "conv2dUnpadded2: impossible pattern needlessly required"+      in rindex @_ @0 arrAt [0, iBw, iBw, 0]+    _ -> error "conv2dSame2: impossible pattern needlessly required"  slicez2   :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)@@ -1295,43 +1357,38 @@  rmaximum2 :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)          => target (TKR 4 r) -> target (TKR 0 r)-rmaximum2 t0 = tlet t0 $ \t -> rindex0 t [0, 0, 0, 0]+rmaximum2 t0 = tlet t0 $ \t -> rindex t [0, 0, 0, 0] -{- TODO: divergent result; bring back when GHC 9.10 dropped: testCNNOPP3 :: Assertion testCNNOPP3 = do   resetVarCounter   let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK+      blackGlyph = AstFromPrimal $ rreplicateN (3 :$: 3 :$: 3 :$: 3 :$: ZSR)                        (rconcrete $ Nested.rscalar 7                         :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))       afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = maxPool2dUnpadded3 $ conv2dUnpadded3 blackGlyph+      afcnn2T = maxPool2dUnpadded3 $ conv2dSame3 blackGlyph   printAstPretty (simplifyInlineContract afcnn2T)     @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [14.0,0.0,14.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"   printAstPretty afcnn2T-    @?= "rfromS (stranspose @[1,2,0] (sreplicate @2 (let w32 = stranspose @[1,2,3,0] (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @2 (let t44 = sreplicate @2 (sgather (sappend (stranspose @[4,1,2,3,5,0] (sappend (sreplicate @1 (stranspose @[0,1,2,4,3] (sgather (str (sreplicate @1 (sconcrete (sfromListLinear [2] [7.0,0.0])))) (\\[i31, i25, i20, i17] -> [ifH (notB (notB (0 <=. negate i31 + negate i20) &&* notB (1 <=. i31 + i20 &&* (-1) <=. negate i31 + negate i20)) &&* notB (notB (0 <=. negate i25 + negate i17) &&* notB (1 <=. i25 + i17 &&* (-1) <=. negate i25 + negate i17))) 0 1])))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @1 (sreplicate @2 (sscalar 0.0))))))))) (stranspose @[4,1,2,3,5,0] (sappend (sreplicate @1 (stranspose @[0,1,2,4,3] (sgather (str (sreplicate @1 (sconcrete (sfromListLinear [2] [7.0,0.0])))) (\\[i30, i24, i19, i17] -> [ifH (notB (notB (0 <=. negate i30 + negate i19) &&* notB (1 <=. i30 + i19 &&* (-1) <=. negate i30 + negate i19)) &&* notB (notB (0 <=. negate i24 + negate i17) &&* notB (1 <=. i24 + i17 &&* (-1) <=. negate i24 + negate i17))) 0 1])))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @1 (sreplicate @2 (sscalar 0.0)))))))))) (\\[i39, i4] -> [2 * i4, 0, 2 * i4, 0, 0, 2 * i39]) + sgather (sappend (sreplicate @1 (sgather (stranspose @[1,2,3,4,5,0] (sappend (sreplicate @1 (stranspose @[0,2,4,3,1] (sgather (str (sreplicate @1 (sconcrete (sfromListLinear [2] [7.0,0.0])))) (\\[i28, i22, i19, i17] -> [ifH (notB (notB (0 <=. negate i28 + negate i19) &&* notB (1 <=. i28 + i19 &&* (-1) <=. negate i28 + negate i19)) &&* notB (notB (0 <=. negate i22 + negate i17) &&* notB (1 <=. i22 + i17 &&* (-1) <=. negate i22 + negate i17))) 0 1])))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sreplicate @1 (sreplicate @2 (sreplicate @2 (sscalar 0.0))))))))) (\\[i27] -> [i27, i27, 0, 1, 0]))) (str (sreplicate @2 (str (sreplicate @2 (sconcrete (sfromListLinear [1] [0.0]))))))) (\\[i38, i4] -> [2 * i4, 0, 2 * i38])) in sappend (str (sappend (stranspose @[1,2,0] (sappend (sgather t44 (\\[i34, i37, i42] -> [i42, i37, i34])) (sreplicate @1 (sreplicate @1 (sreplicate @1 (sscalar 0.0)))))) (sreplicate @1 (sreplicate @1 (sreplicate @2 (sscalar 0.0)))))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sscalar 0.0)))))))))))) in stranspose @[3,4,5,6,0,1,2] w32 !$ [0, 0, 0, 0])))"--}+    @?= "rfromPlain (rfromS (stranspose @[1, 2, 0] (sreplicate @2 (let t124 = sgather1 @2 (stranspose @[2, 1, 0] (sgather1 @2 (str (sgather1 @2 (sreplicate @2 ((let m67 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m68 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2)))) ; v69 = sconcrete (sreplicate [2] 1) + negate (siota (SNat @2)) ; m70 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m71 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2)))) in sgather @[2, 2, 2] (sconcrete (sfromListLinear [2] [7.0,0.0])) (\\[i107, i94, i81] -> [let i60 = ifH (0 <=. negate i94 &&* (notB (notB (0 <=. m71 `sindex0` [i81, i107] &&* 0 <=. negate (m71 `sindex0` [i81, i107])) &&* 1 <=. negate (m70 `sindex0` [i81, i107])) &&* notB (notB (0 <=. m68 `sindex0` [i107, 0] &&* 0 <=. negate (m68 `sindex0` [i107, 0])) &&* 1 <=. negate (m67 `sindex0` [i107, 0])))) 0 1 in ifH (1 <=. i81) (ifH (1 <=. v69 `sindex0` [1]) 1 i60) i60])) + (let m72 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m73 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2)))) ; v74 = sconcrete (sreplicate [2] 1) + negate (siota (SNat @2)) ; m75 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m76 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2)))) in stranspose @[2, 0, 1] (sgather @[2, 2] (sreplicate @1 (sgather @[2, 2, 2] (sconcrete (sfromListLinear [2] [7.0,0.0])) (\\[i88, i47, i114] -> [let i60 = ifH (0 <=. negate i47 &&* (notB (notB (0 <=. m76 `sindex0` [i88, 1] &&* 0 <=. negate (m76 `sindex0` [i88, 1])) &&* 1 <=. negate (m75 `sindex0` [i88, 1])) &&* notB (notB (0 <=. m73 `sindex0` [i114, i114] &&* 0 <=. negate (m73 `sindex0` [i114, i114])) &&* 1 <=. negate (m72 `sindex0` [i114, i114])))) 0 1 in ifH (1 <=. v74 `sindex0` [1]) 1 i60]))) (\\[i100, i87] -> [i87, i87, i100]))))) (\\i1 -> [2 * i1, 0]))) (\\i2 -> [2 * i2]))) (\\i4 -> [2 * i4]) in sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [t124 `sindex0` [0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)])))))" --- TODO: OOMs-_testCNNOPP3b :: Assertion-_testCNNOPP3b = do+testCNNOPP3b :: Assertion+testCNNOPP3b = do   resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dUnpadded3) (FTKR [3, 3, 3, 3] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 1] + sfromR u1 !$ [0, 1, 1, 1])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 1] + sfromR u1 !$ [0, 1, 1, 1])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dSame3) (FTKR [3, 3, 3, 3] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @2 (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 1] + sfromR u1 `sindex0` [0, 1, 1, 1], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]))))"   printArtifactPretty artifactRev-    @?= "\\dret u1 -> let t80 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) in rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t80)))))) [0, 0, 0, 1] + soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t80)))))) [0, 1, 1, 1])"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (let t80 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) in soneHot (t80 !$ [0, 0, 0]) [0, 0, 0, 1] + soneHot (t80 !$ [0, 0, 0]) [0, 1, 1, 1])"+    @?= "\\dret u1 -> rfromS (let m183 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m184 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2)))) ; m185 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m186 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2)))) ; m188 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2)))) ; i191 = ifH (notB (notB (0 <=. m186 `sindex0` [0, 0] &&* 0 <=. negate (m186 `sindex0` [0, 0])) &&* 1 <=. negate (m185 `sindex0` [0, 0])) &&* notB (notB (0 <=. m184 `sindex0` [0, 0] &&* 0 <=. negate (m184 `sindex0` [0, 0])) &&* 1 <=. negate (m183 `sindex0` [0, 0]))) 0 1 ; i192 = ifH (0 <=. m188 `sindex0` [0, 1] &&* 0 <=. negate (m188 `sindex0` [0, 1])) 0 1 ; t194 = ssum @2 (stranspose @[2, 0, 1] (sfromR dret)) ; m195 = t194 !$ [0] ; v196 = m195 !$ [0] ; x197 = v196 `sindex0` [0] ; v198 = soneHot (sfromK x197) [i192] ; v199 = soneHot (sfromK x197) [i191] in soneHot (sfromK (v199 `sindex0` [0])) [0, 0, 0, 1] + soneHot (sfromK (v198 `sindex0` [0])) [0, 1, 1, 1])"+    -- TODO: was once "\\dret u1 -> rfromS (let t88 = ssum @2 (stranspose @[2, 0, 1] (sfromR dret)) ; m89 = t88 !$ [0] ; v90 = m89 !$ [0] ; x91 = v90 `sindex0` [0] in soneHot (sfromK x91) [0, 0, 0, 1] + soneHot (sfromK x91) [0, 1, 1, 1])"+  printAstPretty (simplifyUserCode $ artDerivativeRev artifactRev)+    @?= "rfromS (let m248 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2)))) ; i249 = ifH (0 <=. m248 `sindex0` [0, 1] &&* 0 <=. negate (m248 `sindex0` [0, 1])) 0 1 ; x250 = ssum0 @[2] (stranspose @[0, 1, 3, 2] (sfromR u193) !$ [0, 0, 0]) in sscatter1 @2 (sfromVector0N @[2] (fromList [x250, sscatter1 @1 (sreplicate @1 (sfromK x250)) (\\i251 -> [i249]) `sindex0` [0]])) (\\i226 -> [0, i226, i226, 1]))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (sscatter1 @2 (let x270 = ssum0 @[2] (stranspose @[0, 1, 3, 2] (sfromR dret) !$ [0, 0, 0]) in sfromVector0N @[2] (fromList [x270, x270])) (\\i200 -> [0, i200, i200, 1]))"  maxPool2dUnpadded3-  :: (ADReady target, GoodScalar r)+  :: (ADReady target, NumScalar r)   => target (TKR 4 r) -> target (TKR 4 r) maxPool2dUnpadded3 arr =   rbuild [2, 2, 2, 2] $ \case@@ -1340,25 +1397,25 @@       in rmaximum3 arrt     _ -> error "maxPool2dUnpadded3: impossible pattern needlessly required" -conv2dUnpadded3-  :: (ADReady target, GoodScalar r)+conv2dSame3+  :: (ADReady target, NumScalar r)   => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded3 arrA =+conv2dSame3 arrA =   let shB = [2, 2, 2, 2]   in rbuild shB $ \case     [iImg, _, iBh, iBw] ->       let arrAt = slicez4 shB arrA [iImg, 0, iBw, 1]-      in rindex0 arrAt [0, iBw, iImg, iBh] + rindex0 arrAt [iImg, 1, iBw + 1, iBh]-    _ -> error "conv2dUnpadded3: impossible pattern needlessly required"+      in rindex @_ @0 arrAt [0, iBw, iImg, iBh] + rindex arrAt [iImg, 1, iBw + 1, iBh]+    _ -> error "conv2dSame3: impossible pattern needlessly required"  slicez3-  :: (ADReady target, GoodScalar r, KnownNat n)+  :: (ADReady target, NumScalar r, KnownNat n)   => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r) slicez3 shOut d ixBase =   rbuild shOut $ \_ -> indexz03 d (ixrZipWith (+) ixBase ixBase)  indexz03-  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)+  :: forall target r n. (ADReady target, NumScalar r, KnownNat n)   => target (TKR n r) -> IxROf target n -> target (TKR 0 r) indexz03 d ix = ifH (within3 @target (rshape @target d) ix) (d ! ix) (rscalar 0) @@ -1373,25 +1430,22 @@  rmaximum3 :: (BaseTensor target, LetTensor target, KnownNat n, GoodScalar r)          => target (TKR n r) -> target (TKR 0 r)-rmaximum3 t0 = tlet t0 $ \t -> rindex0 t [0, 0, 0, 0]+rmaximum3 t0 = tlet t0 $ \t -> rindex t [0, 0, 0, 0]  testCNNOPP4 :: Assertion testCNNOPP4 = do   resetVarCounter   let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK-                   $ AstReplicate (SNat @3) knownSTK+      blackGlyph = AstFromPrimal $ rreplicateN (3 :$: 3 :$: 3 :$: 3 :$: ZSR)                        (rconcrete $ Nested.rscalar 7                         :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))       afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)       afcnn2T = maxPool2dUnpadded4 blackGlyph   printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (let w77 = stranspose @[5,6,7,8,0,1,2,3,4] (sfromVector (fromList [stranspose @[2,3,4,0,5,6,7,1] (sgather (stranspose @[6,0,3,1,4,5,2] (sgather (stranspose @[4,1,3,0,2] (sgather (stranspose @[3,0,4,1,2] (sgather (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i52, i55] -> [i52 + i55]))) (\\[i58, i60] -> [2 + (negate i60 + i58), i60]))) (\\[i62, i64, i67] -> [i62 * i64 + i67]))) (\\[i22, i8] -> [2 * i22 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0] in stranspose @[1,2,3,0] (sappend (sreplicate @1 (str (sappend (sreplicate @1 (stranspose @[0,2,4,1,3] w77 !$ [1, 0, 0])) (sreplicate @1 (stranspose @[0,2,4,1,3] w77 !$ [1, 1, 0]))))) (sreplicate @1 (stranspose @[0,4,1,2,3] w77 !$ [1, 1]))))"+    @?= "rfromPlain (rfromS (sgather @[2, 2, 2, 2] (stranspose @[5, 6, 7, 8, 0, 1, 2, 3, 4] (sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[6, 0, 3, 1, 4, 5, 2] (sgather @[2, 2, 2] (stranspose @[4, 1, 3, 0, 2] (sgather @[2, 2] (stranspose @[3, 0, 4, 1, 2] (sgather @[2, 2] (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i143, i146] -> [i143 + i146]))) (\\[i149, i151] -> [2 + (negate i151 + i149), i151]))) (\\[i153, i155, i158] -> [i153 * i155 + i158]))) (\\[i90, i51] -> [2 * i90 + i51])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0]) (\\[i120, i108, i98, i89] -> [ifH (0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i108, 0] &&* (0 <=. sconcrete (sfromListLinear [2,2] [-1,-2,0,-1]) `sindex0` [i108, 0] &&* (notB (notB (0 <=. sconcrete (sfromListLinear [2,2] [1,0,-1,-2]) `sindex0` [i89, 0] &&* 0 <=. negate (sconcrete (sfromListLinear [2,2] [1,0,-1,-2]) `sindex0` [i89, 0])) &&* 1 <=. negate (sconcrete (sfromListLinear [2,2] [0,-1,-2,-3]) `sindex0` [i89, 0])) &&* notB (notB (0 <=. sconcrete (sfromListLinear [2,2,2] [1,0,1,0,1,0,0,-1]) `sindex0` [i120, i98, 0] &&* 0 <=. negate (sconcrete (sfromListLinear [2,2,2] [1,0,1,0,1,0,0,-1]) `sindex0` [i120, i98, 0])) &&* 1 <=. negate (sconcrete (sfromListLinear [2,2,2] [0,-1,0,-1,0,-1,-1,-2]) `sindex0` [i120, i98, 0]))))) 0 1, i120, i108, i98, i89])))"       -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"   printAstPretty afcnn2T-    @?= "rfromS (let w19 = sgather (sfromVector (fromList [stranspose @[2,3,4,0,5,6,7,1] (sgather (stranspose @[6,0,3,1,4,5,2] (sgather (stranspose @[3,0,2,1] (sgather (stranspose @[0,2,1] (sgather (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i32, i5] -> [i32 + i5]))) (\\[i31, i6] -> [i31, 2 + (negate i31 + i6)]))) (\\[i36, i26, i7] -> [i36 * i26 + i7]))) (\\[i22, i8] -> [2 * i22 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) (\\[i34, i28, i24, i21, i15, i12, i10, i9] -> [ifH (0 <=. negate i28 + negate i15 &&* (1 <=. i28 + negate i12 &&* (notB (notB (0 <=. negate i34 * i24 + negate i10) &&* notB (1 <=. i34 * i24 + i10 &&* (-1) <=. negate i34 * i24 + negate i10)) &&* notB (notB (0 <=. (-2) * i21 + negate i9) &&* notB (1 <=. 2 * i21 + i9 &&* (-1) <=. (-2) * i21 + negate i9))))) 0 1, i34, i28, i24, i21, i15, i12, i10, i9]) in stranspose @[4,5,6,7,0,1,2,3] w19 !$ [0, 0, 0, 0])"+    @?= "rfromPlain (rfromS (let w81 = let m40 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m41 = sconcrete (sreplicate [2,2] (-1)) + (str (sreplicate @2 (siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2)))) ; t42 = stranspose @[1, 2, 0] (sreplicate @2 (str (sreplicate @2 (negate (siota (SNat @2)))) * sreplicate @2 (siota (SNat @2)))) + sreplicateN @[2, 2] (negate (siota (SNat @2))) ; t43 = sconcrete (sreplicate [2,2,2] 1) + (stranspose @[1, 2, 0] (sreplicate @2 (str (sreplicate @2 (negate (siota (SNat @2)))) * sreplicate @2 (siota (SNat @2)))) + sreplicateN @[2, 2] (negate (siota (SNat @2)))) ; m44 = str (sreplicate @2 (sconcrete (sreplicate [2] (-2)) * siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2))) ; m45 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (sconcrete (sreplicate [2] (-2)) * siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2)))) in sgather @[2, 2, 2, 2, 2, 2, 2, 2] (stranspose @[1, 2, 3, 4, 5, 6, 7, 8, 0] (sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[6, 0, 3, 1, 4, 5, 2] (sgather @[2, 2, 2] (stranspose @[3, 0, 2, 1] (sgather @[2, 2] (stranspose @[0, 2, 1] (sgather @[2, 2] (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i112, i75] -> [i112 + i75]))) (\\[i111, i66] -> [i111, 2 + (negate i111 + i66)]))) (\\[i122, i100, i58] -> [i122 * i100 + i58]))) (\\[i90, i51] -> [2 * i90 + i51])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)]))) (\\[i120, i108, i98, i89, i76, i67, i59, i52] -> [i120, i108, i98, i89, i76, i67, i59, i52, ifH (0 <=. m40 `sindex0` [i108, i76] &&* (0 <=. m41 `sindex0` [i108, i67] &&* (notB (notB (0 <=. m45 `sindex0` [i89, i52] &&* 0 <=. negate (m45 `sindex0` [i89, i52])) &&* 1 <=. negate (m44 `sindex0` [i89, i52])) &&* notB (notB (0 <=. t43 `sindex0` [i120, i98, i59] &&* 0 <=. negate (t43 `sindex0` [i120, i98, i59])) &&* 1 <=. negate (t42 `sindex0` [i120, i98, i59]))))) 0 1]) in stranspose @[4, 5, 6, 7, 0, 1, 2, 3] w81 !$ [0, 0, 0, 0]))"  -- In this test primal is trivial but gradient is not, so we know how much -- scatters should be able to simplify in the future.@@ -1399,53 +1453,53 @@ testCNNOPP4b = do   resetVarCounter   let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded4 (FTKR [3, 3, 3, 3] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let w120 = stranspose @[5,6,7,8,0,1,2,3,4] (sfromVector (fromList [stranspose @[2,3,4,0,5,6,7,1] (sgather (stranspose @[4,7,0,3,1,5,6,2] (sgather (stranspose @[3,4,7,1,5,6,0,2] (sgather (stranspose @[6,0,7,4,3,2,1,5] (sgather (sslice (SNat @1) (SNat @2) (stranspose @[0,2,1] (sfromR u1))) (\\[i241, i243, i244, i245, i246] -> [i241 + i243]))) (\\[i252, i255] -> [2 + (negate i255 + i252), i255]))) (\\[i258, i260, i263] -> [i258, i260, i258 * i260 + i263]))) (\\[i118, i119] -> [i118, 2 * i118 + i119])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0] in stranspose @[1,2,3,0] (sappend (sreplicate @1 (str (sappend (sreplicate @1 (stranspose @[0,2,4,1,3] w120 !$ [1, 0, 0])) (sreplicate @1 (stranspose @[0,2,4,1,3] w120 !$ [1, 1, 0]))))) (sreplicate @1 (stranspose @[0,4,1,2,3] w120 !$ [1, 1]))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (sgather @[2, 2, 2, 2] (stranspose @[5, 6, 7, 8, 0, 1, 2, 3, 4] (sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[4, 7, 0, 3, 1, 5, 6, 2] (sgather @[2, 2, 2] (stranspose @[3, 4, 7, 1, 5, 6, 0, 2] (sgather @[2, 2] (stranspose @[6, 0, 7, 4, 3, 2, 1, 5] (sgather @[2, 2, 2, 2, 2] (stranspose @[0, 2, 1] (sfromR u1)) (\\[i292, i294, i295, i296, i297] -> [1 + sconcrete (sfromListLinear [2,2] [0,1,1,2]) `sindex0` [i292, i294]]))) (\\[i303, i306] -> [2 + (negate i306 + i303), i306]))) (\\[i309, i311, i314] -> [i309, i311, i309 * i311 + i314]))) (\\[i223, i224] -> [i223, 2 * i223 + i224])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0]) (\\[i225, i226, i227, i228] -> [ifH (0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i226, 0] &&* (0 <=. sconcrete (sfromListLinear [2,2] [-1,-2,0,-1]) `sindex0` [i226, 0] &&* (notB (notB (0 <=. sconcrete (sfromListLinear [2,2] [1,0,-1,-2]) `sindex0` [i228, 0] &&* 0 <=. negate (sconcrete (sfromListLinear [2,2] [1,0,-1,-2]) `sindex0` [i228, 0])) &&* 1 <=. negate (sconcrete (sfromListLinear [2,2] [0,-1,-2,-3]) `sindex0` [i228, 0])) &&* notB (notB (0 <=. sconcrete (sfromListLinear [2,2,2] [1,0,1,0,1,0,0,-1]) `sindex0` [i225, i227, 0] &&* 0 <=. negate (sconcrete (sfromListLinear [2,2,2] [1,0,1,0,1,0,0,-1]) `sindex0` [i225, i227, 0])) &&* 1 <=. negate (sconcrete (sfromListLinear [2,2,2] [0,-1,0,-1,0,-1,-1,-2]) `sindex0` [i225, i227, 0]))))) 0 1, i225, i226, i227, i228]))"       -- TODO: was once "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"   printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w120 = stranspose @[5,6,7,8,0,1,2,3,4] (sfromVector (fromList [stranspose @[2,3,4,0,5,6,7,1] (sgather (stranspose @[4,7,0,3,1,5,6,2] (sgather (stranspose @[3,4,7,1,5,6,0,2] (sgather (stranspose @[6,0,7,4,3,2,1,5] (sgather (sslice (SNat @1) (SNat @2) (stranspose @[0,2,1] (sfromR u1))) (\\[i108, i109, i110, i111, i112] -> [i108 + i109]))) (\\[i113, i114] -> [2 + (negate i114 + i113), i114]))) (\\[i115, i116, i117] -> [i115, i116, i115 * i116 + i117]))) (\\[i118, i119] -> [i118, 2 * i118 + i119])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0] in rfromS (stranspose @[1,2,3,0] (sappend (sreplicate @1 (str (sappend (sreplicate @1 (stranspose @[0,2,4,1,3] w120 !$ [1, 0, 0])) (sreplicate @1 (stranspose @[0,2,4,1,3] w120 !$ [1, 1, 0]))))) (sreplicate @1 (stranspose @[0,4,1,2,3] w120 !$ [1, 1]))))"+    @?= "\\u1 -> rfromS (let m206 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m207 = sconcrete (sreplicate [2,2] (-1)) + (str (sreplicate @2 (siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2)))) ; t208 = stranspose @[1, 2, 0] (sreplicate @2 (str (sreplicate @2 (negate (siota (SNat @2)))) * sreplicate @2 (siota (SNat @2)))) + sreplicateN @[2, 2] (negate (siota (SNat @2))) ; t209 = sconcrete (sreplicate [2,2,2] 1) + (stranspose @[1, 2, 0] (sreplicate @2 (str (sreplicate @2 (negate (siota (SNat @2)))) * sreplicate @2 (siota (SNat @2)))) + sreplicateN @[2, 2] (negate (siota (SNat @2)))) ; m210 = str (sreplicate @2 (sconcrete (sreplicate [2] (-2)) * siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2))) ; m211 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (sconcrete (sreplicate [2] (-2)) * siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2)))) in sgather @[2, 2, 2, 2] (stranspose @[5, 6, 7, 8, 0, 1, 2, 3, 4] (let m212 = str (sreplicate @2 (siota (SNat @2))) + sreplicate @2 (siota (SNat @2)) in sfromVector @2 (fromList [stranspose @[2, 3, 4, 0, 5, 6, 7, 1] (sgather @[2, 2] (stranspose @[4, 7, 0, 3, 1, 5, 6, 2] (sgather @[2, 2, 2] (stranspose @[3, 4, 7, 1, 5, 6, 0, 2] (sgather @[2, 2] (stranspose @[6, 0, 7, 4, 3, 2, 1, 5] (sgather @[2, 2, 2, 2, 2] (stranspose @[0, 2, 1] (sfromR u1)) (\\[i213, i214, i215, i216, i217] -> [1 + m212 `sindex0` [i213, i214]]))) (\\[i218, i219] -> [2 + (negate i219 + i218), i219]))) (\\[i220, i221, i222] -> [i220, i221, i220 * i221 + i222]))) (\\[i223, i224] -> [i223, 2 * i223 + i224])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0]) (\\[i225, i226, i227, i228] -> [ifH (0 <=. m206 `sindex0` [i226, 0] &&* (0 <=. m207 `sindex0` [i226, 0] &&* (notB (notB (0 <=. m211 `sindex0` [i228, 0] &&* 0 <=. negate (m211 `sindex0` [i228, 0])) &&* 1 <=. negate (m210 `sindex0` [i228, 0])) &&* notB (notB (0 <=. t209 `sindex0` [i225, i227, 0] &&* 0 <=. negate (t209 `sindex0` [i225, i227, 0])) &&* 1 <=. negate (t208 `sindex0` [i225, i227, 0]))))) 0 1, i225, i226, i227, i228]))"       -- TODO: was once "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"   printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w122 = stranspose @[4,5,6,7,8,0,1,2,3] (soneHot (stranspose @[0,3,1,4,2] (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (str (ssum @1 (sslice (SNat @0) (SNat @1) (stranspose @[3,0,1,2] (sfromR dret))))))) [1, 0, 0]) + (stranspose @[0,3,1,4,2] (soneHot (ssum @1 (sslice (SNat @1) (SNat @1) (str (ssum @1 (sslice (SNat @0) (SNat @1) (stranspose @[3,0,1,2] (sfromR dret))))))) [1, 1, 0]) + stranspose @[0,2,3,4,1] (soneHot (ssum @1 (sslice (SNat @1) (SNat @1) (stranspose @[3,0,1,2] (sfromR dret)))) [1, 1]))) [0, 0, 0, 0]) in rfromS (stranspose @[0,2,1] (sappend (sconcrete (sreplicate [1,3,3,3] 0.0)) (sappend (sscatter (stranspose @[1,6,5,4,3,7,0,2] (sscatter (stranspose @[6,3,7,0,1,4,5,2] (sscatter (stranspose @[2,4,7,3,0,5,6,1] (sscatter (stranspose @[3,7,0,1,2,4,5,6] (w122 !$ [0])) (\\[i123, i124] -> [i123, 2 * i123 + i124]))) (\\[i125, i126, i127] -> [i125, i126, i125 * i126 + i127]))) (\\[i128, i129] -> [2 + (negate i129 + i128), i129]))) (\\[i130, i131, i132, i133, i134] -> [i130 + i131])) (sconcrete (sfromListLinear [0,3,3,3] [])))))"+    @?= "\\dret u1 -> rfromS (stranspose @[0, 2, 1] (let m206 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m207 = sconcrete (sreplicate [2,2] (-1)) + (str (sreplicate @2 (siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2)))) ; t208 = stranspose @[1, 2, 0] (sreplicate @2 (str (sreplicate @2 (negate (siota (SNat @2)))) * sreplicate @2 (siota (SNat @2)))) + sreplicateN @[2, 2] (negate (siota (SNat @2))) ; t209 = sconcrete (sreplicate [2,2,2] 1) + (stranspose @[1, 2, 0] (sreplicate @2 (str (sreplicate @2 (negate (siota (SNat @2)))) * sreplicate @2 (siota (SNat @2)))) + sreplicateN @[2, 2] (negate (siota (SNat @2)))) ; m210 = str (sreplicate @2 (sconcrete (sreplicate [2] (-2)) * siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2))) ; m211 = sconcrete (sreplicate [2,2] 1) + (str (sreplicate @2 (sconcrete (sreplicate [2] (-2)) * siota (SNat @2))) + sreplicate @2 (negate (siota (SNat @2)))) ; m212 = str (sreplicate @2 (siota (SNat @2))) + sreplicate @2 (siota (SNat @2)) in sscatter @[2, 2, 2, 2, 2] (stranspose @[1, 6, 5, 4, 3, 7, 0, 2] (sscatter @[2, 2] (stranspose @[6, 3, 7, 0, 1, 4, 5, 2] (sscatter @[2, 2, 2] (stranspose @[2, 4, 7, 3, 0, 5, 6, 1] (sscatter @[2, 2] (stranspose @[3, 7, 0, 1, 2, 4, 5, 6] (stranspose @[4, 5, 6, 7, 8, 0, 1, 2, 3] (soneHot (sscatter @[2, 2, 2, 2] (sfromR dret) (\\[i230, i231, i232, i233] -> [ifH (0 <=. m206 `sindex0` [i231, 0] &&* (0 <=. m207 `sindex0` [i231, 0] &&* (notB (notB (0 <=. m211 `sindex0` [i233, 0] &&* 0 <=. negate (m211 `sindex0` [i233, 0])) &&* 1 <=. negate (m210 `sindex0` [i233, 0])) &&* notB (notB (0 <=. t209 `sindex0` [i230, i232, 0] &&* 0 <=. negate (t209 `sindex0` [i230, i232, 0])) &&* 1 <=. negate (t208 `sindex0` [i230, i232, 0]))))) 0 1, i230, i231, i232, i233])) [0, 0, 0, 0]) !$ [0])) (\\[i235, i236] -> [i235, 2 * i235 + i236]))) (\\[i237, i238, i239] -> [i237, i238, i237 * i238 + i239]))) (\\[i240, i241] -> [2 + (negate i241 + i240), i241]))) (\\[i242, i243, i244, i245, i246] -> [1 + m212 `sindex0` [i242, i243]])))"       -- TODO: was once "\\dret u1 -> rfromS (sconcrete (sreplicate [3,3,3,3] 0.0))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (sappend (sconcrete (sreplicate [1,3,3,3] 0.0)) (stranspose @[0,2,1] (sscatter (stranspose @[1,6,5,4,3,7,0,2] (sscatter (stranspose @[6,3,7,0,1,4,5,2] (sscatter (stranspose @[2,4,7,3,0,5,6,1] (sscatter (stranspose @[4,8,3,5,6,7,0,1,2] (soneHot (stranspose @[0,3,1,4,2] (soneHot (stranspose @[1,3,0,2] (sfromR dret) !$ [0, 0]) [1, 0, 0]) + (stranspose @[0,3,1,4,2] (soneHot (stranspose @[1,3,0,2] (sfromR dret) !$ [1, 0]) [1, 1, 0]) + stranspose @[0,2,3,4,1] (soneHot (stranspose @[3,0,1,2] (sfromR dret) !$ [1]) [1, 1]))) [0, 0, 0, 0]) !$ [0]) (\\[i123, i124] -> [i123, 2 * i123 + i124]))) (\\[i125, i126, i127] -> [i125, i126, i125 * i126 + i127]))) (\\[i128, i129] -> [2 + (negate i129 + i128), i129]))) (\\[i130, i131, i132, i133, i134] -> [i130 + i131]))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[0, 2, 1] (sscatter @[2, 2, 2, 2, 2] (stranspose @[1, 6, 5, 4, 3, 7, 0, 2] (sscatter @[2, 2] (stranspose @[6, 3, 7, 0, 1, 4, 5, 2] (sscatter @[2, 2, 2] (stranspose @[2, 4, 7, 3, 0, 5, 6, 1] (sscatter @[2, 2] (sscatter @[2, 2, 2, 2] (sfromR dret) (\\[i230, i231, i232, i233] -> [ifH (0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i231, 0] &&* (0 <=. sconcrete (sfromListLinear [2,2] [-1,-2,0,-1]) `sindex0` [i231, 0] &&* (notB (notB (0 <=. sconcrete (sfromListLinear [2,2] [1,0,-1,-2]) `sindex0` [i233, 0] &&* 0 <=. negate (sconcrete (sfromListLinear [2,2] [1,0,-1,-2]) `sindex0` [i233, 0])) &&* 1 <=. negate (sconcrete (sfromListLinear [2,2] [0,-1,-2,-3]) `sindex0` [i233, 0])) &&* notB (notB (0 <=. sconcrete (sfromListLinear [2,2,2] [1,0,1,0,1,0,0,-1]) `sindex0` [i230, i232, 0] &&* 0 <=. negate (sconcrete (sfromListLinear [2,2,2] [1,0,1,0,1,0,0,-1]) `sindex0` [i230, i232, 0])) &&* 1 <=. negate (sconcrete (sfromListLinear [2,2,2] [0,-1,0,-1,0,-1,-1,-2]) `sindex0` [i230, i232, 0]))))) 0 1, i233, 0, i230, i231, i232, 0, 0, 0]) !$ [0]) (\\[i235, i236] -> [i235, 2 * i235 + i236]))) (\\[i237, i238, i239] -> [i237, i238, i237 * i238 + i239]))) (\\[i240, i241] -> [2 + (negate i241 + i240), i241]))) (\\[i242, i243, i244, i245, i246] -> [1 + sconcrete (sfromListLinear [2,2] [0,1,1,2]) `sindex0` [i242, i243]])))"       -- TODO: was once "\\dret u1 -> rfromS (sconcrete (sreplicate [3,3,3,3] 0.0))"  testCNNOPP5 :: Assertion testCNNOPP5 = do   resetVarCounter   let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @6) knownSTK-                   $ AstReplicate (SNat @6) knownSTK-                   $ AstReplicate (SNat @6) knownSTK-                   $ AstReplicate (SNat @6) knownSTK+      blackGlyph = AstFromPrimal $ rreplicateN (6 :$: 6 :$: 6 :$: 6 :$: ZSR)                        (rconcrete $ Nested.rscalar 7                         :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))       afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = conv2dUnpadded4 blackGlyph+      afcnn2T = conv2dSame4 blackGlyph   printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (sconcrete (sfromListLinear [1,1,2,2] [7.0,0.0,0.0,0.0]))"+    @?= "rfromPlain (rfromS (sreplicateN @[1, 1] (sgather @[2, 2] (sconcrete (sfromListLinear [2] [7.0,0.0])) (\\[i35, i32] -> [ifH (0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i35, 0] &&* 0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i32, 0]) 0 1]))))"+      -- TODO: was once "rfromS (sconcrete (sfromListLinear [1,1,2,2] [7.0,0.0,0.0,0.0]))"   printAstPretty afcnn2T-    @?= "rfromS (sconcrete (sfromListLinear [1,1,2,2] [7.0,0.0,0.0,0.0]))"+    @?= "rfromPlain (rfromS (sreplicateN @[1, 1] (let m29 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m30 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) in sgather @[2, 2] (sconcrete (sfromListLinear [2] [7.0,0.0])) (\\[i35, i32] -> [ifH (0 <=. m29 `sindex0` [i35, 0] &&* 0 <=. m30 `sindex0` [i32, 0]) 0 1]))))"  -- In this test primal is trivial but gradient is not, so we know how much -- scatters should be able to simplify in the future. testCNNOPP5b :: Assertion testCNNOPP5b = do   resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dUnpadded4 (FTKR [5, 5, 5, 5] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0))))))"+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dSame4 (FTKR [5, 5, 5, 5] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (sreplicateN @[1, 1] (let m48 = sgather @[2, 2] (sconcrete (sfromListLinear [2] [0,1])) (\\[i46, i47] -> [ifH (0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i46, 0] &&* 0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i47, 0]) 0 1]) in sgather @[2, 2] (sfromVector @2 (fromList [str (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (sfromR u1 !$ [0, 0])))), sconcrete (sreplicate [2,2] 0.0)])) (\\[i49, i50] -> [m48 `sindex0` [i49, i50], i49, i50])))"+      -- TODO: was once "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sfromVector (fromList [sfromVector (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)])))))"   printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0))))))"+    @?= "\\u1 -> rfromS (sreplicateN @[1, 1] (let m44 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m45 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m48 = sgather @[2, 2] (sconcrete (sfromListLinear [2] [0,1])) (\\[i46, i47] -> [ifH (0 <=. m44 `sindex0` [i46, 0] &&* 0 <=. m45 `sindex0` [i47, 0]) 0 1]) in sgather @[2, 2] (sfromVector @2 (fromList [str (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (sfromR u1 !$ [0, 0])))), sconcrete (sreplicate [2,2] 0.0)])) (\\[i49, i50] -> [m48 `sindex0` [i49, i50], i49, i50])))"   printArtifactPretty artifactRev-    @?= "\\dret u1 -> rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (str (ssum @1 (ssum @1 (sfromR dret)))))))) [0, 0, 0, 0])"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (soneHot (sfromR dret !$ [0, 0, 0, 0]) [0, 0, 0, 0])"+    @?= "\\dret u1 -> rfromS (soneHot (let m44 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m45 = str (sreplicate @2 (negate (siota (SNat @2)))) + sreplicate @2 (negate (siota (SNat @2))) ; m48 = sgather @[2, 2] (sconcrete (sfromListLinear [2] [0,1])) (\\[i46, i47] -> [ifH (0 <=. m44 `sindex0` [i46, 0] &&* 0 <=. m45 `sindex0` [i47, 0]) 0 1]) ; t54 = sscatter @[2, 2] (ssumN @[1, 1] (sfromR dret)) (\\[i52, i53] -> [m48 `sindex0` [i52, i53], i52, i53]) in sappend (sconcrete (sfromListLinear [0,5] [])) (sappend (str (sappend (sconcrete (sfromListLinear [0,2] [])) (sappend (str (t54 !$ [0])) (sconcrete (sreplicate [3,2] 0.0))))) (sconcrete (sreplicate [3,5] 0.0)))) [0, 0])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (let m61 = sgather @[2, 2] (sconcrete (sfromListLinear [2] [0,1])) (\\[i46, i47] -> [ifH (0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i46, 0] &&* 0 <=. sconcrete (sfromListLinear [2,2] [0,-1,-1,-2]) `sindex0` [i47, 0]) 0 1]) in sappend (str (sappend (sscatter @[2, 2] (sfromR dret !$ [0, 0]) (\\[i52, i53] -> [m61 `sindex0` [i52, i53], i53, i52]) !$ [0]) (sconcrete (sreplicate [3,2] 0.0)))) (sconcrete (sreplicate [3,5] 0.0))) [0, 0])"+      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sfromK (sfromR dret `sindex0` [0, 0, 0, 0])) [0, 0, 0, 0])"  maxPool2dUnpadded4-  :: (ADReady target, GoodScalar r)+  :: (ADReady target, NumScalar r)   => target (TKR 4 r) -> target (TKR 4 r) maxPool2dUnpadded4 arr =   rbuild [2, 2, 2, 2] $ \case@@ -1454,19 +1508,19 @@       in rmaximum3 arrt     _ -> error "maxPool2dUnpadded4: impossible pattern needlessly required" -conv2dUnpadded4-  :: (ADReady target, GoodScalar r)+conv2dSame4+  :: (ADReady target, NumScalar r)   => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded4 arrA =+conv2dSame4 arrA =   let shB = [1, 1, 2, 2]   in rbuild shB $ \case     [iImg, _, iBh, iBw] ->       let arrAt = slicez4 shB arrA [iImg, 0, iBh, iBw]-      in rindex0 arrAt [0, 0, 0, 0]-    _ -> error "conv2dUnpadded4: impossible pattern needlessly required"+      in rindex @_ @0 arrAt [0, 0, 0, 0]+    _ -> error "conv2dSame4: impossible pattern needlessly required"  slicez4-  :: (ADReady target, GoodScalar r, KnownNat n)+  :: (ADReady target, NumScalar r, KnownNat n)   => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r) slicez4 shOut d ixBase =   rbuild shOut $ \ixResult -> indexz03 d (ixrZipWith (+) ixBase ixResult)@@ -1475,80 +1529,69 @@ testCNNOPP6 = do   resetVarCounter   let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK+      blackGlyph = AstFromPrimal $ rreplicateN (2 :$: 2 :$: 2 :$: 2 :$: ZSR)                        (rconcrete $ Nested.rscalar 7                         :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))       afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = maxPool2dUnpadded3 $ conv2dUnpadded3z blackGlyph+      afcnn2T = maxPool2dUnpadded3 $ conv2dSame3z blackGlyph   printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (let t45 = str (sfromVector (fromList [stranspose @[1,2,0] (sgather (stranspose @[2,1,0] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sfromListLinear [2,2,2] [7.0,7.0,0.0,0.0,7.0,7.0,0.0,0.0])) (\\[i48] -> [2 * i48]))) (\\[i51] -> [2 * i51]))) (\\[i4] -> [2 * i4])), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0] in sappend (sreplicate @1 (sappend (sreplicate @1 (t45 !$ [0, 0, 0])) (sreplicate @1 (t45 !$ [1, 0, 1])))) (sreplicate @1 (t45 !$ [1, 1])))) (sconcrete (sreplicate [1,2,2] 0.0)))))"-      -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"   printAstPretty afcnn2T-    @?= "rfromS (stranspose @[1,2,0] (sreplicate @2 (sgather (sfromVector (fromList [stranspose @[1,2,0] (sgather (stranspose @[2,1,0] (sgather (str (sgather (sreplicate @2 (str (sreplicate @2 (sgather (sfromVector (fromList [str (sgather (str (sgather (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i9] -> [2 * i9, 2 * i9, 2 * i9]))) (\\[i12] -> [2 * i12])), sconcrete (sreplicate [2,2] 0.0)])) (\\[i19, i17] -> [let i18 = ifH (notB (notB (0 <=. negate i19) &&* notB (0 <=. negate i19)) &&* (notB (notB (0 <=. negate i19) &&* notB (0 <=. negate i19)) &&* notB (notB (0 <=. negate i17) &&* notB (0 <=. negate i17)))) 0 1 in ifH (0 <=. negate i19) i18 (ifH (0 <=. negate i19) i18 1), i19, i17]))))) (\\[i1] -> [2 * i1, 0]))) (\\[i2] -> [2 * i2]))) (\\[i4] -> [2 * i4])), sconcrete (sreplicate [2,2,2] 0.0)])) (\\[i26, i24, i22] -> [let i23 = ifH (notB (notB (0 <=. negate i24) &&* notB (0 <=. negate i24)) &&* notB (notB (0 <=. negate i22) &&* notB (0 <=. negate i22))) 0 1 in ifH (0 <=. negate i26) i23 (ifH (0 <=. negate i26) i23 1), i26, i24, i22]))))"+    @?= "rfromPlain (rfromS (stranspose @[1, 2, 0] (sreplicate @2 (let t31 = sgather1 @2 (stranspose @[2, 1, 0] (sgather1 @2 (stranspose @[1, 2, 0] (sgather1 @2 (sreplicateN @[2, 2] (let m21 = sgather1 @2 (str (sgather1 @2 (sreplicate0N @[2, 2, 2, 2] 7.0) (\\i9 -> [2 * i9, 2 * i9, 2 * i9]))) (\\i12 -> [2 * i12]) in sfromVector @2 (fromList [sfromVector0N @[2] (fromList [m21 `sindex0` [0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]))) (\\i1 -> [2 * i1]))) (\\i2 -> [2 * i2, 0]))) (\\i4 -> [2 * i4]) in sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [t31 `sindex0` [0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)])))))"  testCNNOPP6b :: Assertion testCNNOPP6b = do   resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dUnpadded3z) (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let t58 = str (sfromVector (fromList [stranspose @[1,2,0] (sgather (stranspose @[2,1,0] (sgather (stranspose @[2,0,1] (sgather (stranspose @[1,2,0] (sreplicate @2 (sreplicate @2 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))))) (\\[i93] -> [2 * i93]))) (\\[i96] -> [2 * i96]))) (\\[i57] -> [2 * i57])), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0] in stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (t58 !$ [0, 0, 0])) (sreplicate @1 (t58 !$ [1, 0, 1])))) (sreplicate @1 (t58 !$ [1, 1])))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dSame3z) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @2 (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]))))"   printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let t58 = str (sfromVector (fromList [stranspose @[1,2,0] (sgather (stranspose @[2,1,0] (sgather (stranspose @[2,0,1] (sgather (stranspose @[1,2,0] (sreplicate @2 (sreplicate @2 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))))) (\\[i55] -> [2 * i55]))) (\\[i56] -> [2 * i56]))) (\\[i57] -> [2 * i57])), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0] in rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (t58 !$ [0, 0, 0])) (sreplicate @1 (t58 !$ [1, 0, 1])))) (sreplicate @1 (t58 !$ [1, 1])))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sreplicate @2 (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]))))"   printArtifactPretty artifactRev-    @?= "\\dret u1 -> let t60 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) ; u61 = str (soneHot (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t60)))))) [0, 0, 0] + (soneHot (ssum @1 (sslice (SNat @1) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t60)))))) [1, 0, 1] + soneHot (ssum @1 (sslice (SNat @1) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t60)))) [1, 1])) [0]) ; v65 = ssum @2 (ssum @2 (stranspose @[2,0,1] (sscatter (stranspose @[1,2,0] (sscatter (stranspose @[2,1,0] (sscatter (stranspose @[2,0,1] (u61 !$ [0])) (\\[i62] -> [2 * i62]))) (\\[i63] -> [2 * i63]))) (\\[i64] -> [2 * i64])))) in rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) v65)) [0, 0, 0, 0])"--  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (let t60 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) in soneHot (ssum0 (stranspose @[1,2,0] (sscatter (stranspose @[1,2,0] (sscatter (stranspose @[2,1,0] (sscatter (stranspose @[1,3,0,2] (soneHot (soneHot (t60 !$ [0, 0, 0]) [0, 0, 0] + (soneHot (t60 !$ [0, 0, 1]) [1, 0, 1] + soneHot (t60 !$ [0, 1]) [1, 1])) [0]) !$ [0]) (\\[i62] -> [2 * i62]))) (\\[i63] -> [2 * i63]))) (\\[i64] -> [2 * i64])) !$ [0])) [0, 0, 0, 0])"-      -- TODO: was once "\\dret u1 -> rfromS (soneHot (ssum0 (stranspose @[0,1,3,2] (sfromR dret) !$ [0, 0, 0])) [0, 0, 0, 0])"+    @?= "\\dret u1 -> rfromS (soneHot (sfromK (((ssum @2 (stranspose @[2, 0, 1] (sfromR dret)) !$ [0]) !$ [0]) `sindex0` [0])) [0, 0, 0, 0])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (sfromK (ssum0 @[2] (stranspose @[0, 1, 3, 2] (sfromR dret) !$ [0, 0, 0]))) [0, 0, 0, 0])" -conv2dUnpadded3z-  :: (ADReady target, GoodScalar r)+conv2dSame3z+  :: (ADReady target, NumScalar r)   => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded3z arrA =+conv2dSame3z arrA =   let shB = [2, 2, 2, 2]   in rbuild shB $ \case     [iImg, _, iBh, iBw] ->       let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBw]-      in rindex0 arrAt [iBh, iBw, iImg, iBh]-    _ -> error "conv2dUnpadded3z: impossible pattern needlessly required"+      in rindex @_ @0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dSame3z: impossible pattern needlessly required"  testCNNOPP7 :: Assertion testCNNOPP7 = do   resetVarCounter   let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK+      blackGlyph = AstFromPrimal $ rreplicateN (2 :$: 2 :$: 2 :$: 2 :$: ZSR)                        (rconcrete $ Nested.rscalar 7                         :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))       afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)-      afcnn2T = maxPool2dUnpadded3y $ conv2dUnpadded3y blackGlyph+      afcnn2T = maxPool2dUnpadded3y $ conv2dSame3y blackGlyph   printAstPretty (simplifyInlineContract afcnn2T)-    @?= "rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sgather (stranspose @[3,2,1,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sgather (stranspose @[3,2,0,1] (sgather (stranspose @[2,3,0,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,0.0,0.0,7.0,0.0,0.0,0.0,7.0,0.0,0.0,0.0,7.0,0.0,0.0,0.0])) (\\[i49] -> [2 * i49]))) (\\[i52, i53] -> [2 * i53, 2 * i52]))) (\\[i4] -> [2 * i4])), sconcrete (sreplicate [2,2,2,2] 0.0)])) !$ [0]) (\\[i29, i26, i22] -> [i26, i29, ifH (0 <=. negate i29) (ifH (notB (notB (0 <=. negate i26) &&* notB (0 <=. negate i26)) &&* notB (notB (0 <=. negate i22) &&* notB (0 <=. negate i22))) 0 1) (ifH (0 <=. negate i29) (ifH (notB (notB (0 <=. negate i26) &&* notB (0 <=. negate i26)) &&* notB (notB (0 <=. negate i22) &&* notB (0 <=. negate i22))) 0 1) 1), i22]))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"-      -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"   printAstPretty afcnn2T-    @?= "rfromS (sgather (sfromVector (fromList [stranspose @[1,2,3,0] (sgather (stranspose @[3,2,0,1] (sgather (stranspose @[1,2,0] (sgather (sreplicate @2 (stranspose @[1,2,0] (sreplicate @2 (sgather (sfromVector (fromList [str (sgather (str (sgather (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i9] -> [2 * i9, 2 * i9, 2 * i9]))) (\\[i11] -> [2 * i11])), sconcrete (sreplicate [2,2] 0.0)])) (\\[i19, i17] -> [let i18 = ifH (notB (notB (0 <=. negate i19) &&* notB (0 <=. negate i19)) &&* (notB (notB (0 <=. negate i19) &&* notB (0 <=. negate i19)) &&* notB (notB (0 <=. negate i17) &&* notB (0 <=. negate i17)))) 0 1 in ifH (0 <=. negate i19) i18 (ifH (0 <=. negate i19) i18 1), i19, i17]))))) (\\[i1] -> [2 * i1]))) (\\[i28, i3] -> [2 * i3, 2 * i28]))) (\\[i4] -> [2 * i4])), sconcrete (sreplicate [2,2,2,2] 0.0)])) (\\[i29, i26, i24, i22] -> [let i23 = ifH (notB (notB (0 <=. negate i29) &&* notB (0 <=. negate i29)) &&* (notB (notB (0 <=. negate i26) &&* notB (0 <=. negate i26)) &&* notB (notB (0 <=. negate i22) &&* notB (0 <=. negate i22)))) 0 1 in ifH (0 <=. negate i24) i23 (ifH (0 <=. negate i24) i23 1), i29, i26, i24, i22]))"+    @?= "rfromPlain (rfromS (let u28 = sgather1 @2 (stranspose @[3, 2, 0, 1] (sgather @[2, 2] (stranspose @[1, 2, 0] (sgather1 @2 (sreplicate @2 (stranspose @[1, 2, 0] (sreplicate @2 (let m21 = sgather1 @2 (str (sgather1 @2 (sreplicate0N @[2, 2, 2, 2] 7.0) (\\i9 -> [2 * i9, 2 * i9, 2 * i9]))) (\\i11 -> [2 * i11]) in sfromVector @2 (fromList [sfromVector0N @[2] (fromList [m21 `sindex0` [0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]))))) (\\i1 -> [2 * i1]))) (\\[i32, i3] -> [2 * i3, 2 * i32]))) (\\i4 -> [2 * i4]) in stranspose @[1, 2, 0] (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [u28 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]), sconcrete (sreplicate [2,2,2] 0.0)]))))"  testCNNOPP7b :: Assertion testCNNOPP7b = do   resetVarCounter-  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3y . conv2dUnpadded3y) (FTKR [2, 2, 2, 2] (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sgather (stranspose @[3,2,1,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sgather (stranspose @[3,2,0,1] (sgather (stranspose @[2,3,0,1] (sgather (sreplicate @2 (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0))))) (\\[i119] -> [2 * i119]))) (\\[i122, i123] -> [2 * i123, 2 * i122]))) (\\[i62] -> [2 * i62])), sconcrete (sreplicate [2,2,2,2] 0.0)])) !$ [0]) (\\[i63, i64, i65] -> [i64, i63, ifH (0 <=. negate i63) (ifH (notB (notB (0 <=. negate i64) &&* notB (0 <=. negate i64)) &&* notB (notB (0 <=. negate i65) &&* notB (0 <=. negate i65))) 0 1) (ifH (0 <=. negate i63) (ifH (notB (notB (0 <=. negate i64) &&* notB (0 <=. negate i64)) &&* notB (notB (0 <=. negate i65) &&* notB (0 <=. negate i65))) 0 1) 1), i65]))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3y . conv2dSame3y) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]), sconcrete (sreplicate [2,2,2] 0.0)])))"   printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sgather (stranspose @[3,2,1,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sgather (stranspose @[3,2,0,1] (sgather (stranspose @[2,3,0,1] (sgather (sreplicate @2 (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0))))) (\\[i59] -> [2 * i59]))) (\\[i60, i61] -> [2 * i61, 2 * i60]))) (\\[i62] -> [2 * i62])), sconcrete (sreplicate [2,2,2,2] 0.0)])) !$ [0]) (\\[i63, i64, i65] -> [i64, i63, ifH (0 <=. negate i63) (ifH (notB (notB (0 <=. negate i64) &&* notB (0 <=. negate i64)) &&* notB (notB (0 <=. negate i65) &&* notB (0 <=. negate i65))) 0 1) (ifH (0 <=. negate i63) (ifH (notB (notB (0 <=. negate i64) &&* notB (0 <=. negate i64)) &&* notB (notB (0 <=. negate i65) &&* notB (0 <=. negate i65))) 0 1) 1), i65]))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"+    @?= "\\u1 -> rfromS (stranspose @[1, 2, 0] (sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector @2 (fromList [sfromVector0N @[2] (fromList [sfromR u1 `sindex0` [0, 0, 0, 0], 0.0]), sconcrete (sreplicate [2] 0.0)]), sconcrete (sreplicate [2,2] 0.0)]), sconcrete (sreplicate [2,2,2] 0.0)])))"   printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w70 = stranspose @[3,2,1,0] (soneHot (sscatter (ssum @1 (sslice (SNat @0) (SNat @1) (stranspose @[2,0,1] (sfromR dret)))) (\\[i67, i68, i69] -> [i68, i67, ifH (0 <=. negate i67) (ifH (notB (notB (0 <=. negate i68) &&* notB (0 <=. negate i68)) &&* notB (notB (0 <=. negate i69) &&* notB (0 <=. negate i69))) 0 1) (ifH (0 <=. negate i67) (ifH (notB (notB (0 <=. negate i68) &&* notB (0 <=. negate i68)) &&* notB (notB (0 <=. negate i69) &&* notB (0 <=. negate i69))) 0 1) 1), i69])) [0]) ; m75 = ssum @2 (ssum @2 (sscatter (stranspose @[2,3,0,1] (sscatter (stranspose @[2,3,1,0] (sscatter (stranspose @[3,0,1,2] (w70 !$ [0])) (\\[i71] -> [2 * i71]))) (\\[i72, i73] -> [2 * i73, 2 * i72]))) (\\[i74] -> [2 * i74]))) in rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) m75)))) [0, 0, 0, 0])"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret u1 -> rfromS (soneHot (ssum0 (stranspose @[2,3,0,1] (sscatter (stranspose @[2,3,0,1] (sscatter (stranspose @[2,3,1,0] (sscatter (stranspose @[3,4,2,1,0] (soneHot (sscatter (stranspose @[2,0,1] (sfromR dret) !$ [0]) (\\[i67, i68, i69] -> [i68, i67, ifH (0 <=. negate i67) (ifH (notB (notB (0 <=. negate i68) &&* notB (0 <=. negate i68)) &&* notB (notB (0 <=. negate i69) &&* notB (0 <=. negate i69))) 0 1) (ifH (0 <=. negate i67) (ifH (notB (notB (0 <=. negate i68) &&* notB (0 <=. negate i68)) &&* notB (notB (0 <=. negate i69) &&* notB (0 <=. negate i69))) 0 1) 1), i69])) [0]) !$ [0]) (\\[i71] -> [2 * i71]))) (\\[i72, i73] -> [2 * i73, 2 * i72]))) (\\[i74] -> [2 * i74])) !$ [0, 0])) [0, 0, 0, 0])"-      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sfromR dret !$ [0, 0, 0, 0]) [0, 0, 0, 0])"+    @?= "\\dret u1 -> rfromS (soneHot (sfromK ((((stranspose @[2, 0, 1] (sfromR dret) !$ [0]) !$ [0]) !$ [0]) `sindex0` [0])) [0, 0, 0, 0])"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (sfromK (sfromR dret `sindex0` [0, 0, 0, 0])) [0, 0, 0, 0])"  maxPool2dUnpadded3y-  :: (ADReady target, GoodScalar r)+  :: (ADReady target, NumScalar r)   => target (TKR 4 r) -> target (TKR 4 r) maxPool2dUnpadded3y arr =   rbuild [2, 2, 2, 2] $ \case@@ -1557,48 +1600,59 @@       in rmaximum3 arrt     _ -> error "maxPool2dUnpadded3y: impossible pattern needlessly required" -conv2dUnpadded3y-  :: (ADReady target, GoodScalar r)+conv2dSame3y+  :: (ADReady target, NumScalar r)   => target (TKR 4 r) -> target (TKR 4 r)-conv2dUnpadded3y arrA =+conv2dSame3y arrA =   let shB = [2, 2, 2, 2]   in rbuild shB $ \case     [iImg, _, iBh, iBw] ->       let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBh]-      in rindex0 arrAt [iBh, iBw, iImg, iBh]-    _ -> error "conv2dUnpadded3y: impossible pattern needlessly required"+      in rindex @_ @0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dSame3y: impossible pattern needlessly required"  -- This test uses a disastrous version of smaximum, but shows how--- smaxIndex gets non-trivially vectorized, preserving sharing, too.+-- sargMax gets non-trivially vectorized, preserving sharing, too. testCNNOPP4bU :: Assertion testCNNOPP4bU = do   resetVarCounter   let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpaddedS4 @4 @2) (FTKS (SNat @31 :$$ SNat @31 :$$ SNat @31 :$$ SNat @31 :$$ ZSS) (FTKScalar @Double))-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> let w52 = sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] u1) (\\[i99, i100] -> [2 * i99 + i100]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = smaxIndex (sreshape @[31,31,15,15,16] (stranspose @[2,3,4,0,5,1] w52)) in sgather w52 (\\[i54, i55, i56, i57] -> [i57, remH (kfromS (u53 !$ [i54, i55, i56, i57])) 4, i54, i55, i56, remH (quotH (kfromS (u53 !$ [i54, i55, i56, i57])) 4) 4])"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> let w52 = sgather @[15, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[15, 4] (stranspose @[2, 0, 1] u1) (\\[i107, i108] -> [2 * i107 + i108]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = sargMax (sreshape @[31, 31, 15, 15, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (splainPart w52))))) in sgather @[31, 31, 15, 15] w52 (\\[i54, i55, i56, i57] -> [i57, remH (u53 `sindex0` [i54, i55, i56, i57]) 4, i54, i55, i56, remH (quotH (u53 `sindex0` [i54, i55, i56, i57]) 4) 4])"   printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w52 = sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] u1) (\\[i48, i49] -> [2 * i48 + i49]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = smaxIndex (sreshape @[31,31,15,15,16] (stranspose @[2,3,4,0,5,1] w52)) in sgather w52 (\\[i54, i55, i56, i57] -> [i57, remH (kfromS (u53 !$ [i54, i55, i56, i57])) 4, i54, i55, i56, remH (quotH (kfromS (u53 !$ [i54, i55, i56, i57])) 4) 4])"+    @?= "\\u1 -> let w52 = sgather @[15, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[15, 4] (stranspose @[2, 0, 1] u1) (\\[i48, i49] -> [2 * i48 + i49]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = sargMax (sreshape @[31, 31, 15, 15, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (splainPart w52))))) in sgather @[31, 31, 15, 15] w52 (\\[i54, i55, i56, i57] -> [i57, remH (u53 `sindex0` [i54, i55, i56, i57]) 4, i54, i55, i56, remH (quotH (u53 `sindex0` [i54, i55, i56, i57]) 4) 4])"   printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w52 = sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] u1) (\\[i48, i49] -> [2 * i48 + i49]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = smaxIndex (sreshape @[31,31,15,15,16] (stranspose @[2,3,4,0,5,1] w52)) in stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (sscatter dret (\\[i59, i60, i61, i62] -> [i62, remH (kfromS (u53 !$ [i59, i60, i61, i62])) 4, i59, i60, i61, remH (quotH (kfromS (u53 !$ [i59, i60, i61, i62])) 4) 4])) (\\[i63, i64] -> [2 * i63 + i64]))) (\\[i65, i66] -> [2 * i65 + i66]))"--- TODO: different test result with GHC 9.10:   printArtifactPretty (simplifyArtifact artifactRev)---    @?= "\\dret u1 -> let u53 = smaxIndex (sreshape @[31,31,15,15,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] u1) (\\[i82, i83] -> [2 * i82 + i83]))) (\\[i50, i51] -> [2 * i50 + i51])))) in stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (sscatter dret (\\[i59, i60, i61, i62] -> [i62, remH (kfromS (u53 !$ [i59, i60, i61, i62])) 4, i59, i60, i61, remH (quotH (kfromS (u53 !$ [i59, i60, i61, i62])) 4) 4])) (\\[i63, i64] -> [2 * i63 + i64]))) (\\[i65, i66] -> [2 * i65 + i66]))"+    @?= "\\dret u1 -> stranspose @[1, 2, 0] (sscatter @[15, 4] (stranspose @[3, 4, 1, 2, 0] (sscatter @[15, 4] (let w52 = sgather @[15, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[15, 4] (stranspose @[2, 0, 1] u1) (\\[i48, i49] -> [2 * i48 + i49]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = sargMax (sreshape @[31, 31, 15, 15, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (splainPart w52))))) in sscatter @[31, 31, 15, 15] dret (\\[i59, i60, i61, i62] -> [i62, remH (u53 `sindex0` [i59, i60, i61, i62]) 4, i59, i60, i61, remH (quotH (u53 `sindex0` [i59, i60, i61, i62]) 4) 4])) (\\[i63, i64] -> [2 * i63 + i64]))) (\\[i65, i66] -> [2 * i65 + i66]))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> stranspose @[1, 2, 0] (sscatter @[15, 4] (stranspose @[3, 4, 1, 2, 0] (sscatter @[15, 4] (let u129 = sargMax (sreshape @[31, 31, 15, 15, 16] (stranspose @[2, 3, 4, 5, 0, 1] (sreplicateN @[1, 1] (stranspose @[2, 3, 4, 0, 5, 1] (sgather @[15, 4] (stranspose @[4, 2, 3, 0, 1] (sgather @[15, 4] (stranspose @[2, 0, 1] (splainPart u1)) (\\[i126, i127] -> [2 * i126 + i127]))) (\\[i50, i51] -> [2 * i50 + i51])))))) in sscatter @[31, 31, 15, 15] dret (\\[i59, i60, i61, i62] -> [i62, remH (u129 `sindex0` [i59, i60, i61, i62]) 4, i59, i60, i61, remH (quotH (u129 `sindex0` [i59, i60, i61, i62]) 4) 4])) (\\[i63, i64] -> [2 * i63 + i64]))) (\\[i65, i66] -> [2 * i65 + i66]))" -smaximum4 :: forall r sh target. (ADReady target, GoodScalar r, KnownShS sh)+smaximum4 :: forall r sh target. (ADReady target, NumScalar r, KnownShS sh)           => target (TKS sh r) -> target (TKS '[] r)-smaximum4 t0 =+smaximum4 t0 | Refl <- lemAppNil @sh =   tlet t0 $ \t ->-  ttletPrimal (tprimalPart $ kfromS $ smaxIndex (sflatten t)) $ \maxIndex ->-    sindex0 t-    $ fromLinearIdxS (tprimalPart @target . kconcrete . fromIntegral)-                     (sshape t)-                     maxIndex+  tletPrimal (tprimalPart $ kfromS $ sargMax (sflatten t)) $ \argMax ->+    sindex t+    $ fromLinearIdxS (sshape t)+                     (kplainPart $ kfromPrimal @target argMax) +fromLinearIdxS :: forall sh j. IntegralH j+               => ShS sh -> j -> IxS sh j+fromLinearIdxS = \sh lin -> case go sh lin of (_, ix) -> ix+  where+    go :: ShS sh1 -> j -> (j, IxS sh1 j)+    go ZSS !n = (n, ZIS)+    go ((:$$) n sh) lin =+      let (tensLin, idxInTens) = go sh lin+          tensLin' = tensLin `quotH` fromIntegral (fromSNat' n)+          i = tensLin `remH` fromIntegral (fromSNat' n)+      in (tensLin', i :.$ idxInTens)+ maxPool2dUnpaddedS4   :: forall ksize stride batch_size channels h w target r shOut shK1.      ( KnownNat ksize, KnownNat stride, KnownNat batch_size, KnownNat channels      , KnownNat h, KnownNat w      , 1 <= stride  -- wrongly reported as redundant due to plugins-     , ADReady target, GoodScalar r+     , ADReady target, NumScalar r      , shOut ~ '[batch_size, channels, h `Div` stride, w `Div` stride]      , shK1 ~ '[1, 1, ksize, ksize]      )@@ -1606,9 +1660,143 @@   -> target (TKS shOut r) maxPool2dUnpaddedS4 arr =   let stride = valueOf @stride :: Int-  in sbuild @(Rank shOut) $ \case+  in sbuild @shOut $ \case     [iImg, iChan, iBh, iBw] ->       smaximum4 $ slicezS @shK1 arr [ iImg, iChan                                     , fromIntegral stride * iBh                                     , fromIntegral stride * iBw ]     _ -> error "maxPool2dUnpaddedS4: impossible pattern needlessly required"++type IntOr8 = Int8+type Int8OrDouble = Int8++detSquare :: forall target. ADReady target+          => target (TKR 2 Double) -> target (TKScalar Double)+detSquare =+  let fact :: Int -> Int+      fact = factAcc 1+        where factAcc acc i | i <= 1 = acc+              factAcc acc i = factAcc (i * acc) (i - 1)++      fused :: forall s.+               Int -> Int -> VSM.MVector s IntOr8 -> VSM.MVector s Bool+            -> ST s ()+      fused !len !idx0 !perm !freeSpots = do+        let nthFreeSpot :: Int -> Int -> ST s Int+            nthFreeSpot !pos !el = do+              free <- VSM.read freeSpots el+              if pos <= 0 && free+              then return el+              else nthFreeSpot (pos - fromEnum free) (el + 1)+            loop :: Int -> Int -> Int -> ST s ()+            loop _ _ 0 = return ()+            loop !idx !fi i2 = do+              let fi2 = fi `quot` i2+                  (idxDigit, idxRest) = idx `quotRem` fi2+              el <- nthFreeSpot idxDigit 0+              VSM.write perm (len - i2) (fromIntegral el)+              VSM.write freeSpots el False+              loop idxRest fi2 (i2 - 1)+        loop idx0 (fact len) len++      mutated :: forall s. Int -> ST s (VS.Vector IntOr8)+      mutated !len = do+        perms <- VSM.unsafeNew (len * fact len)+        freeSpots <- VSM.unsafeNew len+        let loop :: Int -> ST s ()+            loop (-1) = return ()+            loop i = do+              VSM.set freeSpots True+              fused len i (VSM.slice (i * len) len perms) freeSpots+              loop (i - 1)+        loop (fact len - 1)+        VS.unsafeFreeze perms++      idx_to_perm :: Int -> Nested.Ranked 2 IntOr8+      idx_to_perm n = Nested.rfromVector [fact n, n] $ runST $ mutated n++      inversion_number_from_idx :: Int -> Nested.Ranked 1 Int8OrDouble+      inversion_number_from_idx n =+        let loop :: Int -> Int -> Int -> Int -> Int8OrDouble+            loop s _ _ i | i == 1 = fromIntegral s+            loop s idx fi i =+              let fi2 = fi `quot` i+                  (s1, idx2) = idx `quotRem` fi2+                  s2 = s + s1+              in loop s2 idx2 fi2 (i - 1)+            f idx0 = loop 0 idx0 (fact n) n+        in Nested.rfromVector [fact n] $ VS.generate (fact n) f++      productR :: target (TKR 1 Double) -> target (TKScalar Double)+      productR = kfromR . rfold (*) (rscalar 1)++      det :: target (TKR 2 Double) -> target (TKScalar Double)+      det a =+        let ell = rwidth a+            p :: PlainOf target (TKR 2 IntOr8)+            p = rconcrete $ idx_to_perm ell+            q :: PlainOf target (TKR 1 Int8OrDouble)+            q = rconcrete $ inversion_number_from_idx ell+            f :: IntOf target -> target (TKScalar Double)+            f i = (-1) ** kfromIntegral (kfromPlain (q `rindex0` [i]))+                  * productR (rgather1 ell a $ \i2 ->+                                [i2, kfromIntegral $ p `rindex0` [i, i2]])+        in withSNat (fact ell) $ \ (SNat @k) ->+             ssum0 $ kbuild1 @k f+      {-+      gradient :: Input -> GradientOutput+      gradient a =+        let ell = rwidth a+        in if ell /= 5+             let art = gradArtifact det a+                 artSimp = simplifyArtifactRev art+             in traceShow ("gradient", printArtifactSimple art) $+                traceShow ("gradSimp", printArtifactSimple artSimp) $+                 gradInterpretArtifact artSimp (chunk ell a)+      -}+  in det++testDetSquarePP :: Assertion+testDetSquarePP = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent detSquare+                                     (FTKR [3, 3] (FTKScalar @Double))+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> rfromS (sscatter @[3, 6] (sreverse (tproject2 (tmapAccumLDer (SNat @3) <lambda> <lambda> <lambda> (sfromPlain (sconcrete (sreplicate [6] (-1.0)) ** sfromIntegral (sconcrete (sfromListLinear [6] [0,1,1,2,2,3]))) * sreplicate0N @[6] dret) (let m15 = sgather @[3, 6] (sfromR m1) (\\[i6, i7] -> [i6, kfromIntegral (sconcrete (sfromListLinear [6,3] [0,1,2,0,2,1,1,0,2,1,2,0,2,0,1,2,1,0]) `sindex0` [i7, i6])]) in tpair (sreverse (tproject2 (tmapAccumLDer (SNat @3) <lambda> <lambda> <lambda> (sconcrete (sreplicate [6] 1.0)) m15))) (sreverse m15))))) (\\[i13, i14] -> [i13, kfromIntegral (sconcrete (sfromListLinear [6,3] [0,1,2,0,2,1,1,0,2,1,2,0,2,0,1,2,1,0]) `sindex0` [i14, i13])]))"+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> rfromS (sscatter @[3, 6] (sreverse (tproject2 (let m8 = sgather @[3, 6] (sfromR m1) (\\[i6, i7] -> [i6, kfromIntegral (sconcrete (sfromListLinear [6,3] [0,1,2,0,2,1,1,0,2,1,2,0,2,0,1,2,1,0]) `sindex0` [i7, i6])]) ; m9 = tmapAccumLDer (SNat @3) <lambda> <lambda> <lambda> (sconcrete (sreplicate [6] 1.0)) m8 ; v10 = sconcrete (sreplicate [6] (-1.0)) ** sfromIntegral (sconcrete (sfromListLinear [6] [0,1,1,2,2,3])) in tmapAccumLDer (SNat @3) <lambda> <lambda> <lambda> (sfromPlain v10 * sreplicate0N @[6] dret) (tpair (sreverse (tproject2 m9)) (sreverse m8))))) (\\[i13, i14] -> [i13, kfromIntegral (sconcrete (sfromListLinear [6,3] [0,1,2,0,2,1,1,0,2,1,2,0,2,0,1,2,1,0]) `sindex0` [i14, i13])]))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> sdot0 (sfromPlain (sconcrete (sreplicate [6] (-1.0)) ** sfromIntegral (sconcrete (sfromListLinear [6] [0,1,1,2,2,3])))) (tproject1 (tmapAccumLDer (SNat @3) <lambda> <lambda> <lambda> (sconcrete (sreplicate [6] 1.0)) (sgather @[3, 6] (sfromR m1) (\\[i6, i7] -> [i6, kfromIntegral (sconcrete (sfromListLinear [6,3] [0,1,2,0,2,1,1,0,2,1,2,0,2,0,1,2,1,0]) `sindex0` [i7, i6])]))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> ssum0 @[6] (let m8 = sgather @[3, 6] (sfromR m1) (\\[i6, i7] -> [i6, kfromIntegral (sconcrete (sfromListLinear [6,3] [0,1,2,0,2,1,1,0,2,1,2,0,2,0,1,2,1,0]) `sindex0` [i7, i6])]) ; m9 = tmapAccumLDer (SNat @3) <lambda> <lambda> <lambda> (sconcrete (sreplicate [6] 1.0)) m8 ; v10 = sconcrete (sreplicate [6] (-1.0)) ** sfromIntegral (sconcrete (sfromListLinear [6] [0,1,1,2,2,3])) in sfromPlain v10 * tproject1 (tpair (tproject1 m9) (sconcrete (sreplicate [3] Z1))))"++testDetSquare3 :: Assertion+testDetSquare3 =+  assertEqualUpToEpsilon' 1e-5+    (ringestData [6,6] [-4.869053606042778e-3,5.097805950791719e-4,2.8429105389292052e-3,2.7764540017998364e-2,4.728231081593177e-3,-2.4786176657837146e-2,-2.3469141678115176e-3,1.1152144353350756e-2,-5.653171988563392e-3,-2.9237121329780706e-3,1.9330615364900276e-3,-9.47499409844526e-4,8.025631739575951e-3,1.532840697726606e-3,-7.047121923697899e-3,-4.305940089911722e-3,1.4834586948324607e-4,1.9922535230070225e-3,-8.022565662650867e-3,-1.2603575569596842e-2,8.144948602197404e-3,-9.343418236792108e-3,8.06305133890595e-3,1.7102844846324558e-2,4.888906896881243e-3,-4.884300437195866e-3,4.144990592410334e-4,-4.765756648788264e-3,-4.741976230085896e-3,1.3056977293001447e-2,1.0096879396193574e-2,1.1166653675142208e-2,-1.279846251614077e-4,-7.896243314270949e-3,-9.67019219688691e-3,-5.316279354872351e-3])+    (rev' @Double @0 (rfromK . detSquare)+                     (ringestData [6,6]+          (map (* 0.01)+             [ 40.1,32.1,40.1,32.1,40.1,22.0+             , 9.71,58.8943200001,18.1,29.1,32.1,40.1+             , 55.8943200001, 1, -2, 0.97,58.200001,1+             , 40.1,29.0,53.99432,7.1,58.8943200001,18.1+             , 32.1,40.1,29.0,53.99432,2.971,58.8943200001+             , 40.1,32.0,53.99432,0.97,25.8943200001, 5 ])))++testDetSquare9 :: Assertion+testDetSquare9 =+  assertEqualUpToEpsilon 1e-10+    (ringestData [10,10] [198.10693359375,-208.391845703125,-74.715576171875,-235.9921875,-56.984375,31.765625,20.578125,-214.5625,-41.9375,-181.75,-158.23355102539063,15.689483642578125,-6.07183837890625,-32.855712890625,-18.2464599609375,-13.33154296875,8.1669921875,9.7578125,36.15625,49.59375,4.102272751895559e15,-7.306139851602269e13,1.6344669707400714e16,-2.0387070048459585e15,1.1609737081775948e14,-5.972802872324755e15,-8.292337523172424e14,6.612381495486813e14,1.3074134839774018e15,-4.551880458211714e15,14.17529296875,-9.8779296875,28.02752685546875,-12.56640625,40.626953125,-29.2734375,4.2890625,41.9375,-2.474609375,-6.21875,-54.611328125,-16.91455078125,-39.5654296875,8.5,-6.6484375,35.75,-9.375e-2,-57.0625,19.8515625,-28.75,32.62109375,16.92626953125,94.29736328125,10.3125,-22.953125,-14.3125,-8.9375,9.625,0.259765625,-17.0,-4.10227275189546e15,7.306139851613288e13,-1.6344669707400854e16,2.0387070048459333e15,-1.1609737081777206e14,5.972802872324765e15,8.292337523172635e14,-6.612381495487908e14,-1.3074134839774085e15,4.551880458211696e15,-44.125,63.90625,24.578125,-34.1015625,-16.609375,-12.015625,0.109375,-7.28125,10.8671875,10.1875,-70.75,-21.0,-31.0625,31.78125,-19.25,77.75,-19.125,43.625,5.875,-32.375,51.25,-17.875,26.5,-10.34375,36.1875,23.4375,1.625,-45.625,-20.59375,-68.25])+    (grad detSquare+          (ringestData [10,10]+             [ 7, 0.92, 0.1, -0.2, 13.1, 9, 8, -4, 34, 2.99432+             , 7, 25.8, 8.1,29.1,32.1,40.1,32.1,40.1,292.0,53.99432+             , 7, 40.1,32.1,40.1,89.0,53.99432,97.1,56.8200001,97.1,52.843201+             , 7, 97.1,55.8943200001,97.1,85.894001,97.1,85.801,18.1,29.1,32.1+             , 7, 40.1,32.1,40.1,32.1,40.1,22.0,53.99432,97.1,82.8943200001+             , 7, 97.1,22.8943200001,97.1,58.8943200001,18.1,29.1,32.1,40.1,32.1+             , 7, 40.1,32.1,40.1,89.0,53.99432,97.1,56.8200001,97.1,52.843201+             , 7, 97.1,55.8943200001, 1, -2, 97.1,58.200001,97.1,55.894320,97.1+             , 7, 29.1,32.1,40.1,29.0,53.99432,97.1,58.8943200001,18.1,29.1+             , 7, 32.1,40.1,32.0,53.99432,97.1,25.8943200001, 5, 2, 6 ]))
test/simplified/TestHighRankSimplified.hs view
@@ -6,7 +6,6 @@  import Prelude -import Data.Int (Int64) import GHC.Exts (IsList (..)) import GHC.TypeLits (KnownNat, type (+), type (-), type (<=)) import Test.Tasty@@ -18,7 +17,8 @@ import Data.Array.Nested.Shaped.Shape  import HordeAd-import HordeAd.Core.Ops (tD, tfromPrimal)+import HordeAd.Core.AstFreshId (resetVarCounter)+import HordeAd.Core.Ops (tD)  import CrossTesting import EqEpsilon@@ -27,7 +27,7 @@ testTrees =   [ testCase "3foo" testFoo   , testCase "3bar" testBar---  , testCase "3barS" testBarS+  , testCase "3barS" testBarS   , testCase "3fooD T Double [1.1, 2.2, 3.3]" testFooD   , testCase "3fooBuild0" testFooBuild0   , testCase "3fooBuildOut" testFooBuildOut@@ -51,9 +51,9 @@   , testCase "3nestedBuildMap1" testNestedBuildMap1   , testCase "3nestedBuildMap10" testNestedBuildMap10   , testCase "3nestedBuildMap11" testNestedBuildMap11---  , testCase "3nestedBuildMap7" testNestedBuildMap7+-- TODO: too slow due to primitive piplines:  , testCase "3nestedBuildMap7" testNestedBuildMap7   , testCase "3nestedSumBuild1" testNestedSumBuild1---  , testCase "3nestedSumBuild5" testNestedSumBuild5+-- TODO: too slow due to primitive piplines:  , testCase "3nestedSumBuild5" testNestedSumBuild5   , testCase "3nestedSumBuildB" testNestedSumBuildB   , testCase "3nestedBuildIndex" testNestedBuildIndex   , testCase "3barReluADValDt" testBarReluADValDt@@ -68,6 +68,8 @@   , testCase "3concatBuild1" testConcatBuild1   , testCase "3concatBuild0m" testConcatBuild0m   , testCase "3concatBuild1m" testConcatBuild1m+  , testCase "3concatBuild2m" testConcatBuild2m+  , testCase "3concatBuild2mPP" testConcatBuild2mPP   , testCase "3concatBuild2" testConcatBuild2   , testCase "3concatBuild22" testConcatBuild22   , testCase "3concatBuild3" testConcatBuild3@@ -82,6 +84,11 @@   , testCase "3logisticC0" testLogisticC0   ] +rfromIndex0 :: forall r target.+               (BaseTensor target, ConvertTensor target, NumScalar r)+            => IntOf target -> target (TKR 0 r)+rfromIndex0 = rfromIntegral . rfromK . tfromPlain STKScalar+ foo :: RealFloatH a => (a,a,a) -> a foo (x,y,z) =   let w = x * sin y@@ -96,15 +103,15 @@ testFoo =   assertEqualUpToEpsilon 1e-3     (ringestData [2,2,1, 2,2] [-4.6947093,1.5697206,-1.6332961,0.34882763,1.5697206,-1.0,-0.9784988,-0.9158946,6.6326222,3.6699238,7.85237,-2.9069107,17.976654,0.3914159,32.98194,19.807974], ringestData [2,2,1, 2,2] [6.943779,-1.436789,33.67549,0.22397964,-1.436789,-1.0,-0.975235,-0.90365005,147.06645,-73.022705,-9.238474,-10.042692,-980.2843,-7.900571,-14.451739,436.9084], ringestData [2,2,1, 2,2] [-4.8945336,2.067469,-1.7196897,1.3341143,2.067469,1.0,0.99846554,0.99536234,6.6943173,3.7482092,7.977362,-3.1475093,18.000969,0.48736274,33.01224,19.845064])-    (grad (kfromR . rsum0 @5 @(TKScalar Float) . foo) (t16, t16, t16))+    (grad (rsum0 @5 @Float . foo) (t16, t16, t16))  bar :: forall a. RealFloatH a => (a, a) -> a bar (x, y) =   let w = foo (x, y, x) * sin y   in atan2H x w + y * w -_barF :: forall a. RealFloatH a => (a, a) -> a-_barF (x, y) =+barF :: forall a. RealFloatH a => (a, a) -> a+barF (x, y) =   let w = _fooF (x, y, x) * sin y   in atan2H x w + y * w @@ -112,21 +119,19 @@ testBar =   assertEqualUpToEpsilon 1e-5     (Concrete $ Nested.rfromListLinear [3,1,2,2,1,2,2] [304.13867,914.9335,823.0187,1464.4688,5264.3306,1790.0055,1535.4309,3541.6572,304.13867,914.9335,823.0187,1464.4688,6632.4355,6047.113,1535.4309,1346.6815,45.92141,6.4903135,5.5406737,1.4242969,6.4903135,1.1458766,4.6446533,2.3550234,88.783676,27.467598,125.27507,18.177452,647.1915,0.3878851,2177.6152,786.1792,6.4903135,6.4903135,6.4903135,6.4903135,2.3550234,2.3550234,2.3550234,2.3550234,21.783596,2.3550234,2.3550234,2.3550234,21.783596,21.783596,21.783596,21.783596],Concrete $ Nested.rfromListLinear [3,1,2,2,1,2,2] [-5728.7617,24965.113,32825.07,-63505.953,-42592.203,145994.88,-500082.5,-202480.06,-5728.7617,24965.113,32825.07,-63505.953,49494.473,-2446.7632,-500082.5,-125885.58,-43.092484,-1.9601002,-98.97709,2.1931143,-1.9601002,1.8243169,-4.0434446,-1.5266153,2020.9731,-538.0603,-84.28137,62.963814,-34986.996,-9.917454,135.30023,17741.998,-1.9601002,-1.9601002,-1.9601002,-1.9601002,-1.5266153,-1.5266153,-1.5266153,-1.5266153,-4029.1775,-1.5266153,-1.5266153,-1.5266153,-4029.1775,-4029.1775,-4029.1775,-4029.1775])-    (cgrad (kfromR . rsum0 . bar @(ADVal Concrete (TKR 7 Float))) (t48, t48))+    (cgrad (rsum0 . bar @(ADVal Concrete (TKR 7 Float))) (t48, t48)) -{- TODO: divergent result; bring back when GHC 9.10 dropped: testBarS :: Assertion testBarS =   assertEqualUpToEpsilon 1e-5-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[3, 1, 2, 2, 1, 2, 2] knownShS [304.13867,914.9335,823.0187,1464.4688,5264.3306,1790.0055,1535.4309,3541.6572,304.13867,914.9335,823.0187,1464.4688,6632.4355,6047.113,1535.4309,1346.6815,45.92141,6.4903135,5.5406737,1.4242969,6.4903135,1.1458766,4.6446533,2.3550234,88.783676,27.467598,125.27507,18.177452,647.1915,0.3878851,2177.6152,786.1792,6.4903135,6.4903135,6.4903135,6.4903135,2.3550234,2.3550234,2.3550234,2.3550234,21.783596,2.3550234,2.3550234,2.3550234,21.783596,21.783596,21.783596,21.783596], sconcrete $ Nested.sfromListPrimLinear @_ @'[3, 1, 2, 2, 1, 2, 2] knownShS [-5728.761,24965.113,32825.074,-63505.957,-42592.203,145994.89,-500082.5,-202480.05,-5728.761,24965.113,32825.074,-63505.957,49494.473,-2446.7632,-500082.5,-125885.58,-43.092484,-1.9601007,-98.97708,2.1931143,-1.9601007,1.8243167,-4.0434446,-1.5266151,2020.9731,-538.06036,-84.28139,62.963818,-34986.992,-9.917454,135.3003,17741.996,-1.9601007,-1.9601007,-1.9601007,-1.9601007,-1.5266151,-1.5266151,-1.5266151,-1.5266151,-4029.1775,-1.5266151,-1.5266151,-1.5266151,-4029.1775,-4029.1775,-4029.1775,-4029.1775])-    (cgrad (kfromS . ssum0 . barF @(ADVal Concrete (TKS '[3, 1, 2, 2, 1, 2, 2] Float))) (sfromR t48, sfromR t48))--}+    (sconcrete $ Nested.sfromListLinear @'[3, 1, 2, 2, 1, 2, 2] knownShS [304.13867,914.9335,823.0187,1464.4688,5264.3306,1790.0055,1535.4309,3541.6572,304.13867,914.9335,823.0187,1464.4688,6632.4355,6047.113,1535.4309,1346.6815,45.92141,6.4903135,5.5406737,1.4242969,6.4903135,1.1458766,4.6446533,2.3550234,88.783676,27.467598,125.27507,18.177452,647.1915,0.3878851,2177.6152,786.1792,6.4903135,6.4903135,6.4903135,6.4903135,2.3550234,2.3550234,2.3550234,2.3550234,21.783596,2.3550234,2.3550234,2.3550234,21.783596,21.783596,21.783596,21.783596], sconcrete $ Nested.sfromListLinear @'[3, 1, 2, 2, 1, 2, 2] knownShS [-5728.761,24965.113,32825.074,-63505.957,-42592.203,145994.89,-500082.5,-202480.05,-5728.761,24965.113,32825.074,-63505.957,49494.473,-2446.7632,-500082.5,-125885.58,-43.092484,-1.9601007,-98.97708,2.1931143,-1.9601007,1.8243167,-4.0434446,-1.5266151,2020.9731,-538.06036,-84.28139,62.963818,-34986.992,-9.917454,135.3003,17741.996,-1.9601007,-1.9601007,-1.9601007,-1.9601007,-1.5266151,-1.5266151,-1.5266151,-1.5266151,-4029.1775,-1.5266151,-1.5266151,-1.5266151,-4029.1775,-4029.1775,-4029.1775,-4029.1775])+    (cgrad (ssum0 . barF @(ADVal Concrete (TKS '[3, 1, 2, 2, 1, 2, 2] Float))) (sfromR t48, sfromR t48))  -- A dual-number and list-based version of a function that goes -- from `R^3` to `R`. fooD :: forall r n. (RealFloatH (ADVal Concrete (TKR n r)))-     => ListR 3 (ADVal Concrete (TKR n r)) -> ADVal Concrete (TKR n r)-fooD (x ::: y ::: z ::: ZR) =+     => IxR 3 (ADVal Concrete (TKR n r)) -> ADVal Concrete (TKR n r)+fooD (x :.: y :.: z :.: ZIR) =   let w = x * sin y   in atan2H z w + z * w @@ -134,11 +139,11 @@ testFooD =   assertEqualUpToEpsilon 1e-10     (fromList [ringestData [1,2,2,1,2,2,2,2,2,1] [18.73108960474591,20.665204824764675,25.821775835995922,18.666613887422585,34.775664100213014,62.54884873632415,37.93303229694526,11.635186977032971,18.73108960474591,20.665204824764675,25.821775835995922,20.600738734367262,34.775664100213014,62.54884873632415,16.663997008808924,3.1300339898598155,1.060799258653783,3.78942741815228,0.1889454555944933,-1.060799258653783,62.54884873632415,37.93303229694526,62.54884873632415,35.99996432769119,62.54884873632415,37.93303229694526,11.635186977032971,18.73108960474591,20.665204824764675,25.821775835995922,20.665204824764675,20.665204824764675,25.821775835995922,34.134947381491145,34.775664100213014,45527.22315787758,-4.488300547708207,2.1475176207684497,8.404498097344806,5.747373381623309,5.096832468946128,-2.4630526910399646,18.666613887422585,1.7769486222994448,-215.8115662030395,16.73214939773215,1.060799258653783,1.060799258653783,1.060799258653783,1.060799258653783,2.1475176207684497,2.1475176207684497,2.1475176207684497,2.1475176207684497,16.08742477551077,16.08742477551077,16.08742477551077,16.08742477551077,2.1475176207684497,2.1475176207684497,2.1475176207684497,2.1475176207684497,16.08742477551077,16.08742477551077,16.08742477551077,16.08742477551077,25.821775835995922,5.096832468946128,7.045006174919766,-1.7808956511653404,16.663997008744435,18.533999054066836,-25.177267779903083,16.60317012020362,25.821775835995922,5.096832468946128,7.045006174919766,-1.7808956511653404,16.663997008744435,18.533999054066836,-12.280721583745471,16.60317012020362,5.161956818274285,18.73108960474591,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,188.11000552192755,34.775664100213014,62.54884873632415,35.99996432769119,62.54884873632415,55.32933980086011,62.54884873632415,55.32933980086011,11.635186977032971,18.73108960474591,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,14.152094926881784,34.775664100213014,62.54884873632415,53.39649491503442,62.54884873632415,14.72904006548922,62.54884873632415,37.93303229694526,11.635186977032971,18.73108960474591,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,57.33025874582143,34.775664100213014,62.54884873632415,36.64432517917614,62.54884873632415,34.06684929392724,62.54884873632415,35.99996432769119], ringestData [1,2,2,1,2,2,2,2,2,1] [647.1354943759653,787.5605199613974,1229.333367336918,642.6917612678424,2229.2701397674327,7210.705208776531,2652.3459120285806,250.02943073785886,647.1354943759653,787.5605199613974,1229.333367336918,782.6578815409038,2229.2701397674327,7210.705208776531,512.2982591657892,18.580536443699742,2.518850510725482,26.993800503829114,0.2243239488720164,2.518850510725482,7210.705208776531,2652.3459120285806,7210.705208776531,2388.9603285490866,7210.705208776531,2652.3459120285806,250.02943073785886,647.1354943759653,787.5605199613974,1229.333367336918,787.5605199613974,787.5605199613974,1229.333367336918,2147.9011858437157,2229.2701397674327,-0.5405182383359878,-0.5328698165396271,-0.5099245509210925,130.7140495214786,61.4116989316311,48.40938174779479,11.696956758139343,642.6917612678424,6.317020301049852,85833.87394976329,516.4928003659018,2.518850510725482,2.518850510725482,2.518850510725482,2.518850510725482,-0.5099245509210925,-0.5099245509210925,-0.5099245509210925,-0.5099245509210925,477.4973215160379,477.4973215160379,477.4973215160379,477.4973215160379,-0.5099245509210925,-0.5099245509210925,-0.5099245509210925,-0.5099245509210925,477.4973215160379,477.4973215160379,477.4973215160379,477.4973215160379,1229.333367336918,48.40938174779479,92.00538642301063,6.3430614471479245,512.2982591618282,633.5999783697488,1168.7578661039847,508.56903530563443,1229.333367336918,48.40938174779479,92.00538642301063,6.3430614471479245,512.2982591618282,633.5999783697488,278.48156010484087,508.56903530563443,49.64077766932281,647.1354943759653,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,65212.963738386214,2229.2701397674327,7210.705208776531,2388.9603285490866,7210.705208776531,5642.338335044463,7210.705208776531,5642.338335044463,250.02943073785886,647.1354943759653,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,369.6431004072799,2229.2701397674327,7210.705208776531,5255.048317224881,7210.705208776531,400.3514287686239,7210.705208776531,2652.3459120285806,250.02943073785886,647.1354943759653,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,6057.774447242021,2229.2701397674327,7210.705208776531,2475.225838667682,7210.705208776531,2139.3419044407133,7210.705208776531,2388.9603285490866], ringestData [1,2,2,1,2,2,2,2,2,1] [18.76237979248771,20.69357069589509,25.8444826804669,18.698011972363496,34.7925278085306,62.558226125235436,37.948492946856575,11.685493300971446,18.76237979248771,20.69357069589509,25.8444826804669,20.629193248844963,34.7925278085306,62.558226125235436,16.699160877305292,3.3121428825170947,1.516071490296981,3.9411848287000124,1.0994899188808887,-1.516071490296981,62.558226125235436,37.948492946856575,62.558226125235436,36.01625479268449,62.558226125235436,37.948492946856575,11.685493300971446,18.76237979248771,20.69357069589509,25.8444826804669,20.69357069589509,20.69357069589509,25.8444826804669,34.1521274657041,34.7925278085306,-45527.22317076194,4.617144085155745,-2.4052046956635262,8.474005308282699,5.84854498865513,5.210650526856928,-2.6906888068615635,18.698011972363496,2.0810391881996813,-215.8142842462135,16.767170338627782,1.516071490296981,1.516071490296981,1.516071490296981,1.516071490296981,-2.4052046956635262,-2.4052046956635262,-2.4052046956635262,-2.4052046956635262,16.123846116986126,16.123846116986126,16.123846116986126,16.123846116986126,-2.4052046956635262,-2.4052046956635262,-2.4052046956635262,-2.4052046956635262,16.123846116986126,16.123846116986126,16.123846116986126,16.123846116986126,25.8444826804669,5.210650526856928,7.127782944309438,-2.0844104722608057,16.69916087724094,18.565621417897145,-25.200555362084323,16.638462541261234,25.8444826804669,5.210650526856928,7.127782944309438,-2.0844104722608057,16.69916087724094,18.565621417897145,-12.328394068734287,16.638462541261234,5.2743697149763085,18.76237979248771,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,188.113123824884,34.7925278085306,62.558226125235436,36.01625479268449,62.558226125235436,55.33994055377702,62.558226125235436,55.33994055377702,11.685493300971446,18.76237979248771,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,14.193483311576621,34.7925278085306,62.558226125235436,53.40747931617656,62.558226125235436,14.768811697198851,62.558226125235436,37.948492946856575,11.685493300971446,18.76237979248771,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,57.34048958248757,34.7925278085306,62.558226125235436,36.660329315674915,62.558226125235436,34.08406370302229,62.558226125235436,36.01625479268449]])-    (cgrad (kfromR . rsum0 . fooD) (fromList [ t128-               , rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (rscalar (0.7 :: Double))+    (cgrad (rsum0 . fooD) (fromList [ t128+               , rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (0.7 :: Concrete (TKScalar Double))                , t128 ])) -fooBuild0 :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)+fooBuild0 :: forall target r n. (ADReady target, NumScalar r, KnownNat n)           => target (TKR (1 + n) r) -> target (TKR (1 + n) r) fooBuild0 v =   let r = rsum v@@ -151,7 +156,7 @@     (rev' @Double @5 fooBuild0 t16)  fooBuildOut-  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)+  :: forall target r n. (ADReady target, NumScalar r, KnownNat n)   => target (TKR (1 + n) r) -> target (TKR (1 + n) r) fooBuildOut v =   rbuild1 2 $ \ix -> ifH (ix ==. 0)@@ -166,35 +171,35 @@  fooBuild2   :: forall target r n.-     (ADReady target, GoodScalar r, KnownNat n, Floating (target (TKR n r)), RealFloat r)+     (ADReady target, NumScalar r, KnownNat n, Floating (target (TKR n r)), Differentiable r)   => target (TKR (1 + n) r) -> target (TKR (1 + n) r) fooBuild2 v =-  rbuild1 2 $ \ix' -> let ix :: PrimalOf target (TKS '[] Int64)+  rbuild1 2 $ \ix' -> let ix :: PlainOf target (TKS '[] Int)                           ix = sfromR $ rfromK ix' in-    ifH (ix - (sprimalPart . sfloor . sfromR) (rsum0  @5 @(TKScalar r)+    ifH (ix - (tplainPart . sfloor . sfromK) (rsum0 @5 @r                       $ rreplicate0N [5,12,11,9,4] (rsum0 v)) - sscalar 10001 >=. sscalar 0-         &&* ix - (sprimalPart . sfloor . sfromR) (rsum0 @5 @(TKScalar r) @target+         &&* ix - (tplainPart . sfloor . sfromK) (rsum0 @5 @r @target                           $ rreplicate0N [5,12,11,9,4] (rsum0 v)) - sscalar 10001 <=. sscalar 1)-        (rindex v [kfromR $ rfromS $ ix - (sprimalPart . sfloor . sfromR) (rsum0  @5 @(TKScalar r) @target+        (rindex v [kfromR $ rfromS $ ix - (tplainPart . sfloor . sfromK) (rsum0 @5 @r @target                                 $ rreplicate0N [5,12,11,9,4] (rsum0 v)) - sscalar 10001])            -- index out of bounds; also fine         (sqrt $ abs $ rindex v [kfromS-                                $ let rr = (ix - (sfromR . rprimalPart . rfloor) (rsum0 v) - sscalar 10001) `remH` sscalar 2+                                $ let rr = (ix - (sfromK . tplainPart . kfloor) (rsum0 v) - sscalar 10001) `remH` sscalar 2                                   in ifH (signum rr ==. negate (signum $ sscalar 2))                                      (rr + sscalar 2)                                      rr])  fooBuild2L   :: forall k target r n.-     (ADReady target, GoodScalar r, KnownNat n, Floating (target (TKR n r)), RealFloat r)-  => ListR k (target (TKR (1 + n) r)) -> target (TKR (1 + n) r)+     (ADReady target, NumScalar r, KnownNat n, Floating (target (TKR n r)), Differentiable r)+  => IxR k (target (TKR (1 + n) r)) -> target (TKR (1 + n) r) fooBuild2L = foldr1 (+) . fmap fooBuild2  testFooBuild91 :: Assertion testFooBuild91 =   assertEqualUpToEpsilon 1e-8     (fromList $ map (ringestData [2]) [[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299]])-    (cgrad (kfromR . rsum0 @1 . fooBuild2L @50 @(ADVal Concrete) @Double @0)+    (cgrad (rsum0 @1 . fooBuild2L @50 @(ADVal Concrete) @Double @0)        (fromList $ map (ringestData [2]) [[0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9]]))  testFooBuild92 :: Assertion@@ -202,7 +207,7 @@   assertEqualUpToEpsilon 1e-8     (fromList $ map (ringestData [2]) [[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299]])     (grad-       (kfromR . rsum0 @1 . fooBuild2L @50 @(AstTensor AstMethodLet FullSpan) @Double @0)+       (rsum0 @1 . fooBuild2L @50 @(AstTensor AstMethodLet FullSpan) @Double @0)        (fromList $ map (ringestData [2]) [[0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9]]))  testFooBuild21 :: Assertion@@ -219,19 +224,19 @@  fooBuild2S   :: forall k sh target r.-     (ADReady target, GoodScalar r, KnownNat k, Floating (target (TKS sh r)), RealFloat r, KnownShS sh)+     (ADReady target, NumScalar r, Floating (target (TKS sh r)), Differentiable r, KnownShS sh)   => target (TKS (k : sh) r) -> target (TKR (1 + Rank sh) r) fooBuild2S v = rfromS $-  sbuild1 @2 $ \ix' -> let ix :: PrimalOf target (TKS '[] Int64)+  sbuild1 @2 $ \ix' -> let ix :: PlainOf target (TKS '[] Int)                            ix = sfromR $ rfromK ix' in-    ifH (ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+    ifH (ix - (tplainPart . sfloor) (sfromK $ ssum0 @[5,12,11,9,4] @r              $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001 >=. srepl 0-         &&* ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+         &&* ix - (tplainPart . sfloor) (sfromK $ ssum0 @[5,12,11,9,4] @r              $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001 <=. srepl 1)-        (sindex v ((kfromS $ ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+        (sindex v ((kfromS $ ix - (tplainPart . sfloor) (sfromK $ ssum0 @[5,12,11,9,4] @r              $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001) :.$ ZIS ))            -- index out of bounds; also fine-        (sqrt $ abs $ sindex v ((kfromR $ rfromS $ let rr = (ix - (sprimalPart . sfloor) (ssum0 v) - srepl 10001) `remH` srepl 2+        (sqrt $ abs $ sindex v ((kfromR $ rfromS $ let rr = (ix - (tplainPart . sfloor) (sfromK $ ssum0 v) - srepl 10001) `remH` srepl 2                                 in ifH (signum rr ==. negate (signum $ srepl 2))                                    (rr + srepl 2)                                    rr) :.$ ZIS))@@ -250,21 +255,21 @@  fooBuildNest2S   :: forall k sh target r.-     (ADReady target, GoodScalar r, KnownNat k, Floating (target (TKS sh r)), RealFloat r, KnownShS sh)+     (ADReady target, NumScalar r, Floating (target (TKS sh r)), Differentiable r, KnownShS sh)   => target (TKS (k : sh) r) -> target (TKR (1 + Rank sh) r) fooBuildNest2S v = rfromS $-  sbuild1 @2 $ \ix' -> let ix :: PrimalOf target (TKS '[] Int64)+  sbuild1 @2 $ \ix' -> let ix :: PlainOf target (TKS '[] Int)                            ix = sfromR $ rfromK ix' in-    ifH (ix - (sunNest @_ @'[] @'[] . sprimalPart . snest knownShS . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+    ifH (ix - (sunNest @_ @'[] @'[] . tplainPart . snest knownShS . sfloor) (sfromK $ ssum0 @[5,12,11,9,4] @r              $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001 >=. srepl 0-         &&* ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+         &&* ix - (tplainPart . sfloor) (sfromK $ ssum0 @[5,12,11,9,4] @r              $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001 <=. srepl 1)--- TODO:        (sindex v (ShapedList.singletonIndex (ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @r  $ sunNest $ treplicate (SNat @5) knownSTK $ snest (knownShS @[12,11])-        (sindex v ((kfromR $ rfromS $ ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r) @target $ sunNest $ tproject2 $ tfromPrimal knownSTK $ tpair tunit (sprimalPart $ snest (knownShS @[5,12,11])+-- TODO:        (sindex v (ShapedList.singletonIndex (ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @r $ sunNest $ treplicate (SNat @5) knownSTK $ snest (knownShS @[12,11])+        (sindex v ((kfromR $ rfromS $ ix - (tplainPart . sfloor) (sfromK $ ssum0 @[5,12,11,9,4] @r @target $ sunNest $ tproject2 $ tfromPrimal knownSTK $ tpair tunit (sprimalPart $ snest (knownShS @[5,12,11])              $ sreplicate0N @[5,12,11,9,4] (ssum0 v))) - srepl 10001) :.$ ZIS))            -- index out of bounds; also fine -- TODO:        (sunNest @_ @'[] @sh $ tlet (snest (knownShS @'[]) $ (sfromPrimal ix - sfloor (ssum0 v) - srepl 10001) `remH` srepl 2) $ \rr -> snest (knownShS @'[]) $ sqrt $ abs $ sindex v (ShapedList.singletonIndex (ifH (signum (sprimalPart (sunNest rr)) ==. negate (signum $ srepl 2)) (sprimalPart (sunNest rr) + srepl 2) (sprimalPart (sunNest rr)))))-        (sunNest @_ @'[] @sh $ tlet ((sfromPrimal ix - sfloor (ssum0 v) - srepl 10001) `remH` srepl 2) $ \rr -> snest (knownShS @'[]) $ sqrt $ abs $ sindex v ((kfromS $ ifH (signum (sprimalPart rr) ==. negate (signum $ srepl 2)) (sprimalPart rr + srepl 2) (sprimalPart rr)) :.$ ZIS))+        (sunNest @_ @'[] @sh $ tlet ((tfromPlain knownSTK ix - sfloor (sfromK $ ssum0 v) - srepl 10001) `remH` srepl 2) $ \rr -> snest (knownShS @'[]) $ sqrt $ abs $ sindex v ((kfromS $ ifH (signum (tplainPart rr) ==. negate (signum $ srepl 2)) (tplainPart rr + srepl 2) (tplainPart rr)) :.$ ZIS))  testFooBuildNest21S :: Assertion testFooBuildNest21S =@@ -279,11 +284,11 @@     (rev' @Double @5 (fooBuildNest2S @2 @[2, 1, 2, 2] . sfromR) t16)  fooBuild3 :: forall target r n.-             ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+             ( ADReady target, NumScalar r, KnownNat n, RealFloatH (target (TKR n r)) )           => target (TKR (1 + n) r) -> target (TKR (1 + n) r) fooBuild3 v =   rbuild1 22 $ \ix ->-    bar ( rreplicate0N (shrTail $ rshape v) (rscalar 1)+    bar ( rreplicate0N (shrTail $ rshape v) 1         , rindex v [minH 1 (ix + 1)] )  -- index not out of bounds  testFooBuild3 :: Assertion@@ -293,13 +298,13 @@     (rev' @Double @5 fooBuild3 t16)  fooBuild5 :: forall target r n.-             ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+             ( ADReady target, NumScalar r, KnownNat n, RealFloatH (target (TKR n r)) )           => target (TKR (1 + n) r) -> target (TKR (1 + n) r) fooBuild5 v =   let r = rsum v       v' = rreplicate0N (shrTail $ rshape v) $ rminimum $ rflatten v   in rbuild1 2 $ \ix ->-       r * foo ( rreplicate0N (shrTail $ rshape v) (rscalar 3)+       r * foo ( rreplicate0N (shrTail $ rshape v) 3                , rrepl (rshape r) 5 * r                , r * v')        + bar (r, rindex v [minH 1 (ix + 1)])  -- index not out of bounds@@ -309,14 +314,14 @@   assertEqualUpToEpsilon 1e-5     (rconcrete $ Nested.rfromListPrimLinear [2,2,1,2,2] [1.1033568028244503e7,74274.22833989389,-5323238.2765011545,253074.03394016018,4.14744804041263e7,242643.98750578283,-1.922371592087736e7,2.730274503834733e7,1.135709425204681e7,6924.195066252549,-5345004.080027547,255679.51406100337,3.8870981856703006e7,241810.92121468345,-1.9380955730171032e7,2.877024321777493e7])     (vjp @_ @(TKR 5 Double)-           fooBuild5 t16 (rreplicate0N [2, 2, 1, 2, 2] (rscalar 42)))+           fooBuild5 t16 (rreplicate0N [2, 2, 1, 2, 2] 42))  testFooBuildDt2 :: Assertion testFooBuildDt2 =   assertEqualUpToEpsilon 1e-5     (rconcrete $ Nested.rfromListPrimLinear [2,2,1,2,2] [2.206713605648901e7,148548.45667978778,-1.0646476553002307e7,506148.0678803204,8.294896080825263e7,485287.9750115657,-3.844743184175473e7,5.460549007669466e7,2.271418850409362e7,13848.390132505112,-1.0690008160055092e7,511359.0281220066,7.774196371340603e7,483621.8424293669,-3.876191146034207e7,5.754048643554987e7])     (vjp @_ @(TKProduct (TKR 5 Double) (TKR 5 Double))-           (\x -> let y = fooBuild5 x in tpair y y) t16 (let dt = rreplicate0N [2, 2, 1, 2, 2] (rscalar 42) in tpair dt dt))+           (\x -> let y = fooBuild5 x in tpair y y) t16 (let dt = rreplicate0N [2, 2, 1, 2, 2] 42 in tpair dt dt))  testFooBuild5 :: Assertion testFooBuild5 =@@ -325,12 +330,12 @@     (rev' @Double @7 fooBuild5 t48)  fooBuild1 :: forall target r n.-             ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+             ( ADReady target, NumScalar r, KnownNat n, RealFloatH (target (TKR n r)) )           => target (TKR (1 + n) r) -> target (TKR (1 + n) r) fooBuild1 v =   let r = rsum v-      tk = rreplicate0N (shrTail $ rshape v)-      v' = tk $ rminimum $ rflatten v+      tk = rreplicate0N (shrTail $ rshape v) . kfromR+      v' = tk $ rfromK $ rminimum $ rflatten v   in rbuild1 3 $ \ix ->        r * foo ( tk (rscalar 3)                , tk (rscalar 5) * r@@ -343,15 +348,15 @@     (ringestData [2,2,1,2,2] [394056.00100873224,2652.651012139068,-190115.65273218407,9038.358355005721,1481231.4430045108,8665.8566966351,-686561.2828884773,975098.0370838332,405610.50900167174,247.29268093759174,-190893.00285812665,9131.411216464405,1388249.3520251075,8636.104329095837,-692176.9903632513,1027508.6863491047])     (rev' @Double @5 fooBuild1 t16) -fooMap1 :: (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+fooMap1 :: (ADReady target, NumScalar r, KnownNat n, Differentiable r)         => IShR (1 + n) -> target (TKR 0 r) -> target (TKR (1 + n) r) fooMap1 sh r =-  let v = fooBuild1 $ rreplicate0N sh (r * r)-  in rmap0N (\x -> x * r + rscalar 5) v+  let v = fooBuild1 $ rreplicate0N sh (kfromR $ r * r)+  in rmap0N (\x -> x * kfromR r + 5) v  testFooMap :: Assertion testFooMap =-  assertEqualUpToEpsilon' 1e-3+  assertEqualUpToEpsilon' 1e-2     (rscalar 2.7518227)     (rev' @Float @1 (fooMap1 [130]) (rscalar 0.1)) @@ -360,21 +365,21 @@ testFooMap1 =   assertEqualUpToEpsilon 1e-6     (rscalar 3901.312463734578)-    (grad (kfromR @_ @Double . rsum0 @7 . fooMap1 [4, 3, 2, 3, 4, 5, 3]) (rscalar 0.1))+    (grad (rsum0 @7 @Double . fooMap1 [4, 3, 2, 3, 4, 5, 3]) (rscalar 0.1))  fooNoGo :: forall target r n.-           ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+           ( ADReady target, NumScalar r, KnownNat n, Differentiable r )         => target (TKR (1 + n) r) -> target (TKR (1 + n) r) fooNoGo v =   let r = rsum v       r0 = rsum0 v       shTail = shrTail (rshape v)   in rbuild1 3 (\ix ->-       bar ( rreplicate0N shTail (rscalar 3.14)+       bar ( rreplicate0N shTail 3.14            , bar ( rrepl shTail 3.14                  , rindex v [ix]) )-       + ifH (rindex v (ix * 2 :.: ZIR) <=. rreplicate0N shTail (rscalar 0) &&* 6 >. abs ix)-               r (rreplicate0N shTail (rscalar 5) * r))+       + ifH (rindex v (ix * 2 :.: ZIR) <=. rreplicate0N shTail 0 &&* 6 >. abs ix)+               r (rreplicate0N shTail 5 * r))      / rslice 1 3 (rmap0N (\x -> ifH (x >. r0) r0 x) v)      * rbuild1 3 (const $ rrepl shTail 1) @@ -389,45 +394,53 @@ testFooNoGo10 =   assertEqualUpToEpsilon 1e-10     (ringestData [5, 3, 1, 2, 2, 1, 2, 2] [8.096867407436072e-8,9.973025492756426e-8,9.976696178938985e-8,5.614458707681111e-8,-1.8338500573636686e-7,-2.144970334428336e-7,7.354143606421902e-7,-1.8140041785503643e-7,8.096867407436072e-8,9.973025492756426e-8,9.976696178938985e-8,5.614458707681111e-8,-2.01381292700262e-7,-2.221588091014473e-7,7.354143606421902e-7,-1.9951065225263367e-7,1.7230532848112822e-7,4.5426218104870796e-7,1.430886696893587e-7,9.354993295163118e-7,-5.225515010723883e-7,1.019433073376504e-6,9.64067025472343e-6,-4.872227980305747e-6,8.089200625992941e-8,9.924319994964371e-8,1.092480101004153e-7,-2.8478802468285825e-7,9.641049518625974e-8,2.9624147815716037e-7,-1.950868158558337e-7,9.547754822865364e-8,4.5426218104870796e-7,4.5426218104870796e-7,4.5426218104870796e-7,4.5426218104870796e-7,-4.872227980305747e-6,-4.872227980305747e-6,-4.872227980305747e-6,-4.872227980305747e-6,9.361277121832246e-8,-4.872227980305747e-6,-4.872227980305747e-6,-4.872227980305747e-6,9.361277121832246e-8,9.361277121832246e-8,9.361277121832246e-8,9.361277121832246e-8,-5.488572216677945e-7,-1.8496203182958057e-7,-1.4603644180845103e-7,-1.2145268106051633e-7,-2.817402689957553e-7,-2.9913537180597976e-7,6.272804203945257e-7,-2.3697344464172694e-7,-5.488572216677945e-7,-1.8496203182958057e-7,-1.4603644180845103e-7,-1.2145268106051633e-7,-2.613973017956691e-7,-3.0013408634207794e-7,6.272804203945257e-7,-2.916736028401805e-7,-7.0114505846358575e-6,-4.303381366239431e-5,-4.897282418246382e-6,-1.710952247892854e-4,-4.2040039667393255e-5,-2.0204742564752248e-4,-1.7017980671040968e-2,-4.247008401789142e-3,-1.056090348050961e-6,-2.210187184450231e-6,-2.7842041329045203e-6,-1.0402806498987974e-5,-1.2967382896879757e-7,-1.9315601705070884e-5,-2.40087090725031e-7,-2.4419692405172046e-7,-4.303381366239431e-5,-4.303381366239431e-5,-4.303381366239431e-5,-4.303381366239431e-5,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-2.683138631810477e-7,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-2.683138631810477e-7,-2.683138631810477e-7,-2.683138631810477e-7,-2.683138631810477e-7,-5.488572216677945e-7,-1.8496203182958057e-7,-1.4603644180845103e-7,-1.2145268106051633e-7,-2.817402689957553e-7,-2.9913537180597976e-7,6.272804203945257e-7,-2.3697344464172694e-7,-5.488572216677945e-7,-1.8496203182958057e-7,-1.4603644180845103e-7,-1.2145268106051633e-7,-2.613973017956691e-7,-3.0013408634207794e-7,6.272804203945257e-7,-2.916736028401805e-7,-7.0114505846358575e-6,-4.303381366239431e-5,-4.897282418246382e-6,-1.710952247892854e-4,-4.2040039667393255e-5,-2.0204742564752248e-4,-1.7017980671040968e-2,-4.247008401789142e-3,-1.056090348050961e-6,-2.210187184450231e-6,-2.7842041329045203e-6,-1.0402806498987974e-5,-1.2967382896879757e-7,-1.9315601705070884e-5,-2.40087090725031e-7,-2.4419692405172046e-7,-4.303381366239431e-5,-4.303381366239431e-5,-4.303381366239431e-5,-4.303381366239431e-5,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-2.683138631810477e-7,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-2.683138631810477e-7,-2.683138631810477e-7,-2.683138631810477e-7,-2.683138631810477e-7,-5.469529675653596e-7,-2.331458950045675e-7,-1.9907443163522408e-7,-1.4019078434680374e-7,-6.95091094132346e-8,-5.685763846730528e-8,-9.268594848659335e-8,-3.010367762029461e-8,-5.469529675653596e-7,-2.331458950045675e-7,-1.9907443163522408e-7,-1.4019078434680374e-7,-3.415394012988984e-8,-5.069973314807702e-8,-9.268594848659335e-8,-6.380451815099858e-8,-6.883755913116986e-6,-4.273807584344302e-5,-4.79037108793574e-6,-1.705307241188017e-4,-4.2267488166320864e-5,-2.0143642393829028e-4,-1.701262134129569e-2,-4.2496361738088365e-3,-1.0224785375169973e-6,-2.1427637177332083e-6,-2.705952143004936e-6,-1.0493018474305117e-5,-1.819666770962338e-7,-1.911089472080586e-5,-9.045482032374276e-8,-2.819821645880664e-7,-4.273807584344302e-5,-4.273807584344302e-5,-4.273807584344302e-5,-4.273807584344302e-5,-4.2496361738088365e-3,-4.2496361738088365e-3,-4.2496361738088365e-3,-4.2496361738088365e-3,-3.019273543907303e-7,-4.2496361738088365e-3,-4.2496361738088365e-3,-4.2496361738088365e-3,-3.019273543907303e-7,-3.019273543907303e-7,-3.019273543907303e-7,-3.019273543907303e-7,8.287292817679557e-8,5.154639175257732e-8,4.672897196261682e-8,3.740648379052369e-8,2.884615384615385e-8,2.7780699895840894e-8,1.5447991761071065e-8,2.546934916639589e-8,8.287292817679557e-8,5.154639175257732e-8,4.672897196261682e-8,3.740648379052369e-8,2.5862068965517245e-8,2.7275544092553562e-8,1.5447991761071065e-8,2.8358432436548274e-8,3.0000000000000004e-7,7.500000000000001e-7,2.5000000000000004e-7,1.5000000000000002e-6,-7.500000000000001e-7,1.6304347826086957e-6,1.5000000000000002e-5,-7.500000000000001e-6,1.1450381679389314e-7,1.6666666666666668e-7,1.8750000000000003e-7,-3.7500000000000006e-7,4.411764705882353e-8,5.00948462422186e-7,-4.545454545454546e-8,5.76923076923077e-8,7.500000000000001e-7,7.500000000000001e-7,7.500000000000001e-7,7.500000000000001e-7,-7.500000000000001e-6,-7.500000000000001e-6,-7.500000000000001e-6,-7.500000000000001e-6,5.999928000863991e-8,-7.500000000000001e-6,-7.500000000000001e-6,-7.500000000000001e-6,5.999928000863991e-8,5.999928000863991e-8,5.999928000863991e-8,5.999928000863991e-8])-   (grad (kfromR @_ @Double . rsum0 @8 . rmap0N (* rscalar 0.000000001) . fooNoGo) (rmap0N (* rscalar 0.01) $ rreplicate 5 t48))+   (grad (rsum0 @8 @Double . rmap0N (* 0.000000001) . fooNoGo) (rmap0N (* 0.01) $ rreplicate 5 t48))  nestedBuildMap :: forall target n r.-                  (ADReady target, GoodScalar r, n <= 6, KnownNat n, Differentiable r)+                  (ADReady target, NumScalar r, n <= 6, KnownNat n, Differentiable r)                => target (TKR 0 r) -> target (TKR (1 + n) r) nestedBuildMap r =-  let w x = rreplicate0N [4] x :: target (TKR 1 r)-      v' = rreplicate0N (177 :$: ZSR) r-      nestedMap x = rmap0N (x /) (w x)+  let w x = rreplicate0N [4] (kfromR x) :: target (TKR 1 r)+      v' = rreplicate0N (177 :$: ZSR) (kfromR r)+      nestedMap x = rmap0N (kfromR x /) (w x)       variableLengthBuild iy = rbuild1 7 (\ix ->         rindex v' (ix + iy :.: ZIR))       doublyBuild =         rbuild1 3 (rreplicate0N (shrTake @n @(6 - n)                              $ 2 :$: 4 :$: 2 :$: 1 :$: 3 :$: 2 :$: ZSR)                    . rminimum . variableLengthBuild)-  in rmap0N (\x -> x * rsum0-                         (rbuild1 3 (\ix -> bar (x, rindex v' [ix]))-                          + fooBuild1 (nestedMap x)-                          / fooMap1 [3] x)+  in rmap0N (\x -> x * (rsum0+                         (rbuild1 3 (\ix -> bar (rfromK x, rindex v' [ix]))+                          + fooBuild1 (nestedMap (rfromK x))+                          / fooMap1 [3] (rfromK x)))             ) doublyBuild +{- one of the exotic pipelines makes this super-slow with AstIndexK testNestedBuildMap1 :: Assertion testNestedBuildMap1 =   assertEqualUpToEpsilon' 1e-8     (rscalar 22.673212907588812)-    (rev' @Double @1 nestedBuildMap (rscalar 0.6))+    (rev' @Double @1 nestedBuildMap (rscalar 0.6)) -} +testNestedBuildMap1 :: Assertion+testNestedBuildMap1 =+  assertEqualUpToEpsilon 1e-8+    (rscalar 27.207855489106613)+    (vjp (nestedBuildMap @_ @1 @Double) (rscalar 0.6) (rrepl [3, 2] 0.6))+ testNestedBuildMap10 :: Assertion testNestedBuildMap10 =   assertEqualUpToEpsilon 1e-8     (map rscalar [109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257, 109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257, 109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257])-    (map (cgrad (kfromR . rsum0 @1 @(TKScalar Double) . nestedBuildMap))+    (map (cgrad @_ @_ @_ @Concrete+                (rsum0 @1 @Double . nestedBuildMap))          (map (Concrete . Nested.rscalar) $ [0.1, 0.2 .. 1] ++ [0.1, 0.2 .. 1] ++ [0.1, 0.2 .. 1]))  testNestedBuildMap11 :: Assertion testNestedBuildMap11 =   assertEqualUpToEpsilon 1e-8     (map rscalar [109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257, 109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257, 109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257])-    (map (grad (kfromR . rsum0 @1 @(TKScalar Double) . nestedBuildMap))+    (map (grad (rsum0 @1 @Double . nestedBuildMap))          (map (Concrete . Nested.rscalar) $ [0.1, 0.2 .. 1] ++ [0.1, 0.2 .. 1] ++ [0.1, 0.2 .. 1]))  {-@@ -442,13 +455,14 @@ -- to nestedBuildMap doesn't help. nestedSumBuild   :: forall target n r.-     (ADReady target, GoodScalar r, n <= 4, KnownNat n, Differentiable r)+     (ADReady target, NumScalar r, n <= 4, KnownNat n, Differentiable r)   => target (TKR n r) -> target (TKR (2 + n) r) nestedSumBuild v =-  rbuild1 13 $ \ix1 -> rbuild1 4 $ \ix2 ->+  rbuild @2 @n (13 :$: 4 :$: ZSR) $ \ (ix1 :.: ix2 :.: ZIR) ->+--  rbuild1 13 $ \ix1 -> rbuild1 4 $ \ix2 ->     ifH (ix2 >. ix1)-        (rmap0N ((* rscalar (-0.00000003)) . sqrt . abs)-         $ nestedBuildMap (rsum0 v)+        (rmap0N ((* (-0.00000003)) . sqrt . abs)+         $ nestedBuildMap (rfromK $ rsum0 v)            `rindex` (ix2 `remH` 3 :.: minH 1 ix1 :.: minH ix1 3 :.: ZIR))         (nestedBuildMap (rscalar 0.00042)          `rindex` (ix2 `remH` 3 :.: minH 1 ix1 :.: minH ix1 3 :.: ZIR))@@ -457,7 +471,7 @@ testNestedSumBuild1 =   assertEqualUpToEpsilon 1e-6     (ringestData [5] [5.738943380972744e-6,5.738943380972744e-6,5.738943380972744e-6,5.738943380972744e-6,5.738943380972744e-6])-    (grad (kfromR . rsum0 @3 @(TKScalar Double) . nestedSumBuild) (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))+    (grad (rsum0 @3 @Double . nestedSumBuild) (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))  {- testNestedSumBuild5 :: Assertion@@ -467,10 +481,10 @@     (rev' @Double @5 nestedSumBuild (rsum (rsum t16))) -} -nestedSumBuildB :: forall target n r. (ADReady target, GoodScalar r, KnownNat n)+nestedSumBuildB :: forall target n r. (ADReady target, NumScalar r, KnownNat n)                 => target (TKR (1 + n) r) -> target (TKR 3 r) nestedSumBuildB v =-  rbuild @2 [13, 4, 2] $ \case+  rbuild @2 [13, 4] $ \case     [ix, ix2] ->       flip rindex [ix2]         (rfromList@@ -479,7 +493,7 @@              , rindex v (fromList                          $ replicate (rlength v - 1)                              (maxH 0 $ minH 1 $ ix2 `quotH` 2 + ix `quotH` 4 - 1))-             , rbuild1 2 (\_ -> rsum0 v)+             , rbuild1 2 (\_ -> rfromK $ rsum0 v)              , rsum (rbuild1 7 (\ix7 ->                  rreplicate 2 (rfromIndex0 ix7)))              ])@@ -503,9 +517,9 @@     (rev' @Double @3 nestedBuildIndex t16)  barRelu-  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+  :: ( ADReady target, NumScalar r, KnownNat n, Differentiable r )   => target (TKR n r) -> target (TKR n r)-barRelu x = let t = rreplicate0N (rshape x) (rscalar 0.001) * x+barRelu x = let t = rreplicate0N (rshape x) 0.001 * x             in relu $ bar (t, relu t)  testBarReluADValDt :: Assertion@@ -513,7 +527,7 @@   assertEqualUpToEpsilon 1e-6     (rconcrete $ Nested.rfromListPrimLinear [2,2,1,2,2] [1.2916050471365906e-2,1.2469757606504572e-2,1.3064120086501589e-2,1.2320300700062944e-2,0.0,1.217049789428711e-2,1.2185494267265312e-2,0.0,1.4105363649830907e-2,1.3506236503127638e-2,1.3359213691150671e-2,0.0,1.7066665416485535e-2,1.2618022646204737e-2,0.0,1.595161947206668e-2])     (vjp @_ @(TKR 5 Double)-           barRelu t16 (rreplicate0N [ 2 , 2 , 1 , 2 , 2 ] (rscalar 42.2)))+           barRelu t16 (rreplicate0N [ 2 , 2 , 1 , 2 , 2 ] 42.2))  testBarReluADValDt2 :: Assertion testBarReluADValDt2 =@@ -524,7 +538,7 @@                   (TKProduct (TKR 5 Float)                              (TKS [2,2,1,2,2] Double)))            (\x -> tpair (rsum x) (tpair (rcast $ barRelu x) (sfromR $ barRelu x))) t16-                  (let dt = rreplicate0N [ 2 , 2 , 1 , 2 , 2 ] (rscalar 42.2)+                  (let dt = rreplicate0N [ 2 , 2 , 1 , 2 , 2 ] 42.2                    in tpair (rsum dt) (tpair (rcast dt) (sfromR dt))))  testBarReluADVal :: Assertion@@ -538,20 +552,19 @@   assertEqualUpToEpsilon' 1e-10     (ringestData [3,1,2,2,1,2,2] [2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885852309100301e-4,2.885923176600045e-4,2.887454843457817e-4,2.886097295122454e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.8860655161315664e-4,2.88595871110374e-4,2.887454843457817e-4,2.885884088500461e-4,2.884182085399516e-4,2.884075468755327e-4,2.8842176240868867e-4,2.8840399312321096e-4,0.0,2.8840370860416445e-4,2.884007943794131e-4,0.0,2.884469945274759e-4,2.8843242392031246e-4,2.884288700806792e-4,0.0,2.885212670262263e-4,2.884110805753153e-4,0.0,2.8849283778617973e-4,2.884075468755327e-4,2.884075468755327e-4,2.884075468755327e-4,2.884075468755327e-4,0.0,0.0,0.0,0.0,2.884892851579934e-4,0.0,0.0,0.0,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4])     (rev' @Double @7 barRelu-         (rmap0N (* rscalar 0.001) t48))+         (rmap0N (* 0.001) t48)) -braidedBuilds :: forall target n r. (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+braidedBuilds :: forall target n r. (ADReady target, NumScalar r, Differentiable r)               => target (TKR (1 + n) r) -> target (TKR 2 r) braidedBuilds r =-  rbuild1 3 (\ix1 ->-    rbuild1 4 (\ix2 -> rindex (rfromList-      [rfromIndex0 ix2, rscalar 7, rsum0 (rslice 1 1 r), rscalar (-0.2)]) (ix1 :.: ZIR)))+  rbuild [3, 4] (\ [ix1, ix2] -> rindex (rfromList+      [rfromIndex0 ix2, rscalar 7, rfromK $ rsum0 (rslice 1 1 r), rscalar (-0.2)]) (ix1 :.: ZIR))  testBraidedBuilds :: Assertion testBraidedBuilds =   assertEqualUpToEpsilon' 1e-10     (ringestData [4] [0.0,4.0,0.0,0.0])-    (rev' @Double @2 (braidedBuilds @_ @0) (rreplicate0N [4] (rscalar 3.4)))+    (rev' @Double @2 (braidedBuilds @_ @0) (rreplicate0N [4] 3.4))  testBraidedBuilds1 :: Assertion testBraidedBuilds1 =@@ -559,40 +572,41 @@     (ringestData [2,2,1,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0])     (rev' @Double @2 braidedBuilds t16) -recycled :: (ADReady target, GoodScalar r, KnownNat n)+recycled :: (ADReady target, NumScalar r, KnownNat n)          => target (TKR n r) -> target (TKR 7 r) recycled r =-  rbuild1 2 $ \_ -> rbuild1 4 $ \_ -> rbuild1 2 $ \_ -> rbuild1 3 $ \_ ->-    nestedSumBuildB (rreplicate 4 r)+  rbuild [2, 4, 2, 3] $ \_ -> nestedSumBuildB (rreplicate 4 r)+--  rbuild1 2 $ \_ -> rbuild1 4 $ \_ -> rbuild1 2 $ \_ -> rbuild1 3 $ \_ ->+--    nestedSumBuildB (rreplicate 4 r)  testRecycled :: Assertion testRecycled =   assertEqualUpToEpsilon' 1e-6     (rrepl [2] 5616)-    (rev' @Double @7 (recycled @_ @_ @1) (rreplicate0N [2] (rscalar 1.0001)))+    (rev' @Double @7 (recycled @_ @_ @1) (rreplicate0N [2] 1.0001))  {- testRecycled1 :: Assertion testRecycled1 =   assertEqualUpToEpsilon' 1e-6     (ringestData [5, 4, 2] [5184.0,5184.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,5424.0,5424.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0])-    (rev' @Double @7 (recycled @_ @_ @3) (rreplicate0N [5, 4, 2] (rscalar 0.0002)))+    (rev' @Double @7 (recycled @_ @_ @3) (rreplicate0N [5, 4, 2] 0.0002)) -}  concatBuild :: forall target r n.-               (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+               (ADReady target, NumScalar r, KnownNat n, Differentiable r)             => target (TKR (1 + n) r) -> target (TKR (3 + n) r) concatBuild r =   rbuild1 7 (\i ->     rconcat [ rbuild1 5 (const r)-            , rbuild1 1 (\j -> rmap0N (* rfromIndex0 (j - i)) r)+            , rbuild1 1 (\j -> rmap0N (* (kfromR . rfromIndex0) (j - i)) r)             , rbuild1 11 (\j ->-                rmap0N (* (rfromIndex0-                  (kfromR (rprimalPart @target (rscalar 125)) * (j `remH` (abs (signum i + abs i) + 1))-                   + maxH j (i `quotH` (j + 1)) * (kfromR . rprimalPart . rfloor) (rsum0 r)+                rmap0N (* ((kfromR . rfromIndex0)+                  (kplainPart @target (kfromR (rscalar 125)) * (j `remH` (abs (signum i + abs i) + 1))+                   + maxH j (i `quotH` (j + 1)) * (kplainPart . kfloor) (rsum0 r)                    - ifH (r <=. r &&* i <. j)-                         (kfromR $ rprimalPart $ rminIndex (rflatten r))-                         ((kfromR . rprimalPart . rfloor) $ rsum0 $ r ! ((i * j) `remH` 7 :.: ZIR))))) r)+                         (kplainPart $ kfromR $ rargMin (rflatten r))+                         ((kplainPart . kfloor) $ rsum0 $ r ! ((i * j) `remH` 7 :.: ZIR))))) r)             , rbuild1 13 (\_k ->                 rsum $ rtr $ rreplicate (rwidth r) (rslice 0 1 r)) ]) @@ -607,15 +621,15 @@ testConcatBuild1 =   assertEqualUpToEpsilon 1e-10     (ringestData [3,1,2,2,1,2,2] [1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3])-    (grad (kfromR . rsum0 @9 @(TKScalar Double) . concatBuild . rmap0N (* rscalar 1e-7)) t48)+    (grad (rsum0 @9 @Double . concatBuild . rmap0N (* 1e-7)) t48)  concatBuildm :: forall target r n.-                (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+                (ADReady target, NumScalar r, KnownNat n, Differentiable r)              => target (TKR (1 + n) r) -> target (TKR (2 + n) r) concatBuildm r =   rbuild1 7 (\i ->-    rmap0N (* (rfromIndex0-      ((kfromR . rprimalPart . rfloor) $ rsum0 $ r ! (i :.: ZIR)))) r)+    rmap0N (* ((kfromR . rfromIndex0)+      ((kplainPart . kfloor) $ rsum0 $ r ! (i :.: ZIR)))) r)  testConcatBuild0m :: Assertion testConcatBuild0m =@@ -627,14 +641,31 @@ testConcatBuild1m :: Assertion testConcatBuild1m =   assertEqualUpToEpsilon' 1e-10+    (ringestData [3,1,2,2,1,2,2] [962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0,962.0])+    (rev' @Double @8 concatBuildm t48)++testConcatBuild2m :: Assertion+testConcatBuild2m =+  assertEqualUpToEpsilon 1e-10     (ringestData [3,1,2,2,1,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])-    (rev' @Double @8 (concatBuildm . rmap0N (* rscalar 1e-7)) t48)+    (vjp (concatBuildm @_ @Double @6 . rmap0N (* 1e-7)) t48 (rrepl [7,3,1,2,2,1,2,2] 1.0)) -concatBuild2 :: (ADReady target, GoodScalar r, KnownNat n)+testConcatBuild2mPP :: Assertion+testConcatBuild2mPP = do+  resetVarCounter+  let t = concatBuildm @(AstTensor AstMethodLet FullSpan) @Double @6+          . rmap0N (* 1e-7)+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent t+                         (FTKR @7 [3,1,2,2,1,2,2] FTKScalar)+  printArtifactSimple (simplifyArtifactRev artifactRev)+    @?= "\\dret w1 -> rfromS (ttranspose (makePerm @[2, 0, 3, 4, 1]) (sreplicateN @[1, 1] (sdot1In (tfromPlain (STKS [3,2,2,2,2,7] STKScalar) (sreplicateN @[3, 2, 2, 2, 2] (sconcrete (sreplicate [7] 1.0e-7) * sfromIntegral (sfloor (ssumN @[1, 2, 2, 1, 2, 2] (ttranspose (makePerm @[1, 2, 3, 4, 5, 6, 0]) (sappend (str (sreplicate @1 (ttranspose (makePerm @[1, 2, 3, 0]) (sreplicate @1 (sconcrete (sreplicate [3,2,2,2,2] 1.0e-7) * ttranspose (makePerm @[1, 4, 0, 2, 3]) (tplainPart (sfromR w1)) !$ [0, 0]))))) (sconcrete (sreplicate [4,1,2,2,1,2,2] 0.0))))))))) (ttranspose (makePerm @[2, 5, 1, 3, 4, 6, 7, 0]) (sfromR dret) !$ [0, 0]))))"++concatBuild2 :: (ADReady target, NumScalar r, KnownNat n)              => target (TKR (1 + n) r) -> target (TKR (3 + n) r) concatBuild2 r =   rbuild1 5 (\i ->-    rbuild1 2 (\j -> rmap0N (* rfromIndex0 (maxH j (i `quotH` (j + 1)))) r))+    rbuild1 2 (\j -> rmap0N (* (kfromR . rfromIndex0) (maxH j (i `quotH` (j + 1)))) r))  testConcatBuild2 :: Assertion testConcatBuild2 =@@ -648,11 +679,11 @@     (ringestData [3,1,2,2,1,2,2] [16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0])     (rev' @Double @9 concatBuild2 t48) -concatBuild3 :: (ADReady target, GoodScalar r)+concatBuild3 :: (ADReady target, NumScalar r)              => target (TKR 1 r) -> target (TKR 2 r) concatBuild3 _r =-  rbuild1 5 (\i ->-    rbuild1 2 (\j -> rfromIndex0 (maxH j (i `quotH` (j + 1)))))+  rbuild [5, 2] (\ [i, j] ->+    rfromIndex0 (maxH j (i `quotH` (j + 1))))  testConcatBuild3 :: Assertion testConcatBuild3 =@@ -680,7 +711,7 @@  logisticOld :: forall target r n.             ( BaseTensor target, LetTensor target-            , BaseTensor (PrimalOf target), KnownNat n, GoodScalar r+            , BaseTensor (PrimalOf target), KnownNat n, NumScalar r             , Floating (PrimalOf target (TKR n r)) )          => target (TKR n r) -> target (TKR n r) logisticOld d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart@@ -711,7 +742,7 @@  logisticA :: forall target r n.             ( BaseTensor target, LetTensor target-            , BaseTensor (PrimalOf target), KnownNat n, GoodScalar r+            , BaseTensor (PrimalOf target), KnownNat n, NumScalar r             , Floating (PrimalOf target (TKR n r)) )          => target (TKR n r) -> target (TKR n r) logisticA d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart@@ -731,7 +762,7 @@  logisticB :: forall target r n.             ( BaseTensor target, LetTensor target-            , BaseTensor (PrimalOf target), KnownNat n, GoodScalar r+            , BaseTensor (PrimalOf target), KnownNat n, NumScalar r             , Floating (PrimalOf target (TKR n r)) )          => target (TKR n r) -> target (TKR n r) logisticB d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart
test/simplified/TestMnistCNNR.hs view
@@ -11,6 +11,7 @@ import Prelude  import Control.Monad (foldM, unless)+import Data.Proxy (Proxy (Proxy)) import System.IO (hPutStrLn, stderr) import System.Random import Test.Tasty@@ -46,12 +47,13 @@ -- which side-steps vectorization. mnistTestCaseCNNRA   :: forall r.-     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> r   -> TestTree mnistTestCaseCNNRA prefix epochs maxBatches khInt kwInt c_outInt n_hiddenInt-                   miniBatchSize totalBatchSize expected =+                   miniBatchSize totalBatchSize expected+                   | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat khInt $ \(_khSNat :: SNat kh) ->   withSNat kwInt $ \(_kwSNat :: SNat kw) ->   withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->@@ -139,7 +141,7 @@   -> TestTree #-}  tensorADValMnistTestsCNNRA :: TestTree-tensorADValMnistTestsCNNRA = testGroup "CNNR ADVal MNIST tests"+tensorADValMnistTestsCNNRA = inOrderTestGroup "CNNR ADVal MNIST tests"   [ mnistTestCaseCNNRA "CNNRA 1 epoch, 1 batch" 1 1 4 4 8 16 1 1                        (1 :: Double)   , mnistTestCaseCNNRA "CNNRA artificial 1 2 3 4 5" 1 1 2 3 4 5 1 10@@ -154,12 +156,13 @@ -- but differentiated anew in each gradient descent iteration. mnistTestCaseCNNRI   :: forall r.-     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> r   -> TestTree mnistTestCaseCNNRI prefix epochs maxBatches khInt kwInt c_outInt n_hiddenInt-                   miniBatchSize totalBatchSize expected =+                   miniBatchSize totalBatchSize expected+                   | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat khInt $ \(_khSNat :: SNat kh) ->   withSNat kwInt $ \(_kwSNat :: SNat kw) ->   withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->@@ -192,7 +195,7 @@                    <$> loadMnistData testGlyphsPath testLabelsPath        let testDataR = mkMnistDataBatchR testData            ftk = tftk @Concrete (knownSTK @(XParams r)) targetInit-       (_, _, var, varAst2) <- funToAstRevIO ftk+       (_, var, varAst2) <- funToAstRevIO ftk        (varGlyph, astGlyph) <-          funToAstIO (FTKR (miniBatchSize                            :$: sizeMnistHeightInt@@ -203,7 +206,7 @@                            :$: sizeMnistLabelInt                            :$: ZSR) FTKScalar) id        let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)-           ast = simplifyInline+           ast = simplifyUserCode                  $ MnistCnnRanked2.convMnistLossFusedR                      miniBatchSize (astGlyph, astLabel)                      (fromTarget varAst2)@@ -263,7 +266,7 @@   -> TestTree #-}  tensorADValMnistTestsCNNRI :: TestTree-tensorADValMnistTestsCNNRI = testGroup "CNNR Intermediate MNIST tests"+tensorADValMnistTestsCNNRI = inOrderTestGroup "CNNR Intermediate MNIST tests"   [ mnistTestCaseCNNRI "CNNRI 1 epoch, 1 batch" 1 1 4 4 8 16 1 1                        (1 :: Double)   , mnistTestCaseCNNRI "CNNRI artificial 1 2 3 4 5" 1 1 2 3 4 5 1 10@@ -279,8 +282,7 @@ -- descent iteration. mnistTestCaseCNNRO   :: forall r.-     ( Differentiable r, GoodScalar r-     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> r   -> TestTree@@ -330,7 +332,7 @@              MnistCnnRanked2.convMnistLossFusedR                miniBatchSize (rprimalPart glyphR, rprimalPart labelR) pars            artRaw = gradArtifact f (fromTarget targetInit, dataInit)-           art = simplifyArtifactGradient artRaw+           art = simplifyArtifactRev artRaw            go :: [MnistDataBatchR r]               -> (Concrete (XParams r), StateAdam (XParams r))               -> (Concrete (XParams r), StateAdam (XParams r))@@ -339,7 +341,7 @@              let parametersAndInput =                    tpair parameters (tpair (rconcrete glyph) (rconcrete label))                  gradient =-                   tproject1 $ fst+                   tproject1 $ snd                    $ revInterpretArtifact art parametersAndInput Nothing              in go rest (updateWithGradientAdam                            @(XParams r)@@ -394,7 +396,7 @@   -> TestTree #-}  tensorADValMnistTestsCNNRO :: TestTree-tensorADValMnistTestsCNNRO = testGroup "CNNR Once MNIST tests"+tensorADValMnistTestsCNNRO = inOrderTestGroup "CNNR Once MNIST tests"   [ mnistTestCaseCNNRO "CNNRO 1 epoch, 1 batch" 1 1 4 4 8 16 1 1                        (1 :: Double)   , mnistTestCaseCNNRO "CNNRO artificial 1 2 3 4 5" 1 1 2 3 4 5 1 10
test/simplified/TestMnistCNNS.hs view
@@ -12,6 +12,7 @@ import Prelude  import Control.Monad (foldM, unless)+import Data.Proxy (Proxy (Proxy)) import GHC.TypeLits (KnownNat, type (<=)) import System.IO (hPutStrLn, stderr) import System.Random@@ -20,6 +21,7 @@ import Text.Printf  import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat')  import HordeAd import HordeAd.Core.Adaptor@@ -51,12 +53,13 @@ mnistTestCaseCNNSA   :: forall kh kw r.      ( 1 <= kh, 1 <= kw-     , Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r )+     , Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r )   => String   -> Int -> Int -> SNat kh -> SNat kw -> Int -> Int -> Int -> Int -> r   -> TestTree mnistTestCaseCNNSA prefix epochs maxBatches kh@SNat kw@SNat c_outInt n_hiddenInt-                   miniBatchSizeInt totalBatchSize expected =+                   miniBatchSizeInt totalBatchSize expected+                   | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->   withSNat n_hiddenInt $ \(_n_hiddenSNat :: SNat n_hidden) ->   withSNat miniBatchSizeInt $ \(miniBatchSize :: SNat miniBatchSize) ->@@ -68,7 +71,7 @@                 0.4 (mkStdGen 44)       name = prefix ++ ": "              ++ unwords [ show epochs, show maxBatches-                        , show (sNatValue kh), show (sNatValue kw)+                        , show (fromSNat' kh), show (fromSNat' kw)                         , show c_outInt, show n_hiddenInt                         , show miniBatchSizeInt                         , show $ widthSTK $ knownSTK @(XParams kh kw c_out n_hidden r)@@ -141,7 +144,7 @@        testErrorFinal @?~ expected  tensorADValMnistTestsCNNSA :: TestTree-tensorADValMnistTestsCNNSA = testGroup "CNNS ADVal MNIST tests"+tensorADValMnistTestsCNNSA = inOrderTestGroup "CNNS ADVal MNIST tests"   [ mnistTestCaseCNNSA "CNNSA 1 epoch, 1 batch"                        1 1 (SNat @4) (SNat @4) 8 16 1 1                        (1 :: Double)@@ -161,12 +164,13 @@ mnistTestCaseCNNSI   :: forall kh kw r.      ( 1 <= kh, 1 <= kw-     , Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r )+     , Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r )   => String   -> Int -> Int -> SNat kh -> SNat kw -> Int -> Int -> Int -> Int -> r   -> TestTree mnistTestCaseCNNSI prefix epochs maxBatches kh@SNat kw@SNat c_outInt n_hiddenInt-                   miniBatchSizeInt totalBatchSize expected =+                   miniBatchSizeInt totalBatchSize expected+                   | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->   withSNat n_hiddenInt $ \(_n_hiddenSNat :: SNat n_hidden) ->   withSNat miniBatchSizeInt $ \(miniBatchSize :: SNat miniBatchSize) ->@@ -178,7 +182,7 @@                 0.4 (mkStdGen 44)       name = prefix ++ ": "              ++ unwords [ show epochs, show maxBatches-                        , show (sNatValue kh), show (sNatValue kw)+                        , show (fromSNat' kh), show (fromSNat' kw)                         , show c_outInt, show n_hiddenInt                         , show miniBatchSizeInt                         , show $ widthSTK $ knownSTK @(XParams kh kw c_out n_hidden r)@@ -200,7 +204,7 @@       withSNat (totalBatchSize * maxBatches) $ \(SNat @lenTestData) -> do        let testDataS = mkMnistDataBatchS testData            ftk = tftk @Concrete (knownSTK @(XParams kh kw c_out n_hidden r)) targetInit-       (_, _, var, varAst2) <- funToAstRevIO ftk+       (_, var, varAst2) <- funToAstRevIO ftk        (varGlyph, astGlyph) <-          funToAstIO (FTKS (miniBatchSize                            :$$ sizeMnistHeight@@ -211,7 +215,7 @@                            :$$ sizeMnistLabel                            :$$ ZSS) FTKScalar) id        let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)-           ast = simplifyInline+           ast = simplifyUserCode                  $ MnistCnnShaped2.convMnistLossFusedS                      kh kw (SNat @c_out) (SNat @n_hidden)                      miniBatchSize (astGlyph, astLabel)@@ -267,7 +271,7 @@        testErrorFinal @?~ expected  tensorADValMnistTestsCNNSI :: TestTree-tensorADValMnistTestsCNNSI = testGroup "CNNS Intermediate MNIST tests"+tensorADValMnistTestsCNNSI = inOrderTestGroup "CNNS Intermediate MNIST tests"   [ mnistTestCaseCNNSI "CNNSI 1 epoch, 1 batch"                        1 1 (SNat @4) (SNat @4) 8 16 1 1                        (1 :: Double)@@ -288,8 +292,7 @@ mnistTestCaseCNNSO   :: forall kh kw r.      ( 1 <= kh, 1 <= kw-     , Differentiable r, GoodScalar r-     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+     , Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r )   => String   -> Int -> Int -> SNat kh -> SNat kw -> Int -> Int -> Int -> Int -> r   -> TestTree@@ -306,7 +309,7 @@                 0.4 (mkStdGen 44)       name = prefix ++ ": "              ++ unwords [ show epochs, show maxBatches-                        , show (sNatValue kh), show (sNatValue kw)+                        , show (fromSNat' kh), show (fromSNat' kw)                         , show c_outInt, show n_hiddenInt                         , show miniBatchSizeInt                         , show $ widthSTK@@ -344,7 +347,7 @@                kh kw (SNat @c_out) (SNat @n_hidden)                miniBatchSize (sprimalPart glyphR, sprimalPart labelR) pars            artRaw = gradArtifact f (fromTarget targetInit, dataInit)-           art = simplifyArtifactGradient artRaw+           art = simplifyArtifactRev artRaw            go :: [MnistDataBatchS miniBatchSize r]               -> ( Concrete (XParams kh kw c_out n_hidden r)                  , StateAdam (XParams kh kw c_out n_hidden r) )@@ -355,7 +358,7 @@              let parametersAndInput =                    tpair parameters (tpair (sconcrete glyph) (sconcrete label))                  gradient =-                   tproject1 $ fst+                   tproject1 $ snd                    $ revInterpretArtifact art parametersAndInput Nothing              in go rest (updateWithGradientAdam                            @(XParams kh kw c_out n_hidden r)@@ -406,7 +409,7 @@        assertEqualUpToEpsilon 1e-1 expected testErrorFinal  tensorADValMnistTestsCNNSO :: TestTree-tensorADValMnistTestsCNNSO = testGroup "CNNS Once MNIST tests"+tensorADValMnistTestsCNNSO = inOrderTestGroup "CNNS Once MNIST tests"   [ mnistTestCaseCNNSO "CNNSO 1 epoch, 1 batch"                        1 1 (SNat @4) (SNat @4) 8 16 1 1                        (1 :: Double)
test/simplified/TestMnistFCNNR.hs view
@@ -9,16 +9,17 @@  import Control.Arrow ((***)) import Control.Monad (foldM, unless)-import Data.Bifunctor (first) import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:))+import GHC.TypeLits (type (<=?)) import System.IO (hPutStrLn, stderr) import System.Random import Test.Tasty import Test.Tasty.HUnit hiding (assert)-import Test.Tasty.QuickCheck hiding (label, shuffle) import Text.Printf  import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)  import HordeAd import HordeAd.Core.Adaptor@@ -54,7 +55,7 @@         MnistFcnnRanked2.mnistTrainBench2VTOGradientX           @Double (Proxy @Float) IgnoreIncomingCotangent           1 (mkStdGen 44) 1500 500-      blackGlyph = rreplicate sizeMnistGlyphInt $ rscalar 7+      blackGlyph = rreplicate0N [sizeMnistGlyphInt] 7       ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))                  targetInit       f :: forall target r. (ADReady target, r ~ Double)@@ -64,10 +65,10 @@                         (tpair (rconcrete $ unConcrete blackGlyph) label)             env = extendEnv (artVarDomainRev art) val emptyEnv         in tproject1 $ tproject2-           $ interpretAst @target env (artDerivativeRev art)+           $ interpretAstFull @target env (artDerivativeRev art)   in assertEqualUpToEpsilon' 1e-10        (ringestData [10] [6.922657834114052e-2,-3.2210167235305924e-5,0.12334696753032606,-4.892729845753193e-3,3.010762414514606e-2,2.0344986964700877e-2,-3.78339785604896e-2,5.77360835535866e-2,0.10761507003315526,-7.909016076299641e-2])-       (rev' f (rreplicate sizeMnistLabelInt $ rscalar 8))+       (rev' f (rreplicate0N [sizeMnistLabelInt] 8))   -- * Using lists of vectors, which is rank 1@@ -80,19 +81,22 @@ -- which side-steps vectorization. mnistTestCase1VTA   :: forall r.-     ( Differentiable r, GoodScalar r+     ( Differentiable r, NumScalar r      , PrintfArg r, AssertEqualUpToEpsilon r )   => String   -> Int -> Int -> Int -> Int -> Double -> Int -> r   -> TestTree mnistTestCase1VTA prefix epochs maxBatches widthHiddenInt widthHidden2Int-                  gamma batchSize expected =+                  gamma batchSize expected+                  | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->   withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+    (stkOfIxR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+    (stkOfIxR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden) :~: True) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden2) :~: True) $   let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters                     Concrete widthHidden widthHidden2 r       valsInit = fst $ randomValue 1 (mkStdGen 44)@@ -168,7 +172,7 @@   -> TestTree #-}  tensorADValMnistTests :: TestTree-tensorADValMnistTests = testGroup "Ranked ADVal MNIST tests"+tensorADValMnistTests = inOrderTestGroup "Ranked ADVal MNIST tests"   [ mnistTestCase1VTA "VTA1 1 epoch, 1 batch" 1 1 300 100 0.02 5000                       (0.2146 :: Double)   , mnistTestCase1VTA "VTA1 artificial 1 2 3 4 5" 1 2 3 4 5 5000@@ -181,19 +185,22 @@ -- but differentiated anew in each gradient descent iteration. mnistTestCase1VTI   :: forall r.-     ( Differentiable r, GoodScalar r+     ( Differentiable r, NumScalar r      , PrintfArg r, AssertEqualUpToEpsilon r )   => String   -> Int -> Int -> Int -> Int -> Double -> Int -> r   -> TestTree mnistTestCase1VTI prefix epochs maxBatches widthHiddenInt widthHidden2Int-                  gamma batchSize expected =+                  gamma batchSize expected+                  | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->   withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+    (stkOfIxR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+    (stkOfIxR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden) :~: True) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden2) :~: True) $   let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters                     Concrete widthHidden widthHidden2 r       valsInit = fst $ randomValue 1 (mkStdGen 44)@@ -220,13 +227,13 @@     trainData <- loadMnistData trainGlyphsPath trainLabelsPath     testData <- map mkMnistDataLinearR . take (batchSize * maxBatches)                 <$> loadMnistData testGlyphsPath testLabelsPath-    (_, _, var, varAst) <- funToAstRevIO ftk+    (_, var, varAst) <- funToAstRevIO ftk     (varGlyph, astGlyph) <-       funToAstIO (FTKR (sizeMnistGlyphInt :$: ZSR) FTKScalar) id     (varLabel, astLabel) <-       funToAstIO (FTKR (sizeMnistLabelInt :$: ZSR) FTKScalar) id     let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)-        ast = simplifyInline+        ast = simplifyUserCode               $ MnistFcnnRanked1.afcnnMnistLoss1                   widthHiddenSNat widthHidden2SNat (astGlyph, astLabel)                   (fromTarget varAst)@@ -280,7 +287,7 @@   -> TestTree #-}  tensorIntermediateMnistTests :: TestTree-tensorIntermediateMnistTests = testGroup "Ranked Intermediate MNIST tests"+tensorIntermediateMnistTests = inOrderTestGroup "Ranked Intermediate MNIST tests"   [ mnistTestCase1VTI "VTI1 1 epoch, 1 batch" 1 1 300 100 0.02 5000                       (0.2116 :: Double)   , mnistTestCase1VTI "VTI1 artificial 1 2 3 4 5" 1 2 3 4 5 5000@@ -294,8 +301,7 @@ -- descent iteration. mnistTestCase1VTO   :: forall r.-     ( Differentiable r, GoodScalar r, ADTensorScalar r ~ r-     , PrintfArg r, AssertEqualUpToEpsilon r)+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> Int -> Int -> Double -> Int -> r   -> TestTree@@ -304,9 +310,11 @@   withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->   withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+    (stkOfIxR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $   withKnownSTK-    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+    (stkOfIxR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden) :~: True) $+  gcastWith (unsafeCoerceRefl :: (1 <=? widthHidden2) :~: True) $   let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters                     Concrete widthHidden widthHidden2 r       valsInit = fst $ randomValue 1 (mkStdGen 44)@@ -345,7 +353,7 @@             widthHiddenSNat widthHidden2SNat             (glyphR, labelR) pars         artRaw = gradArtifact f (valsInit, dataInit)-        art = simplifyArtifactGradient artRaw+        art = simplifyArtifactRev artRaw         go :: [MnistDataLinearR r]            -> Concrete (XParams widthHidden widthHidden2 r)            -> Concrete (XParams widthHidden widthHidden2 r)@@ -353,7 +361,7 @@         go ((glyph, label) : rest) !parameters =           let parametersAndInput =                 tpair parameters (tpair (rconcrete glyph) (rconcrete label))-              gradient = tproject1 $ fst+              gradient = tproject1 $ snd                          $ revInterpretArtifact art parametersAndInput Nothing           in go rest (updateWithGradient gamma knownSTK parameters gradient)     let runBatch :: Concrete (XParams widthHidden widthHidden2 r)@@ -398,7 +406,7 @@   -> TestTree #-}  tensorADOnceMnistTests :: TestTree-tensorADOnceMnistTests = testGroup "Ranked Once MNIST tests"+tensorADOnceMnistTests = inOrderTestGroup "Ranked Once MNIST tests"   [ mnistTestCase1VTO "VTO1 1 epoch, 1 batch" 1 1 300 100 0.02 5000                       (0.2116 :: Double)   , mnistTestCase1VTO "VTO1 artificial 1 2 3 4 5" 1 2 3 4 5 5000@@ -414,13 +422,14 @@ -- which side-steps vectorization. mnistTestCase2VTA   :: forall r.-     ( Differentiable r, GoodScalar r+     ( Differentiable r, NumScalar r      , PrintfArg r, AssertEqualUpToEpsilon r )   => String   -> Int -> Int -> Int -> Int -> Double -> Int -> r   -> TestTree mnistTestCase2VTA prefix epochs maxBatches widthHidden widthHidden2-                  gamma batchSize expected =+                  gamma batchSize expected+                  | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat widthHidden $ \(SNat @widthHidden) ->   withSNat widthHidden2 $ \(SNat @widthHidden2) ->   let targetInit =@@ -490,7 +499,7 @@   -> TestTree #-}  tensorADValMnistTests2 :: TestTree-tensorADValMnistTests2 = testGroup "Ranked2 ADVal MNIST tests"+tensorADValMnistTests2 = inOrderTestGroup "Ranked2 ADVal MNIST tests"   [ mnistTestCase2VTA "VTA2 1 epoch, 1 batch" 1 1 300 100 0.02 5000                        (0.21299999999999997 :: Double)   , mnistTestCase2VTA "VTA2 artificial 1 2 3 4 5" 1 2 3 4 5 5000@@ -505,13 +514,14 @@ -- but differentiated anew in each gradient descent iteration. mnistTestCase2VTI   :: forall r.-     ( Differentiable r, GoodScalar r+     ( Differentiable r, NumScalar r      , PrintfArg r, AssertEqualUpToEpsilon r )   => String   -> Int -> Int -> Int -> Int -> Double -> Int -> r   -> TestTree mnistTestCase2VTI prefix epochs maxBatches widthHidden widthHidden2-                  gamma batchSize expected =+                  gamma batchSize expected+                  | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat widthHidden $ \(SNat @widthHidden) ->   withSNat widthHidden2 $ \(SNat @widthHidden2) ->   let targetInit =@@ -534,13 +544,13 @@     testData <- map mkMnistDataLinearR . take (batchSize * maxBatches)                 <$> loadMnistData testGlyphsPath testLabelsPath     let ftk = tftk @Concrete (knownSTK @(XParams2 r Float)) targetInit-    (_, _, var, varAst) <- funToAstRevIO ftk+    (_, var, varAst) <- funToAstRevIO ftk     (varGlyph, astGlyph) <-       funToAstIO (FTKR (sizeMnistGlyphInt :$: ZSR) FTKScalar) id     (varLabel, astLabel) <-       funToAstIO (FTKR (sizeMnistLabelInt :$: ZSR) FTKScalar) id     let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)-        ast = simplifyInline+        ast = simplifyUserCode               $ MnistFcnnRanked2.afcnnMnistLoss2                   (astGlyph, astLabel)                   (fromTarget varAst)@@ -594,7 +604,7 @@   -> TestTree #-}  tensorIntermediateMnistTests2 :: TestTree-tensorIntermediateMnistTests2 = testGroup "Ranked2 Intermediate MNIST tests"+tensorIntermediateMnistTests2 = inOrderTestGroup "Ranked2 Intermediate MNIST tests"   [ mnistTestCase2VTI "VTI2 1 epoch, 1 batch" 1 1 300 100 0.02 5000                        (0.20779999999999998 :: Double)   , mnistTestCase2VTI "VTI2 artificial 1 2 3 4 5" 1 2 3 4 5 5000@@ -610,8 +620,7 @@ -- descent iteration. mnistTestCase2VTO   :: forall r.-     ( Differentiable r, GoodScalar r-     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> Int -> Int -> Double -> Int -> r   -> TestTree@@ -621,7 +630,7 @@         MnistFcnnRanked2.mnistTrainBench2VTOGradient           @r (Proxy @Float) IgnoreIncomingCotangent           1 (mkStdGen 44) widthHidden widthHidden2-      !art = simplifyArtifactGradient artRaw+      !art = simplifyArtifactRev artRaw       name = prefix ++ ": "              ++ unwords [ show epochs, show maxBatches                         , show widthHidden, show widthHidden2@@ -641,7 +650,7 @@         go ((glyph, label) : rest) !parameters =           let parametersAndInput =                 tpair parameters (tpair (rconcrete glyph) (rconcrete label))-              gradient = tproject1 $ fst+              gradient = tproject1 $ snd                          $ revInterpretArtifact art parametersAndInput Nothing           in go rest (updateWithGradient gamma knownSTK parameters gradient)     let runBatch :: Concrete (XParams2 r Float) -> (Int, [MnistDataLinearR r])@@ -687,97 +696,13 @@   -> TestTree #-}  tensorADOnceMnistTests2 :: TestTree-tensorADOnceMnistTests2 = testGroup "Ranked2 Once MNIST tests"+tensorADOnceMnistTests2 = inOrderTestGroup "Ranked2 Once MNIST tests"   [ mnistTestCase2VTO "VTO2 1 epoch, 1 batch" 1 1 300 100 0.02 5000                        (0.20779999999999998 :: Double)   , mnistTestCase2VTO "VTO2 artificial 1 2 3 4 5" 1 2 3 4 5 5000                        (0.9108 :: Float)   , mnistTestCase2VTO "VTO2 artificial 5 4 3 2 1" 5 4 3 2 1 5000-                       (0.8129 :: Double)+                       (0.817 :: Double)   , mnistTestCase2VTO "VTO2 1 epoch, 0 batch" 1 0 300 100 0.02 5000                        (1 :: Float)-  , testProperty "VTO2 grad vs fwd" $-    \seed0 ->-    forAllShrink (chooseInt (0, 600)) shrinkIntegral $ \width1Hidden ->-    forAllShrink (chooseInt (0, 200)) shrinkIntegral $ \width1Hidden2 ->-    forAllShrink (chooseInt (0, 5)) shrinkIntegral $ \simp ->-    forAll (choose (0.01, 1)) $ \range ->-    forAll (choose (0.01, 1)) $ \range2 ->-    forAll (choose (0.5, 1.5)) $ \dt ->-    forAll (choose (0, 1e-7)) $ \(perturbation :: Double) ->-    withSNat (1 + width1Hidden) $ \(SNat @widthHidden) ->-    withSNat (1 + width1Hidden2) $ \(SNat @widthHidden2) ->-    let (glyph0, seed2) = randomValue @(Concrete (TKS '[SizeMnistGlyph] Double))-                                      0.5 (mkStdGen seed0)-        (label0, seed3) = randomValue @(Concrete (TKS '[SizeMnistLabel] Double))-                                      5 seed2-        (glyph, label) = ( rmap1 (rscalar 0.5 +) $ forgetShape glyph0-                         , rmap1 (rscalar 5 + ) $ forgetShape label0 )-        ds :: Concrete (XParams2 Double Double)-        (ds, seed4) = first forgetShape $-          randomValue-            @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped-                             Concrete widthHidden widthHidden2 Double Double)))-            range seed3-        (targetInit, artRaw) =-          MnistFcnnRanked2.mnistTrainBench2VTOGradient-            @Double (Proxy @Double) UseIncomingCotangent-            range2 seed4 (1 + width1Hidden) (1 + width1Hidden2)-        art = iterate simplifyArtifactGradient artRaw !! simp-        stk = knownSTK @(XParams2 Double Double)-        ftk = tftk @Concrete stk targetInit-        parametersAndInput = tpair targetInit (tpair glyph label)-        (_gradient0, value0) = first tproject1 $-          revInterpretArtifact art parametersAndInput Nothing-        (gradient1, value1) = first tproject1 $-          revInterpretArtifact art parametersAndInput (Just $ kconcrete dt)-        f :: ADVal Concrete (XParams2 Double Double)-          -> ADVal Concrete (TKScalar Double)-        f adinputs =-          MnistFcnnRanked2.afcnnMnistLoss2-            (rfromPrimal glyph, rfromPrimal label) (fromTarget adinputs)-        (derivative2, value2) = cfwdBoth f targetInit ds---        goodDt :: forall r. GoodScalar r => r---        goodDt = ifDifferentiable @r (realToFrac dt) 0---        targetDt :: Concrete (XParams2 Double Double)---        targetDt = replTarget goodDt ftk-        goodPerturbation :: forall r. GoodScalar r => r-        goodPerturbation = ifDifferentiable @r (realToFrac perturbation) 0-        targetPerturbed :: Concrete (XParams2 Double Double)-        targetPerturbed = treplTarget goodPerturbation ftk-        targetInitPerturbed :: Concrete (XParams2 Double Double)-        targetInitPerturbed = taddTarget stk targetInit targetPerturbed-        (derivative3, value3) = cfwdBoth f targetInit targetPerturbed-        value4 :: Concrete (TKScalar Double)-        value4 = MnistFcnnRanked2.afcnnMnistLoss2-                   (rfromPrimal glyph, rfromPrimal label)-                   (fromTarget targetInitPerturbed)-    in-      conjoin-        [ counterexample-            ("Objective function value from grad and jvp matches: "-             ++ show (value1, value2, value1 - value2))-            (abs (value1 - value2) < 1e-10)-        , counterexample-            ("Gradient and derivative agrees: "-             ++ show ( dt, derivative2, tdot0Target ftk gradient1 ds-                     , tdot0Target FTKScalar (kconcrete dt) derivative2-                       - tdot0Target ftk gradient1 ds ))-            (abs (tdot0Target FTKScalar (kconcrete dt) derivative2-                  - tdot0Target ftk gradient1 ds) < 1e-10)---        , counterexample  -- this is implied by the other clauses---            "Gradient is a linear function"---            (gradient1 === tmultTarget stk targetDt gradient0)-        , counterexample-            "Objective function value unaffected by incoming cotangent"-            (value0 === value1)-        , counterexample-            "Objective function value unaffected by derivative perturbation"-            (value2 === value3)-        , counterexample-            ("Derivative approximates the perturbation of value: "-             ++ show ( value2, derivative3, value4-                     , (value3 + derivative3) - value4) )-            (abs ((value3 + derivative3) - value4) < 1e-6)-        ]   ]
test/simplified/TestMnistPP.hs view
@@ -1,3 +1,4 @@+{-# OPTIONS_GHC -fno-late-specialise #-} {-# LANGUAGE OverloadedLists #-} -- | Tests of MNIST nns that pretty-print resulting gradient and primal terms. module TestMnistPP@@ -30,6 +31,8 @@ import MnistRnnRanked2 (ADRnnMnistParameters) import MnistRnnRanked2 qualified +import EqEpsilon+ testTrees :: [TestTree] testTrees = [ tensorMnistPPFCNNR             , tensorMnistPPRNNR@@ -43,7 +46,7 @@        Concrete widthHidden widthHidden2 r)  tensorMnistPPFCNNR :: TestTree-tensorMnistPPFCNNR = testGroup "PP and Ast tests for Short Ranked MNIST"+tensorMnistPPFCNNR = inOrderTestGroup "PP and Ast tests for Short Ranked MNIST"   [ testCase "VTO1 PP Lin" testVTOPP   , testCase "VTO1 Ast Lin" testVTOAst   , testCase "VTO1 PP NonLin" testVTOPPNonLin@@ -61,17 +64,17 @@               => MnistFcnnRanked1.ADFcnnMnist1Parameters Concrete 3 4 r valsInitVTOPP =   ( ( fromList (replicate 3 (Concrete-                             $ Nested.sfromListPrim+                             $ Nested.sfromList1Prim                                  (SNat @SizeMnistGlyph)                                  [1 .. fromIntegral sizeMnistGlyphInt]))-    , Concrete $ Nested.sfromListPrim (SNat @3) [1, 2, 3] )-  , ( fromList (replicate 4 (Concrete $ Nested.sfromListPrim+    , Concrete $ Nested.sfromList1Prim (SNat @3) [1, 2, 3] )+  , ( fromList (replicate 4 (Concrete $ Nested.sfromList1Prim                                           (SNat @3) [1, 2, 3]))-    , Concrete $ Nested.sfromListPrim (SNat @4) [1, 2, 3, 4] )+    , Concrete $ Nested.sfromList1Prim (SNat @4) [1, 2, 3, 4] )   , ( fromList (replicate sizeMnistLabelInt-                          (Concrete $ Nested.sfromListPrim+                          (Concrete $ Nested.sfromList1Prim                                         (SNat @4) [1, 2, 3, 4]))-    , Concrete $ Nested.sfromListPrim (SNat @SizeMnistLabel)+    , Concrete $ Nested.sfromList1Prim (SNat @SizeMnistLabel)                                       [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] ) )  testVTOPP :: Assertion@@ -88,20 +91,20 @@       ftk = tftk @Concrete (knownSTK @(XParams 3 4 Float))                            (toTarget @Concrete valsInitVTOPP)       artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2T ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\v1 -> rfromS (let v4 = sfromVector (fromList [sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject1 (tproject1 (tproject1 v1)))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v4, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v4]) + tproject2 (tproject2 (tproject1 v1)) in sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v5, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v5]) + tproject2 (tproject2 v1))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\v1 -> rfromS (let v4 = sfromVector0N @[3] (fromList [7.0 * ssum0 @[784] (tproject1 (tproject1 (tproject1 (tproject1 v1)))), 7.0 * ssum0 @[784] (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), 7.0 * ssum0 @[784] (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector0N @[4] (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v4, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v4]) + tproject2 (tproject2 (tproject1 v1)) in sfromVector0N @[10] (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v5, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v5]) + tproject2 (tproject2 v1))"   printArtifactPrimalPretty artifactRev-    @?= "\\v1 -> let v4 = sfromVector (fromList [ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject1 (tproject1 (tproject1 v1)))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector (fromList [ssum @3 (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v4), ssum @3 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v4), ssum @3 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v4), ssum @3 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v4)]) + tproject2 (tproject2 (tproject1 v1)) in rfromS (sfromVector (fromList [ssum @4 (tproject1 (tproject1 (tproject2 v1)) * v5), ssum @4 (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v5)]) + tproject2 (tproject2 v1))"+    @?= "\\v1 -> rfromS (let v4 = sfromVector0N @[3] (fromList [ssum0 @[784] (tproject1 (tproject1 (tproject1 (tproject1 v1))) * sfromPlain (sreplicate0N @[784] 7.0)), ssum0 @[784] (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1)))) * sfromPlain (sreplicate0N @[784] 7.0)), ssum0 @[784] (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))) * sfromPlain (sreplicate0N @[784] 7.0))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector0N @[4] (fromList [ssum0 @[3] (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v4), ssum0 @[3] (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v4), ssum0 @[3] (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v4), ssum0 @[3] (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v4)]) + tproject2 (tproject2 (tproject1 v1)) in sfromVector0N @[10] (fromList [ssum0 @[4] (tproject1 (tproject1 (tproject2 v1)) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v5), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v5)]) + tproject2 (tproject2 v1))"   printArtifactPretty artifactRev-    @?= "\\dret v1 -> let v4 = sfromVector (fromList [ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject1 (tproject1 (tproject1 v1)))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector (fromList [ssum @3 (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v4), ssum @3 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v4), ssum @3 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v4), ssum @3 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v4)]) + tproject2 (tproject2 (tproject1 v1)) ; v7 = sreplicate @4 (sfromR dret !$ [9]) ; v8 = sreplicate @4 (sfromR dret !$ [8]) ; v9 = sreplicate @4 (sfromR dret !$ [7]) ; v10 = sreplicate @4 (sfromR dret !$ [6]) ; v11 = sreplicate @4 (sfromR dret !$ [5]) ; v12 = sreplicate @4 (sfromR dret !$ [4]) ; v13 = sreplicate @4 (sfromR dret !$ [3]) ; v14 = sreplicate @4 (sfromR dret !$ [2]) ; v15 = sreplicate @4 (sfromR dret !$ [1]) ; v16 = sreplicate @4 (sfromR dret !$ [0]) ; v17 = tproject1 (tproject1 (tproject2 v1)) * v16 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v15 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v14 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v13 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v12 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v11 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v10 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v9 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v8 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v7)))))))) ; v18 = sreplicate @3 (v17 !$ [3]) ; v19 = sreplicate @3 (v17 !$ [2]) ; v20 = sreplicate @3 (v17 !$ [1]) ; v21 = sreplicate @3 (v17 !$ [0]) ; v22 = tproject1 (tproject1 (tproject2 (tproject1 v1))) * v21 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v20 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v19 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v18)) in tpair (tpair (tpair (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [0])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [1])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [2])) Z1))) v22) (tpair (tpair (v4 * v21) (tpair (v4 * v20) (tpair (v4 * v19) (tpair (v4 * v18) Z1)))) v17)) (tpair (tpair (v5 * v16) (tpair (v5 * v15) (tpair (v5 * v14) (tpair (v5 * v13) (tpair (v5 * v12) (tpair (v5 * v11) (tpair (v5 * v10) (tpair (v5 * v9) (tpair (v5 * v8) (tpair (v5 * v7) Z1)))))))))) (sfromR dret))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret v1 -> let v4 = sfromVector (fromList [sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject1 (tproject1 (tproject1 v1)))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v4, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v4]) + tproject2 (tproject2 (tproject1 v1)) ; x7 = sfromR dret !$ [9] ; x8 = sfromR dret !$ [8] ; x9 = sfromR dret !$ [7] ; x10 = sfromR dret !$ [6] ; x11 = sfromR dret !$ [5] ; x12 = sfromR dret !$ [4] ; x13 = sfromR dret !$ [3] ; x14 = sfromR dret !$ [2] ; x15 = sfromR dret !$ [1] ; x16 = sfromR dret !$ [0] ; v17 = tproject1 (tproject1 (tproject2 v1)) * sreplicate @4 x16 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * sreplicate @4 x15 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * sreplicate @4 x14 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * sreplicate @4 x13 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * sreplicate @4 x12 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * sreplicate @4 x11 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * sreplicate @4 x10 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * sreplicate @4 x9 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * sreplicate @4 x8 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * sreplicate @4 x7)))))))) ; x18 = v17 !$ [3] ; x19 = v17 !$ [2] ; x20 = v17 !$ [1] ; x21 = v17 !$ [0] ; v22 = tproject1 (tproject1 (tproject2 (tproject1 v1))) * sreplicate @3 x21 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * sreplicate @3 x20 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * sreplicate @3 x19 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * sreplicate @3 x18)) in tpair (tpair (tpair (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [0])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [1])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [2])) Z1))) v22) (tpair (tpair (v4 * sreplicate @3 x21) (tpair (v4 * sreplicate @3 x20) (tpair (v4 * sreplicate @3 x19) (tpair (v4 * sreplicate @3 x18) Z1)))) v17)) (tpair (tpair (v5 * sreplicate @4 x16) (tpair (v5 * sreplicate @4 x15) (tpair (v5 * sreplicate @4 x14) (tpair (v5 * sreplicate @4 x13) (tpair (v5 * sreplicate @4 x12) (tpair (v5 * sreplicate @4 x11) (tpair (v5 * sreplicate @4 x10) (tpair (v5 * sreplicate @4 x9) (tpair (v5 * sreplicate @4 x8) (tpair (v5 * sreplicate @4 x7) Z1)))))))))) (sfromR dret))"+    @?= "\\dret v1 -> let v4 = sfromVector0N @[3] (fromList [ssum0 @[784] (tproject1 (tproject1 (tproject1 (tproject1 v1))) * sfromPlain (sreplicate0N @[784] 7.0)), ssum0 @[784] (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1)))) * sfromPlain (sreplicate0N @[784] 7.0)), ssum0 @[784] (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))) * sfromPlain (sreplicate0N @[784] 7.0))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector0N @[4] (fromList [ssum0 @[3] (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v4), ssum0 @[3] (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v4), ssum0 @[3] (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v4), ssum0 @[3] (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v4)]) + tproject2 (tproject2 (tproject1 v1)) ; v7 = sreplicate0N @[4] (sfromR dret `sindex0` [9]) ; v8 = sreplicate0N @[4] (sfromR dret `sindex0` [8]) ; v9 = sreplicate0N @[4] (sfromR dret `sindex0` [7]) ; v10 = sreplicate0N @[4] (sfromR dret `sindex0` [6]) ; v11 = sreplicate0N @[4] (sfromR dret `sindex0` [5]) ; v12 = sreplicate0N @[4] (sfromR dret `sindex0` [4]) ; v13 = sreplicate0N @[4] (sfromR dret `sindex0` [3]) ; v14 = sreplicate0N @[4] (sfromR dret `sindex0` [2]) ; v15 = sreplicate0N @[4] (sfromR dret `sindex0` [1]) ; v16 = sreplicate0N @[4] (sfromR dret `sindex0` [0]) in tpair (let v17 = tproject1 (tproject1 (tproject2 v1)) * v16 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v15 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v14 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v13 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v12 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v11 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v10 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v9 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v8 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v7)))))))) ; v18 = sreplicate0N @[3] (v17 `sindex0` [3]) ; v19 = sreplicate0N @[3] (v17 `sindex0` [2]) ; v20 = sreplicate0N @[3] (v17 `sindex0` [1]) ; v21 = sreplicate0N @[3] (v17 `sindex0` [0]) in tpair (let v22 = tproject1 (tproject1 (tproject2 (tproject1 v1))) * v21 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v20 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v19 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v18)) in tpair (tpair (sfromPlain (sreplicate0N @[784] 7.0) * sreplicate0N @[784] (v22 `sindex0` [0])) (tpair (sfromPlain (sreplicate0N @[784] 7.0) * sreplicate0N @[784] (v22 `sindex0` [1])) (tpair (sfromPlain (sreplicate0N @[784] 7.0) * sreplicate0N @[784] (v22 `sindex0` [2])) Z1))) v22) (tpair (tpair (v4 * v21) (tpair (v4 * v20) (tpair (v4 * v19) (tpair (v4 * v18) Z1)))) v17)) (tpair (tpair (v5 * v16) (tpair (v5 * v15) (tpair (v5 * v14) (tpair (v5 * v13) (tpair (v5 * v12) (tpair (v5 * v11) (tpair (v5 * v10) (tpair (v5 * v9) (tpair (v5 * v8) (tpair (v5 * v7) Z1)))))))))) (sfromR dret))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret v1 -> let v4 = sfromVector0N @[3] (fromList [7.0 * ssum0 @[784] (tproject1 (tproject1 (tproject1 (tproject1 v1)))), 7.0 * ssum0 @[784] (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), 7.0 * ssum0 @[784] (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector0N @[4] (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v4, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v4]) + tproject2 (tproject2 (tproject1 v1)) ; x7 = sfromR dret `sindex0` [9] ; x8 = sfromR dret `sindex0` [8] ; x9 = sfromR dret `sindex0` [7] ; x10 = sfromR dret `sindex0` [6] ; x11 = sfromR dret `sindex0` [5] ; x12 = sfromR dret `sindex0` [4] ; x13 = sfromR dret `sindex0` [3] ; x14 = sfromR dret `sindex0` [2] ; x15 = sfromR dret `sindex0` [1] ; x16 = sfromR dret `sindex0` [0] in tpair (let v17 = tproject1 (tproject1 (tproject2 v1)) * sreplicate0N @[4] x16 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * sreplicate0N @[4] x15 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * sreplicate0N @[4] x14 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * sreplicate0N @[4] x13 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * sreplicate0N @[4] x12 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * sreplicate0N @[4] x11 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * sreplicate0N @[4] x10 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * sreplicate0N @[4] x9 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * sreplicate0N @[4] x8 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * sreplicate0N @[4] x7)))))))) ; x18 = v17 `sindex0` [3] ; x19 = v17 `sindex0` [2] ; x20 = v17 `sindex0` [1] ; x21 = v17 `sindex0` [0] in tpair (let v22 = tproject1 (tproject1 (tproject2 (tproject1 v1))) * sreplicate0N @[3] x21 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * sreplicate0N @[3] x20 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * sreplicate0N @[3] x19 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * sreplicate0N @[3] x18)) in tpair (tpair (sreplicate0N @[784] (7.0 * v22 `sindex0` [0])) (tpair (sreplicate0N @[784] (7.0 * v22 `sindex0` [1])) (tpair (sreplicate0N @[784] (7.0 * v22 `sindex0` [2])) Z1))) v22) (tpair (tpair (v4 * sreplicate0N @[3] x21) (tpair (v4 * sreplicate0N @[3] x20) (tpair (v4 * sreplicate0N @[3] x19) (tpair (v4 * sreplicate0N @[3] x18) Z1)))) v17)) (tpair (tpair (v5 * sreplicate0N @[4] x16) (tpair (v5 * sreplicate0N @[4] x15) (tpair (v5 * sreplicate0N @[4] x14) (tpair (v5 * sreplicate0N @[4] x13) (tpair (v5 * sreplicate0N @[4] x12) (tpair (v5 * sreplicate0N @[4] x11) (tpair (v5 * sreplicate0N @[4] x10) (tpair (v5 * sreplicate0N @[4] x9) (tpair (v5 * sreplicate0N @[4] x8) (tpair (v5 * sreplicate0N @[4] x7) Z1)))))))))) (sfromR dret))"  testVTOAst :: Assertion testVTOAst = do   let ftk = tftk @Concrete (knownSTK @(XParams 3 4 Float))                  (toTarget @Concrete valsInitVTOPP)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan (XParams 3 4 Float)       var = AstVar varName       vals = toTarget @Concrete valsInitVTOPP@@ -117,7 +120,7 @@   interpretAstFull @Concrete env afcnn1     @?= afcnn2 valsInitVTOPP   interpretAstFull @Concrete env-                   (simplifyInline @(TKR 1 Float) afcnn1)+                   (simplifyUserCode @(TKR 1 Float) afcnn1)     @?= afcnn2 valsInitVTOPP   interpretAstFull @Concrete env                    (simplifyInlineContract @(TKR 1 Float) afcnn1)@@ -138,20 +141,20 @@                            (toTarget @Concrete valsInitVTOPP)       artifactRevnonLin =         revArtifactAdapt UseIncomingCotangent afcnn2TnonLin ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRevnonLin)-    @?= "\\v1 -> rfromS (let v15 = scast (recip (sconcrete (sreplicate [3] 1.0) + exp (negate (sfromVector (fromList [sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject1 (tproject1 (tproject1 v1)))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))])) + negate (tproject2 (tproject1 (tproject1 v1)))))) ; v19 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (scast (sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v15, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v15, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v15, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v15]))) + negate (tproject2 (tproject2 (tproject1 v1))))) ; v22 = exp (sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v19, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v19]) + tproject2 (tproject2 v1)) in sreplicate @10 (recip (ssum0 v22)) * v22)"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRevnonLin)+    @?= "\\v1 -> rfromS (let v13 = recip (scast (sconcrete (sreplicate [3] 1.0) + exp (negate (sfromVector0N @[3] (fromList [7.0 * ssum0 @[784] (tproject1 (tproject1 (tproject1 (tproject1 v1)))), 7.0 * ssum0 @[784] (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), 7.0 * ssum0 @[784] (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))])) + negate (tproject2 (tproject1 (tproject1 v1)))))) ; v16 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (scast (sfromVector0N @[4] (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v13, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v13, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v13, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v13]))) + negate (tproject2 (tproject2 (tproject1 v1))))) ; v19 = exp (sfromVector0N @[10] (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v16, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v16]) + tproject2 (tproject2 v1)) in sreplicate0N @[10] (recip (ssum0 @[10] v19)) * v19)"   printArtifactPrimalPretty artifactRevnonLin-    @?= "\\v1 -> let v9 = sfromVector (fromList [ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject1 (tproject1 (tproject1 v1)))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v10 = exp (negate v9) ; v11 = sconcrete (sreplicate [3] 1.0) + v10 ; v12 = recip v11 ; v15 = scast v12 ; v16 = scast (sfromVector (fromList [ssum @3 (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v15), ssum @3 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v15), ssum @3 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v15), ssum @3 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v15)])) + tproject2 (tproject2 (tproject1 v1)) ; v17 = exp (negate v16) ; v18 = sconcrete (sreplicate [4] 1.0) + v17 ; v19 = recip v18 ; v22 = exp (sfromVector (fromList [ssum @4 (tproject1 (tproject1 (tproject2 v1)) * v19), ssum @4 (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v19)]) + tproject2 (tproject2 v1)) ; x23 = ssum @10 v22 ; v24 = sreplicate @10 (recip x23) in rfromS (v24 * v22)"+    @?= "\\v1 -> rfromS (let v9 = sfromVector0N @[3] (fromList [ssum0 @[784] (tproject1 (tproject1 (tproject1 (tproject1 v1))) * sfromPlain (sreplicate0N @[784] 7.0)), ssum0 @[784] (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1)))) * sfromPlain (sreplicate0N @[784] 7.0)), ssum0 @[784] (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))) * sfromPlain (sreplicate0N @[784] 7.0))]) + tproject2 (tproject1 (tproject1 v1)) ; v10 = exp (negate v9) ; v11 = recip (sconcrete (sreplicate [3] 1.0) + v10) ; v13 = scast (v11 + sconcrete (sreplicate [3] 0.0)) ; v14 = scast (sfromVector0N @[4] (fromList [ssum0 @[3] (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v13), ssum0 @[3] (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v13), ssum0 @[3] (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v13), ssum0 @[3] (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v13)])) + tproject2 (tproject2 (tproject1 v1)) ; v15 = exp (negate v14) ; v16 = recip (sconcrete (sreplicate [4] 1.0) + v15) ; v18 = v16 + sconcrete (sreplicate [4] 0.0) ; v19 = exp (sfromVector0N @[10] (fromList [ssum0 @[4] (tproject1 (tproject1 (tproject2 v1)) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v18)]) + tproject2 (tproject2 v1)) ; x20 = ssum0 @[10] v19 ; v21 = sreplicate0N @[10] (recip x20) in v21 * v19)"   printArtifactPretty artifactRevnonLin-    @?= "\\dret v1 -> let v9 = sfromVector (fromList [ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject1 (tproject1 (tproject1 v1)))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v10 = exp (negate v9) ; v11 = sconcrete (sreplicate [3] 1.0) + v10 ; v12 = recip v11 ; v13 = sconcrete (sreplicate [3] 1.0) + negate v12 ; v14 = v12 * v13 ; v15 = scast v12 ; v16 = scast (sfromVector (fromList [ssum @3 (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v15), ssum @3 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v15), ssum @3 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v15), ssum @3 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v15)])) + tproject2 (tproject2 (tproject1 v1)) ; v17 = exp (negate v16) ; v18 = sconcrete (sreplicate [4] 1.0) + v17 ; v19 = recip v18 ; v20 = sconcrete (sreplicate [4] 1.0) + negate v19 ; v21 = v19 * v20 ; v22 = exp (sfromVector (fromList [ssum @4 (tproject1 (tproject1 (tproject2 v1)) * v19), ssum @4 (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v19)]) + tproject2 (tproject2 v1)) ; x23 = ssum @10 v22 ; v24 = sreplicate @10 (recip x23) ; v26 = v22 * (sreplicate @10 (negate (recip (x23 * x23)) * ssum @10 (v22 * sfromR dret)) + v24 * sfromR dret) ; v27 = sreplicate @4 (v26 !$ [9]) ; v28 = sreplicate @4 (v26 !$ [8]) ; v29 = sreplicate @4 (v26 !$ [7]) ; v30 = sreplicate @4 (v26 !$ [6]) ; v31 = sreplicate @4 (v26 !$ [5]) ; v32 = sreplicate @4 (v26 !$ [4]) ; v33 = sreplicate @4 (v26 !$ [3]) ; v34 = sreplicate @4 (v26 !$ [2]) ; v35 = sreplicate @4 (v26 !$ [1]) ; v36 = sreplicate @4 (v26 !$ [0]) ; v37 = v21 * (tproject1 (tproject1 (tproject2 v1)) * v36 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v35 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v34 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v33 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v32 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v31 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v30 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v29 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v28 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v27))))))))) ; v38 = scast v37 ; v39 = sreplicate @3 (v38 !$ [3]) ; v40 = sreplicate @3 (v38 !$ [2]) ; v41 = sreplicate @3 (v38 !$ [1]) ; v42 = sreplicate @3 (v38 !$ [0]) ; v43 = v14 * scast (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v42 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v41 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v40 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v39))) in tpair (tpair (tpair (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [0])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [1])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [2])) Z1))) v43) (tpair (tpair (v15 * v42) (tpair (v15 * v41) (tpair (v15 * v40) (tpair (v15 * v39) Z1)))) v37)) (tpair (tpair (v19 * v36) (tpair (v19 * v35) (tpair (v19 * v34) (tpair (v19 * v33) (tpair (v19 * v32) (tpair (v19 * v31) (tpair (v19 * v30) (tpair (v19 * v29) (tpair (v19 * v28) (tpair (v19 * v27) Z1)))))))))) v26)"-  printArtifactPretty (simplifyArtifact artifactRevnonLin)-    @?= "\\dret v1 -> let v12 = recip (sconcrete (sreplicate [3] 1.0) + exp (negate (sfromVector (fromList [sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject1 (tproject1 (tproject1 v1)))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))])) + negate (tproject2 (tproject1 (tproject1 v1))))) ; v15 = scast v12 ; v19 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (scast (sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v15, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v15, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v15, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v15]))) + negate (tproject2 (tproject2 (tproject1 v1))))) ; v22 = exp (sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v19, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v19]) + tproject2 (tproject2 v1)) ; x23 = ssum0 v22 ; v26 = v22 * (sreplicate @10 (negate (recip (x23 * x23)) * sdot0 v22 (sfromR dret)) + sreplicate @10 (recip x23) * sfromR dret) ; x27 = v26 !$ [9] ; x28 = v26 !$ [8] ; x29 = v26 !$ [7] ; x30 = v26 !$ [6] ; x31 = v26 !$ [5] ; x32 = v26 !$ [4] ; x33 = v26 !$ [3] ; x34 = v26 !$ [2] ; x35 = v26 !$ [1] ; x36 = v26 !$ [0] ; v37 = (v19 * (sconcrete (sreplicate [4] 1.0) + negate v19)) * (tproject1 (tproject1 (tproject2 v1)) * sreplicate @4 x36 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * sreplicate @4 x35 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * sreplicate @4 x34 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * sreplicate @4 x33 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * sreplicate @4 x32 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * sreplicate @4 x31 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * sreplicate @4 x30 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * sreplicate @4 x29 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * sreplicate @4 x28 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * sreplicate @4 x27))))))))) ; v38 = scast v37 ; x39 = v38 !$ [3] ; x40 = v38 !$ [2] ; x41 = v38 !$ [1] ; x42 = v38 !$ [0] ; v43 = (v12 * (sconcrete (sreplicate [3] 1.0) + negate v12)) * scast (tproject1 (tproject1 (tproject2 (tproject1 v1))) * sreplicate @3 x42 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * sreplicate @3 x41 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * sreplicate @3 x40 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * sreplicate @3 x39))) in tpair (tpair (tpair (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [0])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [1])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [2])) Z1))) v43) (tpair (tpair (v15 * sreplicate @3 x42) (tpair (v15 * sreplicate @3 x41) (tpair (v15 * sreplicate @3 x40) (tpair (v15 * sreplicate @3 x39) Z1)))) v37)) (tpair (tpair (v19 * sreplicate @4 x36) (tpair (v19 * sreplicate @4 x35) (tpair (v19 * sreplicate @4 x34) (tpair (v19 * sreplicate @4 x33) (tpair (v19 * sreplicate @4 x32) (tpair (v19 * sreplicate @4 x31) (tpair (v19 * sreplicate @4 x30) (tpair (v19 * sreplicate @4 x29) (tpair (v19 * sreplicate @4 x28) (tpair (v19 * sreplicate @4 x27) Z1)))))))))) v26)"+    @?= "\\dret v1 -> let v9 = sfromVector0N @[3] (fromList [ssum0 @[784] (tproject1 (tproject1 (tproject1 (tproject1 v1))) * sfromPlain (sreplicate0N @[784] 7.0)), ssum0 @[784] (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1)))) * sfromPlain (sreplicate0N @[784] 7.0)), ssum0 @[784] (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))) * sfromPlain (sreplicate0N @[784] 7.0))]) + tproject2 (tproject1 (tproject1 v1)) ; v10 = exp (negate v9) ; v11 = recip (sconcrete (sreplicate [3] 1.0) + v10) ; v12 = v11 * (sconcrete (sreplicate [3] 1.0) + negate v11) ; v13 = scast (v11 + sconcrete (sreplicate [3] 0.0)) ; v14 = scast (sfromVector0N @[4] (fromList [ssum0 @[3] (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v13), ssum0 @[3] (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v13), ssum0 @[3] (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v13), ssum0 @[3] (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v13)])) + tproject2 (tproject2 (tproject1 v1)) ; v15 = exp (negate v14) ; v16 = recip (sconcrete (sreplicate [4] 1.0) + v15) ; v17 = v16 * (sconcrete (sreplicate [4] 1.0) + negate v16) ; v18 = v16 + sconcrete (sreplicate [4] 0.0) ; v19 = exp (sfromVector0N @[10] (fromList [ssum0 @[4] (tproject1 (tproject1 (tproject2 v1)) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v18), ssum0 @[4] (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v18)]) + tproject2 (tproject2 v1)) ; x20 = ssum0 @[10] v19 ; v21 = sreplicate0N @[10] (recip x20) ; v23 = v19 * (sreplicate0N @[10] (negate (recip (x20 * x20)) * ssum0 @[10] (v19 * sfromR dret)) + v21 * sfromR dret) ; v24 = sreplicate0N @[4] (v23 `sindex0` [9]) ; v25 = sreplicate0N @[4] (v23 `sindex0` [8]) ; v26 = sreplicate0N @[4] (v23 `sindex0` [7]) ; v27 = sreplicate0N @[4] (v23 `sindex0` [6]) ; v28 = sreplicate0N @[4] (v23 `sindex0` [5]) ; v29 = sreplicate0N @[4] (v23 `sindex0` [4]) ; v30 = sreplicate0N @[4] (v23 `sindex0` [3]) ; v31 = sreplicate0N @[4] (v23 `sindex0` [2]) ; v32 = sreplicate0N @[4] (v23 `sindex0` [1]) ; v33 = sreplicate0N @[4] (v23 `sindex0` [0]) in tpair (let v34 = v17 * (tproject1 (tproject1 (tproject2 v1)) * v33 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v32 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v31 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v30 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v29 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v28 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v27 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v26 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v25 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v24))))))))) ; v35 = scast v34 ; v36 = sreplicate0N @[3] (v35 `sindex0` [3]) ; v37 = sreplicate0N @[3] (v35 `sindex0` [2]) ; v38 = sreplicate0N @[3] (v35 `sindex0` [1]) ; v39 = sreplicate0N @[3] (v35 `sindex0` [0]) in tpair (let v40 = v12 * scast (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v39 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v38 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v37 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v36))) in tpair (tpair (sfromPlain (sreplicate0N @[784] 7.0) * sreplicate0N @[784] (v40 `sindex0` [0])) (tpair (sfromPlain (sreplicate0N @[784] 7.0) * sreplicate0N @[784] (v40 `sindex0` [1])) (tpair (sfromPlain (sreplicate0N @[784] 7.0) * sreplicate0N @[784] (v40 `sindex0` [2])) Z1))) v40) (tpair (tpair (v13 * v39) (tpair (v13 * v38) (tpair (v13 * v37) (tpair (v13 * v36) Z1)))) v34)) (tpair (tpair (v18 * v33) (tpair (v18 * v32) (tpair (v18 * v31) (tpair (v18 * v30) (tpair (v18 * v29) (tpair (v18 * v28) (tpair (v18 * v27) (tpair (v18 * v26) (tpair (v18 * v25) (tpair (v18 * v24) Z1)))))))))) v23)"+  printArtifactPretty (simplifyArtifactRev artifactRevnonLin)+    @?= "\\dret v1 -> let v11 = recip (sconcrete (sreplicate [3] 1.0) + exp (negate (sfromVector0N @[3] (fromList [7.0 * ssum0 @[784] (tproject1 (tproject1 (tproject1 (tproject1 v1)))), 7.0 * ssum0 @[784] (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), 7.0 * ssum0 @[784] (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))])) + negate (tproject2 (tproject1 (tproject1 v1))))) ; v13 = scast v11 ; v16 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (scast (sfromVector0N @[4] (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v13, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v13, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v13, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v13]))) + negate (tproject2 (tproject2 (tproject1 v1))))) ; v19 = exp (sfromVector0N @[10] (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v16, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v16, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v16]) + tproject2 (tproject2 v1)) ; x20 = ssum0 @[10] v19 ; v23 = v19 * (sreplicate0N @[10] (negate (recip (x20 * x20)) * sdot0 v19 (sfromR dret)) + sreplicate0N @[10] (recip x20) * sfromR dret) ; x24 = v23 `sindex0` [9] ; x25 = v23 `sindex0` [8] ; x26 = v23 `sindex0` [7] ; x27 = v23 `sindex0` [6] ; x28 = v23 `sindex0` [5] ; x29 = v23 `sindex0` [4] ; x30 = v23 `sindex0` [3] ; x31 = v23 `sindex0` [2] ; x32 = v23 `sindex0` [1] ; x33 = v23 `sindex0` [0] in tpair (let v34 = v16 * ((sconcrete (sreplicate [4] 1.0) + negate v16) * (tproject1 (tproject1 (tproject2 v1)) * sreplicate0N @[4] x33 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * sreplicate0N @[4] x32 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * sreplicate0N @[4] x31 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * sreplicate0N @[4] x30 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * sreplicate0N @[4] x29 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * sreplicate0N @[4] x28 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * sreplicate0N @[4] x27 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * sreplicate0N @[4] x26 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * sreplicate0N @[4] x25 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * sreplicate0N @[4] x24)))))))))) ; v35 = scast v34 ; x36 = v35 `sindex0` [3] ; x37 = v35 `sindex0` [2] ; x38 = v35 `sindex0` [1] ; x39 = v35 `sindex0` [0] in tpair (let v40 = v11 * ((sconcrete (sreplicate [3] 1.0) + negate v11) * scast (tproject1 (tproject1 (tproject2 (tproject1 v1))) * sreplicate0N @[3] x39 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * sreplicate0N @[3] x38 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * sreplicate0N @[3] x37 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * sreplicate0N @[3] x36)))) in tpair (tpair (sreplicate0N @[784] (7.0 * v40 `sindex0` [0])) (tpair (sreplicate0N @[784] (7.0 * v40 `sindex0` [1])) (tpair (sreplicate0N @[784] (7.0 * v40 `sindex0` [2])) Z1))) v40) (tpair (tpair (v13 * sreplicate0N @[3] x39) (tpair (v13 * sreplicate0N @[3] x38) (tpair (v13 * sreplicate0N @[3] x37) (tpair (v13 * sreplicate0N @[3] x36) Z1)))) v34)) (tpair (tpair (v16 * sreplicate0N @[4] x33) (tpair (v16 * sreplicate0N @[4] x32) (tpair (v16 * sreplicate0N @[4] x31) (tpair (v16 * sreplicate0N @[4] x30) (tpair (v16 * sreplicate0N @[4] x29) (tpair (v16 * sreplicate0N @[4] x28) (tpair (v16 * sreplicate0N @[4] x27) (tpair (v16 * sreplicate0N @[4] x26) (tpair (v16 * sreplicate0N @[4] x25) (tpair (v16 * sreplicate0N @[4] x24) Z1)))))))))) v23)"  testVTOAstNonLin :: Assertion testVTOAstNonLin = do   let ftk = tftk @Concrete (knownSTK @(XParams 3 4 Double))                  (toTarget @Concrete valsInitVTOPP)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan (XParams 3 4 Double)       var = AstVar varName       vals = toTarget @Concrete valsInitVTOPP@@ -167,7 +170,7 @@   interpretAstFull @Concrete env afcnn1     @?= afcnn2 valsInitVTOPP   interpretAstFull @Concrete env-                   (simplifyInline @(TKR 1 Double) afcnn1)+                   (simplifyUserCode @(TKR 1 Double) afcnn1)     @?= afcnn2 valsInitVTOPP   interpretAstFull @Concrete env                    (simplifyInlineContract @(TKR 1 Double) afcnn1)@@ -198,19 +201,19 @@                            (toTarget @Concrete valsInitVT2OPP)       artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2T ftk   printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> let m5 = str (sreplicate @5 (scast (ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m6 = str (sreplicate @2 (scast (ssum @4 (m5 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) in rfromS (ssum @5 (m6 * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1)))"+    @?= "\\m1 -> rfromS (let v5 = ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) ; m6 = str (sreplicate @5 (scast (sconcrete (sreplicate [4] 7.0) * v5 + sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m7 = str (sreplicate @2 (scast (ssum @4 (m6 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) in ssum @5 (m7 * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1)))"   printArtifactPretty artifactRev-    @?= "\\dret m1 -> let m5 = str (sreplicate @5 (scast (ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m6 = str (sreplicate @2 (scast (ssum @4 (m5 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) ; v8 = ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 (sfromR dret))) ; m9 = sreplicate @4 (scast v8) ; v10 = scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m9))) in tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [3,4] 7.0) * sreplicate @3 v10))) (rfromS v10)) (tpair (rfromS (str (m5 * m9))) (rfromS v8))) (tpair (rfromS (str (m6 * sreplicate @5 (sfromR dret)))) dret)"+    @?= "\\dret m1 -> let v5 = ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) ; m6 = str (sreplicate @5 (scast (sconcrete (sreplicate [4] 7.0) * v5 + sfromR (tproject2 (tproject1 (tproject1 m1)))))) in tpair (let v9 = ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 (sfromR dret))) ; m10 = sreplicate @4 (scast v9) in tpair (let v11 = scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m10))) in tpair (rfromS (str (sreplicate @3 (sconcrete (sreplicate [4] 7.0) * v11)))) (rfromS v11)) (tpair (rfromS (str (m6 * m10))) (rfromS v9))) (tpair (rfromS (str (let m7 = str (sreplicate @2 (scast (ssum @4 (m6 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) in m7 * sreplicate @5 (sfromR dret)))) (rfromS (sfromR dret)))"   printArtifactPretty artifactRev-    @?= "\\dret m1 -> let m5 = str (sreplicate @5 (scast (ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m6 = str (sreplicate @2 (scast (ssum @4 (m5 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) ; v8 = ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 (sfromR dret))) ; m9 = sreplicate @4 (scast v8) ; v10 = scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m9))) in tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [3,4] 7.0) * sreplicate @3 v10))) (rfromS v10)) (tpair (rfromS (str (m5 * m9))) (rfromS v8))) (tpair (rfromS (str (m6 * sreplicate @5 (sfromR dret)))) dret)"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4] FTKScalar)) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,5] FTKScalar)) ConvSX)) ConvId)) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKR (SNat @1) STKScalar))) (let v5 = scast (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1)))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; v8 = sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 (sfromR dret)) ; v9 = scast v8 ; v10 = scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v9)) in tpair (tpair (tpair (sconcrete (sreplicate [4,3] 7.0) * str (sreplicate @3 v10)) v10) (tpair (sreplicate @5 v5 * str (sreplicate @4 v9)) v8)) (tpair (sreplicate @2 (scast (sdot1In (sreplicate @5 v5) (sfromR (tproject1 (tproject2 (tproject1 m1))))) + sfromR (tproject2 (tproject2 (tproject1 m1)))) * str (sreplicate @5 (sfromR dret))) dret))"+    @?= "\\dret m1 -> let v5 = ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) ; m6 = str (sreplicate @5 (scast (sconcrete (sreplicate [4] 7.0) * v5 + sfromR (tproject2 (tproject1 (tproject1 m1)))))) in tpair (let v9 = ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 (sfromR dret))) ; m10 = sreplicate @4 (scast v9) in tpair (let v11 = scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m10))) in tpair (rfromS (str (sreplicate @3 (sconcrete (sreplicate [4] 7.0) * v11)))) (rfromS v11)) (tpair (rfromS (str (m6 * m10))) (rfromS v9))) (tpair (rfromS (str (let m7 = str (sreplicate @2 (scast (ssum @4 (m6 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) in m7 * sreplicate @5 (sfromR dret)))) (rfromS (sfromR dret)))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKR (SNat @1) STKScalar))) (let v6 = scast (sconcrete (sreplicate [4] 7.0) * ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) in tpair (let v9 = sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 (sfromR dret)) ; v10 = scast v9 in tpair (let v11 = scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v10)) in tpair (str (sreplicate @3 (sconcrete (sreplicate [4] 7.0) * v11))) v11) (tpair (sreplicate @5 v6 * str (sreplicate @4 v10)) v9)) (tpair (sreplicate @2 (scast (sdot1In (sreplicate @5 v6) (sfromR (tproject1 (tproject2 (tproject1 m1))))) + sfromR (tproject2 (tproject2 (tproject1 m1)))) * str (sreplicate @5 (sfromR dret))) dret))"  testVT2OAst :: Assertion testVT2OAst = do   let ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))                  (toTarget @Concrete valsInitVT2OPP)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan (XParams2 Double Float)       var = AstVar varName       vals = toTarget @Concrete valsInitVT2OPP@@ -226,7 +229,7 @@   interpretAstFull @Concrete env afcnn1     @?= afcnn2 valsInitVT2OPP   interpretAstFull @Concrete env-                   (simplifyInline @(TKR 1 Double) afcnn1)+                   (simplifyUserCode @(TKR 1 Double) afcnn1)     @?= afcnn2 valsInitVT2OPP   interpretAstFull @Concrete env                    (simplifyInlineContract @(TKR 1 Double) afcnn1)@@ -249,12 +252,13 @@              , fromPrimal $ rcast $ rconcrete $ unConcrete a4 )            , ( rcast $ fromPrimal $ rconcrete $ unConcrete a5              , fromPrimal $ rcast $ rconcrete $ unConcrete a6 ) )-      ast3 = fun1ToAst (FTKR (0 :$: ZSR) (FTKScalar @Float))-                       (const $ afcnn2TnonLin constant)+      ast3 = snd $ funToAst @FullSpan+                            (FTKR (0 :$: ZSR) (FTKScalar @Float))+                            (const $ afcnn2TnonLin constant)   "\\dummy" ++ " -> " ++ printAstSimple ast3-    @?= "\\dummy -> rfromS (tlet (exp (ssum @5 (str (sreplicate @2 (tlet (ssum @4 (str (sreplicate @5 (tlet (tfromPrimal (STKS [4] STKScalar) (ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (scast (sconcrete (sfromListLinear [4,3] [1.0,2.0,3.0,1.0,2.0,3.0,1.0,2.0,3.0,1.0,2.0,3.0])))) + scast (sconcrete (sfromListLinear [4] [1.0,2.0,3.0,4.0])))) (\\v5 -> ttletPrimal (recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sfromR (tprimalPart (rfromS v5)))))) (\\v6 -> tfromPrimal (STKS [4] STKScalar) v6 + sfromR (tfromDual (tdualPart (STKR (SNat @1) STKScalar) (rfromS (tfromPrimal (STKS [4] STKScalar) (v6 * (sconcrete (sreplicate [4] 1.0) + negate v6)) * sfromR (tfromDual (tdualPart (STKR (SNat @1) STKScalar) (rfromS v5))))))))))) * tfromPrimal (STKS [4,5] STKScalar) (sconcrete (sfromListLinear [4,5] [1.0,1.0,1.0,1.0,1.0,2.0,2.0,2.0,2.0,2.0,3.0,3.0,3.0,3.0,3.0,4.0,4.0,4.0,4.0,4.0]))) + tfromPrimal (STKS [5] STKScalar) (scast (sconcrete (sfromListLinear [5] [1.0,2.0,3.0,4.0,5.0])))) (\\v7 -> ttletPrimal (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (sfromR (tprimalPart (rfromS v7)))))) (\\v8 -> tfromPrimal (STKS [5] STKScalar) v8 + sfromR (tfromDual (tdualPart (STKR (SNat @1) STKScalar) (rfromS (tfromPrimal (STKS [5] STKScalar) (v8 * (sconcrete (sreplicate [5] 1.0) + negate v8)) * sfromR (tfromDual (tdualPart (STKR (SNat @1) STKScalar) (rfromS v7))))))))))) * tfromPrimal (STKS [5,2] STKScalar) (str (scast (sconcrete (sfromListLinear [2,5] [1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0]))))) + tfromPrimal (STKS [2] STKScalar) (scast (sconcrete (sfromListLinear [2] [1.0,2.0]))))) (\\v9 -> sreplicate @2 (recip (ssum @2 v9)) * v9))"+    @?= "\\dummy -> tfromPlain (STKR (SNat @1) STKScalar) (rfromS (tlet (exp (ssum @5 (str (sreplicate @2 (tlet (ssum @4 (sconcrete (sfromListLinear [4,5] [1.0,1.0,1.0,1.0,1.0,2.0,2.0,2.0,2.0,2.0,3.0,3.0,3.0,3.0,3.0,4.0,4.0,4.0,4.0,4.0]) * str (sreplicate @5 (tlet (sconcrete (sreplicate [4] 7.0) * ssum @3 (str (scast (sconcrete (sfromListLinear [4,3] [1.0,2.0,3.0,1.0,2.0,3.0,1.0,2.0,3.0,1.0,2.0,3.0])))) + scast (sconcrete (sfromListLinear [4] [1.0,2.0,3.0,4.0]))) (\\v7 -> recip (sconcrete (sreplicate [4] 1.0) + exp (negate v7)))))) + scast (sconcrete (sfromListLinear [5] [1.0,2.0,3.0,4.0,5.0]))) (\\v10 -> recip (sconcrete (sreplicate [5] 1.0) + exp (negate v10))))) * str (scast (sconcrete (sfromListLinear [2,5] [1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0])))) + scast (sconcrete (sfromListLinear [2] [1.0,2.0])))) (\\v12 -> sreplicate0N @[2] (recip (ssum0 @[2] v12)) * v12)))"   "\\dummy" ++ " -> " ++ printAstSimple (simplifyInlineContract ast3)-    @?= "\\dummy -> rfromS (tlet (exp (sdot1In (sreplicate @2 (tlet (sdot1In (sreplicate @5 (ttletPrimal (recip (sconcrete (sreplicate [4] 1.0) + exp (sconcrete (sfromListLinear [4] [-43.0,-44.0,-45.0,-46.0])))) (\\v6 -> tfromPrimal (STKS [4] STKScalar) v6 + tfromDual (tdualPart (STKS [4] STKScalar) (tfromPrimal (STKS [4] STKScalar) (v6 * (sconcrete (sreplicate [4] 1.0) + negate v6)) * tfromDual (tdualPart (STKS [4] STKScalar) (tfromPrimal (STKS [4] STKScalar) (sconcrete (sreplicate [4] 0.0))))))))) (tfromPrimal (STKS [5,4] STKScalar) (sconcrete (sfromListLinear [5,4] [1.0,2.0,3.0,4.0,1.0,2.0,3.0,4.0,1.0,2.0,3.0,4.0,1.0,2.0,3.0,4.0,1.0,2.0,3.0,4.0]))) + tfromPrimal (STKS [5] STKScalar) (sconcrete (sfromListLinear [5] [1.0,2.0,3.0,4.0,5.0]))) (\\v7 -> ttletPrimal (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (tprimalPart v7)))) (\\v8 -> tfromPrimal (STKS [5] STKScalar) v8 + tfromDual (tdualPart (STKS [5] STKScalar) (tfromPrimal (STKS [5] STKScalar) (v8 * (sconcrete (sreplicate [5] 1.0) + negate v8)) * tfromDual (tdualPart (STKS [5] STKScalar) v7))))))) (tfromPrimal (STKS [2,5] STKScalar) (sconcrete (sfromListLinear [2,5] [1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0]))) + tfromPrimal (STKS [2] STKScalar) (sconcrete (sfromListLinear [2] [1.0,2.0])))) (\\v9 -> sreplicate @2 (recip (ssum0 v9)) * v9))"+    @?= "\\dummy -> tfromPlain (STKR (SNat @1) STKScalar) (rfromS (sconcrete (sfromListLinear [2] [0.26894158,0.73105836])))"  testVT2OPPNonLin2 :: Assertion testVT2OPPNonLin2 = do@@ -269,20 +273,20 @@                            (toTarget @Concrete valsInitVT2OPP)       artifactRevnonLin =         revArtifactAdapt UseIncomingCotangent afcnn2TnonLin ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRevnonLin)-    @?= "\\m1 -> rfromS (let v23 = exp (sdot1In (sreplicate @2 (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 (scast (recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1))))))))) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))))) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) in sreplicate @2 (recip (ssum0 v23)) * v23)"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRevnonLin)+    @?= "\\m1 -> rfromS (let v21 = exp (sdot1In (sreplicate @2 (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 (recip (scast (sconcrete (sreplicate [4] 1.0) + exp (sconcrete (sreplicate [4] (-7.0)) * ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1))))))))) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))))) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) in sreplicate0N @[2] (recip (ssum0 @[2] v21)) * v21)"   printArtifactPrimalPretty artifactRevnonLin-    @?= "\\m1 -> let v10 = ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v11 = exp (negate v10) ; v12 = sconcrete (sreplicate [4] 1.0) + v11 ; v13 = recip v12 ; m16 = str (sreplicate @5 (scast v13)) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = sconcrete (sreplicate [5] 1.0) + v18 ; v20 = recip v19 ; v23 = exp (ssum @5 (str (sreplicate @2 v20) * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum @2 v23 ; v25 = sreplicate @2 (recip x24) in rfromS (v25 * v23)"+    @?= "\\m1 -> rfromS (let v10 = ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) ; v11 = sconcrete (sreplicate [4] 7.0) * v10 + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v12 = exp (negate v11) ; v13 = recip (sconcrete (sreplicate [4] 1.0) + v12) ; m15 = str (sreplicate @5 (scast (v13 + sconcrete (sreplicate [4] 0.0)))) ; v16 = scast (ssum @4 (m15 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v17 = exp (negate v16) ; v18 = recip (sconcrete (sreplicate [5] 1.0) + v17) ; m20 = str (sreplicate @2 (v18 + sconcrete (sreplicate [5] 0.0))) ; v21 = exp (ssum @5 (m20 * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x22 = ssum0 @[2] v21 ; v23 = sreplicate0N @[2] (recip x22) in v23 * v21)"   printArtifactPretty artifactRevnonLin-    @?= "\\dret m1 -> let v10 = ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v11 = exp (negate v10) ; v12 = sconcrete (sreplicate [4] 1.0) + v11 ; v13 = recip v12 ; v14 = sconcrete (sreplicate [4] 1.0) + negate v13 ; v15 = v13 * v14 ; m16 = str (sreplicate @5 (scast v13)) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = sconcrete (sreplicate [5] 1.0) + v18 ; v20 = recip v19 ; v21 = sconcrete (sreplicate [5] 1.0) + negate v20 ; v22 = v20 * v21 ; v23 = exp (ssum @5 (str (sreplicate @2 v20) * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum @2 v23 ; v25 = sreplicate @2 (recip x24) ; v27 = v23 * (sreplicate @2 (negate (recip (x24 * x24)) * ssum @2 (v23 * sfromR dret)) + v25 * sfromR dret) ; v28 = v22 * ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 v27)) ; m29 = sreplicate @4 (scast v28) ; v30 = v15 * scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m29))) in tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [3,4] 7.0) * sreplicate @3 v30))) (rfromS v30)) (tpair (rfromS (str (m16 * m29))) (rfromS v28))) (tpair (rfromS (str (str (sreplicate @2 v20) * sreplicate @5 v27))) (rfromS v27))"-  printArtifactPretty (simplifyArtifact artifactRevnonLin)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4] FTKScalar)) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2] FTKScalar)) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKS [2] STKScalar))) (let v13 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v16 = scast v13 ; v20 = recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 v16) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; v23 = exp (sdot1In (sreplicate @2 v20) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum0 v23 ; v27 = v23 * (sreplicate @2 (negate (recip (x24 * x24)) * sdot0 v23 (sfromR dret)) + sreplicate @2 (recip x24) * sfromR dret) ; v28 = (v20 * (sconcrete (sreplicate [5] 1.0) + negate v20)) * sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 v27) ; v29 = scast v28 ; v30 = (v13 * (sconcrete (sreplicate [4] 1.0) + negate v13)) * scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v29)) in tpair (tpair (tpair (sconcrete (sreplicate [4,3] 7.0) * str (sreplicate @3 v30)) v30) (tpair (sreplicate @5 v16 * str (sreplicate @4 v29)) v28)) (tpair (sreplicate @2 v20 * str (sreplicate @5 v27)) v27))"+    @?= "\\dret m1 -> let v10 = ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) ; v11 = sconcrete (sreplicate [4] 7.0) * v10 + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v12 = exp (negate v11) ; v13 = recip (sconcrete (sreplicate [4] 1.0) + v12) ; v14 = v13 * (sconcrete (sreplicate [4] 1.0) + negate v13) ; m15 = str (sreplicate @5 (scast (v13 + sconcrete (sreplicate [4] 0.0)))) ; v16 = scast (ssum @4 (m15 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v17 = exp (negate v16) ; v18 = recip (sconcrete (sreplicate [5] 1.0) + v17) ; v19 = v18 * (sconcrete (sreplicate [5] 1.0) + negate v18) ; m20 = str (sreplicate @2 (v18 + sconcrete (sreplicate [5] 0.0))) ; v21 = exp (ssum @5 (m20 * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x22 = ssum0 @[2] v21 ; v23 = sreplicate0N @[2] (recip x22) ; v25 = v21 * (sreplicate0N @[2] (negate (recip (x22 * x22)) * ssum0 @[2] (v21 * sfromR dret)) + v23 * sfromR dret) in tpair (let v26 = v19 * ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 v25)) ; m27 = sreplicate @4 (scast v26) in tpair (let v28 = v14 * scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m27))) in tpair (rfromS (str (sreplicate @3 (sconcrete (sreplicate [4] 7.0) * v28)))) (rfromS v28)) (tpair (rfromS (str (m15 * m27))) (rfromS v26))) (tpair (rfromS (str (m20 * sreplicate @5 v25))) (rfromS v25))"+  printArtifactPretty (simplifyArtifactRev artifactRevnonLin)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKS [2] STKScalar))) (let v13 = recip (sconcrete (sreplicate [4] 1.0) + exp (sconcrete (sreplicate [4] (-7.0)) * ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v15 = scast v13 ; v18 = recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 v15) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; v21 = exp (sdot1In (sreplicate @2 v18) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) ; x22 = ssum0 @[2] v21 ; v25 = v21 * (sreplicate0N @[2] (negate (recip (x22 * x22)) * sdot0 v21 (sfromR dret)) + sreplicate0N @[2] (recip x22) * sfromR dret) in tpair (let v26 = v18 * ((sconcrete (sreplicate [5] 1.0) + negate v18) * sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 v25)) ; v27 = scast v26 in tpair (let v28 = v13 * ((sconcrete (sreplicate [4] 1.0) + negate v13) * scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v27))) in tpair (str (sreplicate @3 (sconcrete (sreplicate [4] 7.0) * v28))) v28) (tpair (sreplicate @5 v15 * str (sreplicate @4 v27)) v26)) (tpair (sreplicate @2 v18 * str (sreplicate @5 v25)) v25))"  testVT2OAstNonLin2 :: Assertion testVT2OAstNonLin2 = do   let ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))                  (toTarget @Concrete valsInitVT2OPP)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan (XParams2 Double Float)       var = AstVar varName       vals = toTarget @Concrete valsInitVT2OPP@@ -298,7 +302,7 @@   interpretAstFull @Concrete env afcnn1     @?= afcnn2 valsInitVT2OPP   interpretAstFull @Concrete env-                   (simplifyInline @(TKR 1 Double) afcnn1)+                   (simplifyUserCode @(TKR 1 Double) afcnn1)     @?= afcnn2 valsInitVT2OPP   interpretAstFull @Concrete env                    (simplifyInlineContract @(TKR 1 Double) afcnn1)@@ -319,20 +323,20 @@                            (toTarget @Concrete valsInitVT2OPP)       artifactRevnonLin =         revArtifactAdapt UseIncomingCotangent afcnn2TnonLin ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRevnonLin)-    @?= "\\m1 -> let v23 = exp (sdot1In (sreplicate @2 (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 (scast (recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1))))))))) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))))) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) in kfromS (negate (sdot0 (sconcrete (sreplicate [2] 8.0)) (log (sreplicate @2 (recip (ssum0 v23)) * v23))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRevnonLin)+    @?= "\\m1 -> let v22 = exp (sdot1In (sreplicate @2 (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 (recip (scast (sconcrete (sreplicate [4] 1.0) + exp (sconcrete (sreplicate [4] (-7.0)) * ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1))))))))) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))))) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) in (-8.0) * ssum0 @[2] (log (sreplicate0N @[2] (recip (ssum0 @[2] v22)) * v22))"   printArtifactPrimalPretty artifactRevnonLin-    @?= "\\m1 -> let v10 = ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v11 = exp (negate v10) ; v12 = sconcrete (sreplicate [4] 1.0) + v11 ; v13 = recip v12 ; m16 = str (sreplicate @5 (scast v13)) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = sconcrete (sreplicate [5] 1.0) + v18 ; v20 = recip v19 ; v23 = exp (ssum @5 (str (sreplicate @2 v20) * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum @2 v23 ; v25 = sreplicate @2 (recip x24) ; v26 = v25 * v23 ; v27 = log v26 in kfromS (negate (ssum @2 (sconcrete (sreplicate [2] 8.0) * v27)))"+    @?= "\\m1 -> let v11 = ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) ; v12 = sconcrete (sreplicate [4] 7.0) * v11 + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v13 = exp (negate v12) ; v14 = recip (sconcrete (sreplicate [4] 1.0) + v13) ; m16 = str (sreplicate @5 (scast (v14 + sconcrete (sreplicate [4] 0.0)))) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = recip (sconcrete (sreplicate [5] 1.0) + v18) ; m21 = str (sreplicate @2 (v19 + sconcrete (sreplicate [5] 0.0))) ; v22 = exp (ssum @5 (m21 * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x23 = ssum0 @[2] v22 ; v24 = sreplicate0N @[2] (recip x23) ; v25 = v24 * v22 ; x26 = ssum0 @[2] (log v25) in (-8.0) * x26"   printArtifactPretty artifactRevnonLin-    @?= "\\dret m1 -> let v10 = ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v11 = exp (negate v10) ; v12 = sconcrete (sreplicate [4] 1.0) + v11 ; v13 = recip v12 ; v14 = sconcrete (sreplicate [4] 1.0) + negate v13 ; v15 = v13 * v14 ; m16 = str (sreplicate @5 (scast v13)) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = sconcrete (sreplicate [5] 1.0) + v18 ; v20 = recip v19 ; v21 = sconcrete (sreplicate [5] 1.0) + negate v20 ; v22 = v20 * v21 ; v23 = exp (ssum @5 (str (sreplicate @2 v20) * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum @2 v23 ; v25 = sreplicate @2 (recip x24) ; v26 = v25 * v23 ; v29 = sconcrete (sreplicate [2] 8.0) * (recip v26 * sreplicate @2 (sscalar (-1.0) * sfromK dret)) ; v30 = v23 * (sreplicate @2 (negate (recip (x24 * x24)) * ssum @2 (v23 * v29)) + v25 * v29) ; v31 = v22 * ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 v30)) ; m32 = sreplicate @4 (scast v31) ; v33 = v15 * scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m32))) in tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [3,4] 7.0) * sreplicate @3 v33))) (rfromS v33)) (tpair (rfromS (str (m16 * m32))) (rfromS v31))) (tpair (rfromS (str (str (sreplicate @2 v20) * sreplicate @5 v30))) (rfromS v30))"-  printArtifactPretty (simplifyArtifact artifactRevnonLin)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4] FTKScalar)) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2] FTKScalar)) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKS [2] STKScalar))) (let v13 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v16 = scast v13 ; v20 = recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 v16) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; v23 = exp (sdot1In (sreplicate @2 v20) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum0 v23 ; x25 = recip x24 ; v29 = sconcrete (sreplicate [2] 8.0) * (recip (sreplicate @2 x25 * v23) * sreplicate @2 (sscalar (-1.0) * sfromK dret)) ; v30 = v23 * (sreplicate @2 (negate (recip (x24 * x24)) * sdot0 v23 v29) + sreplicate @2 x25 * v29) ; v31 = (v20 * (sconcrete (sreplicate [5] 1.0) + negate v20)) * sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 v30) ; v32 = scast v31 ; v33 = (v13 * (sconcrete (sreplicate [4] 1.0) + negate v13)) * scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v32)) in tpair (tpair (tpair (sconcrete (sreplicate [4,3] 7.0) * str (sreplicate @3 v33)) v33) (tpair (sreplicate @5 v16 * str (sreplicate @4 v32)) v31)) (tpair (sreplicate @2 v20 * str (sreplicate @5 v30)) v30))"+    @?= "\\dret m1 -> let v11 = ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) ; v12 = sconcrete (sreplicate [4] 7.0) * v11 + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v13 = exp (negate v12) ; v14 = recip (sconcrete (sreplicate [4] 1.0) + v13) ; v15 = v14 * (sconcrete (sreplicate [4] 1.0) + negate v14) ; m16 = str (sreplicate @5 (scast (v14 + sconcrete (sreplicate [4] 0.0)))) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = recip (sconcrete (sreplicate [5] 1.0) + v18) ; v20 = v19 * (sconcrete (sreplicate [5] 1.0) + negate v19) ; m21 = str (sreplicate @2 (v19 + sconcrete (sreplicate [5] 0.0))) ; v22 = exp (ssum @5 (m21 * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x23 = ssum0 @[2] v22 ; v24 = sreplicate0N @[2] (recip x23) ; v25 = v24 * v22 ; v28 = recip v25 * sreplicate0N @[2] ((-8.0) * dret) ; v29 = v22 * (sreplicate0N @[2] (negate (recip (x23 * x23)) * ssum0 @[2] (v22 * v28)) + v24 * v28) in tpair (let v30 = v20 * ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 v29)) ; m31 = sreplicate @4 (scast v30) in tpair (let v32 = v15 * scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m31))) in tpair (rfromS (str (sreplicate @3 (sconcrete (sreplicate [4] 7.0) * v32)))) (rfromS v32)) (tpair (rfromS (str (m16 * m31))) (rfromS v30))) (tpair (rfromS (str (m21 * sreplicate @5 v29))) (rfromS v29))"+  printArtifactPretty (simplifyArtifactRev artifactRevnonLin)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKS [2] STKScalar))) (let v14 = recip (sconcrete (sreplicate [4] 1.0) + exp (sconcrete (sreplicate [4] (-7.0)) * ssum @3 (str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v16 = scast v14 ; v19 = recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 v16) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; v22 = exp (sdot1In (sreplicate @2 v19) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) ; x23 = ssum0 @[2] v22 ; x24 = recip x23 ; v28 = sreplicate0N @[2] ((-8.0) * dret) / (sreplicate0N @[2] x24 * v22) ; v29 = v22 * (sreplicate0N @[2] (negate (recip (x23 * x23)) * sdot0 v22 v28) + sreplicate0N @[2] x24 * v28) in tpair (let v30 = v19 * ((sconcrete (sreplicate [5] 1.0) + negate v19) * sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 v29)) ; v31 = scast v30 in tpair (let v32 = v14 * ((sconcrete (sreplicate [4] 1.0) + negate v14) * scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v31))) in tpair (str (sreplicate @3 (sconcrete (sreplicate [4] 7.0) * v32))) v32) (tpair (sreplicate @5 v16 * str (sreplicate @4 v31)) v30)) (tpair (sreplicate @2 v19 * str (sreplicate @5 v29)) v29))"  testVT2OAstNonLin3 :: Assertion testVT2OAstNonLin3 = do   let ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))                  (toTarget @Concrete valsInitVT2OPP)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan (XParams2 Double Float)       var = AstVar varName       vals = toTarget @Concrete valsInitVT2OPP@@ -349,7 +353,7 @@   interpretAstFull @Concrete env afcnn1     @?= afcnn2 valsInitVT2OPP   interpretAstFull @Concrete env-                   (simplifyInline @(TKScalar Double) afcnn1)+                   (simplifyUserCode @(TKScalar Double) afcnn1)     @?= afcnn2 valsInitVT2OPP   interpretAstFull @Concrete env                    (simplifyInlineContract @(TKScalar Double) afcnn1)@@ -359,7 +363,7 @@ -- * RNNR tests  tensorMnistPPRNNR :: TestTree-tensorMnistPPRNNR = testGroup "PP and Ast tests for RNNR MNIST"+tensorMnistPPRNNR = inOrderTestGroup "PP and Ast tests for RNNR MNIST"   [ testCase "RNNO PP" testRNNOPP   , testCase "RNNO Ast" testRNNOAst   , testCase "RNNO PP 2" testRNNOPP2@@ -400,9 +404,7 @@   let batch_size = 1       sizeMnistHeightI = 1       blackGlyph :: AstTensor AstMethodLet PrimalSpan (TKR 3 Double)-      blackGlyph = AstReplicate (SNat @1) knownSTK-                   $ AstReplicate (SNat @1) knownSTK-                   $ AstReplicate (SNat @1) knownSTK+      blackGlyph = rreplicateN (1 :$: 1 :$: 1 :$: ZSR)                        (rconcrete $ Nested.rscalar 7                         :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))       afcnn2T :: ADRnnMnistParameters (AstTensor AstMethodLet FullSpan)@@ -413,14 +415,14 @@                  (knownSTK @(X (ADRnnMnistParameters Concrete Double)))                  (toTarget @Concrete $ valsInitRNNOPP 1 sizeMnistHeightI)       artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2T ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (str (sreplicate @1 (str (sfromR (tproject1 (tproject2 m1))) !$ [0] * sreplicate @10 (tanh (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) !$ [0, 0] * tanh (sscalar 7.0 * sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) !$ [0, 0] + sfromR (tproject2 (tproject1 (tproject1 m1))) !$ [0]) + sfromR (tproject2 (tproject2 (tproject1 m1))) !$ [0])))) + str (sreplicate @1 (sfromR (tproject2 (tproject2 m1)))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (str (sreplicate @1 (str (sfromR (tproject1 (tproject2 m1))) !$ [0] * sreplicate0N @[10] (tanh (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) `sindex0` [0, 0] * tanh (7.0 * sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) `sindex0` [0, 0] + sfromR (tproject2 (tproject1 (tproject1 m1))) `sindex0` [0]) + sfromR (tproject2 (tproject2 (tproject1 m1))) `sindex0` [0])))) + str (sreplicate @1 (sfromR (tproject2 (tproject2 m1)))))"   printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> let x16 = sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) !$ [0, 0] ; x18 = tanh (sscalar 7.0 * x16 + sfromR (tproject2 (tproject1 (tproject1 m1))) !$ [0]) ; x19 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) !$ [0, 0] ; x21 = tanh (x19 * x18 + sfromR (tproject2 (tproject2 (tproject1 m1))) !$ [0]) ; v22 = str (sfromR (tproject1 (tproject2 m1))) !$ [0] in rfromS (str (sreplicate @1 (v22 * sreplicate @10 x21)) + str (sreplicate @1 (sfromR (tproject2 (tproject2 m1)))))"+    @?= "\\m1 -> rfromS (let x16 = sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) `sindex0` [0, 0] ; x17 = tanh (7.0 * x16 + sfromR (tproject2 (tproject1 (tproject1 m1))) `sindex0` [0]) ; x18 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) `sindex0` [0, 0] ; x19 = tanh (x18 * x17 + sfromR (tproject2 (tproject2 (tproject1 m1))) `sindex0` [0]) ; v20 = str (sfromR (tproject1 (tproject2 m1))) !$ [0] in str (sreplicate @1 (v20 * sreplicate0N @[10] x19)) + str (sreplicate @1 (sfromR (tproject2 (tproject2 m1)))))"   printArtifactPretty artifactRev-    @?= "\\dret m1 -> let x16 = sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) !$ [0, 0] ; x18 = tanh (sscalar 7.0 * x16 + sfromR (tproject2 (tproject1 (tproject1 m1))) !$ [0]) ; x19 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) !$ [0, 0] ; x21 = tanh (x19 * x18 + sfromR (tproject2 (tproject2 (tproject1 m1))) !$ [0]) ; v22 = str (sfromR (tproject1 (tproject2 m1))) !$ [0] ; x24 = (sscalar 1.0 + negate x21 * x21) * ssum @10 (v22 * ssum @1 (str (sfromR dret))) ; x25 = (sscalar 1.0 + negate x18 * x18) * (x19 * x24) in tpair (tpair (tpair (tpair (rfromS (soneHot (sscalar 7.0 * x25) [0, 0])) (rfromS (soneHot (sscalar 0.0) [0, 0]))) (rfromS (soneHot x25 [0]))) (tpair (tpair (rfromS (soneHot (x18 * x24) [0, 0])) (rfromS (soneHot (sscalar 0.0) [0, 0]))) (rfromS (soneHot x24 [0])))) (tpair (rfromS (str (soneHot (sreplicate @10 x21 * ssum @1 (str (sfromR dret))) [0]))) (rfromS (ssum @1 (str (sfromR dret)))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1,1] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1,1] FTKScalar)) ConvSX))) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1] FTKScalar)) ConvSX))) (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1,1] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1,1] FTKScalar)) ConvSX))) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1] FTKScalar)) ConvSX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [10,1] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [10] FTKScalar)) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKProduct (STKS [1,1] STKScalar) (STKS [1,1] STKScalar)) (STKS [1] STKScalar)) (STKProduct (STKProduct (STKS [1,1] STKScalar) (STKS [1,1] STKScalar)) (STKS [1] STKScalar))) (STKProduct (STKS [10,1] STKScalar) (STKS [10] STKScalar))) (let x18 = tanh (sscalar 7.0 * sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) !$ [0, 0] + sfromR (tproject2 (tproject1 (tproject1 m1))) !$ [0]) ; x19 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) !$ [0, 0] ; x21 = tanh (x19 * x18 + sfromR (tproject2 (tproject2 (tproject1 m1))) !$ [0]) ; x24 = (sscalar 1.0 + negate x21 * x21) * sdot0 (str (sfromR (tproject1 (tproject2 m1))) !$ [0]) (str (sfromR dret) !$ [0]) ; x25 = (sscalar 1.0 + negate x18 * x18) * (x19 * x24) in tpair (tpair (tpair (tpair (sreplicate @1 (sreplicate @1 (sscalar 7.0 * x25))) (sconcrete (sfromListLinear [1,1] [0.0]))) (sreplicate @1 x25)) (tpair (tpair (sreplicate @1 (sreplicate @1 (x18 * x24))) (sconcrete (sfromListLinear [1,1] [0.0]))) (sreplicate @1 x24))) (tpair (str (sreplicate @1 (sreplicate @10 x21 * str (sfromR dret) !$ [0]))) (str (sfromR dret) !$ [0])))"+    @?= "\\dret m1 -> let x16 = sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) `sindex0` [0, 0] ; x17 = tanh (7.0 * x16 + sfromR (tproject2 (tproject1 (tproject1 m1))) `sindex0` [0]) ; x18 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) `sindex0` [0, 0] ; x19 = tanh (x18 * x17 + sfromR (tproject2 (tproject2 (tproject1 m1))) `sindex0` [0]) ; v20 = str (sfromR (tproject1 (tproject2 m1))) !$ [0] ; v22 = ssum @1 (str (sfromR dret)) in tpair (let x23 = (1.0 + negate (x19 * x19)) * ssum0 @[10] (v20 * v22) in tpair (let x24 = (1.0 + negate (x17 * x17)) * (x18 * x23) in tpair (tpair (rfromS (soneHot (sfromK (7.0 * x24)) [0, 0])) (rfromS (sconcrete (sfromListLinear [1,1] [0.0])))) (rfromS (soneHot (sfromK x24) [0]))) (tpair (tpair (rfromS (soneHot (sfromK (x17 * x23)) [0, 0])) (rfromS (sconcrete (sfromListLinear [1,1] [0.0])))) (rfromS (soneHot (sfromK x23) [0])))) (tpair (rfromS (str (soneHot (sreplicate0N @[10] x19 * v22) [0]))) (rfromS (ssum @1 (str (sfromR dret)))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKProduct (STKProduct (STKProduct (STKS [1,1] STKScalar) (STKS [1,1] STKScalar)) (STKS [1] STKScalar)) (STKProduct (STKProduct (STKS [1,1] STKScalar) (STKS [1,1] STKScalar)) (STKS [1] STKScalar))) (STKProduct (STKS [10,1] STKScalar) (STKS [10] STKScalar))) (let x17 = tanh (7.0 * sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) `sindex0` [0, 0] + sfromR (tproject2 (tproject1 (tproject1 m1))) `sindex0` [0]) ; x18 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) `sindex0` [0, 0] ; x19 = tanh (x18 * x17 + sfromR (tproject2 (tproject2 (tproject1 m1))) `sindex0` [0]) ; v22 = str (sfromR dret) !$ [0] in tpair (let x23 = (1.0 + negate x19 * x19) * sdot0 (str (sfromR (tproject1 (tproject2 m1))) !$ [0]) v22 in tpair (let x24 = (1.0 + negate x17 * x17) * (x18 * x23) in tpair (tpair (sreplicate0N @[1, 1] (7.0 * x24)) (sconcrete (sfromListLinear [1,1] [0.0]))) (sreplicate0N @[1] x24)) (tpair (tpair (sreplicate0N @[1, 1] (x17 * x23)) (sconcrete (sfromListLinear [1,1] [0.0]))) (sreplicate0N @[1] x23))) (tpair (str (sreplicate @1 (sreplicate0N @[10] x19 * v22))) (str (sfromR dret) !$ [0])))"  testRNNOAst :: Assertion testRNNOAst = do@@ -429,7 +431,7 @@       ftk = tftk @Concrete                  (knownSTK @(X (ADRnnMnistParameters Concrete Double)))                  (toTarget @Concrete $ valsInitRNNOPP 1 sizeMnistHeightI)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan                        (X (ADRnnMnistParameters Concrete Double))       var = AstVar varName@@ -447,7 +449,7 @@   interpretAstFull @Concrete env afcnn1     @?= afcnn2 (valsInitRNNOPP 1 sizeMnistHeightI)   interpretAstFull @Concrete env-                   (simplifyInline @(TKR 2 Double) afcnn1)+                   (simplifyUserCode @(TKR 2 Double) afcnn1)     @?= afcnn2 (valsInitRNNOPP 1 sizeMnistHeightI)   interpretAstFull @Concrete env                    (simplifyInlineContract @(TKR 2 Double) afcnn1)@@ -459,9 +461,7 @@   let batch_size = 2       sizeMnistHeightI = 2       blackGlyph :: AstTensor AstMethodLet PrimalSpan (TKR 3 Double)-      blackGlyph = AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK-                   $ AstReplicate (SNat @2) knownSTK+      blackGlyph = rreplicateN (2 :$: 2 :$: 2 :$: ZSR)                        (rconcrete $ Nested.rscalar 7                         :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))       afcnn2T :: ADRnnMnistParameters (AstTensor AstMethodLet FullSpan)@@ -472,14 +472,14 @@                  (knownSTK @(X (ADRnnMnistParameters Concrete Double)))                  (toTarget @Concrete $ valsInitRNNOPP 2 sizeMnistHeightI)       artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2T ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\m1 -> rfromS (let m40 = sappend (tanh (str (sreplicate @2 (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))))) (tanh (str (sreplicate @2 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @2 (tanh (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1))))))) in smatmul2 (sfromR (tproject1 (tproject2 m1))) (tanh ((smatmul2 (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (tanh ((str (sreplicate @2 (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + smatmul2 (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))) (sslice (SNat @0) (SNat @2) m40)) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))))) + smatmul2 (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))) (sslice (SNat @2) (SNat @2) m40)) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 m1)))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\m1 -> rfromS (let m43 = sappend (tanh (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))))) (tanh (str (sreplicate @2 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @2 (tanh (sconcrete (sreplicate [2] 7.0) * ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1))))))) in smatmul2 (sfromR (tproject1 (tproject2 m1))) (tanh (smatmul2 (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (tanh (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))))) + (smatmul2 (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))) (sslice (SNat @0) (SNat @2) m43) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))))) + (smatmul2 (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))) (sslice (SNat @2) (SNat @2) m43) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 m1)))))"   printArtifactPrimalPretty artifactRev-    @?= "\\m1 -> let m37 = tanh ((str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 0.0))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v38 = tanh ((ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) + ssum @2 (sconcrete (sreplicate [2,2] 0.0))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m39 = tanh ((str (sreplicate @2 (ssum @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * str (sreplicate @2 v38)))) + str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 0.0))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m40 = sappend m37 m39 ; m41 = tanh ((sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + ssum @2 (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * stranspose @[2,1,0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m40))))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))) ; m42 = tanh ((ssum @2 (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2,0,1] (sreplicate @2 m41)) + ssum @2 (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2,0,1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m40))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) in rfromS (ssum @2 (stranspose @[2,1,0] (sreplicate @2 (sfromR (tproject1 (tproject2 m1)))) * str (sreplicate @10 m42)) + str (sreplicate @2 (sfromR (tproject2 (tproject2 m1)))))"+    @?= "\\m1 -> rfromS (let v38 = ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) ; m39 = tanh (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * v38)) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v40 = ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) ; v41 = tanh (sconcrete (sreplicate [2] 7.0) * v40 + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m42 = tanh (str (sreplicate @2 (ssum @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * str (sreplicate @2 v41)))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m43 = sappend m39 m42 ; v44 = ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) ; m45 = tanh (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * v44) + (ssum @2 (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * stranspose @[2, 1, 0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m43)))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m46 = tanh (ssum @2 (stranspose @[1, 2, 0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2, 0, 1] (sreplicate @2 m45)) + (ssum @2 (stranspose @[1, 2, 0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2, 0, 1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m43)))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1))))))) in ssum @2 (stranspose @[2, 1, 0] (sreplicate @2 (sfromR (tproject1 (tproject2 m1)))) * str (sreplicate @10 m46)) + str (sreplicate @2 (sfromR (tproject2 (tproject2 m1)))))"   printArtifactPretty artifactRev-    @?= "\\dret m1 -> let m37 = tanh ((str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 0.0))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v38 = tanh ((ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) + ssum @2 (sconcrete (sreplicate [2,2] 0.0))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m39 = tanh ((str (sreplicate @2 (ssum @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * str (sreplicate @2 v38)))) + str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 0.0))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m40 = sappend m37 m39 ; m41 = tanh ((sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + ssum @2 (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * stranspose @[2,1,0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m40))))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))) ; m42 = tanh ((ssum @2 (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2,0,1] (sreplicate @2 m41)) + ssum @2 (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2,0,1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m40))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m44 = (sconcrete (sreplicate [2,2] 1.0) + negate m42 * m42) * ssum @10 (str (stranspose @[2,1,0] (sreplicate @2 (sfromR (tproject1 (tproject2 m1)))) * sreplicate @2 (sfromR dret))) ; m45 = (sconcrete (sreplicate [2,2] 1.0) + negate m41 * m41) * ssum @2 (stranspose @[1,2,0] (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * sreplicate @2 m44)) ; m46 = sappend (sconcrete (sfromListLinear [0,2] [])) (sappend (str (ssum @2 (stranspose @[2,1,0] (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * sreplicate @2 m45)))) (sconcrete (sreplicate [2,2] 0.0))) + sappend (sconcrete (sreplicate [2,2] 0.0)) (sappend (str (ssum @2 (stranspose @[1,2,0] (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * sreplicate @2 m44)))) (sconcrete (sfromListLinear [0,2] []))) ; m47 = (sconcrete (sreplicate [2,2] 1.0) + negate m39 * m39) * sslice (SNat @2) (SNat @2) m46 ; m48 = sreplicate @2 (ssum @2 (str m47)) ; v49 = (sconcrete (sreplicate [2] 1.0) + negate v38 * v38) * ssum @2 (str (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * m48)) ; m50 = (sconcrete (sreplicate [2,2] 1.0) + negate m37 * m37) * sslice (SNat @0) (SNat @2) m46 in tpair (tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [2,2] 7.0) * sreplicate @2 (ssum @2 (str m50))) + (str (sconcrete (sreplicate [2,2] 7.0) * sreplicate @2 v49) + str (sconcrete (sreplicate [2,2] 7.0) * sreplicate @2 (ssum @2 m45))))) (rfromS (str (sconcrete (sreplicate [2,2] 0.0)) + (str (sconcrete (sreplicate [2,2] 0.0)) + str (ssum @2 (str (stranspose @[2,1,0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m40))) * sreplicate @2 m45))))))) (rfromS (ssum @2 (str m50) + (v49 + ssum @2 m45)))) (tpair (tpair (rfromS (str (str (sreplicate @2 v38) * m48) + str (ssum @2 (stranspose @[2,0,1] (stranspose @[2,0,1] (sreplicate @2 m41) * sreplicate @2 m44))))) (rfromS (str (sconcrete (sreplicate [2,2] 0.0)) + str (ssum @2 (stranspose @[2,0,1] (stranspose @[2,0,1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m40))) * sreplicate @2 m44)))))) (rfromS (ssum @2 (str m47) + ssum @2 (str m44))))) (tpair (rfromS (ssum @2 (stranspose @[2,1,0] (str (sreplicate @10 m42) * sreplicate @2 (sfromR dret))))) (rfromS (ssum @2 (str (sfromR dret)))))"-  printArtifactPretty (simplifyArtifact artifactRev)-    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [10,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [10] FTKScalar)) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2] STKScalar)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2] STKScalar))) (STKProduct (STKS [10,2] STKScalar) (STKS [10] STKScalar))) (let m37 = tanh (str (sreplicate @2 (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v38 = tanh (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m39 = tanh (str (sreplicate @2 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @2 v38))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m40 = sappend m37 m39 ; m41 = tanh ((sreplicate @2 (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) + smatmul2 (str (sslice (SNat @0) (SNat @2) m40)) (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))) ; m42 = tanh ((smatmul2 (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (str m41) + smatmul2 (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))) (sslice (SNat @2) (SNat @2) m40)) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m44 = (sconcrete (sreplicate [2,2] 1.0) + negate m42 * m42) * smatmul2 (str (sfromR (tproject1 (tproject2 m1)))) (sfromR dret) ; m45 = (sconcrete (sreplicate [2,2] 1.0) + negate m41 * m41) * smatmul2 (str m44) (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) ; m46 = sappend (smatmul2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1)))))) (str m45)) (sconcrete (sreplicate [2,2] 0.0)) + sappend (sconcrete (sreplicate [2,2] 0.0)) (smatmul2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1)))))) m44) ; m47 = (sconcrete (sreplicate [2,2] 1.0) + negate m39 * m39) * sslice (SNat @2) (SNat @2) m46 ; v48 = ssum @2 (str m47) ; v49 = (sconcrete (sreplicate [2] 1.0) + negate v38 * v38) * sdot1In (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))))) (sreplicate @2 v48) ; m50 = (sconcrete (sreplicate [2,2] 1.0) + negate m37 * m37) * sslice (SNat @0) (SNat @2) m46 in tpair (tpair (tpair (tpair (sconcrete (sreplicate [2,2] 7.0) * str (sreplicate @2 (ssum @2 (str m50))) + (sconcrete (sreplicate [2,2] 7.0) * str (sreplicate @2 v49) + sconcrete (sreplicate [2,2] 7.0) * str (sreplicate @2 (ssum @2 m45)))) (smatmul2 (str m45) (str (sslice (SNat @0) (SNat @2) m40)))) (ssum @2 (str m50) + (v49 + ssum @2 m45))) (tpair (tpair (sreplicate @2 v38 * str (sreplicate @2 v48) + smatmul2 m44 m41) (smatmul2 m44 (str (sslice (SNat @2) (SNat @2) m40)))) (ssum @2 (str m47) + ssum @2 (str m44)))) (tpair (smatmul2 (sfromR dret) (str m42)) (ssum @2 (str (sfromR dret)))))"+    @?= "\\dret m1 -> let v38 = ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) ; m39 = tanh (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * v38)) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v40 = ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) ; v41 = tanh (sconcrete (sreplicate [2] 7.0) * v40 + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m42 = tanh (str (sreplicate @2 (ssum @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * str (sreplicate @2 v41)))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m43 = sappend m39 m42 ; v44 = ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) ; m45 = tanh (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * v44) + (ssum @2 (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * stranspose @[2, 1, 0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m43)))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m46 = tanh (ssum @2 (stranspose @[1, 2, 0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2, 0, 1] (sreplicate @2 m45)) + (ssum @2 (stranspose @[1, 2, 0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2, 0, 1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m43)))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1))))))) in tpair (let m48 = (sconcrete (sreplicate [2,2] 1.0) + negate (m46 * m46)) * ssum @10 (str (stranspose @[2, 1, 0] (sreplicate @2 (sfromR (tproject1 (tproject2 m1)))) * sreplicate @2 (sfromR dret))) ; m49 = (sconcrete (sreplicate [2,2] 1.0) + negate (m45 * m45)) * ssum @2 (stranspose @[1, 2, 0] (stranspose @[1, 2, 0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * sreplicate @2 m48)) ; m50 = sappend (sconcrete (sfromListLinear [0,2] [])) (sappend (str (ssum @2 (stranspose @[2, 1, 0] (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * sreplicate @2 m49)))) (sconcrete (sreplicate [2,2] 0.0))) + sappend (sconcrete (sreplicate [2,2] 0.0)) (sappend (str (ssum @2 (stranspose @[1, 2, 0] (stranspose @[1, 2, 0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * sreplicate @2 m48)))) (sconcrete (sfromListLinear [0,2] []))) ; m51 = (sconcrete (sreplicate [2,2] 1.0) + negate (m42 * m42)) * sslice (SNat @2) (SNat @2) m50 ; m52 = sreplicate @2 (ssum @2 (str m51)) in tpair (let v53 = (sconcrete (sreplicate [2] 1.0) + negate (v41 * v41)) * ssum @2 (str (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * m52)) ; m54 = (sconcrete (sreplicate [2,2] 1.0) + negate (m39 * m39)) * sslice (SNat @0) (SNat @2) m50 in tpair (tpair (rfromS (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * ssum @2 (str m54))) + (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * v53)) + str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * ssum @2 m49))))) (rfromS (str (ssum @2 (str (stranspose @[2, 1, 0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m43))) * sreplicate @2 m49)))))) (rfromS (ssum @2 (str m54) + (v53 + ssum @2 m49)))) (tpair (tpair (rfromS (str (str (sreplicate @2 v41) * m52) + str (ssum @2 (stranspose @[2, 0, 1] (stranspose @[2, 0, 1] (sreplicate @2 m45) * sreplicate @2 m48))))) (rfromS (str (ssum @2 (stranspose @[2, 0, 1] (stranspose @[2, 0, 1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m43))) * sreplicate @2 m48)))))) (rfromS (ssum @2 (str m51) + ssum @2 (str m48))))) (tpair (rfromS (ssum @2 (stranspose @[2, 1, 0] (str (sreplicate @10 m46) * sreplicate @2 (sfromR dret))))) (rfromS (ssum @2 (str (sfromR dret)))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKProduct (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2] STKScalar)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2] STKScalar))) (STKProduct (STKS [10,2] STKScalar) (STKS [10] STKScalar))) (let m39 = tanh (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v41 = tanh (sconcrete (sreplicate [2] 7.0) * ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m42 = tanh (str (sreplicate @2 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @2 v41))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m43 = sappend m39 m42 ; m45 = tanh (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * ssum @2 (str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + (smatmul2 (str (sslice (SNat @0) (SNat @2) m43)) (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1)))))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m46 = tanh (smatmul2 (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (str m45) + (smatmul2 (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))) (sslice (SNat @2) (SNat @2) m43) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1))))))) in tpair (let m48 = (sconcrete (sreplicate [2,2] 1.0) + negate m46 * m46) * smatmul2 (str (sfromR (tproject1 (tproject2 m1)))) (sfromR dret) ; m49 = (sconcrete (sreplicate [2,2] 1.0) + negate m45 * m45) * smatmul2 (str m48) (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) ; m50 = sappend (smatmul2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1)))))) (str m49)) (sconcrete (sreplicate [2,2] 0.0)) + sappend (sconcrete (sreplicate [2,2] 0.0)) (smatmul2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1)))))) m48) ; m51 = (sconcrete (sreplicate [2,2] 1.0) + negate m42 * m42) * sslice (SNat @2) (SNat @2) m50 ; v52 = ssum @2 (str m51) in tpair (let v53 = (sconcrete (sreplicate [2] 1.0) + negate v41 * v41) * sdot1In (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))))) (sreplicate @2 v52) ; m54 = (sconcrete (sreplicate [2,2] 1.0) + negate m39 * m39) * sslice (SNat @0) (SNat @2) m50 in tpair (tpair (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * ssum @2 (str m54))) + (str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * v53)) + str (sreplicate @2 (sconcrete (sreplicate [2] 7.0) * ssum @2 m49)))) (smatmul2 (str m49) (str (sslice (SNat @0) (SNat @2) m43)))) (ssum @2 (str m54) + (v53 + ssum @2 m49))) (tpair (tpair (sreplicate @2 v41 * str (sreplicate @2 v52) + smatmul2 m48 m45) (smatmul2 m48 (str (sslice (SNat @2) (SNat @2) m43)))) (ssum @2 (str m51) + ssum @2 (str m48)))) (tpair (smatmul2 (sfromR dret) (str m46)) (ssum @2 (str (sfromR dret)))))"  testRNNOAst2 :: Assertion testRNNOAst2 = do@@ -488,7 +488,7 @@       ftk = tftk @Concrete                  (knownSTK @(X (ADRnnMnistParameters Concrete Double)))                  (toTarget @Concrete $ valsInitRNNOPP 2 sizeMnistHeightI)-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan                        (X (ADRnnMnistParameters Concrete Double))       var = AstVar varName@@ -506,7 +506,7 @@   interpretAstFull @Concrete env afcnn1     @?= afcnn2 (valsInitRNNOPP 2 sizeMnistHeightI)   interpretAstFull @Concrete env-                   (simplifyInline @(TKR 2 Double) afcnn1)+                   (simplifyUserCode @(TKR 2 Double) afcnn1)     @?= afcnn2 (valsInitRNNOPP 2 sizeMnistHeightI)   interpretAstFull @Concrete env                    (simplifyInlineContract @(TKR 2 Double) afcnn1)@@ -516,7 +516,7 @@ -- * CNNR tests  tensorMnistCNNRPP :: TestTree-tensorMnistCNNRPP = testGroup "Ast tests for CNNR MNIST"+tensorMnistCNNRPP = inOrderTestGroup "Ast tests for CNNR MNIST"   [ testCase "CNNO PP 1" testCNNOPP1   , testCase "CNNO Ast 1" testCNNOAst1   , testCase "CNNO PP 2" testCNNOPP2@@ -553,14 +553,10 @@                  sizeMnistHeightI sizeMnistWidthI batch_size                  (rconcrete $ unConcrete blackGlyph)       artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2 ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let t189 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[7,9,4] (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [7,9] 7.0)) (\\[i358, i359] -> [i358 + i359]))) (\\[i185, i186] -> [i185 + i186])))))) * sreplicate @7 (sreplicate @9 (sreplicate @1 (sreplicate @1 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0])))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @9 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; t204 = sreshape @[3,4,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i192, i193, i194, i195] -> [ifH (sscalar -0.0 <=. negate (t189 !$ [0, kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i192, i194]), kfromS (sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) !$ [i193, i195])])) 0 1]) * sgather (t189 !$ [0]) (\\[i198, i199, i200, i201] -> [kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i198, i200]), kfromS (sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) !$ [i199, i201])])) ; t213 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[3,4,4] (stranspose @[2,0,3,1] (sgather (stranspose @[2,3,4,0,1] (sgather (stranspose @[3,5,0,4,1,2] (sgather (stranspose @[3,2,1,6,5,4,0] (sreplicate @2 (stranspose @[5,4,3,0,1,2] (sreplicate @4 (sreplicate @3 (sreplicate @2 (stranspose @[2,1,0] t204))))))) (\\[i346, i350] -> [kfromS (smaxIndex (t204 !$ [i346, i350])), i350, i346]))) (\\[i355, i356] -> [i355, i356, i355 + i356]))) (\\[i209, i210] -> [i209, i209 + i210, i210])) * sreplicate @3 (sreplicate @4 (sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0])))))) + stranspose @[2,0,1] (sreplicate @3 (sreplicate @4 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m223 = sreshape @[2,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i215, i216, i217] -> [ifH (sscalar -0.0 <=. negate (t213 !$ [0, i216, kfromS (sconcrete (sfromListLinear [2,2] [0,1,2,3]) !$ [i215, i217])])) 0 1]) * stranspose @[0,2,1] (sgather (str (sslice (SNat @0) (SNat @2) (t213 !$ [0]))) (\\[i219, i220] -> [kfromS (sconcrete (sfromListLinear [2,2] [0,1,2,3]) !$ [i219, i220])]))) ; m227 = sreplicate @1 (sreplicate @5 (sdot0 (sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0]) (sgather m223 (\\[i224] -> [i224, kfromS (smaxIndex (m223 !$ [i224]))])))) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) in str (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) * sreplicate @10 (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i228] -> [ifH (sscalar -0.0 <=. negate (m227 !$ [0, i228])) 0 1]) * m227 !$ [0]) + str (sreplicate @5 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w187 = stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [7,9] 7.0)) (\\[i183, i184] -> [i183 + i184]))) (\\[i185, i186] -> [i185 + i186])))))) ; w188 = sreplicate @7 (sreplicate @9 (sreplicate @1 (sreplicate @1 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0])))) ; t189 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[7,9,4] (w187 * w188)))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @9 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; m190 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m191 = str (sreplicate @2 (sconcrete (sreplicate [4] 2) * siota (SNat @4))) + sreplicate @4 (siota (SNat @2)) ; u202 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i192, i193, i194, i195] -> [let x196 = m190 !$ [i192, i194] ; x197 = m191 !$ [i193, i195] in ifH (sscalar -0.0 <=. negate (t189 !$ [0, kfromS x196, kfromS x197])) 0 1]) ; u203 = sgather (t189 !$ [0]) (\\[i198, i199, i200, i201] -> [kfromS (m190 !$ [i198, i200]), kfromS (m191 !$ [i199, i201])]) ; t204 = sreshape @[3,4,4] (u202 * u203) ; u211 = stranspose @[2,0,3,1] (sgather (stranspose @[2,3,4,0,1] (sgather (stranspose @[3,5,0,4,1,2] (sgather (stranspose @[3,2,1,6,5,4,0] (sreplicate @2 (stranspose @[5,4,3,0,1,2] (sreplicate @4 (sreplicate @3 (sreplicate @2 (stranspose @[2,1,0] t204))))))) (\\[i205, i206] -> [kfromS (smaxIndex (t204 !$ [i205, i206])), i206, i205]))) (\\[i207, i208] -> [i207, i208, i207 + i208]))) (\\[i209, i210] -> [i209, i209 + i210, i210])) ; u212 = sreplicate @3 (sreplicate @4 (sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0])) ; t213 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[3,4,4] (u211 * u212)))) + stranspose @[2,0,1] (sreplicate @3 (sreplicate @4 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m214 = str (sreplicate @2 (sconcrete (sreplicate [2] 2) * siota (SNat @2))) + sreplicate @2 (siota (SNat @2)) ; t221 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i215, i216, i217] -> [let x218 = m214 !$ [i215, i217] in ifH (sscalar -0.0 <=. negate (t213 !$ [0, i216, kfromS x218])) 0 1]) ; t222 = stranspose @[0,2,1] (sgather (str (sslice (SNat @0) (SNat @2) (t213 !$ [0]))) (\\[i219, i220] -> [kfromS (m214 !$ [i219, i220])])) ; m223 = sreshape @[2,4] (t221 * t222) ; v225 = sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0] ; v226 = sgather m223 (\\[i224] -> [i224, kfromS (smaxIndex (m223 !$ [i224]))]) ; m227 = sreplicate @1 (sreplicate @5 (ssum @2 (v225 * v226))) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; v229 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i228] -> [ifH (sscalar -0.0 <=. negate (m227 !$ [0, i228])) 0 1]) ; v230 = m227 !$ [0] ; m231 = str (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) ; m232 = sreplicate @10 (v229 * v230) in rfromS (m231 * m232 + str (sreplicate @5 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w187 = stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [7,9] 7.0)) (\\[i183, i184] -> [i183 + i184]))) (\\[i185, i186] -> [i185 + i186])))))) ; w188 = sreplicate @7 (sreplicate @9 (sreplicate @1 (sreplicate @1 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0])))) ; t189 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[7,9,4] (w187 * w188)))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @9 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; m190 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m191 = str (sreplicate @2 (sconcrete (sreplicate [4] 2) * siota (SNat @4))) + sreplicate @4 (siota (SNat @2)) ; u202 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i192, i193, i194, i195] -> [let x196 = m190 !$ [i192, i194] ; x197 = m191 !$ [i193, i195] in ifH (sscalar -0.0 <=. negate (t189 !$ [0, kfromS x196, kfromS x197])) 0 1]) ; u203 = sgather (t189 !$ [0]) (\\[i198, i199, i200, i201] -> [kfromS (m190 !$ [i198, i200]), kfromS (m191 !$ [i199, i201])]) ; t204 = sreshape @[3,4,4] (u202 * u203) ; u211 = stranspose @[2,0,3,1] (sgather (stranspose @[2,3,4,0,1] (sgather (stranspose @[3,5,0,4,1,2] (sgather (stranspose @[3,2,1,6,5,4,0] (sreplicate @2 (stranspose @[5,4,3,0,1,2] (sreplicate @4 (sreplicate @3 (sreplicate @2 (stranspose @[2,1,0] t204))))))) (\\[i205, i206] -> [kfromS (smaxIndex (t204 !$ [i205, i206])), i206, i205]))) (\\[i207, i208] -> [i207, i208, i207 + i208]))) (\\[i209, i210] -> [i209, i209 + i210, i210])) ; u212 = sreplicate @3 (sreplicate @4 (sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0])) ; t213 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[3,4,4] (u211 * u212)))) + stranspose @[2,0,1] (sreplicate @3 (sreplicate @4 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m214 = str (sreplicate @2 (sconcrete (sreplicate [2] 2) * siota (SNat @2))) + sreplicate @2 (siota (SNat @2)) ; t221 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i215, i216, i217] -> [let x218 = m214 !$ [i215, i217] in ifH (sscalar -0.0 <=. negate (t213 !$ [0, i216, kfromS x218])) 0 1]) ; t222 = stranspose @[0,2,1] (sgather (str (sslice (SNat @0) (SNat @2) (t213 !$ [0]))) (\\[i219, i220] -> [kfromS (m214 !$ [i219, i220])])) ; m223 = sreshape @[2,4] (t221 * t222) ; v225 = sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0] ; v226 = sgather m223 (\\[i224] -> [i224, kfromS (smaxIndex (m223 !$ [i224]))]) ; m227 = sreplicate @1 (sreplicate @5 (ssum @2 (v225 * v226))) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; v229 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i228] -> [ifH (sscalar -0.0 <=. negate (m227 !$ [0, i228])) 0 1]) ; v230 = m227 !$ [0] ; m231 = str (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) ; m232 = sreplicate @10 (v229 * v230) ; m234 = soneHot (v229 * ssum @10 (m231 * sfromR dret)) [0] ; v235 = sreplicate @2 (ssum @5 (ssum @1 m234)) ; t239 = soneHot (sappend (sconcrete (sfromListLinear [0,4] [])) (sappend (str (sscatter (stranspose @[0,2,1] (t221 * sreshape @[2,2,2] (sscatter (v225 * v235) (\\[i236] -> [i236, kfromS (smaxIndex (m223 !$ [i236]))])))) (\\[i237, i238] -> [kfromS (m214 !$ [i237, i238])]))) (sconcrete (sreplicate [1,4] 0.0)))) [0] ; u240 = sreshape @[3,4,2,2] (stranspose @[1,2,0] (sreplicate @4 (ssum @1 t239))) ; t251 = soneHot (sscatter (u202 * sreshape @[3,4,2,2] (stranspose @[2,1,0] (ssum @2 (ssum @3 (ssum @4 (stranspose @[3,4,5,2,1,0] (ssum @2 (stranspose @[6,2,1,0,5,4,3] (sscatter (stranspose @[2,4,5,0,3,1] (sscatter (stranspose @[3,4,0,1,2] (sscatter (stranspose @[1,3,0,2] (u212 * u240)) (\\[i241, i242] -> [i241, i241 + i242, i242]))) (\\[i243, i244] -> [i243, i244, i243 + i244]))) (\\[i245, i246] -> [kfromS (smaxIndex (t204 !$ [i245, i246])), i246, i245])))))))))) (\\[i247, i248, i249, i250] -> [kfromS (m190 !$ [i247, i249]), kfromS (m191 !$ [i248, i250])])) [0] in tpair (tpair (tpair (rfromS (soneHot (ssum @1 (ssum @1 (ssum @9 (ssum @7 (w187 * sreshape @[7,9,1,1,2,2] (stranspose @[1,2,0] (sreplicate @4 (ssum @1 t251)))))))) [0, 0])) (rfromS (ssum @9 (ssum @7 (stranspose @[1,2,0] t251))))) (tpair (rfromS (soneHot (ssum @4 (ssum @3 (u211 * u240))) [0, 0])) (rfromS (ssum @4 (ssum @3 (stranspose @[1,2,0] t239)))))) (tpair (tpair (rfromS (soneHot (v226 * v235) [0])) (rfromS (ssum @5 (str m234)))) (tpair (rfromS (str (soneHot (ssum @5 (str (m232 * sfromR dret))) [0]))) (rfromS (ssum @5 (str (sfromR dret))))))"--- TODO: different test result with GHC 9.10:   printArtifactPretty (simplifyArtifact artifactRev)---    @?= "\\dret u1 -> tfromS (STKProduct (STKProduct (STKProduct (STKR (SNat @4) STKScalar) (STKR (SNat @1) STKScalar)) (STKProduct (STKR (SNat @4) STKScalar) (STKR (SNat @1) STKScalar))) (STKProduct (STKProduct (STKR (SNat @2) STKScalar) (STKR (SNat @1) STKScalar)) (STKProduct (STKR (SNat @2) STKScalar) (STKR (SNat @1) STKScalar)))) (let u181 = sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [7,9] 7.0)) (\\[i296, i297] -> [i296 + i297]))) (\\[i179, i180] -> [i179 + i180]) ; t183 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[7,9,4] (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] u181)))) * sreplicate @7 (sreplicate @9 (sreplicate @1 (sreplicate @1 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0])))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @9 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; u192 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i184, i185, i186, i187] -> [ifH (sscalar -0.0 <=. negate (t183 !$ [0, 2 * i184 + i186, 2 * i185 + i187])) 0 1]) ; t194 = sreshape @[3,4,4] (u192 * sgather (t183 !$ [0]) (\\[i188, i189, i190, i191] -> [2 * i188 + i190, 2 * i189 + i191])) ; u201 = sgather (stranspose @[2,3,4,0,1] (sgather (stranspose @[3,5,0,4,1,2] (sgather (stranspose @[3,2,1,6,5,4,0] (sreplicate @2 (stranspose @[5,4,3,0,1,2] (sreplicate @4 (sreplicate @3 (sreplicate @2 (stranspose @[2,1,0] t194))))))) (\\[i284, i288] -> [kfromS (smaxIndex (t194 !$ [i284, i288])), i288, i284]))) (\\[i293, i294] -> [i293, i294, i293 + i294]))) (\\[i199, i200] -> [i199, i199 + i200, i200]) ; m202 = sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0] ; t203 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[3,4,4] (stranspose @[2,0,3,1] u201 * sreplicate @3 (sreplicate @4 m202))))) + stranspose @[2,0,1] (sreplicate @3 (sreplicate @4 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; t209 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i204, i205, i206] -> [ifH (sscalar -0.0 <=. negate (t203 !$ [0, i205, 2 * i204 + i206])) 0 1]) ; m211 = sreshape @[2,4] (t209 * stranspose @[0,2,1] (sgather (str (sslice (SNat @0) (SNat @2) (t203 !$ [0]))) (\\[i207, i208] -> [2 * i207 + i208]))) ; v213 = sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0] ; v214 = sgather m211 (\\[i212] -> [i212, kfromS (smaxIndex (m211 !$ [i212]))]) ; m215 = sreplicate @1 (sreplicate @5 (sdot0 v213 v214)) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; v217 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i216] -> [ifH (sscalar -0.0 <=. negate (m215 !$ [0, i216])) 0 1]) ; v222 = v217 * sdot1In (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) (str (sfromR dret)) ; x223 = ssum0 v222 ; m227 = sappend (str (sscatter (stranspose @[0,2,1] t209 * stranspose @[0,2,1] (sreshape @[2,2,2] (sscatter (v213 * sreplicate @2 x223) (\\[i224] -> [i224, kfromS (smaxIndex (m211 !$ [i224]))])))) (\\[i225, i226] -> [2 * i225 + i226]))) (sconcrete (sreplicate [1,4] 0.0)) ; u228 = sreshape @[3,4,2,2] (stranspose @[1,2,0] (sreplicate @4 m227)) ; m239 = sscatter (u192 * sreshape @[3,4,2,2] (ssum @2 (ssum @3 (ssum @4 (ssum @2 (stranspose @[6,5,4,3,2,1,0] (sscatter (stranspose @[2,4,5,0,3,1] (sscatter (stranspose @[3,4,0,1,2] (sscatter (stranspose @[1,3,0,2] (sreplicate @3 (sreplicate @4 m202)) * stranspose @[1,3,0,2] u228) (\\[i229, i230] -> [i229, i229 + i230, i230]))) (\\[i231, i232] -> [i231, i232, i231 + i232]))) (\\[i233, i234] -> [kfromS (smaxIndex (t194 !$ [i233, i234])), i234, i233])))))))) (\\[i235, i236, i237, i238] -> [2 * i235 + i237, 2 * i236 + i238]) in tpair (tpair (tpair (sreplicate @1 (sreplicate @1 (ssum @9 (sdot1In (stranspose @[0,3,1,2] u181) (stranspose @[2,3,1,4,5,0] (sreshape @[7,9,1,1,2,2] (stranspose @[1,2,0] (sreplicate @4 m239))) !$ [0, 0]))))) (ssum @9 (ssum @7 (stranspose @[1,2,0] (sreplicate @1 m239))))) (tpair (sreplicate @1 (sreplicate @1 (ssum @4 (sdot1In (stranspose @[0,3,1,2] u201) (stranspose @[1,2,3,0] u228))))) (ssum @4 (ssum @3 (stranspose @[1,2,0] (sreplicate @1 m227)))))) (tpair (tpair (sreplicate @1 (v214 * sreplicate @2 x223)) (sreplicate @1 (ssum0 v222))) (tpair (str (sreplicate @1 (sdot1In (sreplicate @10 (v217 * m215 !$ [0])) (sfromR dret)))) (ssum @5 (str (sfromR dret))))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (let t172 = sreplicate @1 (ssum @2 (sdot1In (sfromPlain (stranspose @[1, 2, 0] (sgather @[9, 2] (stranspose @[2, 0, 1] (sgather @[7, 2] (sconcrete (sreplicate [7,9] 7.0)) (\\[i304, i305] -> [i304 + i305]))) (\\[i166, i167] -> [i166 + i167])))) (stranspose @[2, 0, 3, 1] (sreplicate @7 (stranspose @[1, 2, 0] (sreplicate @9 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0]))))))) + stranspose @[2, 0, 1] (sreplicateN @[7, 9] (sfromR (tproject2 (tproject1 (tproject1 u1))))) ; t185 = sreshape @[3, 4, 4] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (sfromPlain (sgather @[3, 4, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i175, i176, i177, i178] -> [ifH (0.0 <=. negate (splainPart t172 `sindex0` [0, sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i175, i177], sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) `sindex0` [i176, i178]])) 0 1])) * sgather @[3, 4, 2, 2] (t172 !$ [0]) (\\[i179, i180, i181, i182] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i179, i181], sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) `sindex0` [i180, i182]])))) ; t194 = sreplicate @1 (ssum @2 (sdot1In (sgather @[2, 3, 4, 2] (sgather @[3, 4] (stranspose @[6, 5, 4, 0, 1, 2, 3] (sreplicateN @[3, 4, 2, 2] (stranspose @[2, 1, 0] t185))) (\\[i186, i187] -> [i186, i187, kargMax (splainPart t185 !$ [i186, i187])])) (\\[i189, i190, i191, i188] -> [i190 + i188, i191 + i189, i190, i191, i188, i189])) (stranspose @[2, 0, 3, 1] (sreplicate @3 (stranspose @[1, 2, 0] (sreplicate @4 (sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0]))))))) + stranspose @[2, 0, 1] (sreplicateN @[3, 4] (sfromR (tproject2 (tproject2 (tproject1 u1))))) ; m203 = sreshape @[2, 4] (stranspose @[2, 0, 1] (sreplicateN @[1, 1] (sfromPlain (sgather @[2, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i196, i197, i198] -> [ifH (0.0 <=. negate (splainPart t194 `sindex0` [0, i197, sconcrete (sfromListLinear [2,2] [0,1,2,3]) `sindex0` [i196, i198]])) 0 1])) * stranspose @[0, 2, 1] (sgather @[2, 2] (str (sslice (SNat @0) (SNat @2) (t194 !$ [0]))) (\\[i199, i200] -> [sconcrete (sfromListLinear [2,2] [0,1,2,3]) `sindex0` [i199, i200]]))))) ; m207 = sreplicate0N @[1, 5] (sdot0 (sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0]) (sreshape @[2] (sreplicateN @[1, 1] (sgather1 @2 m203 (\\i204 -> [i204, kargMax (splainPart m203 !$ [i204])]))))) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) in str (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) * sreplicate @10 (sfromPlain (sgather1 @5 (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\i208 -> [ifH (0.0 <=. negate (splainPart m207 `sindex0` [0, i208])) 0 1])) * m207 !$ [0]) + str (sreplicate @5 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKS [1,1,2,2] STKScalar) (STKS [1] STKScalar)) (STKProduct (STKS [1,1,2,2] STKScalar) (STKS [1] STKScalar))) (STKProduct (STKProduct (STKS [1,2] STKScalar) (STKS [1] STKScalar)) (STKProduct (STKS [10,1] STKScalar) (STKS [10] STKScalar)))) (let u170 = sgather @[9, 2] (stranspose @[2, 0, 1] (sgather @[7, 2] (sconcrete (sreplicate [7,9] 7.0)) (\\[i350, i351] -> [i350 + i351]))) (\\[i166, i167] -> [i166 + i167]) ; t172 = sreplicate @1 (ssum @2 (sdot1In (sfromPlain (stranspose @[1, 2, 0] u170)) (stranspose @[2, 0, 3, 1] (sreplicate @7 (stranspose @[1, 2, 0] (sreplicate @9 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0]))))))) + stranspose @[2, 0, 1] (sreplicateN @[7, 9] (sfromR (tproject2 (tproject1 (tproject1 u1))))) ; u183 = sgather @[3, 4, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i175, i176, i177, i178] -> [ifH (0.0 <=. negate (splainPart t172 `sindex0` [0, sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i175, i177], sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) `sindex0` [i176, i178]])) 0 1]) ; t185 = sreshape @[3, 4, 4] (stranspose @[2, 3, 0, 1] (sreplicateN @[1, 1] (sfromPlain u183 * sgather @[3, 4, 2, 2] (t172 !$ [0]) (\\[i179, i180, i181, i182] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i179, i181], sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) `sindex0` [i180, i182]])))) ; u192 = sgather @[2, 2, 3, 4] (sgather @[3, 4] (stranspose @[6, 5, 4, 0, 1, 2, 3] (sreplicateN @[3, 4, 2, 2] (stranspose @[2, 1, 0] t185))) (\\[i186, i187] -> [i186, i187, kargMax (splainPart t185 !$ [i186, i187])])) (\\[i188, i189, i190, i191] -> [i190 + i188, i191 + i189, i190, i191, i188, i189]) ; m193 = sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0] ; t194 = sreplicate @1 (ssum @2 (sdot1In (stranspose @[1, 2, 3, 0] u192) (stranspose @[2, 0, 3, 1] (sreplicate @3 (stranspose @[1, 2, 0] (sreplicate @4 m193)))))) + stranspose @[2, 0, 1] (sreplicateN @[3, 4] (sfromR (tproject2 (tproject2 (tproject1 u1))))) ; t201 = sgather @[2, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i196, i197, i198] -> [ifH (0.0 <=. negate (splainPart t194 `sindex0` [0, i197, sconcrete (sfromListLinear [2,2] [0,1,2,3]) `sindex0` [i196, i198]])) 0 1]) ; m203 = sreshape @[2, 4] (stranspose @[2, 0, 1] (sreplicateN @[1, 1] (sfromPlain t201 * stranspose @[0, 2, 1] (sgather @[2, 2] (str (sslice (SNat @0) (SNat @2) (t194 !$ [0]))) (\\[i199, i200] -> [sconcrete (sfromListLinear [2,2] [0,1,2,3]) `sindex0` [i199, i200]]))))) ; v205 = sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0] ; v206 = sreshape @[2] (sreplicateN @[1, 1] (sgather1 @2 m203 (\\i204 -> [i204, kargMax (splainPart m203 !$ [i204])]))) ; m207 = sreplicate0N @[1, 5] (sdot0 v205 v206) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; v209 = sgather1 @5 (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\i208 -> [ifH (0.0 <=. negate (splainPart m207 `sindex0` [0, i208])) 0 1]) ; v214 = sfromPlain v209 * sdot1In (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) (str (sfromR dret)) ; x215 = ssum0 @[5] v214 in tpair (let m219 = sappend (str (sscatter @[2, 2] (sfromPlain (stranspose @[0, 2, 1] t201) * stranspose @[1, 2, 0, 4, 3] (sreshape @[2, 1, 1, 2, 2] (sscatter1 @2 (sreshape @[1, 1, 2] (v205 * sreplicate0N @[2] x215) !$ [0, 0]) (\\i216 -> [i216, kargMax (splainPart m203 !$ [i216])]))) !$ [0, 0]) (\\[i217, i218] -> [sconcrete (sfromListLinear [2,2] [0,1,2,3]) `sindex0` [i217, i218]]))) (sconcrete (sreplicate [1,4] 0.0)) in tpair (let m231 = sscatter @[3, 4, 2, 2] (sfromPlain u183 * stranspose @[2, 3, 0, 1] (sreshape @[3, 4, 1, 1, 2, 2] (ssumN @[3, 4, 2, 2] (stranspose @[3, 4, 5, 6, 0, 1, 2] (sscatter @[3, 4] (sscatter @[2, 2, 3, 4] (stranspose @[1, 2, 0] (sreplicate @3 (stranspose @[1, 2, 0] (sreplicate @4 m193))) * sreplicateN @[2, 2] m219) (\\[i221, i222, i223, i224] -> [i223 + i221, i224 + i222, i223, i224, i221, i222])) (\\[i225, i226] -> [i225, i226, kargMax (splainPart t185 !$ [i225, i226])]))))) !$ [0, 0]) (\\[i227, i228, i229, i230] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i227, i229], sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) `sindex0` [i228, i230]]) in tpair (sreplicateN @[1, 1] (ssum @9 (sdot1In (sfromPlain (stranspose @[0, 3, 1, 2] u170)) (stranspose @[3, 0, 1, 2] (sreplicateN @[2, 2] m231))))) (ssumN @[7, 9] (stranspose @[1, 2, 0] (sreplicate @1 m231)))) (tpair (sreplicateN @[1, 1] (ssum @4 (sdot1In (stranspose @[3, 0, 1, 2] u192) (stranspose @[3, 0, 1, 2] (sreplicateN @[2, 2] m219))))) (ssumN @[3, 4] (stranspose @[1, 2, 0] (sreplicate @1 m219))))) (tpair (tpair (sreplicate @1 (v206 * sreplicate0N @[2] x215)) (ssum @5 (str (sreplicate @1 v214)))) (tpair (str (sreplicate @1 (sdot1In (sreplicate @10 (sfromPlain v209 * m207 !$ [0])) (sfromR dret)))) (ssum @5 (str (sfromR dret))))))"  testCNNOAst1 :: Assertion testCNNOAst1 = do@@ -571,7 +567,7 @@                  (knownSTK @(X (MnistCnnRanked2.ADCnnMnistParameters                                   Concrete Double)))                  vals-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan                        (X (MnistCnnRanked2.ADCnnMnistParameters                              Concrete Double))@@ -599,7 +595,7 @@   interpretAstFull @Concrete env afcnn1     @?= afcnn2 valsInit   interpretAstFull @Concrete env-                   (simplifyInline @(TKR 2 Double) afcnn1)+                   (simplifyUserCode @(TKR 2 Double) afcnn1)     @?= afcnn2 valsInit   interpretAstFull @Concrete env                    (simplifyInlineContract @(TKR 2 Double) afcnn1)@@ -634,14 +630,10 @@                  sizeMnistHeightI sizeMnistWidthI batch_size                  (rconcrete $ unConcrete blackGlyph)       artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2 ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> rfromS (let t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i473, i474] -> [i473 + i474]))) (\\[i251, i252] -> [i251 + i252]))))))) * str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0]))))))))))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; u271 = sreshape @[4,7,11,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS (sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) !$ [i259, i261]), kfromS (sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) !$ [i260, i262])])) 0 1]) * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) !$ [i265, i267]), kfromS (sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) !$ [i266, i268])]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i458, i460, i461] -> [kfromS (smaxIndex (u271 !$ [i461, i458, i460])), i460, i458, i461]))) (\\[i466, i468] -> [i466, i468, i466 + i468]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) * str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; u296 = sreshape @[4,3,5,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i284, i286]), kfromS (sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) !$ [i285, i287])])) 0 1]) * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i290, i292]), kfromS (sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) !$ [i291, i293])]))) ; m301 = str (sreplicate @7 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) in smatmul2 (sfromR (tproject1 (tproject2 (tproject2 u1)))) (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) * m301) + str (sreplicate @7 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"-  printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w253 = sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i249, i250] -> [i249 + i250]))) (\\[i251, i252] -> [i251 + i252]))))))) ; w254 = str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0])))))))) ; t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (w253 * w254))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; m256 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; m257 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; w269 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [let x263 = m256 !$ [i259, i261] ; x264 = m257 !$ [i260, i262] in ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS x263, kfromS x264])) 0 1]) ; w270 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (m256 !$ [i265, i267]), kfromS (m257 !$ [i266, i268])])) ; u271 = sreshape @[4,7,11,4] (w269 * w270) ; w279 = sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i272, i273, i274] -> [kfromS (smaxIndex (u271 !$ [i274, i272, i273])), i273, i272, i274]))) (\\[i275, i276] -> [i275, i276, i275 + i276]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (w279 * str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m281 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m282 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; w294 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [let x288 = m281 !$ [i284, i286] ; x289 = m282 !$ [i285, i287] in ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS x288, kfromS x289])) 0 1]) ; w295 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (m281 !$ [i290, i292]), kfromS (m282 !$ [i291, i293])])) ; u296 = sreshape @[4,3,5,4] (w294 * w295) ; m300 = str (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))) ; m301 = str (sreplicate @7 (ssum @60 (str (sfromR (tproject1 (tproject1 (tproject2 u1)))) * m300))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; m304 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) ; t305 = str (sreplicate @10 (m304 * m301)) in rfromS (ssum @5 (stranspose @[2,1,0] (sreplicate @7 (sfromR (tproject1 (tproject2 (tproject2 u1))))) * t305) + str (sreplicate @7 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"-  printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w253 = sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i249, i250] -> [i249 + i250]))) (\\[i251, i252] -> [i251 + i252]))))))) ; w254 = str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0])))))))) ; t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (w253 * w254))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; m256 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; m257 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; w269 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [let x263 = m256 !$ [i259, i261] ; x264 = m257 !$ [i260, i262] in ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS x263, kfromS x264])) 0 1]) ; w270 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (m256 !$ [i265, i267]), kfromS (m257 !$ [i266, i268])])) ; u271 = sreshape @[4,7,11,4] (w269 * w270) ; w279 = sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i272, i273, i274] -> [kfromS (smaxIndex (u271 !$ [i274, i272, i273])), i273, i272, i274]))) (\\[i275, i276] -> [i275, i276, i275 + i276]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (w279 * str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m281 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m282 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; w294 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [let x288 = m281 !$ [i284, i286] ; x289 = m282 !$ [i285, i287] in ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS x288, kfromS x289])) 0 1]) ; w295 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (m281 !$ [i290, i292]), kfromS (m282 !$ [i291, i293])])) ; u296 = sreshape @[4,3,5,4] (w294 * w295) ; m300 = str (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))) ; m301 = str (sreplicate @7 (ssum @60 (str (sfromR (tproject1 (tproject1 (tproject2 u1)))) * m300))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; m304 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) ; t305 = str (sreplicate @10 (m304 * m301)) ; m307 = m304 * ssum @10 (str (stranspose @[2,1,0] (sreplicate @7 (sfromR (tproject1 (tproject2 (tproject2 u1))))) * sreplicate @5 (sfromR dret))) ; m308 = sreplicate @60 (ssum @7 (str m307)) ; t316 = stranspose @[2,0,1] (sscatter (stranspose @[1,2,3,4,0] (w294 * sreshape @[4,3,5,2,2] (sscatter (sreshape @[4,3,5] (ssum @5 (str (str (sfromR (tproject1 (tproject1 (tproject2 u1)))) * m308)))) (\\[i309, i310, i311] -> [i309, i310, i311, kfromS (smaxIndex (u296 !$ [i309, i310, i311]))])))) (\\[i312, i313, i314, i315] -> [kfromS (m281 !$ [i312, i314]), kfromS (m282 !$ [i313, i315])])) ; w317 = sreshape @[4,7,11,4,3,4] (stranspose @[1,2,3,0] (sreplicate @48 t316)) ; t329 = stranspose @[2,0,1] (sscatter (stranspose @[1,2,3,4,0] (w269 * sreshape @[4,7,11,2,2] (stranspose @[3,2,1,0] (ssum @4 (ssum @7 (ssum @11 (stranspose @[4,5,6,3,2,1,0] (ssum @3 (stranspose @[7,3,2,1,0,6,5,4] (sscatter (stranspose @[2,5,4,6,0,3,1] (sscatter (stranspose @[3,5,0,4,1,2] (sscatter (stranspose @[1,4,0,2,3] (ssum @4 (str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1))))))) * w317))) (\\[i318, i319] -> [i318, i318 + i319, i319]))) (\\[i320, i321] -> [i320, i321, i320 + i321]))) (\\[i322, i323, i324] -> [kfromS (smaxIndex (u271 !$ [i324, i322, i323])), i323, i322, i324]))))))))))) (\\[i325, i326, i327, i328] -> [kfromS (m256 !$ [i325, i327]), kfromS (m257 !$ [i326, i328])])) in tpair (tpair (tpair (rfromS (str (soneHot (ssum @1 (str (ssum @1 (str (ssum @23 (str (ssum @14 (str (w253 * sreshape @[4,14,23,1,1,3,4] (stranspose @[1,2,3,0] (sreplicate @12 t329))))))))))) [0]))) (rfromS (ssum @23 (ssum @14 (stranspose @[1,2,0] t329))))) (tpair (rfromS (ssum @11 (str (ssum @7 (str (w279 * w317)))))) (rfromS (ssum @11 (ssum @7 (stranspose @[1,2,0] t316)))))) (tpair (tpair (rfromS (str (m300 * m308))) (rfromS (ssum @7 (str m307)))) (tpair (rfromS (ssum @7 (stranspose @[2,1,0] (t305 * sreplicate @5 (sfromR dret))))) (rfromS (ssum @7 (str (sfromR dret))))))"--- TODO: different test result with GHC 9.10:   printArtifactPretty (simplifyArtifact artifactRev)---    @?= "\\dret u1 -> tfromS (STKProduct (STKProduct (STKProduct (STKR (SNat @4) STKScalar) (STKR (SNat @1) STKScalar)) (STKProduct (STKR (SNat @4) STKScalar) (STKR (SNat @1) STKScalar))) (STKProduct (STKProduct (STKR (SNat @2) STKScalar) (STKR (SNat @1) STKScalar)) (STKProduct (STKR (SNat @2) STKScalar) (STKR (SNat @1) STKScalar)))) (let u247 = sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i397, i398] -> [i397 + i398]))) (\\[i245, i246] -> [i245 + i246]) ; t249 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] u247))))) * str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0]))))))))))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; w259 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i250, i251, i252, i253, i254] -> [ifH (sscalar -0.0 <=. negate (t249 !$ [i250, 2 * i251 + i253, 2 * i252 + i254])) 0 1]) ; u261 = sreshape @[4,7,11,4] (w259 * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t249) (\\[i255, i256, i257, i258] -> [2 * i255 + i257, 2 * i256 + i258]))) ; w269 = sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u261))))))) (\\[i380, i382, i383] -> [kfromS (smaxIndex (u261 !$ [i383, i380, i382])), i382, i380, i383]))) (\\[i388, i390] -> [i388, i390, i388 + i390]))) (\\[i267, i268] -> [i267, i267 + i268, i268]) ; t270 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (sreplicate @4 (stranspose @[2,0,3,4,1] w269) * str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; w280 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i271, i272, i273, i274, i275] -> [ifH (sscalar -0.0 <=. negate (t270 !$ [i271, 2 * i272 + i274, 2 * i273 + i275])) 0 1]) ; u282 = sreshape @[4,3,5,4] (w280 * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t270) (\\[i276, i277, i278, i279] -> [2 * i276 + i278, 2 * i277 + i279]))) ; v286 = sreshape @[60] (sgather u282 (\\[i283, i284, i285] -> [i283, i284, i285, kfromS (smaxIndex (u282 !$ [i283, i284, i285]))])) ; m287 = str (sreplicate @7 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @5 v286))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; m290 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i288, i289] -> [ifH (sscalar -0.0 <=. negate (m287 !$ [i288, i289])) 0 1]) ; m293 = m290 * smatmul2 (str (sfromR (tproject1 (tproject2 (tproject2 u1))))) (sfromR dret) ; v294 = ssum @7 (str m293) ; t302 = sscatter (stranspose @[1,2,3,4,0] w280 * stranspose @[1,2,3,4,0] (sreshape @[4,3,5,2,2] (sscatter (sreshape @[4,3,5] (sdot1In (str (sfromR (tproject1 (tproject1 (tproject2 u1))))) (sreplicate @60 v294))) (\\[i295, i296, i297] -> [i295, i296, i297, kfromS (smaxIndex (u282 !$ [i295, i296, i297]))])))) (\\[i298, i299, i300, i301] -> [2 * i298 + i300, 2 * i299 + i301]) ; w303 = sreshape @[4,7,11,4,3,4] (stranspose @[1,2,3,0] (sreplicate @48 (stranspose @[2,0,1] t302))) ; t315 = sscatter (stranspose @[1,2,3,4,0] w259 * stranspose @[1,2,3,4,0] (sreshape @[4,7,11,2,2] (ssum @4 (ssum @7 (ssum @11 (ssum @3 (stranspose @[7,6,5,4,3,2,1,0] (sscatter (stranspose @[2,5,4,6,0,3,1] (sscatter (stranspose @[3,5,0,4,1,2] (sscatter (sdot1In (stranspose @[2,5,0,3,4,1] (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))) (stranspose @[2,5,1,3,4,0] w303)) (\\[i304, i305] -> [i304, i304 + i305, i305]))) (\\[i306, i307] -> [i306, i307, i306 + i307]))) (\\[i308, i309, i310] -> [kfromS (smaxIndex (u261 !$ [i310, i308, i309])), i309, i308, i310]))))))))) (\\[i311, i312, i313, i314] -> [2 * i311 + i313, 2 * i312 + i314]) in tpair (tpair (tpair (str (sreplicate @1 (ssum @23 (sdot1In (stranspose @[4,0,3,2,1] (sreplicate @4 (stranspose @[2,1,3,0] u247))) (stranspose @[3,4,2,0,5,6,1] (sreshape @[4,14,23,1,1,3,4] (stranspose @[1,2,3,0] (sreplicate @12 (stranspose @[2,0,1] t315)))) !$ [0, 0]))))) (ssum @23 (ssum @14 t315))) (tpair (ssum @11 (sdot1In (stranspose @[2,0,3,4,5,1] (sreplicate @4 (stranspose @[2,0,3,4,1] w269))) (stranspose @[2,0,3,4,5,1] w303))) (ssum @11 (ssum @7 t302)))) (tpair (tpair (sreplicate @5 v286 * str (sreplicate @60 v294)) (ssum @7 (str m293))) (tpair (smatmul2 (sfromR dret) (str m290 * str m287)) (ssum @7 (str (sfromR dret))))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> rfromS (let t236 = ssum @4 (sdot1In (sfromPlain (stranspose @[2, 0, 3, 4, 1] (sreplicate @4 (stranspose @[3, 1, 2, 0] (sgather @[23, 4] (stranspose @[2, 0, 1] (sgather @[14, 3] (sreplicate0N @[14, 23] 7.0) (\\[i434, i435] -> [i434 + i435]))) (\\[i230, i231] -> [i230 + i231])))))) (stranspose @[3, 1, 0, 4, 2] (sreplicate @14 (stranspose @[1, 2, 3, 0] (sreplicate @23 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0])))))) + stranspose @[2, 0, 1] (sreplicateN @[14, 23] (sfromR (tproject2 (tproject1 (tproject1 u1))))) ; u250 = sreshape @[4, 7, 11, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain (sgather @[4, 7, 11, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i239, i240, i241, i242, i243] -> [ifH (0.0 <=. negate (splainPart t236 `sindex0` [i239, sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i240, i242], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i241, i243]])) 0 1])) * stranspose @[4, 0, 1, 2, 3] (sgather @[7, 11, 2, 2] (stranspose @[1, 2, 0] t236) (\\[i244, i245, i246, i247] -> [sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i244, i246], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i245, i247]]))))) ; t259 = ssumN @[3, 4] (sdot1In (stranspose @[2, 3, 0, 4, 5, 1] (sreplicate @4 (stranspose @[4, 2, 3, 0, 1] (sgather @[7, 11, 3, 4] (stranspose @[3, 4, 5, 6, 0, 1, 2] (sgather @[7, 11, 4] (stranspose @[4, 3, 2, 1, 7, 6, 5, 0] (sreplicate @4 (stranspose @[3, 4, 5, 6, 0, 1, 2] (sreplicateN @[3, 11, 7] (stranspose @[3, 2, 1, 0] u250))))) (\\[i427, i428, i429] -> [i429, i427, i428, kargMax (splainPart u250 !$ [i429, i427, i428])]))) (\\[i254, i255, i256, i257] -> [i254, i255, i256, i257, i254 + i256, i255 + i257]))))) (stranspose @[3, 4, 1, 0, 5, 2] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2, 0, 1] (sreplicateN @[7, 11] (sfromR (tproject2 (tproject2 (tproject1 u1))))) ; u273 = sreshape @[4, 3, 5, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain (sgather @[4, 3, 5, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i262, i263, i264, i265, i266] -> [ifH (0.0 <=. negate (splainPart t259 `sindex0` [i262, sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i263, i265], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i264, i266]])) 0 1])) * stranspose @[4, 0, 1, 2, 3] (sgather @[3, 5, 2, 2] (stranspose @[1, 2, 0] t259) (\\[i267, i268, i269, i270] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i267, i269], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i268, i270]]))))) ; m278 = str (sreplicate @7 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @5 (sreshape @[60] (sgather @[4, 3, 5] u273 (\\[i274, i275, i276] -> [i274, i275, i276, kargMax (splainPart u273 !$ [i274, i275, i276])])))))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) in smatmul2 (sfromR (tproject1 (tproject2 (tproject2 u1)))) (sfromPlain (sgather @[5, 7] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i279, i280] -> [ifH (0.0 <=. negate (splainPart m278 `sindex0` [i279, i280])) 0 1])) * m278) + str (sreplicate @7 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKS [4,1,3,4] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [4,4,3,4] STKScalar) (STKS [4] STKScalar))) (STKProduct (STKProduct (STKS [5,60] STKScalar) (STKS [5] STKScalar)) (STKProduct (STKS [10,5] STKScalar) (STKS [10] STKScalar)))) (let u234 = sgather @[23, 4] (stranspose @[2, 0, 1] (sgather @[14, 3] (sreplicate0N @[14, 23] 7.0) (\\[i516, i517] -> [i516 + i517]))) (\\[i230, i231] -> [i230 + i231]) ; t236 = ssum @4 (sdot1In (sfromPlain (stranspose @[2, 0, 3, 4, 1] (sreplicate @4 (stranspose @[3, 1, 2, 0] u234)))) (stranspose @[3, 1, 0, 4, 2] (sreplicate @14 (stranspose @[1, 2, 3, 0] (sreplicate @23 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0])))))) + stranspose @[2, 0, 1] (sreplicateN @[14, 23] (sfromR (tproject2 (tproject1 (tproject1 u1))))) ; w248 = sgather @[4, 7, 11, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i239, i240, i241, i242, i243] -> [ifH (0.0 <=. negate (splainPart t236 `sindex0` [i239, sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i240, i242], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i241, i243]])) 0 1]) ; u250 = sreshape @[4, 7, 11, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain w248 * stranspose @[4, 0, 1, 2, 3] (sgather @[7, 11, 2, 2] (stranspose @[1, 2, 0] t236) (\\[i244, i245, i246, i247] -> [sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i244, i246], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i245, i247]]))))) ; w258 = sgather @[7, 11, 3, 4] (stranspose @[3, 4, 5, 6, 0, 1, 2] (sgather @[7, 11, 4] (stranspose @[4, 3, 2, 1, 7, 6, 5, 0] (sreplicate @4 (stranspose @[3, 4, 5, 6, 0, 1, 2] (sreplicateN @[3, 11, 7] (stranspose @[3, 2, 1, 0] u250))))) (\\[i507, i508, i509] -> [i509, i507, i508, kargMax (splainPart u250 !$ [i509, i507, i508])]))) (\\[i254, i255, i256, i257] -> [i254, i255, i256, i257, i254 + i256, i255 + i257]) ; t259 = ssumN @[3, 4] (sdot1In (stranspose @[2, 3, 0, 4, 5, 1] (sreplicate @4 (stranspose @[4, 2, 3, 0, 1] w258))) (stranspose @[3, 4, 1, 0, 5, 2] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2, 0, 1] (sreplicateN @[7, 11] (sfromR (tproject2 (tproject2 (tproject1 u1))))) ; w271 = sgather @[4, 3, 5, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i262, i263, i264, i265, i266] -> [ifH (0.0 <=. negate (splainPart t259 `sindex0` [i262, sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i263, i265], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i264, i266]])) 0 1]) ; u273 = sreshape @[4, 3, 5, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain w271 * stranspose @[4, 0, 1, 2, 3] (sgather @[3, 5, 2, 2] (stranspose @[1, 2, 0] t259) (\\[i267, i268, i269, i270] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i267, i269], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i268, i270]]))))) ; v277 = sreshape @[60] (sgather @[4, 3, 5] u273 (\\[i274, i275, i276] -> [i274, i275, i276, kargMax (splainPart u273 !$ [i274, i275, i276])])) ; m278 = str (sreplicate @7 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @5 v277))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; m281 = sgather @[5, 7] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i279, i280] -> [ifH (0.0 <=. negate (splainPart m278 `sindex0` [i279, i280])) 0 1]) ; m284 = sfromPlain m281 * smatmul2 (str (sfromR (tproject1 (tproject2 (tproject2 u1))))) (sfromR dret) ; v285 = ssum @7 (str m284) in tpair (let t293 = sscatter @[3, 5, 2, 2] (sfromPlain (stranspose @[1, 2, 3, 4, 0] w271) * stranspose @[3, 4, 1, 2, 5, 6, 0] (sreshape @[4, 3, 5, 1, 1, 2, 2] (sscatter @[4, 3, 5] (sreshape @[4, 3, 5] (sdot1In (str (sfromR (tproject1 (tproject1 (tproject2 u1))))) (sreplicate @60 v285))) (\\[i286, i287, i288] -> [i286, i287, i288, kargMax (splainPart u273 !$ [i286, i287, i288])]))) !$ [0, 0]) (\\[i289, i290, i291, i292] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i289, i291], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i290, i292]]) in tpair (let t305 = sscatter @[7, 11, 2, 2] (sfromPlain (stranspose @[1, 2, 3, 4, 0] w248) * stranspose @[3, 4, 1, 2, 5, 6, 0] (sreshape @[4, 7, 11, 1, 1, 2, 2] (ssumN @[4, 3, 11, 7] (stranspose @[7, 6, 5, 4, 0, 1, 2, 3] (sscatter @[7, 11, 4] (stranspose @[4, 5, 6, 0, 1, 2, 3] (sscatter @[7, 11, 3, 4] (sdot1In (sreplicateN @[7, 11] (stranspose @[2, 3, 1, 0] (sfromR (tproject1 (tproject2 (tproject1 u1)))))) (stranspose @[4, 5, 1, 2, 0, 3] (sreplicateN @[4, 3, 4] (stranspose @[2, 0, 1] t293)))) (\\[i294, i295, i296, i297] -> [i294, i295, i296, i297, i294 + i296, i295 + i297]))) (\\[i298, i299, i300] -> [i300, i298, i299, kargMax (splainPart u250 !$ [i300, i298, i299])]))))) !$ [0, 0]) (\\[i301, i302, i303, i304] -> [sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i301, i303], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i302, i304]]) in tpair (str (sreplicate @1 (ssum @23 (sdot1In (sfromPlain (stranspose @[4, 0, 1, 2, 3] (sreplicate @4 (stranspose @[3, 1, 2, 0] u234)))) (stranspose @[4, 2, 0, 1, 3] (sreplicateN @[3, 4] (stranspose @[2, 0, 1] t305))))))) (ssumN @[14, 23] t305)) (tpair (ssum @11 (sdot1In (stranspose @[5, 0, 1, 2, 3, 4] (sreplicate @4 (stranspose @[4, 2, 3, 0, 1] w258))) (stranspose @[5, 3, 0, 1, 2, 4] (sreplicateN @[4, 3, 4] (stranspose @[2, 0, 1] t293))))) (ssumN @[7, 11] t293))) (tpair (tpair (sreplicate @5 v277 * str (sreplicate @60 v285)) (ssum @7 (str m284))) (tpair (smatmul2 (sfromR dret) (sfromPlain (str m281) * str m278)) (ssum @7 (str (sfromR dret))))))"  testCNNOAst2 :: Assertion testCNNOAst2 = do@@ -652,7 +644,7 @@                  (knownSTK @(X (MnistCnnRanked2.ADCnnMnistParameters                                   Concrete Double)))                  vals-      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      varName = mkAstVarName ftk . intToAstVarId $ 100000000       var :: AstTensor AstMethodLet FullSpan                        (X (MnistCnnRanked2.ADCnnMnistParameters                              Concrete Double))@@ -677,14 +669,17 @@                  sizeMnistHeightI sizeMnistWidthI batch_size                  (rconcrete $ unConcrete blackGlyph)       afcnn1 = afcnn2 $ fromTarget var-  interpretAstFull @Concrete env afcnn1-    @?= afcnn2 valsInit-  interpretAstFull @Concrete env-                   (simplifyInline @(TKR 2 Double) afcnn1)-    @?= afcnn2 valsInit-  interpretAstFull @Concrete env-                   (simplifyInlineContract @(TKR 2 Double) afcnn1)-    @?= afcnn2 valsInit+  assertEqualUpToEpsilon 1e-5+    (interpretAstFull @Concrete env afcnn1)+    (afcnn2 valsInit)+  assertEqualUpToEpsilon 1e-5+    (interpretAstFull @Concrete env+                      (simplifyUserCode @(TKR 2 Double) afcnn1))+    (afcnn2 valsInit)+  assertEqualUpToEpsilon 1e-5+    (interpretAstFull @Concrete env+                      (simplifyInlineContract @(TKR 2 Double) afcnn1))+    (afcnn2 valsInit)  testCNNOPP2S :: Assertion testCNNOPP2S = do@@ -716,11 +711,11 @@                  (SNat @4) (SNat @5) batch_size                  (sconcrete $ unConcrete blackGlyph)       artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2 ftk-  printArtifactPrimalPretty (simplifyArtifact artifactRev)-    @?= "\\u1 -> let t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i473, i474] -> [i473 + i474]))) (\\[i251, i252] -> [i251 + i252]))))))) * str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0]))))))))))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (tproject2 (tproject1 (tproject1 u1))))) ; u271 = sreshape @[4,7,11,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS (sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) !$ [i259, i261]), kfromS (sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) !$ [i260, i262])])) 0 1]) * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) !$ [i265, i267]), kfromS (sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) !$ [i266, i268])]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i458, i460, i461] -> [kfromS (smaxIndex (u271 !$ [i461, i458, i460])), i460, i458, i461]))) (\\[i466, i468] -> [i466, i468, i466 + i468]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) * str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (tproject2 (tproject2 (tproject1 u1))))) ; u296 = sreshape @[4,3,5,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i284, i286]), kfromS (sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) !$ [i285, i287])])) 0 1]) * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i290, i292]), kfromS (sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) !$ [i291, i293])]))) ; m301 = str (sreplicate @7 (sdot1In (tproject1 (tproject1 (tproject2 u1))) (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) in smatmul2 (tproject1 (tproject2 (tproject2 u1))) (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) * m301) + str (sreplicate @7 (tproject2 (tproject2 (tproject2 u1))))"+  printArtifactPrimalPretty (simplifyArtifactRev artifactRev)+    @?= "\\u1 -> let t236 = ssum @4 (sdot1In (sfromPlain (stranspose @[2, 0, 3, 4, 1] (sreplicate @4 (stranspose @[3, 1, 2, 0] (sgather @[23, 4] (stranspose @[2, 0, 1] (sgather @[14, 3] (sreplicate0N @[14, 23] 7.0) (\\[i434, i435] -> [i434 + i435]))) (\\[i230, i231] -> [i230 + i231])))))) (stranspose @[3, 1, 0, 4, 2] (sreplicate @14 (stranspose @[1, 2, 3, 0] (sreplicate @23 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0])))))) + stranspose @[2, 0, 1] (sreplicateN @[14, 23] (tproject2 (tproject1 (tproject1 u1)))) ; u250 = sreshape @[4, 7, 11, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain (sgather @[4, 7, 11, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i239, i240, i241, i242, i243] -> [ifH (0.0 <=. negate (splainPart t236 `sindex0` [i239, sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i240, i242], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i241, i243]])) 0 1])) * stranspose @[4, 0, 1, 2, 3] (sgather @[7, 11, 2, 2] (stranspose @[1, 2, 0] t236) (\\[i244, i245, i246, i247] -> [sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i244, i246], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i245, i247]]))))) ; t259 = ssumN @[3, 4] (sdot1In (stranspose @[2, 3, 0, 4, 5, 1] (sreplicate @4 (stranspose @[4, 2, 3, 0, 1] (sgather @[7, 11, 3, 4] (stranspose @[3, 4, 5, 6, 0, 1, 2] (sgather @[7, 11, 4] (stranspose @[4, 3, 2, 1, 7, 6, 5, 0] (sreplicate @4 (stranspose @[3, 4, 5, 6, 0, 1, 2] (sreplicateN @[3, 11, 7] (stranspose @[3, 2, 1, 0] u250))))) (\\[i427, i428, i429] -> [i429, i427, i428, kargMax (splainPart u250 !$ [i429, i427, i428])]))) (\\[i254, i255, i256, i257] -> [i254, i255, i256, i257, i254 + i256, i255 + i257]))))) (stranspose @[3, 4, 1, 0, 5, 2] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @11 (tproject1 (tproject2 (tproject1 u1)))))))) + stranspose @[2, 0, 1] (sreplicateN @[7, 11] (tproject2 (tproject2 (tproject1 u1)))) ; u273 = sreshape @[4, 3, 5, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain (sgather @[4, 3, 5, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i262, i263, i264, i265, i266] -> [ifH (0.0 <=. negate (splainPart t259 `sindex0` [i262, sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i263, i265], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i264, i266]])) 0 1])) * stranspose @[4, 0, 1, 2, 3] (sgather @[3, 5, 2, 2] (stranspose @[1, 2, 0] t259) (\\[i267, i268, i269, i270] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i267, i269], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i268, i270]]))))) ; m278 = str (sreplicate @7 (sdot1In (tproject1 (tproject1 (tproject2 u1))) (sreplicate @5 (sreshape @[60] (sgather @[4, 3, 5] u273 (\\[i274, i275, i276] -> [i274, i275, i276, kargMax (splainPart u273 !$ [i274, i275, i276])])))))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) in smatmul2 (tproject1 (tproject2 (tproject2 u1))) (sfromPlain (sgather @[5, 7] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i279, i280] -> [ifH (0.0 <=. negate (splainPart m278 `sindex0` [i279, i280])) 0 1])) * m278) + str (sreplicate @7 (tproject2 (tproject2 (tproject2 u1))))"   printArtifactPrimalPretty artifactRev-    @?= "\\u1 -> let w253 = sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i249, i250] -> [i249 + i250]))) (\\[i251, i252] -> [i251 + i252]))))))) ; w254 = str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0])))))))) ; t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (w253 * w254))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (tproject2 (tproject1 (tproject1 u1))))) ; m256 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; m257 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; w269 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [let x263 = m256 !$ [i259, i261] ; x264 = m257 !$ [i260, i262] in ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS x263, kfromS x264])) 0 1]) ; w270 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (m256 !$ [i265, i267]), kfromS (m257 !$ [i266, i268])])) ; u271 = sreshape @[4,7,11,4] (w269 * w270) ; w279 = sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i272, i273, i274] -> [kfromS (smaxIndex (u271 !$ [i274, i272, i273])), i273, i272, i274]))) (\\[i275, i276] -> [i275, i276, i275 + i276]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (w279 * str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (tproject2 (tproject2 (tproject1 u1))))) ; m281 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m282 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; w294 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [let x288 = m281 !$ [i284, i286] ; x289 = m282 !$ [i285, i287] in ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS x288, kfromS x289])) 0 1]) ; w295 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (m281 !$ [i290, i292]), kfromS (m282 !$ [i291, i293])])) ; u296 = sreshape @[4,3,5,4] (w294 * w295) ; m300 = str (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))) ; m301 = str (sreplicate @7 (ssum @60 (str (tproject1 (tproject1 (tproject2 u1))) * m300))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m304 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) ; t305 = str (sreplicate @10 (m304 * m301)) in ssum @5 (stranspose @[2,1,0] (sreplicate @7 (tproject1 (tproject2 (tproject2 u1)))) * t305) + str (sreplicate @7 (tproject2 (tproject2 (tproject2 u1))))"+    @?= "\\u1 -> let w234 = stranspose @[1, 2, 0] (sreplicate @4 (stranspose @[3, 1, 2, 0] (sgather @[23, 4] (stranspose @[2, 0, 1] (sgather @[14, 3] (sreplicate0N @[14, 23] 7.0) (\\[i232, i233] -> [i232 + i233]))) (\\[i230, i231] -> [i230 + i231])))) ; w235 = stranspose @[2, 3, 1, 0] (sreplicate @14 (stranspose @[1, 2, 3, 0] (sreplicate @23 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0])))) ; t236 = ssumN @[3, 4] (sfromPlain w234 * w235) + stranspose @[2, 0, 1] (sreplicateN @[14, 23] (tproject2 (tproject1 (tproject1 u1)))) ; m237 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; m238 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; w248 = sgather @[4, 7, 11, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i239, i240, i241, i242, i243] -> [ifH (0.0 <=. negate (splainPart t236 `sindex0` [i239, m238 `sindex0` [i240, i242], m237 `sindex0` [i241, i243]])) 0 1]) ; w249 = stranspose @[4, 0, 1, 2, 3] (sgather @[7, 11, 2, 2] (stranspose @[1, 2, 0] t236) (\\[i244, i245, i246, i247] -> [m238 `sindex0` [i244, i246], m237 `sindex0` [i245, i247]])) ; u250 = sreshape @[4, 7, 11, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain w248 * w249))) ; w258 = stranspose @[1, 2, 3, 0] (sreplicate @4 (stranspose @[4, 2, 3, 0, 1] (sgather @[7, 11, 3, 4] (stranspose @[3, 4, 5, 6, 0, 1, 2] (sgather @[7, 11, 4] (stranspose @[4, 3, 2, 1, 7, 6, 5, 0] (sreplicate @4 (stranspose @[3, 4, 5, 6, 0, 1, 2] (sreplicateN @[3, 11, 7] (stranspose @[3, 2, 1, 0] u250))))) (\\[i251, i252, i253] -> [i253, i251, i252, kargMax (splainPart u250 !$ [i253, i251, i252])]))) (\\[i254, i255, i256, i257] -> [i254, i255, i256, i257, i254 + i256, i255 + i257])))) ; t259 = ssumN @[4, 3, 4] (w258 * stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))) + stranspose @[2, 0, 1] (sreplicateN @[7, 11] (tproject2 (tproject2 (tproject1 u1)))) ; m260 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; m261 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; w271 = sgather @[4, 3, 5, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i262, i263, i264, i265, i266] -> [ifH (0.0 <=. negate (splainPart t259 `sindex0` [i262, m261 `sindex0` [i263, i265], m260 `sindex0` [i264, i266]])) 0 1]) ; w272 = stranspose @[4, 0, 1, 2, 3] (sgather @[3, 5, 2, 2] (stranspose @[1, 2, 0] t259) (\\[i267, i268, i269, i270] -> [m261 `sindex0` [i267, i269], m260 `sindex0` [i268, i270]])) ; u273 = sreshape @[4, 3, 5, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain w271 * w272))) ; m277 = str (sreplicate @5 (sreshape @[60] (sgather @[4, 3, 5] u273 (\\[i274, i275, i276] -> [i274, i275, i276, kargMax (splainPart u273 !$ [i274, i275, i276])])))) ; m278 = str (sreplicate @7 (ssum @60 (str (tproject1 (tproject1 (tproject2 u1))) * m277))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m281 = sgather @[5, 7] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i279, i280] -> [ifH (0.0 <=. negate (splainPart m278 `sindex0` [i279, i280])) 0 1]) ; t282 = str (sreplicate @10 (sfromPlain m281 * m278)) in ssum @5 (stranspose @[2, 1, 0] (sreplicate @7 (tproject1 (tproject2 (tproject2 u1)))) * t282) + str (sreplicate @7 (tproject2 (tproject2 (tproject2 u1))))"   printArtifactPretty artifactRev-    @?= "\\dret u1 -> let w253 = sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i249, i250] -> [i249 + i250]))) (\\[i251, i252] -> [i251 + i252]))))))) ; w254 = str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0])))))))) ; t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (w253 * w254))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (tproject2 (tproject1 (tproject1 u1))))) ; m256 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; m257 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; w269 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [let x263 = m256 !$ [i259, i261] ; x264 = m257 !$ [i260, i262] in ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS x263, kfromS x264])) 0 1]) ; w270 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (m256 !$ [i265, i267]), kfromS (m257 !$ [i266, i268])])) ; u271 = sreshape @[4,7,11,4] (w269 * w270) ; w279 = sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i272, i273, i274] -> [kfromS (smaxIndex (u271 !$ [i274, i272, i273])), i273, i272, i274]))) (\\[i275, i276] -> [i275, i276, i275 + i276]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (w279 * str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (tproject2 (tproject2 (tproject1 u1))))) ; m281 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m282 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; w294 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [let x288 = m281 !$ [i284, i286] ; x289 = m282 !$ [i285, i287] in ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS x288, kfromS x289])) 0 1]) ; w295 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (m281 !$ [i290, i292]), kfromS (m282 !$ [i291, i293])])) ; u296 = sreshape @[4,3,5,4] (w294 * w295) ; m300 = str (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))) ; m301 = str (sreplicate @7 (ssum @60 (str (tproject1 (tproject1 (tproject2 u1))) * m300))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m304 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) ; t305 = str (sreplicate @10 (m304 * m301)) ; m307 = m304 * ssum @10 (str (stranspose @[2,1,0] (sreplicate @7 (tproject1 (tproject2 (tproject2 u1)))) * sreplicate @5 dret)) ; m308 = sreplicate @60 (ssum @7 (str m307)) ; t316 = stranspose @[2,0,1] (sscatter (stranspose @[1,2,3,4,0] (w294 * sreshape @[4,3,5,2,2] (sscatter (sreshape @[4,3,5] (ssum @5 (str (str (tproject1 (tproject1 (tproject2 u1))) * m308)))) (\\[i309, i310, i311] -> [i309, i310, i311, kfromS (smaxIndex (u296 !$ [i309, i310, i311]))])))) (\\[i312, i313, i314, i315] -> [kfromS (m281 !$ [i312, i314]), kfromS (m282 !$ [i313, i315])])) ; w317 = sreshape @[4,7,11,4,3,4] (stranspose @[1,2,3,0] (sreplicate @48 t316)) ; t329 = stranspose @[2,0,1] (sscatter (stranspose @[1,2,3,4,0] (w269 * sreshape @[4,7,11,2,2] (stranspose @[3,2,1,0] (ssum @4 (ssum @7 (ssum @11 (stranspose @[4,5,6,3,2,1,0] (ssum @3 (stranspose @[7,3,2,1,0,6,5,4] (sscatter (stranspose @[2,5,4,6,0,3,1] (sscatter (stranspose @[3,5,0,4,1,2] (sscatter (stranspose @[1,4,0,2,3] (ssum @4 (str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1)))))) * w317))) (\\[i318, i319] -> [i318, i318 + i319, i319]))) (\\[i320, i321] -> [i320, i321, i320 + i321]))) (\\[i322, i323, i324] -> [kfromS (smaxIndex (u271 !$ [i324, i322, i323])), i323, i322, i324]))))))))))) (\\[i325, i326, i327, i328] -> [kfromS (m256 !$ [i325, i327]), kfromS (m257 !$ [i326, i328])])) in tpair (tpair (tpair (str (soneHot (ssum @1 (str (ssum @1 (str (ssum @23 (str (ssum @14 (str (w253 * sreshape @[4,14,23,1,1,3,4] (stranspose @[1,2,3,0] (sreplicate @12 t329))))))))))) [0])) (ssum @23 (ssum @14 (stranspose @[1,2,0] t329)))) (tpair (ssum @11 (str (ssum @7 (str (w279 * w317))))) (ssum @11 (ssum @7 (stranspose @[1,2,0] t316))))) (tpair (tpair (str (m300 * m308)) (ssum @7 (str m307))) (tpair (ssum @7 (stranspose @[2,1,0] (t305 * sreplicate @5 dret))) (ssum @7 (str dret))))"--- TODO: different test result with GHC 9.10: printArtifactPretty (simplifyArtifact artifactRev)---    @?= "\\dret u1 -> let u247 = sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i397, i398] -> [i397 + i398]))) (\\[i245, i246] -> [i245 + i246]) ; t249 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] u247))))) * str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0]))))))))))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (tproject2 (tproject1 (tproject1 u1))))) ; w259 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i250, i251, i252, i253, i254] -> [ifH (sscalar -0.0 <=. negate (t249 !$ [i250, 2 * i251 + i253, 2 * i252 + i254])) 0 1]) ; u261 = sreshape @[4,7,11,4] (w259 * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t249) (\\[i255, i256, i257, i258] -> [2 * i255 + i257, 2 * i256 + i258]))) ; w269 = sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u261))))))) (\\[i380, i382, i383] -> [kfromS (smaxIndex (u261 !$ [i383, i380, i382])), i382, i380, i383]))) (\\[i388, i390] -> [i388, i390, i388 + i390]))) (\\[i267, i268] -> [i267, i267 + i268, i268]) ; t270 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (sreplicate @4 (stranspose @[2,0,3,4,1] w269) * str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (tproject2 (tproject2 (tproject1 u1))))) ; w280 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i271, i272, i273, i274, i275] -> [ifH (sscalar -0.0 <=. negate (t270 !$ [i271, 2 * i272 + i274, 2 * i273 + i275])) 0 1]) ; u282 = sreshape @[4,3,5,4] (w280 * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t270) (\\[i276, i277, i278, i279] -> [2 * i276 + i278, 2 * i277 + i279]))) ; v286 = sreshape @[60] (sgather u282 (\\[i283, i284, i285] -> [i283, i284, i285, kfromS (smaxIndex (u282 !$ [i283, i284, i285]))])) ; m287 = str (sreplicate @7 (sdot1In (tproject1 (tproject1 (tproject2 u1))) (sreplicate @5 v286))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m290 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i288, i289] -> [ifH (sscalar -0.0 <=. negate (m287 !$ [i288, i289])) 0 1]) ; m293 = m290 * smatmul2 (str (tproject1 (tproject2 (tproject2 u1)))) dret ; v294 = ssum @7 (str m293) ; t302 = sscatter (stranspose @[1,2,3,4,0] w280 * stranspose @[1,2,3,4,0] (sreshape @[4,3,5,2,2] (sscatter (sreshape @[4,3,5] (sdot1In (str (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @60 v294))) (\\[i295, i296, i297] -> [i295, i296, i297, kfromS (smaxIndex (u282 !$ [i295, i296, i297]))])))) (\\[i298, i299, i300, i301] -> [2 * i298 + i300, 2 * i299 + i301]) ; w303 = sreshape @[4,7,11,4,3,4] (stranspose @[1,2,3,0] (sreplicate @48 (stranspose @[2,0,1] t302))) ; t315 = sscatter (stranspose @[1,2,3,4,0] w259 * stranspose @[1,2,3,4,0] (sreshape @[4,7,11,2,2] (ssum @4 (ssum @7 (ssum @11 (ssum @3 (stranspose @[7,6,5,4,3,2,1,0] (sscatter (stranspose @[2,5,4,6,0,3,1] (sscatter (stranspose @[3,5,0,4,1,2] (sscatter (sdot1In (stranspose @[2,5,0,3,4,1] (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))) (stranspose @[2,5,1,3,4,0] w303)) (\\[i304, i305] -> [i304, i304 + i305, i305]))) (\\[i306, i307] -> [i306, i307, i306 + i307]))) (\\[i308, i309, i310] -> [kfromS (smaxIndex (u261 !$ [i310, i308, i309])), i309, i308, i310]))))))))) (\\[i311, i312, i313, i314] -> [2 * i311 + i313, 2 * i312 + i314]) in tpair (tpair (tpair (str (sreplicate @1 (ssum @23 (sdot1In (stranspose @[4,0,3,2,1] (sreplicate @4 (stranspose @[2,1,3,0] u247))) (stranspose @[3,4,2,0,5,6,1] (sreshape @[4,14,23,1,1,3,4] (stranspose @[1,2,3,0] (sreplicate @12 (stranspose @[2,0,1] t315)))) !$ [0, 0]))))) (ssum @23 (ssum @14 t315))) (tpair (ssum @11 (sdot1In (stranspose @[2,0,3,4,5,1] (sreplicate @4 (stranspose @[2,0,3,4,1] w269))) (stranspose @[2,0,3,4,5,1] w303))) (ssum @11 (ssum @7 t302)))) (tpair (tpair (sreplicate @5 v286 * str (sreplicate @60 v294)) (ssum @7 (str m293))) (tpair (smatmul2 dret (str m290 * str m287)) (ssum @7 (str dret))))"+    @?= "\\dret u1 -> let w234 = stranspose @[1, 2, 0] (sreplicate @4 (stranspose @[3, 1, 2, 0] (sgather @[23, 4] (stranspose @[2, 0, 1] (sgather @[14, 3] (sreplicate0N @[14, 23] 7.0) (\\[i232, i233] -> [i232 + i233]))) (\\[i230, i231] -> [i230 + i231])))) ; w235 = stranspose @[2, 3, 1, 0] (sreplicate @14 (stranspose @[1, 2, 3, 0] (sreplicate @23 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0])))) ; t236 = ssumN @[3, 4] (sfromPlain w234 * w235) + stranspose @[2, 0, 1] (sreplicateN @[14, 23] (tproject2 (tproject1 (tproject1 u1)))) ; m237 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; m238 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; w248 = sgather @[4, 7, 11, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i239, i240, i241, i242, i243] -> [ifH (0.0 <=. negate (splainPart t236 `sindex0` [i239, m238 `sindex0` [i240, i242], m237 `sindex0` [i241, i243]])) 0 1]) ; w249 = stranspose @[4, 0, 1, 2, 3] (sgather @[7, 11, 2, 2] (stranspose @[1, 2, 0] t236) (\\[i244, i245, i246, i247] -> [m238 `sindex0` [i244, i246], m237 `sindex0` [i245, i247]])) ; u250 = sreshape @[4, 7, 11, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain w248 * w249))) ; w258 = stranspose @[1, 2, 3, 0] (sreplicate @4 (stranspose @[4, 2, 3, 0, 1] (sgather @[7, 11, 3, 4] (stranspose @[3, 4, 5, 6, 0, 1, 2] (sgather @[7, 11, 4] (stranspose @[4, 3, 2, 1, 7, 6, 5, 0] (sreplicate @4 (stranspose @[3, 4, 5, 6, 0, 1, 2] (sreplicateN @[3, 11, 7] (stranspose @[3, 2, 1, 0] u250))))) (\\[i251, i252, i253] -> [i253, i251, i252, kargMax (splainPart u250 !$ [i253, i251, i252])]))) (\\[i254, i255, i256, i257] -> [i254, i255, i256, i257, i254 + i256, i255 + i257])))) ; t259 = ssumN @[4, 3, 4] (w258 * stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))) + stranspose @[2, 0, 1] (sreplicateN @[7, 11] (tproject2 (tproject2 (tproject1 u1)))) ; m260 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; m261 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; w271 = sgather @[4, 3, 5, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i262, i263, i264, i265, i266] -> [ifH (0.0 <=. negate (splainPart t259 `sindex0` [i262, m261 `sindex0` [i263, i265], m260 `sindex0` [i264, i266]])) 0 1]) ; w272 = stranspose @[4, 0, 1, 2, 3] (sgather @[3, 5, 2, 2] (stranspose @[1, 2, 0] t259) (\\[i267, i268, i269, i270] -> [m261 `sindex0` [i267, i269], m260 `sindex0` [i268, i270]])) ; u273 = sreshape @[4, 3, 5, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain w271 * w272))) ; m277 = str (sreplicate @5 (sreshape @[60] (sgather @[4, 3, 5] u273 (\\[i274, i275, i276] -> [i274, i275, i276, kargMax (splainPart u273 !$ [i274, i275, i276])])))) ; m278 = str (sreplicate @7 (ssum @60 (str (tproject1 (tproject1 (tproject2 u1))) * m277))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m281 = sgather @[5, 7] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i279, i280] -> [ifH (0.0 <=. negate (splainPart m278 `sindex0` [i279, i280])) 0 1]) ; t282 = str (sreplicate @10 (sfromPlain m281 * m278)) ; m284 = sfromPlain m281 * ssum @10 (str (stranspose @[2, 1, 0] (sreplicate @7 (tproject1 (tproject2 (tproject2 u1)))) * sreplicate @5 dret)) ; m285 = sreplicate @60 (ssum @7 (str m284)) in tpair (let t293 = stranspose @[2, 0, 1] (sscatter @[3, 5, 2, 2] (stranspose @[1, 2, 3, 4, 0] (sfromPlain w271 * ssumN @[1, 1] (stranspose @[3, 4, 0, 1, 2] (sreshape @[4, 3, 5, 1, 1, 2, 2] (sscatter @[4, 3, 5] (sreshape @[4, 3, 5] (ssum @5 (str (str (tproject1 (tproject1 (tproject2 u1))) * m285)))) (\\[i286, i287, i288] -> [i286, i287, i288, kargMax (splainPart u273 !$ [i286, i287, i288])])))))) (\\[i289, i290, i291, i292] -> [m261 `sindex0` [i289, i291], m260 `sindex0` [i290, i292]])) in tpair (let t305 = stranspose @[2, 0, 1] (sscatter @[7, 11, 2, 2] (stranspose @[1, 2, 3, 4, 0] (sfromPlain w248 * ssumN @[1, 1] (stranspose @[3, 4, 0, 1, 2] (sreshape @[4, 7, 11, 1, 1, 2, 2] (stranspose @[3, 2, 1, 0] (ssumN @[3, 11, 7] (stranspose @[4, 5, 6, 0, 1, 2, 3] (ssum @4 (stranspose @[7, 3, 2, 1, 0, 6, 5, 4] (sscatter @[7, 11, 4] (stranspose @[4, 5, 6, 0, 1, 2, 3] (sscatter @[7, 11, 3, 4] (stranspose @[3, 4, 1, 2, 0] (ssum @4 (stranspose @[3, 0, 1, 2] (stranspose @[2, 3, 4, 1, 0] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @11 (tproject1 (tproject2 (tproject1 u1)))))) * sreplicateN @[4, 3, 4] t293)))) (\\[i294, i295, i296, i297] -> [i294, i295, i296, i297, i294 + i296, i295 + i297]))) (\\[i298, i299, i300] -> [i300, i298, i299, kargMax (splainPart u250 !$ [i300, i298, i299])]))))))))))) (\\[i301, i302, i303, i304] -> [m238 `sindex0` [i301, i303], m237 `sindex0` [i302, i304]])) in tpair (str (soneHot (ssum @23 (stranspose @[3, 0, 1, 2] (ssum @14 (stranspose @[3, 2, 0, 1] (sfromPlain w234 * sreplicateN @[3, 4] t305))))) [0])) (ssumN @[14, 23] (stranspose @[1, 2, 0] t305))) (tpair (ssum @11 (stranspose @[4, 0, 1, 2, 3] (ssum @7 (stranspose @[4, 3, 0, 1, 2] (w258 * sreplicateN @[4, 3, 4] t293))))) (ssumN @[7, 11] (stranspose @[1, 2, 0] t293)))) (tpair (tpair (str (m277 * m285)) (ssum @7 (str m284))) (tpair (ssum @7 (stranspose @[2, 1, 0] (t282 * sreplicate @5 dret))) (ssum @7 (str dret))))"+  printArtifactPretty (simplifyArtifactRev artifactRev)+    @?= "\\dret u1 -> let u234 = sgather @[23, 4] (stranspose @[2, 0, 1] (sgather @[14, 3] (sreplicate0N @[14, 23] 7.0) (\\[i516, i517] -> [i516 + i517]))) (\\[i230, i231] -> [i230 + i231]) ; t236 = ssum @4 (sdot1In (sfromPlain (stranspose @[2, 0, 3, 4, 1] (sreplicate @4 (stranspose @[3, 1, 2, 0] u234)))) (stranspose @[3, 1, 0, 4, 2] (sreplicate @14 (stranspose @[1, 2, 3, 0] (sreplicate @23 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0])))))) + stranspose @[2, 0, 1] (sreplicateN @[14, 23] (tproject2 (tproject1 (tproject1 u1)))) ; w248 = sgather @[4, 7, 11, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i239, i240, i241, i242, i243] -> [ifH (0.0 <=. negate (splainPart t236 `sindex0` [i239, sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i240, i242], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i241, i243]])) 0 1]) ; u250 = sreshape @[4, 7, 11, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain w248 * stranspose @[4, 0, 1, 2, 3] (sgather @[7, 11, 2, 2] (stranspose @[1, 2, 0] t236) (\\[i244, i245, i246, i247] -> [sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i244, i246], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i245, i247]]))))) ; w258 = sgather @[7, 11, 3, 4] (stranspose @[3, 4, 5, 6, 0, 1, 2] (sgather @[7, 11, 4] (stranspose @[4, 3, 2, 1, 7, 6, 5, 0] (sreplicate @4 (stranspose @[3, 4, 5, 6, 0, 1, 2] (sreplicateN @[3, 11, 7] (stranspose @[3, 2, 1, 0] u250))))) (\\[i507, i508, i509] -> [i509, i507, i508, kargMax (splainPart u250 !$ [i509, i507, i508])]))) (\\[i254, i255, i256, i257] -> [i254, i255, i256, i257, i254 + i256, i255 + i257]) ; t259 = ssumN @[3, 4] (sdot1In (stranspose @[2, 3, 0, 4, 5, 1] (sreplicate @4 (stranspose @[4, 2, 3, 0, 1] w258))) (stranspose @[3, 4, 1, 0, 5, 2] (sreplicate @7 (stranspose @[1, 2, 3, 4, 0] (sreplicate @11 (tproject1 (tproject2 (tproject1 u1)))))))) + stranspose @[2, 0, 1] (sreplicateN @[7, 11] (tproject2 (tproject2 (tproject1 u1)))) ; w271 = sgather @[4, 3, 5, 2, 2] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i262, i263, i264, i265, i266] -> [ifH (0.0 <=. negate (splainPart t259 `sindex0` [i262, sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i263, i265], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i264, i266]])) 0 1]) ; u273 = sreshape @[4, 3, 5, 4] (stranspose @[2, 3, 4, 0, 1] (sreplicateN @[1, 1] (sfromPlain w271 * stranspose @[4, 0, 1, 2, 3] (sgather @[3, 5, 2, 2] (stranspose @[1, 2, 0] t259) (\\[i267, i268, i269, i270] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i267, i269], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i268, i270]]))))) ; v277 = sreshape @[60] (sgather @[4, 3, 5] u273 (\\[i274, i275, i276] -> [i274, i275, i276, kargMax (splainPart u273 !$ [i274, i275, i276])])) ; m278 = str (sreplicate @7 (sdot1In (tproject1 (tproject1 (tproject2 u1))) (sreplicate @5 v277))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m281 = sgather @[5, 7] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i279, i280] -> [ifH (0.0 <=. negate (splainPart m278 `sindex0` [i279, i280])) 0 1]) ; m284 = sfromPlain m281 * smatmul2 (str (tproject1 (tproject2 (tproject2 u1)))) dret ; v285 = ssum @7 (str m284) in tpair (let t293 = sscatter @[3, 5, 2, 2] (sfromPlain (stranspose @[1, 2, 3, 4, 0] w271) * stranspose @[3, 4, 1, 2, 5, 6, 0] (sreshape @[4, 3, 5, 1, 1, 2, 2] (sscatter @[4, 3, 5] (sreshape @[4, 3, 5] (sdot1In (str (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @60 v285))) (\\[i286, i287, i288] -> [i286, i287, i288, kargMax (splainPart u273 !$ [i286, i287, i288])]))) !$ [0, 0]) (\\[i289, i290, i291, i292] -> [sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) `sindex0` [i289, i291], sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) `sindex0` [i290, i292]]) in tpair (let t305 = sscatter @[7, 11, 2, 2] (sfromPlain (stranspose @[1, 2, 3, 4, 0] w248) * stranspose @[3, 4, 1, 2, 5, 6, 0] (sreshape @[4, 7, 11, 1, 1, 2, 2] (ssumN @[4, 3, 11, 7] (stranspose @[7, 6, 5, 4, 0, 1, 2, 3] (sscatter @[7, 11, 4] (stranspose @[4, 5, 6, 0, 1, 2, 3] (sscatter @[7, 11, 3, 4] (sdot1In (sreplicateN @[7, 11] (stranspose @[2, 3, 1, 0] (tproject1 (tproject2 (tproject1 u1))))) (stranspose @[4, 5, 1, 2, 0, 3] (sreplicateN @[4, 3, 4] (stranspose @[2, 0, 1] t293)))) (\\[i294, i295, i296, i297] -> [i294, i295, i296, i297, i294 + i296, i295 + i297]))) (\\[i298, i299, i300] -> [i300, i298, i299, kargMax (splainPart u250 !$ [i300, i298, i299])]))))) !$ [0, 0]) (\\[i301, i302, i303, i304] -> [sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) `sindex0` [i301, i303], sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) `sindex0` [i302, i304]]) in tpair (str (sreplicate @1 (ssum @23 (sdot1In (sfromPlain (stranspose @[4, 0, 1, 2, 3] (sreplicate @4 (stranspose @[3, 1, 2, 0] u234)))) (stranspose @[4, 2, 0, 1, 3] (sreplicateN @[3, 4] (stranspose @[2, 0, 1] t305))))))) (ssumN @[14, 23] t305)) (tpair (ssum @11 (sdot1In (stranspose @[5, 0, 1, 2, 3, 4] (sreplicate @4 (stranspose @[4, 2, 3, 0, 1] w258))) (stranspose @[5, 3, 0, 1, 2, 4] (sreplicateN @[4, 3, 4] (stranspose @[2, 0, 1] t293))))) (ssumN @[7, 11] t293))) (tpair (tpair (sreplicate @5 v277 * str (sreplicate @60 v285)) (ssum @7 (str m284))) (tpair (smatmul2 dret (sfromPlain (str m281) * str m278)) (ssum @7 (str dret))))"
test/simplified/TestMnistRNNR.hs view
@@ -17,6 +17,7 @@ import Prelude  import Control.Monad (foldM, unless)+import Data.Proxy (Proxy (Proxy)) import System.IO (hPutStrLn, stderr) import System.Random import Test.Tasty@@ -51,12 +52,13 @@ -- which side-steps vectorization. mnistTestCaseRNNRA   :: forall r.-     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> Int -> Int -> Int -> r   -> TestTree mnistTestCaseRNNRA prefix epochs maxBatches width miniBatchSize totalBatchSize-                   expected =+                   expected+                   | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat width $ \(SNat @width) ->   let targetInit =         forgetShape $ fst@@ -150,7 +152,7 @@   -> TestTree #-}  tensorADValMnistTestsRNNRA :: TestTree-tensorADValMnistTestsRNNRA = testGroup "RNNR ADVal MNIST tests"+tensorADValMnistTestsRNNRA = inOrderTestGroup "RNNR ADVal MNIST tests"   [ mnistTestCaseRNNRA "RNNRA 1 epoch, 1 batch" 1 1 128 150 5000                        (0.6026 :: Double)   , mnistTestCaseRNNRA "RNNRA artificial 1 2 3 4 5" 2 3 4 5 50@@ -165,12 +167,13 @@ -- but differentiated anew in each gradient descent iteration. mnistTestCaseRNNRI   :: forall r.-     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> Int -> Int -> Int -> r   -> TestTree mnistTestCaseRNNRI prefix epochs maxBatches width miniBatchSize totalBatchSize-                   expected =+                   expected+                   | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   withSNat width $ \(SNat @width) ->   let targetInit =         forgetShape $ fst@@ -201,7 +204,7 @@            ftk = tftk @Concrete                       (knownSTK @(X (ADRnnMnistParameters Concrete r)))                       targetInit-       (_, _, var, varAst) <- funToAstRevIO ftk+       (_, var, varAst) <- funToAstRevIO ftk        (varGlyph, astGlyph) <-          funToAstIO (FTKR (miniBatchSize                            :$: sizeMnistHeightInt@@ -212,7 +215,7 @@                            :$: sizeMnistLabelInt                            :$: ZSR) FTKScalar) id        let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)-           ast = simplifyInline+           ast = simplifyUserCode                  $ MnistRnnRanked2.rnnMnistLossFusedR                      miniBatchSize (astGlyph, astLabel)                      (fromTarget varAst)@@ -280,7 +283,7 @@   -> TestTree #-}  tensorADValMnistTestsRNNRI :: TestTree-tensorADValMnistTestsRNNRI = testGroup "RNNR Intermediate MNIST tests"+tensorADValMnistTestsRNNRI = inOrderTestGroup "RNNR Intermediate MNIST tests"   [ mnistTestCaseRNNRI "RNNRI 1 epoch, 1 batch" 1 1 128 150 5000                        (0.6026 :: Double)   , mnistTestCaseRNNRI "RNNRI artificial 1 2 3 4 5" 2 3 4 5 50@@ -296,8 +299,7 @@ -- descent iteration. mnistTestCaseRNNRO   :: forall r.-     ( Differentiable r, GoodScalar r-     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> Int -> Int -> Int -> r   -> TestTree@@ -342,7 +344,7 @@              MnistRnnRanked2.rnnMnistLossFusedR                miniBatchSize (rprimalPart glyphR, rprimalPart labelR) pars            artRaw = gradArtifact f (fromTarget targetInit, dataInit)-           art = simplifyArtifactGradient artRaw+           art = simplifyArtifactRev artRaw            go :: [MnistDataBatchR r]               -> ( Concrete (X (ADRnnMnistParameters Concrete r))                  , StateAdam (X (ADRnnMnistParameters Concrete r)) )@@ -352,7 +354,7 @@            go ((glyph, label) : rest) (!parameters, !stateAdam) =              let parametersAndInput =                    tpair parameters (tpair (rconcrete glyph) (rconcrete label))-                 gradient = tproject1 $ fst+                 gradient = tproject1 $ snd                             $ revInterpretArtifact                                 art parametersAndInput Nothing              in go rest (updateWithGradientAdam@@ -415,7 +417,7 @@   -> TestTree #-}  tensorADValMnistTestsRNNRO :: TestTree-tensorADValMnistTestsRNNRO = testGroup "RNNR Once MNIST tests"+tensorADValMnistTestsRNNRO = inOrderTestGroup "RNNR Once MNIST tests"   [ mnistTestCaseRNNRO "RNNRO 1 epoch, 1 batch" 1 1 128 150 5000                        (0.6026 :: Double)   , mnistTestCaseRNNRO "RNNRO artificial 1 2 3 4 5" 2 3 4 5 50
test/simplified/TestMnistRNNS.hs view
@@ -26,6 +26,7 @@ import Text.Printf  import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (fromSNat')  import HordeAd import HordeAd.Core.Adaptor@@ -56,18 +57,19 @@ -- which side-steps vectorization. mnistTestCaseRNNSA   :: forall width batch_size r.-     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> SNat width -> SNat batch_size -> Int -> r   -> TestTree mnistTestCaseRNNSA prefix epochs maxBatches width@SNat batch_size@SNat-                   totalBatchSize expected =+                   totalBatchSize expected+                   | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   let targetInit =         fst $ randomValue @(Concrete (XParams width r)) 0.23 (mkStdGen 44)-      miniBatchSize = sNatValue batch_size+      miniBatchSize = fromSNat' batch_size       name = prefix ++ ": "              ++ unwords [ show epochs, show maxBatches-                        , show (sNatValue width), show miniBatchSize+                        , show (fromSNat' width), show miniBatchSize                         , show $ widthSTK                           $ knownSTK @(XParams width r)                         , show (tsize knownSTK targetInit) ]@@ -113,7 +115,7 @@                    ftest @len (mkMnistDataBatchS chunk) parameters2                  testScore = ftest @lenTestData testDataS parameters2                  lenChunk = length chunk-             unless (sNatValue width < 10) $ do+             unless (fromSNat' width < 10) $ do                hPutStrLn stderr $                  printf "\n%s: (Batch %d with %d points)"                         prefix k lenChunk@@ -130,7 +132,7 @@                     -> IO (Concrete (XParams width r))            runEpoch n (params2, _) | n > epochs = return params2            runEpoch n paramsStateAdam@(!_, !_) = do-             unless (sNatValue width < 10) $+             unless (fromSNat' width < 10) $                hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n              let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData                  chunks = take maxBatches@@ -150,7 +152,7 @@   -> TestTree #-}  tensorADValMnistTestsRNNSA :: TestTree-tensorADValMnistTestsRNNSA = testGroup "RNNS ADVal MNIST tests"+tensorADValMnistTestsRNNSA = inOrderTestGroup "RNNS ADVal MNIST tests"   [ mnistTestCaseRNNSA "RNNSA 1 epoch, 1 batch" 1 1 (SNat @128) (SNat @150) 5000                        (0.6026 :: Double)   , mnistTestCaseRNNSA "RNNSA artificial 1 2 3 4 5" 2 3 (SNat @4) (SNat @5) 50@@ -165,18 +167,19 @@ -- but differentiated anew in each gradient descent iteration. mnistTestCaseRNNSI   :: forall width batch_size r.-     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> SNat width -> SNat batch_size -> Int -> r   -> TestTree mnistTestCaseRNNSI prefix epochs maxBatches width@SNat batch_size@SNat-                   totalBatchSize expected =+                   totalBatchSize expected+                   | Dict0 <- lemTKScalarAllNumAD (Proxy @r) =   let targetInit =         fst $ randomValue @(Concrete (XParams width r)) 0.23 (mkStdGen 44)-      miniBatchSize = sNatValue batch_size+      miniBatchSize = fromSNat' batch_size       name = prefix ++ ": "              ++ unwords [ show epochs, show maxBatches-                        , show (sNatValue width), show miniBatchSize+                        , show (fromSNat' width), show miniBatchSize                         , show $ widthSTK                           $ knownSTK @(XParams width r)                         , show (tsize knownSTK targetInit) ]@@ -199,11 +202,11 @@      $ \(SNat @lenTestData) -> do        let testDataS = mkMnistDataBatchS @lenTestData testData            ftk = tftk @Concrete (knownSTK @(XParams width r)) targetInit-       (_, _, var, varAst) <- funToAstRevIO ftk+       (_, var, varAst) <- funToAstRevIO ftk        (varGlyph, astGlyph) <- funToAstIO (FTKS knownShS FTKScalar) id        (varLabel, astLabel) <- funToAstIO (FTKS knownShS FTKScalar) id        let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)-           ast = simplifyInline+           ast = simplifyUserCode                  $ MnistRnnShaped2.rnnMnistLossFusedS                      width batch_size (astGlyph, astLabel)                      (fromTarget varAst)@@ -232,7 +235,7 @@                    ftest @len (mkMnistDataBatchS chunk) parameters2                  testScore = ftest @lenTestData testDataS parameters2                  lenChunk = length chunk-             unless (sNatValue width < 10) $ do+             unless (fromSNat' width < 10) $ do                hPutStrLn stderr $                  printf "\n%s: (Batch %d with %d points)"                         prefix k lenChunk@@ -249,7 +252,7 @@                     -> IO (Concrete (XParams width r))            runEpoch n (params2, _) | n > epochs = return params2            runEpoch n paramsStateAdam@(!_, !_) = do-             unless (sNatValue width < 10) $+             unless (fromSNat' width < 10) $                hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n              let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData                  chunks = take maxBatches@@ -267,7 +270,7 @@   -> TestTree #-}  tensorADValMnistTestsRNNSI :: TestTree-tensorADValMnistTestsRNNSI = testGroup "RNNS Intermediate MNIST tests"+tensorADValMnistTestsRNNSI = inOrderTestGroup "RNNS Intermediate MNIST tests"   [ mnistTestCaseRNNSI "RNNSI 1 epoch, 1 batch" 1 1 (SNat @128) (SNat @150) 5000                        (0.6026 :: Double)   , mnistTestCaseRNNSI "RNNSI artificial 1 2 3 4 5" 2 3 (SNat @4) (SNat @5) 50@@ -283,8 +286,7 @@ -- descent iteration. mnistTestCaseRNNSO   :: forall width batch_size r.-     ( Differentiable r, GoodScalar r-     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+     (Differentiable r, NumScalar r, PrintfArg r, AssertEqualUpToEpsilon r)   => String   -> Int -> Int -> SNat width -> SNat batch_size -> Int -> r   -> TestTree@@ -292,10 +294,10 @@                    totalBatchSize expected =   let targetInit =         fst $ randomValue @(Concrete (XParams width r)) 0.23 (mkStdGen 44)-      miniBatchSize = sNatValue batch_size+      miniBatchSize = fromSNat' batch_size       name = prefix ++ ": "              ++ unwords [ show epochs, show maxBatches-                        , show (sNatValue width), show miniBatchSize+                        , show (fromSNat' width), show miniBatchSize                         , show $ widthSTK                           $ knownSTK @(XParams width r)                         , show (tsize knownSTK targetInit) ]@@ -337,7 +339,7 @@                width batch_size (sprimalPart glyphS, sprimalPart labelS) pars            artRaw = revArtifactAdapt IgnoreIncomingCotangent                                      f (FTKProduct ftk ftkData)-           art = simplifyArtifactGradient artRaw+           art = simplifyArtifactRev artRaw            go :: [MnistDataBatchS batch_size r]               -> ( Concrete (XParams width r)                  , StateAdam (XParams width r) )@@ -347,7 +349,7 @@            go ((glyph, label) : rest) (!parameters, !stateAdam) =              let parametersAndInput =                    tpair parameters (tpair (sconcrete glyph) (sconcrete label))-                 gradient = tproject1 $ fst+                 gradient = tproject1 $ snd                             $ revInterpretArtifact                                 art parametersAndInput Nothing              in go rest (updateWithGradientAdam@@ -368,7 +370,7 @@                    ftest @len (mkMnistDataBatchS chunk) parameters2                  testScore = ftest @lenTestData testDataS parameters2                  lenChunk = length chunk-             unless (sNatValue width < 10) $ do+             unless (fromSNat' width < 10) $ do                hPutStrLn stderr $                  printf "\n%s: (Batch %d with %d points)"                         prefix k lenChunk@@ -385,7 +387,7 @@                     -> IO (Concrete (XParams width r))            runEpoch n (params2, _) | n > epochs = return params2            runEpoch n paramsStateAdam@(!_, !_) = do-             unless (sNatValue width < 10) $+             unless (fromSNat' width < 10) $                hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n              let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData                  chunks = take maxBatches@@ -403,7 +405,7 @@   -> TestTree #-}  tensorADValMnistTestsRNNSO :: TestTree-tensorADValMnistTestsRNNSO = testGroup "RNNS Once MNIST tests"+tensorADValMnistTestsRNNSO = inOrderTestGroup "RNNS Once MNIST tests"   [ mnistTestCaseRNNSO "RNNSO 1 epoch, 1 batch" 1 1 (SNat @128) (SNat @150) 5000                        (0.6026 :: Double)   , mnistTestCaseRNNSO "RNNSO artificial 1 2 3 4 5" 2 3 (SNat @4) (SNat @5) 50
test/simplified/TestRevFwdFold.hs view
@@ -18,7 +18,9 @@ import Data.Array.Nested.Shaped.Shape  import HordeAd+import HordeAd.Core.AstEnv import HordeAd.Core.AstFreshId (resetVarCounter)+import HordeAd.Core.AstInterpret import HordeAd.Core.Ops (tbuild1, treplicate) import HordeAd.Core.OpsConcrete () @@ -146,14 +148,17 @@   , testCase "4S0rmapAccumRD00SCall" testSin0rmapAccumRD00SCall   , testCase "4S0rmapAccumRD00Sall0" testSin0rmapAccumRD00Sall0   , testCase "4S0rmapAccumRD00Sall" testSin0rmapAccumRD00Sall-  , testCase "4S0rmapAccumRD0R" testSin0rmapAccumRD0R+  , testCase "4S0rmapAccumRD0R0" testSin0rmapAccumRD0R0+  , testCase "4S0rmapAccumRD0R1" testSin0rmapAccumRD0R1+  , testCase "4S0rmapAccumRD0R2" testSin0rmapAccumRD0R2+  , testCase "4S0rmapAccumRD0R3" testSin0rmapAccumRD0R3   , testCase "4S0rmapAccumRD01SN" testSin0rmapAccumRD01SN   , testCase "4S0rmapAccumRD01SN3" testSin0rmapAccumRD01SN3   , testCase "4S0rmapAccumRD01SN5" testSin0rmapAccumRD01SN5   , testCase "4S0rmapAccumRD01SN51" testSin0rmapAccumRD01SN51   , testCase "4S0rmapAccumRD01SN531a" testSin0rmapAccumRD01SN531a   , testCase "4S0rmapAccumRD01SN531b0" testSin0rmapAccumRD01SN531b0---  , testCase "4S0rmapAccumRD01SN531b0PP" testSin0rmapAccumRD01SN531b0PP+  , testCase "4S0rmapAccumRD01SN531b0PP" testSin0rmapAccumRD01SN531b0PP   , testCase "4S0rmapAccumRD01SN531b0PPj" testSin0rmapAccumRD01SN531b0PPj   , testCase "4S0rmapAccumRD01SN531bRPPj" testSin0rmapAccumRD01SN531bRPPj   , testCase "4S0rmapAccumRD01SN531c" testSin0rmapAccumRD01SN531c@@ -179,8 +184,12 @@   , testCase "4S0FoldNestedS1" testSin0FoldNestedS1   , testCase "4S0FoldNestedS1PP" testSin0FoldNestedS1PP   , testCase "4S0FoldNestedR1PP" testSin0FoldNestedR1PP+  , testCase "4S0FoldNestedK0PP" testSin0FoldNestedK0PP+  , testCase "4S0FoldNestedS0PP" testSin0FoldNestedS0PP+  , testCase "4S0FoldNestedR0PP" testSin0FoldNestedR0PP+  , testCase "4S0FoldNestedK0LengthPPs" testSin0FoldNestedK0LengthPPs+  , testCase "4S0FoldNestedS0LengthPPs" testSin0FoldNestedS0LengthPPs   , testCase "4S0FoldNestedR0LengthPPs" testSin0FoldNestedR0LengthPPs-  , testCase "4S0FoldNestedR1LengthPPs" testSin0FoldNestedR1LengthPPs   , testCase "4S0FoldNestedR2LengthPPs" testSin0FoldNestedR2LengthPPs   , testCase "4S0FoldNestedR3LengthPPs" testSin0FoldNestedR3LengthPPs -- takes too long:    , testCase "4S0FoldNestedR4LengthPPs" testSin0FoldNestedR4LengthPPs@@ -219,13 +228,16 @@   , testCase "4S0revhV7" testSin0revhV7   ] +sfromIndex0 :: forall r target. (ADReady target, NumScalar r)+            => IntOf target -> target (TKS '[] r)+sfromIndex0 = sfromR . rfromIntegral . tfromPlain knownSTK . rfromK+ foo :: RealFloatH a => (a, a, a) -> a foo (x, y, z) =   let w = x * sin y   in atan2H z w + z * w -fooRgrad :: forall g a.-           (ADReady g, GoodScalar a, Differentiable a, ADTensorScalar a ~ a)+fooRgrad :: forall g a. (ADReady g, NumScalar a, Differentiable a)         => (a, a, a) -> (g (TKR 0 a), g (TKR 0 a), g (TKR 0 a)) fooRgrad (x, y, z) =   let f :: forall f. ADReady f => f (TKProduct (TKProduct (TKR 0 a) (TKR 0 a)) (TKR 0 a)) -> f (TKR 0 a)@@ -256,13 +268,13 @@   resetVarCounter   let (a1, _, _) = fooRgrad @(AstTensor AstMethodLet PrimalSpan) @Double (1.1, 2.2, 3.3)   printAstPretty a1-    @?= "rfromS (let x15 = let x10 = sin (sscalar 2.2) ; x11 = sscalar 1.1 * x10 ; x12 = recip (sscalar 10.889999999999999 + x11 * x11) ; x13 = sin (sscalar 2.2) ; x14 = sscalar (-3.3) * x12 in tpair (tpair (x10 * x14 + sscalar 3.3 * x13) (sscalar 1.1 * (cos (sscalar 2.2) * x14) + sscalar 3.63 * cos (sscalar 2.2))) (x11 * x12 + sscalar 1.1 * x13) in tproject1 (tproject1 x15))"+    @?= "rfromK 2.4396285219055063"  testFooRrevPP2 :: Assertion testFooRrevPP2 = do   let (a1, _, _) = fooRgrad @(AstTensor AstMethodLet PrimalSpan) @Double (1.1, 2.2, 3.3)   printAstSimple (simplifyInlineContract a1)-    @?= "rfromS (sscalar 2.668038132604647 + sscalar (-2.668038132604647) * recip (sscalar 11.680936386336942))"+    @?= "tfromPlain (STKR (SNat @0) STKScalar) (rfromK 2.4396285219055063)"  testFooRrev3 :: Assertion testFooRrev3 = do@@ -286,14 +298,14 @@   resetVarCounter   let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin (rscalar 1.1)   printAstPretty a1-    @?= "rfromS (cos (sscalar 1.1))"+    @?= "rfromK 0.4535961214255773"  testSin0RrevPP2 :: Assertion testSin0RrevPP2 = do   resetVarCounter   let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin (rscalar 1.1)   printAstSimple a1-    @?= "rfromS (cos (sscalar 1.1))"+    @?= "tfromPlain (STKR (SNat @0) STKScalar) (rfromK 0.4535961214255773)"  testSin0Rrev3 :: Assertion testSin0Rrev3 = do@@ -311,8 +323,8 @@ testSin0RrevPP4 :: Assertion testSin0RrevPP4 = do   let a1 = (rrev1 sin . rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin) (rscalar 1.1)-  printAstPretty (simplifyInline a1)-    @?= "rfromS (cos (cos (sscalar 1.1)))"+  printAstPretty (simplifyUserCode a1)+    @?= "rfromK 0.8988770945225438"  testSin0Rrev5 :: Assertion testSin0Rrev5 = do@@ -325,7 +337,7 @@   resetVarCounter   let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 (rrev1 sin) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sscalar (-0.8912073600614354))"+    @?= "rfromK (-0.8912073600614354)"  testSin0Rrev3' :: Assertion testSin0Rrev3' = do@@ -357,15 +369,15 @@   let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double)       a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin (rscalar 1.1)   printAstPretty a1-    @?= "rfromS (cos (sscalar 1.1))"+    @?= "rfromK 0.4535961214255773"  testSin0RfwdPP1 :: Assertion testSin0RfwdPP1 = do   resetVarCounter   let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double)       a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin (rscalar 1.1)-  printAstPretty (simplifyInline a1)-    @?= "rfromS (cos (sscalar 1.1))"+  printAstPretty (simplifyUserCode a1)+    @?= "rfromK 0.4535961214255773"  testSin0RfwdPP1FullUnsimp :: Assertion testSin0RfwdPP1FullUnsimp = do@@ -373,7 +385,7 @@   let a1 :: AstTensor AstMethodLet FullSpan (TKR 0 Double)       a1 = rfwd1 @(AstTensor AstMethodLet FullSpan) @Double @0 @0 sin (rscalar 1.1)   printAstPretty a1-    @?= "rfromS (sscalar 1.0 * cos (sscalar 1.1))"+    @?= "rfromK 0.4535961214255773"  testSin0RfwdPP1Full :: Assertion testSin0RfwdPP1Full = do@@ -381,7 +393,7 @@   let a1 :: AstTensor AstMethodLet FullSpan (TKR 0 Double)       a1 = rfwd1 @(AstTensor AstMethodLet FullSpan) @Double @0 @0 sin (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sscalar 0.4535961214255773)"+    @?= "rfromK 0.4535961214255773"  testSin0Rfwd3 :: Assertion testSin0Rfwd3 = do@@ -401,15 +413,19 @@ testSin0RfwdPP4P = do   let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double)       a1 = (rfwd1 sin . rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin) (rscalar 1.1)-  printAstPretty (simplifyInline a1)-    @?= "rfromS (cos (cos (sscalar 1.1)))"+  interpretAstPrimal @Concrete emptyEnv a1+    @?= rscalar 0.8988770945225438+  printAstPretty (simplifyUserCode a1)+    @?= "rfromK 0.8988770945225438"  testSin0RfwdPP4Dual :: Assertion testSin0RfwdPP4Dual = do   let a1 :: AstTensor AstMethodLet DualSpan (TKR 0 Double)       a1 = (rfwd1 sin . rfwd1 @(AstTensor AstMethodLet DualSpan) @Double @0 @0 sin) (rscalar 1.1)+  interpretAstDual @Concrete emptyEnv a1+    @?= DummyDualTarget (FTKR [] FTKScalar)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sdualPart (sscalar 0.0) * cos (sdualPart (sscalar 0.0) * cos (sdualPart (sscalar 0.0))))"+    @?= "rdualPart (rfromK 0.0)"  testSin0Rfwd5 :: Assertion testSin0Rfwd5 = do@@ -422,7 +438,7 @@   let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double)       a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 (rfwd1 sin) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sscalar (-0.8912073600614354))"+    @?= "rfromK (-0.8912073600614354)"  testSin0Rfwd3' :: Assertion testSin0Rfwd3' = do@@ -452,8 +468,8 @@ testSin0RrevPP5S = do   resetVarCounter   let a1 = srev1 @(AstTensor AstMethodLet PrimalSpan) @Double @'[] @'[] (srev1 sin) (srepl 1.1)-  printAstPretty (simplifyInline a1)-    @?= "negate (sin (sscalar 1.1))"+  printAstPretty (simplifyUserCode a1)+    @?= "sfromK (-0.8912073600614354)"  testSin0Fold0 :: Assertion testSin0Fold0 = do@@ -496,8 +512,8 @@                   f x0 = rfold (\_x _a -> rscalar 7)                            (rscalar 5) (rreplicate 1 x0)               in f) (rscalar 1.1)-  printAstPretty a1-    @?= "rsum (tproject2 (tmapAccumRDer (SNat @1) <lambda> <lambda> <lambda> (tconvert (ConvCmp (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp ConvXS (Conv0X STKScalar))) (STKScalar) 1.0) (tpair (sconcrete (sfromListLinear [1] [Z1])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sscalar 5.0) (sconcrete (sfromListLinear [1] [1.1]))))) (sconcrete (sfromListLinear [1] [1.1]))))))"+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromK (tproject2 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> 1.0 (tpair (sreverse (tproject2 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> 5.0 (sconcrete (sfromListLinear [1] [1.1]))))) (sconcrete (sfromListLinear [1] [1.1])))) `sindex0` [0])"  testSin0FoldB2 :: Assertion testSin0FoldB2 = do@@ -789,7 +805,9 @@                         (sreplicate @1 a0)             in rfromS . f . sfromR) (rscalar 1.1)   printAstPretty a1-    @?= "rfromS (let v6 = tmapAccumRDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sreplicate [5] 1.0)) (tpair (sconcrete (sfromListLinear [1] [Z1])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sreplicate [5] 1.1)) (sconcrete (sfromListLinear [1] [1.1]))))) (sconcrete (sfromListLinear [1] [1.1])))) in ssum @5 (tproject1 v6) + tproject2 v6 !$ [0])"+    @?= "rfromK (let v4 = tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sreplicate [5] 1.0)) (tpair (sreverse (tproject2 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sreplicate [5] 1.1)) (sconcrete (sfromListLinear [1] [1.1]))))) (sconcrete (sfromListLinear [1] [1.1]))) in ssum0 @[5] (sprimalPart (tproject1 v4)) + sprimalPart (tproject2 v4) `sindex0` [0])"+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromK (let v4 = tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sreplicate [5] 1.0)) (tpair (sreverse (tproject2 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sreplicate [5] 1.1)) (sconcrete (sfromListLinear [1] [1.1]))))) (sconcrete (sfromListLinear [1] [1.1]))) in ssum0 @[5] (sprimalPart (tproject1 v4)) + sprimalPart (tproject2 v4) `sindex0` [0])"  testSin0Fold18Sgrad :: Assertion testSin0Fold18Sgrad = do@@ -829,7 +847,7 @@                         (rreplicate 3 a0))   assertEqualUpToEpsilon 1e-10     (rscalar 98.72666469795735)-    (cgrad (kfromR . rsum0 . h) (rscalar 1.1))+    (cgrad (rsum0 . h) (rscalar 1.1))  testSin0Fold8Sjvp :: Assertion testSin0Fold8Sjvp = do@@ -868,7 +886,8 @@ testSin0Fold5Sjvp = do   assertEqualUpToEpsilon 1e-10     (rscalar 1.4291653807319993)-    (cjvp (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+    (cjvp @_ @_ @_ @Concrete+          (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)                f a0 = sfold (let g :: forall f2. ADReady f2                                    => f2 (TKS '[] Double) -> f2 (TKS '[2, 5] Double)                                    -> f2 (TKS '[] Double)@@ -887,7 +906,8 @@ testSin0Fold5Sfwds = do   assertEqualUpToEpsilon 1e-10     (srepl 1.4291653807319993)-    (cjvp (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+    (cjvp @_ @_ @_ @Concrete+          (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)                f a0 = sfold (let g :: forall f2. ADReady f2                                    => f2 (TKS '[] Double) -> f2 (TKS '[2, 5] Double)                                    -> f2 (TKS '[] Double)@@ -939,7 +959,7 @@   assertEqualUpToEpsilon' 1e-10     (ringestData [1,1,1,1,1] [1.360788364276732] :: Concrete (TKR 5 Double))     (rev' (\a0 -> rscan (\_x a -> sin a)-                        (rreplicate0N [1,1,1,1,1] (rscalar 84))+                        (rreplicate0N [1,1,1,1,1] 84)                         (rreplicate 3 a0)) (ringestData [1,1,1,1,1] [1.1]))  testSin0Scan4 :: Assertion@@ -947,7 +967,7 @@   assertEqualUpToEpsilon' 1e-10     (ringestData [1,1,1,1,1] [-0.4458209450295252] :: Concrete (TKR 5 Double))     (rev' (\a0 -> rscan (\x a -> atan2H (sin x) (sin a))-                        (rreplicate0N [1,1,1,1,1] (rscalar 2) * a0)+                        (rreplicate0N [1,1,1,1,1] 2 * a0)                         (rreplicate 3 a0)) (ringestData [1,1,1,1,1] [1.1]))  testSin0Scan5 :: Assertion@@ -958,7 +978,7 @@                                  $ atan2H (sin $ rreplicate 5 x)                                          (rsum $ sin $ rsum                                           $ rtr $ rreplicate 7 a))-                        (rreplicate0N [1,1,1,1] (rscalar 2) * a0)+                        (rreplicate0N [1,1,1,1] 2 * a0)                         (rreplicate 3 (rreplicate 2 (rreplicate 5 a0))))           (ringestData [1,1,1,1] [1.1])) @@ -987,7 +1007,7 @@                                  $ atan2H (rsum (rtr $ sin x))                                          (rreplicate 2                                           $ sin (rsum $ rreplicate 7 a)))-                        (rreplicate 2 (rreplicate 5 (rreplicate0N [1,1,1] (rscalar 2) * a0)))+                        (rreplicate 2 (rreplicate 5 (rreplicate0N [1,1,1] 2 * a0)))                         (rreplicate 3 a0)) (ringestData [1,1,1] [1.1]))  testSin0Scan8grad :: Assertion@@ -1035,15 +1055,15 @@                  (\x0 -> rscan (\x _a -> sin x) x0                            (rrepl @1 @Double [2] 42)) (rscalar 1.1)   printAstPrettyButNested (simplifyInlineContract a1)-    @?= "rfromS (sscalar 1.0 + tproject1 (tmapAccumRDer (SNat @2) (\\x6 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair (cos (sfromR (tproject1 (tproject2 (tproject2 x6)))) * (tproject1 (tproject2 x6) + tproject1 x6)) (sscalar 0.0))) (\\x13 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair ((sfromR (tproject1 (tproject2 (tproject2 (tproject1 x13)))) * negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x13))))))) * (tproject1 (tproject2 (tproject2 x13)) + tproject1 (tproject2 x13)) + (tproject1 (tproject2 (tproject1 x13)) + tproject1 (tproject1 x13)) * cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x13)))))) (sscalar 0.0))) (\\x22 -> tconvert (ConvCmp (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (ConvT2 ConvId (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX)))) (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)))))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)))) (let x30 = cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x22))))) * tproject1 (tproject1 x22) in tpair x30 (tpair x30 (tpair (negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x22)))))) * ((tproject1 (tproject2 (tproject2 x22)) + tproject1 (tproject2 x22)) * tproject1 (tproject1 x22))) (sscalar 0.0))))) (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x31 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar))) (let x37 = sin (tproject1 x31) in tpair x37 (tpair (tproject1 x31) x37))) (\\x38 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar))) (let x49 = tproject1 (tproject1 x38) * cos (tproject1 (tproject2 x38)) in tpair x49 (tpair (tproject1 (tproject1 x38)) x49))) (\\x50 -> tpair (cos (tproject1 (tproject2 x50)) * (sfromR (tproject2 (tproject2 (tproject1 x50))) + tproject1 (tproject1 x50)) + sfromR (tproject1 (tproject2 (tproject1 x50)))) (sscalar 0.0)) (sscalar 1.1) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))"+    @?= "rfromK (1.0 + tproject1 (tmapAccumLDer (SNat @2) (\\x17 -> tpair (cos (tproject1 (tproject2 (tproject2 x17))) * (tproject1 (tproject2 x17) + tproject1 x17)) Z1) (\\x18 -> tpair (negate (tproject1 (tproject2 (tproject2 (tproject1 x18)))) * (sin (tproject1 (tproject2 (tproject2 (tproject2 x18)))) * (tproject1 (tproject2 (tproject2 x18)) + tproject1 (tproject2 x18))) + (tproject1 (tproject2 (tproject1 x18)) + tproject1 (tproject1 x18)) * cos (tproject1 (tproject2 (tproject2 (tproject2 x18))))) Z1) (\\x19 -> let x37 = cos (tproject1 (tproject2 (tproject2 (tproject2 x19)))) * tproject1 (tproject1 x19) in tpair x37 (tpair x37 (tpair (negate (sin (tproject1 (tproject2 (tproject2 (tproject2 x19))))) * ((tproject1 (tproject2 (tproject2 x19)) + tproject1 (tproject2 x19)) * tproject1 (tproject1 x19))) 0.0))) 0.0 (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (sreverse (sprimalPart (sfromR (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x6 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)))) (STKProduct STKScalar (STKProduct STKScalar STKScalar)) (let x40 = sin (tproject1 x6) in tpair x40 (tpair (tproject1 x6) x40))) (\\x7 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)))) (STKProduct STKScalar (STKProduct STKScalar STKScalar)) (let x47 = tproject1 (tproject1 x7) * cos (tproject1 (tproject2 x7)) in tpair x47 (tpair (tproject1 (tproject1 x7)) x47))) (\\x8 -> tpair (cos (tproject1 (tproject2 x8)) * (kfromR (tproject2 (tproject2 (tproject1 x8))) + tproject1 (tproject1 x8)) + kfromR (tproject1 (tproject2 (tproject1 x8)))) 0.0) 1.1 (sconcrete (sreplicate [2] 42.0)))))))) (sconcrete (sreplicate [2] 42.0))))))"  testSin0Scan1RevPPForComparison :: Assertion testSin0Scan1RevPPForComparison = do   resetVarCounter   let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1                  (\x0 -> rfromList [sin (sin x0), sin x0, x0]) (rscalar 1.1)-  printAstPretty (simplifyInline a1)-    @?= "rfromS (sscalar 1.0 + (cos (sscalar 1.1) * cos (sin (sscalar 1.1)) + cos (sscalar 1.1)))"+  printAstPretty a1+    @?= "rfromK 1.738669201145692"  testSin0ScanFwdPP :: Assertion testSin0ScanFwdPP = do@@ -1052,7 +1072,7 @@                  (\x0 -> rscan (\x _a -> sin x) x0                            (rrepl @1 @Double [2] 42)) (rscalar 1.1)   printAstPrettyButNested (simplifyInlineContract a1)-    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) (\\x7 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (let x17 = tproject1 x7 * cos (sfromR (tproject1 (tproject2 (tproject2 x7)))) in tpair x17 x17)) (\\x18 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (let x28 = tproject1 (tproject1 x18) * cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x18))))) + (sfromR (tproject1 (tproject2 (tproject2 (tproject1 x18)))) * negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x18))))))) * tproject1 (tproject2 x18) in tpair x28 x28)) (\\x29 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (ConvT2 ConvId (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)))))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)))) (let x38 = sfromR (tproject2 (tproject1 x29)) + tproject1 (tproject1 x29) in tpair (cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x29))))) * x38) (tpair (sscalar 0.0) (tpair (negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x29)))))) * (tproject1 (tproject2 x29) * x38)) (sscalar 0.0))))) (sscalar 1.0) (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x39 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar))) (let x45 = sin (tproject1 x39) in tpair x45 (tpair (tproject1 x39) x45))) (\\x46 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar))) (let x50 = tproject1 (tproject1 x46) * cos (tproject1 (tproject2 x46)) in tpair x50 (tpair (tproject1 (tproject1 x46)) x50))) (\\x51 -> tpair (cos (tproject1 (tproject2 x51)) * (sfromR (tproject2 (tproject2 (tproject1 x51))) + tproject1 (tproject1 x51)) + sfromR (tproject1 (tproject2 (tproject1 x51)))) (sscalar 0.0)) (sscalar 1.1) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))))"+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sprimalPart (sfromR (tproject2 (tmapAccumLDer (SNat @2) (\\x18 -> tconvert (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (let x33 = tproject1 x18 * cos (kfromR (tproject1 (tproject2 (tproject2 x18)))) in tpair x33 x33)) (\\x19 -> tconvert (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (let x40 = tproject1 (tproject1 x19) * cos (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x19))))) + negate (kfromR (tproject1 (tproject2 (tproject2 (tproject1 x19))))) * (sin (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x19))))) * tproject1 (tproject2 x19)) in tpair x40 x40)) (\\x20 -> tconvert (ConvT2 ConvId (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar))) (let x32 = kfromR (tproject2 (tproject1 x20)) + tproject1 (tproject1 x20) in tpair (cos (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x20))))) * x32) (tpair 0.0 (tpair (negate (sin (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x20)))))) * (tproject1 (tproject2 x20) * x32)) 0.0)))) 1.0 (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x7 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)))) (STKProduct STKScalar (STKProduct STKScalar STKScalar)) (let x43 = sin (tproject1 x7) in tpair x43 (tpair (tproject1 x7) x43))) (\\x8 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)))) (STKProduct STKScalar (STKProduct STKScalar STKScalar)) (let x50 = tproject1 (tproject1 x8) * cos (tproject1 (tproject2 x8)) in tpair x50 (tpair (tproject1 (tproject1 x8)) x50))) (\\x9 -> tpair (cos (tproject1 (tproject2 x9)) * (kfromR (tproject2 (tproject2 (tproject1 x9))) + tproject1 (tproject1 x9)) + kfromR (tproject1 (tproject2 (tproject1 x9)))) 0.0) 1.1 (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0)))))))))"  testSin0ScanFwdPPFull :: Assertion testSin0ScanFwdPPFull = do@@ -1061,16 +1081,16 @@                  (\x0 -> rscan (\x _a -> sin x) x0                            (rrepl @1 @Double [2] 42)) (rscalar 1.1)   printAstPrettyButNested (simplifyInlineContract a1)-    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) (\\x7 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (let x17 = sfromR (tproject1 x7) * cos (sfromR (tproject1 (tproject2 (tproject2 x7)))) in tpair x17 x17)) (\\x18 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (let x28 = sfromR (tproject1 (tproject1 x18)) * cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x18))))) + (sfromR (tproject1 (tproject2 (tproject2 (tproject1 x18)))) * negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x18))))))) * sfromR (tproject1 (tproject2 x18)) in tpair x28 x28)) (\\x29 -> tconvert (ConvCmp (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (ConvT2 ConvId (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)))) (let x38 = tproject2 (tconvert (ConvT2 (ConvCmp ConvXS (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvRX)) (ConvCmp ConvXS (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvRX))) (STKProduct (STKR (SNat @0) STKScalar) (STKR (SNat @0) STKScalar)) (tproject1 x29)) + tproject1 (tconvert (ConvT2 (ConvCmp ConvXS (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvRX)) (ConvCmp ConvXS (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvRX))) (STKProduct (STKR (SNat @0) STKScalar) (STKR (SNat @0) STKScalar)) (tproject1 x29)) in tpair (cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x29))))) * x38) (tpair (sscalar 0.0) (tpair (negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x29)))))) * (sfromR (tproject1 (tproject2 x29)) * x38)) (sscalar 0.0))))) (rfromS (sscalar 1.0)) (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x39 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKProduct (STKR (SNat @0) STKScalar) (STKS [] STKScalar))) (let x45 = sin (sfromR (tproject1 x39)) in tpair x45 (tpair (tproject1 x39) x45))) (\\x46 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKProduct (STKR (SNat @0) STKScalar) (STKS [] STKScalar))) (let x50 = sfromR (tproject1 (tproject1 x46)) * cos (sfromR (tproject1 (tproject2 x46))) in tpair x50 (tpair (tproject1 (tproject1 x46)) x50))) (\\x51 -> tconvert (ConvCmp (ConvT2 ConvId (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair (cos (sfromR (tproject1 (tproject2 x51))) * (sfromR (tproject2 (tproject2 (tproject1 x51))) + sfromR (tproject1 (tproject1 x51))) + sfromR (tproject1 (tproject2 (tproject1 x51)))) (sscalar 0.0))) (rfromS (sscalar 1.1)) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))))"+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) (\\x18 -> tconvert (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (let x33 = tproject1 x18 * cos (kfromR (tproject1 (tproject2 (tproject2 x18)))) in tpair x33 x33)) (\\x19 -> tconvert (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (let x40 = tproject1 (tproject1 x19) * cos (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x19))))) + negate (kfromR (tproject1 (tproject2 (tproject2 (tproject1 x19))))) * (sin (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x19))))) * tproject1 (tproject2 x19)) in tpair x40 x40)) (\\x20 -> tconvert (ConvT2 ConvId (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) ConvId))) (STKProduct STKScalar (STKProduct STKScalar (STKProduct STKScalar STKScalar))) (let x32 = kfromR (tproject2 (tproject1 x20)) + tproject1 (tproject1 x20) in tpair (cos (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x20))))) * x32) (tpair 0.0 (tpair (negate (sin (kfromR (tproject1 (tproject2 (tproject2 (tproject2 x20)))))) * (tproject1 (tproject2 x20) * x32)) 0.0)))) 1.0 (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x7 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)))) (STKProduct STKScalar (STKProduct STKScalar STKScalar)) (let x43 = sin (tproject1 x7) in tpair x43 (tpair (tproject1 x7) x43))) (\\x8 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)) (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar)))) (STKProduct STKScalar (STKProduct STKScalar STKScalar)) (let x50 = tproject1 (tproject1 x8) * cos (tproject1 (tproject2 x8)) in tpair x50 (tpair (tproject1 (tproject1 x8)) x50))) (\\x9 -> tpair (cos (tproject1 (tproject2 x9)) * (kfromR (tproject2 (tproject2 (tproject1 x9))) + tproject1 (tproject1 x9)) + kfromR (tproject1 (tproject2 (tproject1 x9)))) 0.0) 1.1 (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))))"  testSin0Scan1Rev2PP1 :: Assertion testSin0Scan1Rev2PP1 = do   resetVarCounter   let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1                  (\x0 -> rscan (\x a -> sin x - a) x0-                           (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7]))) (rscalar 1.1)+                           (rconcrete (Nested.rfromListPrimLinear @1 @Double [2] [5, 7]))) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sscalar 1.0 + tproject1 (tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0]))))))"+    @?= "rfromK (1.0 + tproject1 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 0.0 (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (sreverse (sprimalPart (sfromR (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.1 (sconcrete (sfromListLinear [2] [5.0,7.0])))))))) (sconcrete (sfromListLinear [2] [7.0,5.0]))))))"  testSin0Scan1Rev2PPA :: Assertion testSin0Scan1Rev2PPA = do@@ -1078,10 +1098,10 @@   let art = revArtifactAdapt                  UseIncomingCotangent                  (\x0 -> rscan @_ @_ @_ @(TKScalar Double) (\x a -> sin x - a) x0-                           (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7])))+                           (rconcrete (Nested.rfromListPrimLinear @1 @Double [2] [5, 7])))                  (FTKR ZSR FTKScalar)-  printArtifactPretty (simplifyArtifact art)-    @?= "\\dret x1 -> rfromS (sfromR dret !$ [0] + tproject1 (tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sslice (SNat @1) (SNat @2) (sfromR dret)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> x1 (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0]))))))"+  printArtifactPretty (simplifyArtifactRev art)+    @?= "\\dret x1 -> rfromK (sfromR dret `sindex0` [0] + tproject1 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 0.0 (tpair (sreverse (sslice (SNat @1) (SNat @2) (sfromR dret))) (tpair (sreverse (sfromR (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> x1 (sconcrete (sfromListLinear [2] [5.0,7.0]))))))) (sconcrete (sfromListLinear [2] [7.0,5.0]))))))"  testSin0Scan1Rev2PPForComparison :: Assertion testSin0Scan1Rev2PPForComparison = do@@ -1090,15 +1110,15 @@                  UseIncomingCotangent                  (\x0 -> rfromList [sin (sin x0 - rscalar 5) - rscalar 7, sin x0 - rscalar 5, x0])                  (FTKR ZSR FTKScalar)-  printArtifactPretty @_ @(TKR 1 Double) (simplifyArtifact art)-    @?= "\\dret x1 -> rfromS (cos (sfromR x1) * (cos (sscalar (-5.0) + sin (sfromR x1)) * sfromR dret !$ [0]) + (cos (sfromR x1) * sfromR dret !$ [1] + sfromR dret !$ [2]))"+  printArtifactPretty @_ @(TKR 1 Double) (simplifyArtifactRev art)+    @?= "\\dret x1 -> rfromK (cos (kfromR x1) * (cos ((-5.0) + sin (kfromR x1)) * sfromR dret `sindex0` [0]) + (cos (kfromR x1) * sfromR dret `sindex0` [1] + sfromR dret `sindex0` [2]))"  testSin0Scan1Rev2 :: Assertion testSin0Scan1Rev2 = do   assertEqualUpToEpsilon' 1e-10     (ringestData [] [1.1961317861865948] :: Concrete (TKR 0 Double))     (rev' (\x0 -> rscan (\x a -> sin x - a) x0-                    (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7]))) (rscalar 1.1))+                    (rconcrete (Nested.rfromListPrimLinear @1 @Double [2] [5, 7]))) (rscalar 1.1))  testSin0Scan1Rev2ForComparison :: Assertion testSin0Scan1Rev2ForComparison = do@@ -1113,16 +1133,15 @@                  (\x0 -> rscan (\x a -> sin x - a) x0                            (rfromList [x0 * rscalar 5, x0 * rscalar 7])) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (let v11 = tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001])))) in sscalar 1.0 + (sscalar 5.0 * tproject2 v11 !$ [0] + (sscalar 7.0 * tproject2 v11 !$ [1] + tproject1 v11)))"-+    @?= "rfromK (let v5 = tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 0.0 (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (sreverse (sfromR (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.1 (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))))) (sconcrete (sfromListLinear [2] [7.700000000000001,5.5])))) in 1.0 + (5.0 * sprimalPart (tproject2 v5) `sindex0` [1] + (7.0 * sprimalPart (tproject2 v5) `sindex0` [0] + kprimalPart (tproject1 v5))))"  testSin0Scan1Rev3PPForComparison :: Assertion testSin0Scan1Rev3PPForComparison = do   resetVarCounter   let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1                  (\x0 -> rfromList [sin (sin x0 - x0 * rscalar 5) - x0 * rscalar 7, sin x0 - x0 * rscalar 5, x0]) (rscalar 1.1)-  printAstPretty (simplifyInline a1)-    @?= "rfromS (let x9 = cos (sscalar (-5.5) + sin (sscalar 1.1)) in sscalar (-11.0) + (cos (sscalar 1.1) * x9 + (sscalar (-5.0) * x9 + cos (sscalar 1.1))))"+  printAstPretty a1+    @?= "rfromK (-10.076255083995068)"  testSin0ScanFwd3PP :: Assertion testSin0ScanFwd3PP = do@@ -1132,7 +1151,7 @@                  (\x0 -> rscan (\x a -> sin x - a) x0                            (rfromList [x0 * rscalar 5, x0 * rscalar 7])) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sfromListLinear [2] [5.0,7.0])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))))))"+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sprimalPart (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.0 (tpair (sconcrete (sfromListLinear [2] [5.0,7.0])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.1 (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001])))))))))"  testSin0Scan1Rev3 :: Assertion testSin0Scan1Rev3 = do@@ -1197,7 +1216,7 @@                         (rreplicate 3 a0))   assertEqualUpToEpsilon 1e-10     (rconcrete $ Nested.rfromListPrimLinear [] [285.95794829475744])-    (cgrad (kfromR . rsum0 . h) (rscalar 1.1))+    (cgrad (rsum0 . h) (rscalar 1.1))  testUnitriangular0PP :: Assertion testUnitriangular0PP = do@@ -1207,30 +1226,30 @@            $ \i -> rbuild1 k            $ \j -> ifH (i <=. j) (rscalar 0) (rscalar 1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3, i2] -> [ifH (0 <=. i2 + negate i3) 0 1]))"+    @?= "rfromPlain (rfromS (sgather @[1000000, 1000000] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3, i2] -> [ifH (i3 <=. i2) 0 1])))" -unitriangular1 :: (KnownNat k, GoodScalar rk, ADReady target)+unitriangular1 :: (KnownNat k, NumScalar rk, ADReady target)                => Int -> IShR k -> target (TKR (2 + k) rk) unitriangular1 k sh =   rbuild1 k $ \i ->     rbuild1 k $ \j ->-      ifH (i <=. j) (rreplicate0N sh (rscalar 0)) (rreplicate0N sh (rscalar 1))+      ifH (i <=. j) (rreplicate0N sh 0) (rreplicate0N sh 1)  testUnitriangular1PP :: Assertion testUnitriangular1PP = do   resetVarCounter   let sh = 2 :$: 3 :$: 6 :$: ZSR       k = 10-      a1 = unitriangular1 @3 @Double @(AstTensor AstMethodLet PrimalSpan) k sh-  printAstPretty (simplifyInline a1)-    @?= "rfromS (sgather (sconcrete (sfromListLinear [2,2,3,6] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])) (\\[i3, i2] -> [ifH (0 <=. i2 + negate i3) 0 1]))"+      a1 = unitriangular1 @3 @Double @(AstTensor AstMethodLet PlainSpan) k sh+  printAstPretty (simplifyUserCode a1)+    @?= "rfromS (stranspose @[3, 4, 0, 1, 2] (sreplicateN @[2, 3, 6] (sgather @[10, 10] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3, i2] -> [ifH (i3 <=. i2) 0 1]))))" -unitriangular2 :: (KnownNat k, GoodScalar rk, ADReady target)+unitriangular2 :: (KnownNat k, NumScalar rk, ADReady target)                => Int -> IShR k -> target (TKR (2 + k) rk) unitriangular2 k sh =-  rgather @_ @_ @1 (k :$: k :$: sh)-          (rfromList [ rreplicate0N sh (rscalar 0)-                     , rreplicate0N sh (rscalar 1) ])+  rgather @_ @_ @1 (k :$: k :$: ZSR)+          (rfromList [ rreplicate0N sh 0+                     , rreplicate0N sh 1 ])           (\(i :.: j :.: ZIR) -> ifH (i <. j) 0 1 :.: ZIR)  testUnitriangular2PP :: Assertion@@ -1238,9 +1257,9 @@   resetVarCounter   let sh = 2 :$: 3 :$: 6 :$: ZSR       k = 10-      a1 = unitriangular2 @3 @Double @(AstTensor AstMethodLet PrimalSpan) k sh-  printAstPretty (simplifyInline a1)-    @?= "rfromS (sgather (sconcrete (sfromListLinear [2,2,3,6] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])) (\\[i1, i2] -> [ifH (0 <=. i1 + negate i2) 1 0]))"+      a1 = unitriangular2 @3 @Double @(AstTensor AstMethodLet PlainSpan) k sh+  printAstPretty (simplifyUserCode a1)+    @?= "rfromS (sgather @[10, 10] (sconcrete (sfromListLinear [2,2,3,6] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])) (\\[i1, i2] -> [ifH (1 <=. negate i1 + i2) 0 1]))"  testSin0rmapAccumRD0S :: Assertion testSin0rmapAccumRD0S = do@@ -1264,7 +1283,8 @@ testSin0rmapAccumRD00SC = do   assertEqualUpToEpsilon 1e-10     (srepl 1)-    (cgrad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+    (cgrad @_ @_ @_ @Concrete+           (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)                 f x0 = tproject1 $ tmapAccumL (Proxy @f) (SNat @0)                           (FTKS ZSS FTKScalar)                           (FTKS ZSS FTKScalar)@@ -1330,13 +1350,14 @@                            in g)                           x0                           (treplicate (SNat @7) stkUnit tunit)-            in kfromS . ssum0 . f) (srepl 1.1))+            in ssum0 . f) (srepl 1.1))  testSin0rmapAccumRD00SCacc0 :: Assertion testSin0rmapAccumRD00SCacc0 = do   assertEqualUpToEpsilon 1e-10     (srepl 0)-    (cvjp (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)+    (cvjp @_ @_ @_ @Concrete+          (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)                f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @0)                           ftkUnit                           ftkUnit@@ -1348,7 +1369,7 @@                            in g)                           tunit                           (srepl 0)-            in f) (srepl 1.1) (sfromList0N []))+            in f) (srepl 1.1) (srepl @'[0] 0))  {- TODO: crashes due to a zero dimension testSin0rmapAccumRD00SCacc0 :: Assertion@@ -1367,14 +1388,15 @@                            in g)                           tunit                           (srepl 0)-             in kfromS . ssum0 . f) (srepl 1.1))+             in ssum0 . f) (srepl 1.1)) -}  testSin0rmapAccumRD00SCacc :: Assertion testSin0rmapAccumRD00SCacc = do   assertEqualUpToEpsilon 1e-10     (srepl 0)-    (cgrad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Z1)+    (cgrad @_ @_ @_ @Concrete+           (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Z1)                 f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @7)                           ftkUnit                           ftkUnit@@ -1386,7 +1408,7 @@                            in g)                           tunit                           (srepl 0)-             in kfromS . ssum0 . f) (srepl 1.1))+             in ssum0 . f) (srepl 1.1))  testSin0rmapAccumRD00Sacc0 :: Assertion testSin0rmapAccumRD00Sacc0 = do@@ -1404,7 +1426,7 @@                            in g)                           tunit                           (srepl 0)-             in kfromS . ssum0 . f) (srepl 1.1))+             in ssum0 . f) (srepl 1.1))  testSin0rmapAccumRD00Sacc :: Assertion testSin0rmapAccumRD00Sacc = do@@ -1422,13 +1444,14 @@                            in g)                           tunit                           (srepl 0)-             in kfromS . ssum0 . f) (srepl 1.1))+             in ssum0 . f) (srepl 1.1))  testSin0rmapAccumRD00SCall0 :: Assertion testSin0rmapAccumRD00SCall0 = do   assertEqualUpToEpsilon 1e-10     (srepl 0)-    (cvjp (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)+    (cvjp @_ @_ @_ @Concrete+          (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)                f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @0)                           ftkUnit                           ftkUnit@@ -1440,7 +1463,7 @@                            in g)                           tunit                           (treplicate (SNat @0) stkUnit tunit)-            in f) (srepl 1.1) (sfromList0N []))+            in f) (srepl 1.1) (srepl @'[0] 0))  {- TODO: crashes due to a zero dimension testSin0rmapAccumRD00SCall00 :: Assertion@@ -1459,13 +1482,14 @@                            in g)                           tunit                           (treplicate (SNat @0) stkUnit tunit)-             in kfromS . ssum0 . f) (srepl 1.1))+             in ssum0 . f) (srepl 1.1)) -} testSin0rmapAccumRD00SCall :: Assertion testSin0rmapAccumRD00SCall = do   assertEqualUpToEpsilon 1e-10     (srepl 0)-    (cgrad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Z1)+    (cgrad @_ @_ @_ @Concrete+           (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Z1)                 f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @7)                           ftkUnit                           ftkUnit@@ -1477,7 +1501,7 @@                            in g)                           tunit                           (treplicate (SNat @7) stkUnit tunit)-            in kfromS . ssum0 . f) (srepl 1.1))+            in ssum0 . f) (srepl 1.1))  testSin0rmapAccumRD00Sall0 :: Assertion testSin0rmapAccumRD00Sall0 = do@@ -1495,7 +1519,7 @@                            in g)                           tunit                           (treplicate (SNat @0) stkUnit tunit)-             in kfromS . ssum0 . f) (srepl 1.1))+             in ssum0 . f) (srepl 1.1))  testSin0rmapAccumRD00Sall :: Assertion testSin0rmapAccumRD00Sall = do@@ -1513,11 +1537,56 @@                            in g)                           tunit                           (treplicate (SNat @7) stkUnit tunit)-             in kfromS . ssum0 . f) (srepl 1.1))+           in ssum0 . f) (srepl 1.1)) -testSin0rmapAccumRD0R :: Assertion-testSin0rmapAccumRD0R = do+testSin0rmapAccumRD0R0 :: Assertion+testSin0rmapAccumRD0R0 = do   assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 1 Double)+               f x0 = rreverse (rreverse (rreplicate 0 x0) + rreplicate 0 x0)+           in f) (rscalar 1.1))++testSin0rmapAccumRD0R1 :: Assertion+testSin0rmapAccumRD0R1 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f _ = tproject1 $ tlet (tmapAccumL (Proxy @f) (SNat @0)+                          (FTKR ZSR FTKScalar)+                          (FTKR ZSR FTKScalar)+                          (FTKR ZSR FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKR 0 Double) -> g (TKR 0 Double)+                                 -> g (TKProduct (TKR 0 Double) (TKR 0 Double))+                               g _x _a = tpair (rscalar 1) (rscalar 2)+                           in g)+                          (rscalar 3)+                          (rreverse $ rrepl [0] 0)) $ \a ->+                            tpair (tproject1 a) (rreverse (tproject2 a))+           in f) (rscalar 1.1))++testSin0rmapAccumRD0R2 :: Assertion+testSin0rmapAccumRD0R2 = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 1 :: Concrete (TKR 0 Double))+    (cgrad (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+                f x0 = tproject1 $ tmapAccumR (Proxy @f) (SNat @0)+                          (FTKR ZSR FTKScalar)+                          (FTKR ZSR FTKScalar)+                          (FTKR ZSR FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKR 0 Double) -> g (TKR 0 Double)+                                 -> g (TKProduct (TKR 0 Double) (TKR 0 Double))+                               g x _a = tpair (sin x) (sin x)+                           in g)+                          x0+                          (rrepl [0] 0)+           in kfromR . f) (rscalar 1.1))++testSin0rmapAccumRD0R3 :: Assertion+testSin0rmapAccumRD0R3 = do+  assertEqualUpToEpsilon' 1e-10     (rscalar 1)     (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)                f x0 = tproject1 $ tmapAccumR (Proxy @f) (SNat @0)@@ -1531,7 +1600,7 @@                            in g)                           x0                           (rrepl [0] 0)-            in f) (rscalar 1.1))+           in f) (rscalar 1.1))  testSin0rmapAccumRD01SN :: Assertion testSin0rmapAccumRD01SN = do@@ -1601,11 +1670,11 @@                                                             (TKS '[3] Double)))                                g x a =                                  tpair (sin x-                                           - smaxIndex+                                           - (sfromIntegral $ sargMax                                                @2 @'[]-                                               (tproject2 $ tproject1 a))+                                               (tproject2 $ tproject1 a)))                                    (tpair (sreplicate @3-                                             (sindex0 @'[2]+                                             (sindex @'[2]                                                        (tproject1 $ tproject2 a) [1]                                               / sin x / srepl 3))                                           (sreplicate @3@@ -1624,8 +1693,8 @@     (rev' (let f :: forall f. ADReady f                  => f (TKS '[] Double) -> f (TKS '[] Double)                f x0 = (\res -> ssum @6 (tproject1 $ tproject1 res)-                               + ssum0 @'[6, 5, 4, 3]-                                   (tproject2 res))+                               + (sfromK $ ssum0 @'[6, 5, 4, 3]+                                   (tproject2 res)))                       $ tbuild1 @f (SNat @6) knownSTK $ \j ->                          tmapAccumR (Proxy @f) (SNat @5)                           (FTKProduct (FTKS ZSS FTKScalar)@@ -1648,11 +1717,11 @@                                      x1 = tproject2 xh                                  in tpair (tpair                                             (sin x-                                           - smaxIndex+                                           - (sfromIntegral $ sargMax                                                @2 @'[]-                                               (tproject2 $ tproject2 a))+                                               (tproject2 $ tproject2 a)))                                             (sreplicate @3-                                             (sindex0 @'[2]+                                             (sindex @'[2]                                                        (tproject1 $ tproject2 a) [1]                                               / sin x / srepl 3)))                                          (sbuild1 @4 $ \i ->@@ -1663,7 +1732,7 @@                           (tpair (x0 / (srepl 1 + sfromIndex0 j))                                  (sreplicate @3 x0))                           (tpair (tpair (srepl 1) (sreplicate0N @'[5, 3]-                                               (sfromIndex0 j)))+                                               (kfromS $ sfromIndex0 j)))                                  (tpair (srepl 3) (srepl 4)))            in rfromS . f . sfromR) (rscalar 1.1)) @@ -1706,17 +1775,17 @@                                              (ssum @1 (tproject1 a))                                            - sin x / srepl 3                                            - sreplicate @3-                                             (sindex0 @'[3]+                                             (sindex @'[3]                                                        (tproject2 a) [1]-                                             - smaxIndex+                                             - (sfromIntegral $ sargMax                                                  @3 @'[]-                                                 (sin x / srepl 3)))+                                                 (sin x / srepl 3))))                            in g)                           (tpair (x0 / (srepl 1 + sreplicate @3 (sfromIndex0 j)))                                  (sreplicate @6 (sfromIndex0 i)                                           - sflatten (sappend x0 x0)))                           (tpair (sfromList [srepl (-0.1), sreshape @'[] @'[1] $ sfromIndex0 j])-                                 ((sfromList0N+                                 ((sreshape $ sfromList                                            [sscalar 0.4, sscalar (-0.01), sscalar (-0.3), sfromIndex0 i, sscalar 0.5, sscalar 1.3]))))            in rfromS . f . sfromR) (ringestData [3] [1.1, 2, 3.14])) @@ -1743,9 +1812,8 @@                           (rconcrete $ Nested.rfromListPrimLinear [0] []))            in f) (rscalar 1.1)) --- Different result with -O1:-_testSin0rmapAccumRD01SN531b0PP :: Assertion-_testSin0rmapAccumRD01SN531b0PP = do+testSin0rmapAccumRD01SN531b0PP :: Assertion+testSin0rmapAccumRD01SN531b0PP = do   resetVarCounter   let f :: forall f. ADReady f                  => f (TKR 0 Double)@@ -1767,7 +1835,7 @@   printAstPrettyButNested     (simplifyInlineContract      $ f @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))-    @?= "rfromS (sreplicate @2 (sreplicate @2 (tproject1 (tmapAccumRDer (SNat @0) (\\x7 -> tpair (tproject1 x7) Z1) (\\x8 -> tpair (tproject1 (tproject1 x8)) Z1) (\\x11 -> tpair (tproject1 (tproject1 x11)) (sscalar 0.0)) (sscalar 1.1) (sconcrete (sfromListLinear [0] []))))))"+    @?= "rfromS (sreplicate0N @[2, 2] (tproject1 (tmapAccumLDer (SNat @0) (\\x8 -> tpair (tproject1 x8) Z1) (\\x9 -> tpair (tproject1 (tproject1 x9)) Z1) (\\x12 -> tpair (tproject1 (tproject1 x12)) 0.0) 1.1 (sconcrete (sfromListLinear [0] [])))))"  testSin0rmapAccumRD01SN531b0PPj :: Assertion testSin0rmapAccumRD01SN531b0PPj = do@@ -1792,7 +1860,7 @@   printAstPretty     (simplifyInlineContract      $ f @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))-    @?= "rfromS (tproject1 (tmapAccumRDer (SNat @0) <lambda> <lambda> <lambda> (sconcrete (sfromListLinear [2,2] [1.1,2.1,2.1,3.1])) (sconcrete (sfromListLinear [0,2,2] []))))"+    @?= "rfromS (tproject1 (tmapAccumLDer (SNat @0) <lambda> <lambda> <lambda> (sconcrete (sfromListLinear [2,2] [1.1,2.1,2.1,3.1])) (sconcrete (sfromListLinear [0,2,2] []))))"  testSin0rmapAccumRD01SN531bRPPj :: Assertion testSin0rmapAccumRD01SN531bRPPj = do@@ -1811,13 +1879,13 @@                                  -> g (TKProduct (TKR 0 Double) TKUnit)                                h xh _a = tpair xh tunit                            in h)-                          (rfromIndex0 (i + j) + x0)+                          (tfromPlain knownSTK (rfromIntegral (rfromK (i + j))) + x0)                           (rconcrete $ Nested.rfromListPrimLinear [0] [])))))                         $ \ !d -> tproject1 d   printAstPretty     (simplifyInlineContract      $ f @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))-    @?= "rfromS (tproject1 (tmapAccumRDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sfromListLinear [2,2] [1.1,2.1,2.1,3.1])) (sconcrete (sfromListLinear [0,2,2] []))))"+    @?= "rfromS (tproject1 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sfromListLinear [2,2] [1.1,2.1,2.1,3.1])) (sconcrete (sfromListLinear [0,2,2] []))))"  testSin0rmapAccumRD01SN531c :: Assertion testSin0rmapAccumRD01SN531c = do@@ -1841,7 +1909,7 @@                                       (srepl 1 - sin x / srepl 3 - a)                            in g)                           (x0 / (srepl 1 + sfromIndex0 j))-                          (sfromList0N [sscalar 0.4, sfromIndex0 i])))+                          (sreshape $ sfromList [sscalar 0.4, sfromIndex0 i])))            in rfromS . f . sfromR) (rscalar 1.1))  testSin0rmapAccumRD01SN531Slice :: Assertion@@ -1890,10 +1958,10 @@                                g x _a =                                  tpair (sin x - x)                                        (tpair (sreplicate @3-                                             (sindex0 @'[3] x [1]-                                              - smaxIndex+                                             (sindex @'[3] x [1]+                                              - (sfromIntegral $ sargMax                                                   @3 @'[]-                                                  (x / sin x / srepl 3)))+                                                  (x / sin x / srepl 3))))                                           (sreplicate @3                                              (ssum @3 x)                                            + sin x / srepl 3))@@ -1937,28 +2005,29 @@                                            (srepl 1 - sreplicate @7                                                  (ssum                                                   $ sin x - tproject2 a)))-                                        (sreplicate0N (sscalar 1)+                                        (sreplicate0N 1                                          - sreplicate @3                                              (ssum @1 (tproject1 a))-                                           - sin x / sreplicate0N (sscalar 3)+                                           - sin x / sreplicate0N 3                                            - sreplicate @3-                                             (sindex0 @'[3]+                                             (sindex @'[3]                                                        (tproject2 a) [1]-                                             - smaxIndex+                                             - (sfromIntegral $ sargMax                                                  @3 @'[]-                                                 (sin x / sreplicate0N (sscalar 3)))))+                                                 (sin x / sreplicate0N 3)))))                            in g)                           tunit                           (tpair (singestData [-0.1, 0.23])-                                 (sfromList0N-                                    [sindex0 x0 [1], sscalar (-0.01), sscalar (-0.3), ssum x0, sscalar 0.5, sscalar 1.3]))+                                 (sreshape $ sfromList+                                    [sindex @'[3] @'[] x0 [1], sscalar (-0.01), sscalar (-0.3), ssum x0, sscalar 0.5, sscalar 1.3]))            in rfromS . f . sfromR) (ringestData [3] [1.1, 2, 3.14]))  testSin0rmapAccumRD01SN58 :: Assertion testSin0rmapAccumRD01SN58 = do   assertEqualUpToEpsilon 1e-10-    (sconcrete $ Nested.sfromListPrimLinear @_ @'[5] knownShS [0,0,0,0,1.1])-    (cjvp (let f :: forall f. ADReady f+    (sconcrete $ Nested.sfromListPrimLinear @'[5] knownShS [0,0,0,0,1.1])+    (cjvp @_ @_ @_ @Concrete+          (let f :: forall f. ADReady f                  => f (TKS '[] Double) -> f (TKS '[5] Double)                f x0 = tproject2                       $ tmapAccumR (Proxy @f) (SNat @5)@@ -2001,12 +2070,12 @@                                                             (TKS '[3] Double)))                                g x a =                                   tpair (sin x-                                           ** smaxIndex+                                           ** (sfromIntegral $ sargMax                                                 @2 @'[]-                                                (tproject2 $ tproject1 a))+                                                (tproject2 $ tproject1 a)))                                     (tpair (sreplicate @3                                              (sin x / srepl 6-                                              + sindex0 @'[2]+                                              + sindex @'[2]                                                         (tproject1 $ tproject2 a) [1]                                                 / sin x / srepl 3))                                        (sreplicate @3@@ -2015,10 +2084,10 @@                            in g)                           x0                           (tpair-                             (tpair (sreplicate0N $ sscalar 0)-                                    (sreplicate0N $ sscalar 0))-                             (tpair (sreplicate0N $ sscalar 0)-                                    (sreplicate0N $ sscalar 0)))+                             (tpair (sreplicate0N 0)+                                    (sreplicate0N 0))+                             (tpair (sreplicate0N 0)+                                    (sreplicate0N 0)))             in tdot0Target (FTKProduct (FTKS ZSS FTKScalar)                              (FTKProduct (FTKS (SNat @1 :$$ SNat @3 :$$ ZSS) FTKScalar)                                          (FTKS (SNat @1 :$$ SNat @3 :$$ ZSS) FTKScalar))) (treplTarget 1 (FTKProduct (FTKS ZSS FTKScalar)@@ -2058,7 +2127,7 @@                                           $ rreplicate 2 $ rreplicate 5                                           $ rsum $ rsum a))                          (FTKR (8 :$: 3 :$: 1 :$: 1 :$: 1 :$: 1 :$: ZSR) FTKScalar)-                         (rreplicate0N [1,1,1,1] (rscalar 2) * a0)+                         (rreplicate0N [1,1,1,1] 2 * a0)                          (rreplicate 3 (rreplicate 8 (rreplicate 3 a0)))                          )           (ringestData [1,1,1,1] [1.1]))@@ -2073,7 +2142,7 @@                                           $ sin (rsum (rreplicate 7 a))))                        (FTKR (1 :$: 1 :$: 1 :$: ZSR) FTKScalar)                        (rreplicate 2 (rreplicate 5-                                        (rreplicate0N [1,1,1] (rscalar 2) * a0)))+                                        (rreplicate0N [1,1,1] 2 * a0)))                        (rreplicate 3 a0))           (ringestData [1,1,1] [1.1])) @@ -2098,7 +2167,7 @@                                 x                     in g)                       (rreplicate 2 (rreplicate 5-                                      (rreplicate0N [1,1,1] (rscalar 2) * a0)))+                                      (rreplicate0N [1,1,1] 2 * a0)))                       (rreplicate 4 a0))        (ringestData [1,1,1] [1.1])) @@ -2138,7 +2207,7 @@                            (FTKR ZSR FTKScalar)                            x0 (rrepl @1 @Double [2] 42)) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sscalar 1.0 + tproject1 (tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))"+    @?= "rfromK (1.0 + tproject1 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 0.0 (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (sreverse (sprimalPart (sfromR (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.1 (sconcrete (sreplicate [2] 42.0)))))))) (sconcrete (sreplicate [2] 42.0))))))"  testSin0ScanDFwdPP :: Assertion testSin0ScanDFwdPP = do@@ -2148,7 +2217,7 @@                            (FTKR ZSR FTKScalar)                            x0 (rrepl @1 @Double [2] 42)) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))))"+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sprimalPart (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.0 (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.1 (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0)))))))))"  testSin0ScanD1Rev2PP :: Assertion testSin0ScanD1Rev2PP = do@@ -2156,9 +2225,9 @@   let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1                  (\x0 -> rscanZip (\x a -> sin x - a)                          (FTKR ZSR FTKScalar)-                         x0 (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7]))) (rscalar 1.1)+                         x0 (rconcrete (Nested.rfromListPrimLinear @1 @Double [2] [5, 7]))) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sscalar 1.0 + tproject1 (tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0]))))))"+    @?= "rfromK (1.0 + tproject1 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 0.0 (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (sreverse (sprimalPart (sfromR (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.1 (sconcrete (sfromListLinear [2] [5.0,7.0])))))))) (sconcrete (sfromListLinear [2] [7.0,5.0]))))))"  testSin0ScanDFwd2PP :: Assertion testSin0ScanDFwd2PP = do@@ -2167,9 +2236,9 @@       a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1                  (\x0 -> rscanZip (\x a -> sin x - a)                          (FTKR ZSR FTKScalar)-                         x0 (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7]))) (rscalar 1.1)+                         x0 (rconcrete (Nested.rfromListPrimLinear @1 @Double [2] [5, 7]))) (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0]))))))))"+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sprimalPart (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.0 (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.1 (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0])))))))))"  testSin0ScanDFwd3PP :: Assertion testSin0ScanDFwd3PP = do@@ -2181,7 +2250,7 @@                                 x0 (rfromList [x0 * rscalar 5, x0 * rscalar 7]))                  (rscalar 1.1)   printAstPretty (simplifyInlineContract a1)-    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sfromListLinear [2] [5.0,7.0])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))))))"+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sprimalPart (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.0 (tpair (sconcrete (sfromListLinear [2] [5.0,7.0])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> 1.1 (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001])))))))))"  testSin0ScanD1jvp :: Assertion testSin0ScanD1jvp = do@@ -2242,7 +2311,7 @@                        (rreplicate 3 a0))   assertEqualUpToEpsilon 1e-10     (rconcrete $ Nested.rfromListPrimLinear [] [285.95794829475744])-    (cgrad (kfromR . rsum0 . h) (rscalar 1.1))+    (cgrad (rsum0 . h) (rscalar 1.1))  testSin0FoldNestedS1 :: Assertion testSin0FoldNestedS1 = do@@ -2267,7 +2336,7 @@       g = kgrad (kfromS . f) (FTKS ZSS FTKScalar)   printAstPretty     (g @(AstTensor AstMethodLet PrimalSpan) (sscalar 1.1))-    @?= "let v6 = tmapAccumRDer (SNat @11) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sreplicate [11] Z1)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @11) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sreplicate [11] 1.1))))) (sconcrete (sreplicate [11] 1.1)))) in ssum @11 (tproject2 v6) + tproject1 v6"+    @?= "sfromK (let v4 = tmapAccumLDer (SNat @11) <lambda> <lambda> <lambda> 1.0 (tpair (sreverse (tproject2 (tmapAccumLDer (SNat @11) <lambda> <lambda> <lambda> 1.1 (sconcrete (sreplicate [11] 1.1))))) (sconcrete (sreplicate [11] 1.1))) in ssum0 @[11] (sprimalPart (tproject2 v4)) + kprimalPart (tproject1 v4))"  testSin0FoldNestedR1PP :: Assertion testSin0FoldNestedR1PP = do@@ -2280,30 +2349,82 @@       g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)       g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)   printAstPretty-    (g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))-    @?= "rfromS (let v6 = tmapAccumRDer (SNat @11) <lambda> <lambda> <lambda> (tconvert (ConvCmp (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp ConvXS (Conv0X STKScalar))) (STKScalar) 1.0) (tpair (sconcrete (sreplicate [11] Z1)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @11) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sreplicate [11] 1.1))))) (sconcrete (sreplicate [11] 1.1)))) in ssum @11 (sfromR (tproject2 v6)) + sfromR (tproject1 v6))"+    (simplifyInlineContract+     $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))+    @?= "rfromK (let v4 = tmapAccumLDer (SNat @11) <lambda> <lambda> <lambda> 1.0 (tpair (sreverse (sfromR (tproject2 (tmapAccumLDer (SNat @11) <lambda> <lambda> <lambda> 1.1 (sconcrete (sreplicate [11] 1.1)))))) (sconcrete (sreplicate [11] 1.1))) in ssum0 @[11] (sprimalPart (tproject2 v4)) + kprimalPart (tproject1 v4))" -testSin0FoldNestedR0LengthPPs :: Assertion-testSin0FoldNestedR0LengthPPs = do+testSin0FoldNestedK0PP :: Assertion+testSin0FoldNestedK0PP = do   resetVarCounter+  let f :: forall f. ADReady f => f (TKScalar Double) -> f (TKScalar Double)+      f z = tfold (SNat @2) STKScalar STKScalar (\x a -> x + a)+                  z (treplicate (SNat @2) STKScalar z)+      g :: forall f. ADReady f => f (TKScalar Double) -> f (TKScalar Double)+      g = kgrad f FTKScalar+  printAstPrettyButNested+    (simplifyInlineContract+     $ g @(AstTensor AstMethodLet PrimalSpan) 1.1)+    @?= "let v4 = tmapAccumLDer (SNat @2) (\\x12 -> tpair (tproject1 x12) (tproject1 x12)) (\\x13 -> tpair (tproject1 (tproject1 x13)) (tproject1 (tproject1 x13))) (\\x14 -> tpair (tproject2 (tproject1 x14) + tproject1 (tproject1 x14)) (tpair 0.0 0.0)) 1.0 (tpair (sreverse (tproject2 (tmapAccumLDer (SNat @2) (\\x6 -> tpair (tproject1 x6 + tproject2 x6) (tproject1 x6)) (\\x7 -> tpair (tproject1 (tproject1 x7) + tproject2 (tproject1 x7)) (tproject1 (tproject1 x7))) (\\x8 -> tpair (tproject1 (tproject1 x8) + tproject2 (tproject1 x8)) (tproject1 (tproject1 x8))) 1.1 (sconcrete (sreplicate [2] 1.1))))) (sconcrete (sreplicate [2] 1.1))) in ssum0 @[2] (sprimalPart (tproject2 v4)) + kprimalPart (tproject1 v4)"++testSin0FoldNestedS0PP :: Assertion+testSin0FoldNestedS0PP = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+      f z = sfold (\x a -> x + a)+                  z (sreplicate @2 z)+      g :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+      g = kgrad (kfromS . f) (FTKS ZSS FTKScalar)+  printAstPrettyButNested+    (simplifyInlineContract+     $ g @(AstTensor AstMethodLet PrimalSpan) (sscalar 1.1))+    @?= "sfromK (let v4 = tmapAccumLDer (SNat @2) (\\x12 -> tconvert (ConvT2 ConvId (ConvCmp ConvXS (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (tpair (tproject1 x12) (tproject1 x12))) (\\x13 -> tconvert (ConvT2 ConvId (ConvCmp ConvXS (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (tpair (tproject1 (tproject1 x13)) (tproject1 (tproject1 x13)))) (\\x14 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp ConvXS (Conv0X STKScalar)) (ConvCmp ConvXS (Conv0X STKScalar)))) (STKProduct STKScalar (STKProduct STKScalar STKScalar)) (tpair (kfromS (tproject2 (tproject1 x14)) + tproject1 (tproject1 x14)) (tpair 0.0 0.0))) 1.0 (tpair (sreverse (tproject2 (tmapAccumLDer (SNat @2) (\\x6 -> tconvert (ConvT2 ConvId (ConvCmp ConvXS (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (tpair (tproject1 x6 + kfromS (tproject2 x6)) (tproject1 x6))) (\\x7 -> tconvert (ConvT2 ConvId (ConvCmp ConvXS (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (tpair (tproject1 (tproject1 x7) + kfromS (tproject2 (tproject1 x7))) (tproject1 (tproject1 x7)))) (\\x8 -> tconvert (ConvT2 ConvId (ConvCmp ConvXS (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (tpair (tproject1 (tproject1 x8) + kfromS (tproject2 (tproject1 x8))) (tproject1 (tproject1 x8)))) 1.1 (sconcrete (sreplicate [2] 1.1))))) (sconcrete (sreplicate [2] 1.1))) in ssum0 @[2] (sprimalPart (tproject2 v4)) + kprimalPart (tproject1 v4))"++testSin0FoldNestedR0PP :: Assertion+testSin0FoldNestedR0PP = do+  resetVarCounter   let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)-      f z = rfold (\x a -> x + tan a)+      f z = rfold (\x a -> x + a)                   z (rreplicate 2 z)       g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)       g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  printAstPrettyButNested+    (simplifyInlineContract+     $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))+    @?= "rfromK (let v4 = tmapAccumLDer (SNat @2) (\\x12 -> tpair (tproject1 x12) (tproject1 x12)) (\\x13 -> tpair (tproject1 (tproject1 x13)) (tproject1 (tproject1 x13))) (\\x14 -> tpair (tproject2 (tproject1 x14) + tproject1 (tproject1 x14)) (tpair 0.0 0.0)) 1.0 (tpair (sreverse (sfromR (tproject2 (tmapAccumLDer (SNat @2) (\\x6 -> tconvert (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (tpair (tproject1 x6 + tproject2 x6) (tproject1 x6))) (\\x7 -> tconvert (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))) (STKProduct STKScalar STKScalar) (tpair (tproject1 (tproject1 x7) + tproject2 (tproject1 x7)) (tproject1 (tproject1 x7)))) (\\x8 -> tpair (tproject1 (tproject1 x8) + kfromR (tproject2 (tproject1 x8))) (tproject1 (tproject1 x8))) 1.1 (sconcrete (sreplicate [2] 1.1)))))) (sconcrete (sreplicate [2] 1.1))) in ssum0 @[2] (sprimalPart (tproject2 v4)) + kprimalPart (tproject1 v4))"++testSin0FoldNestedK0LengthPPs :: Assertion+testSin0FoldNestedK0LengthPPs = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKScalar Double) -> f (TKScalar Double)+      f z = tfold (SNat @2) STKScalar STKScalar (\x a -> x + tan a)+                  z (treplicate (SNat @2) STKScalar z)+      g :: forall f. ADReady f => f (TKScalar Double) -> f (TKScalar Double)+      g = kgrad f FTKScalar   length     (printAstSimple       (simplifyInlineContract-       $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))-    @?= 4454+       $ g @(AstTensor AstMethodLet PrimalSpan) 1.1))+    @?= 1702 -testSin0FoldNestedR1LengthPPs :: Assertion-testSin0FoldNestedR1LengthPPs = do+testSin0FoldNestedS0LengthPPs :: Assertion+testSin0FoldNestedS0LengthPPs = do   resetVarCounter+  let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+      f z = sfold (\x a -> x + tan a)+                  z (sreplicate @2 z)+      g :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+      g = kgrad (kfromS . f) (FTKS ZSS FTKScalar)+  length+    (printAstSimple+      (simplifyInlineContract+       $ g @(AstTensor AstMethodLet PrimalSpan) (sscalar 1.1)))+    @?= 2516++testSin0FoldNestedR0LengthPPs :: Assertion+testSin0FoldNestedR0LengthPPs = do+  resetVarCounter   let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)-      f z = rfold (\x a ->-               rfold (\x2 a2 -> x2 + tan a2)-                     a (rreplicate 2 x))+      f z = rfold (\x a -> x + tan a)                   z (rreplicate 2 z)       g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)       g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)@@ -2311,7 +2432,7 @@     (printAstSimple       (simplifyInlineContract        $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))-    @?= 42039+    @?= 1953  testSin0FoldNestedR2LengthPPs :: Assertion testSin0FoldNestedR2LengthPPs = do@@ -2329,7 +2450,7 @@     (printAstSimple        (simplifyInlineContract         $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))-    @?= 530520+    @?= 300893  testSin0FoldNestedR3LengthPPs :: Assertion testSin0FoldNestedR3LengthPPs = do@@ -2349,7 +2470,7 @@     (printAstSimple        (simplifyInlineContract         $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))-    @?= 8035414+    @?= 4648181  -- Takes 70s, probably due to something (simplification?) forcing all derivs. _testSin0FoldNestedR4LengthPPs :: Assertion@@ -2416,7 +2537,7 @@     (printAstSimple        (simplifyInlineContract         $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))-    @?= 232921+    @?= 127739  testSin0FoldNestedR2Dummy7 :: Assertion testSin0FoldNestedR2Dummy7 = do@@ -2734,7 +2855,7 @@                         rfold (\x2 a2 ->                           rfold (\x3 a3 ->                             rfold (\x4 a4 ->-                              rfold (\_x5 _a5 -> 0)+                              rfold (\_x5 _a5 -> rscalar 0)                                     a4 (rreplicate 0 x4))                                   a3 (rreplicate 0 x3))                                 a2 (rreplicate 0 x2))@@ -2806,7 +2927,7 @@                             a0 (rreplicate 2 a0)            in f) (rscalar 1.1)   length (printAstSimple (simplifyInlineContract a1))-    @?= 67598+    @?= 50116  testSin0revhV :: Assertion testSin0revhV = do@@ -2825,7 +2946,7 @@       f = kgrad @_ @Double @g (kfromR . sin) (FTKR ZSR FTKScalar)   printAstSimple (f @(AstTensor AstMethodLet PrimalSpan)                                     (rscalar 1.1))-    @?= "rfromS (cos (sscalar 1.1))"+    @?= "tfromPlain (STKR (SNat @0) STKScalar) (rfromK 0.4535961214255773)"  testSin0revhV4 :: Assertion testSin0revhV4 = do@@ -2839,7 +2960,7 @@                doms3 x (ringestData [4] [1, 2, 3, 4])   assertEqualUpToEpsilon 1e-10     (rfromList [rscalar 0, rscalar 0, rscalar 0])-    (cgrad (kfromR . rsum0 . f) (rreplicate 3 (rscalar 1.1)))+    (cgrad @_ @_ @_ @Concrete (rsum0 . f) (rreplicate 3 (rscalar 1.1)))  testSin0revhV5 :: Assertion testSin0revhV5 = do@@ -2853,7 +2974,7 @@                 doms3 x (singestData @'[4] [1, 2, 3, 4])   assertEqualUpToEpsilon 1e-10     (sfromList @3 [sscalar 0, sscalar 0, sscalar 0])-    (cgrad (kfromS . ssum0 . f) (sreplicate @3 (sscalar 1.1)))+    (cgrad @_ @_ @_ @Concrete (ssum0 . f) (sreplicate @3 (sscalar 1.1)))  testSin0revhV6 :: Assertion testSin0revhV6 = do@@ -2867,7 +2988,7 @@                doms3 x (ringestData [4] [1, 2, 3, 4])   assertEqualUpToEpsilon 1e-10     (ringestData [3] [4.0,6.0,8.0])-    (cgrad (kfromR . rsum0 . f) (rreplicate 3 (rscalar 1.1)))+    (cgrad @_ @_ @_ @Concrete (rsum0 . f) (rreplicate 3 (rscalar 1.1)))  testSin0revhV7 :: Assertion testSin0revhV7 = do@@ -2881,4 +3002,4 @@                 doms3 x (singestData @'[4] [1, 2, 3, 4])   assertEqualUpToEpsilon 1e-10     (singestData @'[3] [4.0,6.0,8.0])-    (cgrad (kfromS . ssum0 . f) (sreplicate @3 (sscalar 1.1)))+    (cgrad @_ @_ @_ @Concrete (ssum0 . f) (sreplicate @3 (sscalar 1.1)))
test/tool/CrossTesting.hs view
@@ -8,6 +8,7 @@  import Prelude +import Data.Proxy (Proxy (Proxy)) import GHC.Exts (IsList (..)) import GHC.TypeLits (KnownNat) import System.IO.Unsafe (unsafePerformIO)@@ -23,9 +24,9 @@   , revInterpretArtifact   , revProduceArtifactWithoutInterpretation   )-import HordeAd.AstEngine import HordeAd.Core.Adaptor import HordeAd.Core.Ast+import HordeAd.Core.AstEngine import HordeAd.Core.AstEnv import HordeAd.Core.AstFreshId import HordeAd.Core.AstInterpret@@ -33,7 +34,7 @@ import HordeAd.Core.CarriersADVal import HordeAd.Core.CarriersAst import HordeAd.Core.CarriersConcrete-import HordeAd.Core.ConvertTensor (kfromR, kfromS)+import HordeAd.Core.ConvertTensor (rfromK) import HordeAd.Core.Ops import HordeAd.Core.OpsADVal import HordeAd.Core.TensorKind@@ -45,36 +46,35 @@  crevMaybeBoth   :: forall r m f src tgt.-     ( f ~ Concrete, X src ~ X (DValue src), KnownSTK (X src)+     ( GoodScalar r, f ~ Concrete, X src ~ X (DValue src), KnownSTK (X src)      , AdaptableTarget (ADVal Concrete) src-     , AdaptableTarget (ADVal Concrete) tgt      , AdaptableTarget Concrete (DValue src)      , tgt ~ ADVal f (TKR m r) )   => (src -> tgt)   -> DValue src-  -> (f (ADTensorKind (X src)), f (TKR m r))+  -> (f (TKR m r), f (ADTensorKind (X src))) {-# INLINE crevMaybeBoth #-}-crevMaybeBoth f vals =+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 m, KnownNat n, GoodScalar r+        ( 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 PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+          , 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 PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+          , 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@@ -90,14 +90,14 @@   return (resNormalSharing, resTotalSharing)  rev1 :: forall r m n v a w.-        ( KnownNat m, KnownNat n, GoodScalar r+        ( 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 PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+           , 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@@ -110,23 +110,23 @@       g :: ADVal Concrete (TKR n r)         -> ADVal Concrete (TKR m r)       g inputs = f $ fromTarget inputs-      !(!gradient1, !value1) = crevMaybeBoth g vals+      !(!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 PrimalSpan) (TKR n r)-         -> ADVal (AstRaw PrimalSpan) (TKR m r)-      g9 inputs = f @(ADVal (AstRaw PrimalSpan))+      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-      !(!gradient9, !value9) = revInterpretArtifact7 artifactsGradAst9+      !(!value9, !gradient9) = revInterpretArtifact7 artifactsGradAst9       revInterpretArtifact7         :: AstArtifactRev (TKR n r) (TKR m r)-        -> (Concrete (ADTensorKind (TKR n r)), Concrete (TKR m r))+        -> (Concrete (TKR m r), Concrete (ADTensorKind (TKR n r)))       revInterpretArtifact7 a1 = revInterpretArtifact a1 vals Nothing       hGeneral         :: (ADReady fgen, ADReady f1)@@ -136,8 +136,7 @@         -> fgen (TKR n r)         -> fgen (TKR m r)       hGeneral fx1 fx2 gx inputs =-        let (var, ast) =-              funToAst (FTKR (rshape vals) FTKScalar) Nothing (fx1 . f . fx2)+        let (var, ast) = funToAst (FTKR (rshape vals) FTKScalar) (fx1 . f . fx2)             env = extendEnv var inputs emptyEnv         in interpretAstFull env (gx ast)       h :: ADReady f1@@ -149,11 +148,11 @@       h fx1 fx2 gx inputs =         hGeneral @(ADVal Concrete) fx1 fx2 gx                  (fromTarget inputs)-      !(!gradient2, !value2) =+      !(!value2, !gradient2) =         crevMaybeBoth (h id id id) vals-      !(!gradient3, !value3) =+      !(!value3, !gradient3) =         crevMaybeBoth (h id id simplifyInlineContract) vals-      !(!gradient2UnSimp, !value2UnSimp) =+      !(!value2UnSimp, !gradient2UnSimp) =         crevMaybeBoth (h unAstNoSimplify AstNoSimplify id) vals       gradientRrev2UnSimp =         rrev1 @Concrete @r @n @m @r@@ -162,8 +161,8 @@         rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))                 ftkz (rrev1 @_ @r @n @m @r                             (hGeneral unAstNoSimplify AstNoSimplify id)) vals-      !(!gradient3UnSimp, !value3UnSimp) =-        crevMaybeBoth (h (simplifyInline . unAstNoSimplify) AstNoSimplify simplifyInlineContract)+      !(!value3UnSimp, !gradient3UnSimp) =+        crevMaybeBoth (h (simplifyUserCode . unAstNoSimplify) AstNoSimplify simplifyInlineContract)                       vals       gradientRrev3UnSimp =         rrev1 @Concrete @r @n @m @r@@ -172,7 +171,7 @@         rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))                 ftkz (rrev1 @_ @r @n @m @r                             (hGeneral unAstNoSimplify AstNoSimplify simplifyInlineContract)) vals-      !(!gradient4, !value4) =+      !(!value4, !gradient4) =         crevMaybeBoth (h unAstNoVectorize AstNoVectorize id)                       vals           -- use the AstNoVectorize instance that does no vectorization@@ -184,7 +183,7 @@         rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))                 ftkz (rrev1 @_ @r @n @m @r                             (hGeneral unAstNoVectorize AstNoVectorize id)) vals-      !(!gradient5, !value5) =+      !(!value5, !gradient5) =         crevMaybeBoth (h unAstNoVectorize AstNoVectorize simplifyInlineContract)                       vals       gradientRrev5 =@@ -194,74 +193,74 @@         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) Nothing f+      astVectSimp = simplifyInlineContract $ snd $ funToAst (FTKR (rshape vals) FTKScalar) f       astSimp =         simplifyInlineContract $ simplifyInlineContract $ snd  -- builds simplify with difficulty-        $ funToAst (FTKR (rshape vals) FTKScalar) Nothing (unAstNoVectorize . f . AstNoVectorize)+        $ 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 PrimalSpan) (TKR n r)-           -> ADVal (AstRaw PrimalSpan) (TKR m r)+           -> ADVal (AstRaw FullSpan) (TKR n r)+           -> ADVal (AstRaw FullSpan) (TKR m r)       hAst fx1 fx2 gx inputs-        = hGeneral @(ADVal (AstRaw PrimalSpan))+        = 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-      !(!gradient2Ast, !value2Ast) =+      !(!value2Ast, !gradient2Ast) =         revInterpretArtifact7 artifactsGradAst-      !(!gradient2AstS, !value2AstS) =-        revInterpretArtifact7 (simplifyArtifact artifactsGradAst)+      !(!value2AstS, !gradient2AstS) =+        revInterpretArtifact7 (simplifyArtifactRev artifactsGradAst)       artifactsGradAstT =         fst $ revProduceArtifactWithoutInterpretation-                UseIncomingCotangent (hAst simplifyInline id id) ftk-      !(!gradient2AstST, !value2AstST) =-        revInterpretArtifact7 (simplifyArtifact artifactsGradAstT)+                UseIncomingCotangent (hAst simplifyUserCode id id) ftk+      !(!value2AstST, !gradient2AstST) =+        revInterpretArtifact7 (simplifyArtifactRev artifactsGradAstT)       artifactsSimpleAst =         fst $ revProduceArtifactWithoutInterpretation                 IgnoreIncomingCotangent (hAst id id simplifyInlineContract) ftk-      !(!gradient3Ast, !value3Ast) =+      !(!value3Ast, !gradient3Ast) =         revInterpretArtifact7 artifactsSimpleAst-      !(!gradient3AstS, !value3AstS) =-        revInterpretArtifact7 (simplifyArtifact artifactsSimpleAst)+      !(!value3AstS, !gradient3AstS) =+        revInterpretArtifact7 (simplifyArtifactRev artifactsSimpleAst)       artifactsGradAstUnSimp =         fst $ revProduceArtifactWithoutInterpretation                 IgnoreIncomingCotangent (hAst unAstNoSimplify AstNoSimplify id) ftk-      !(!gradient2AstUnSimp, !value2AstUnSimp) =+      !(!value2AstUnSimp, !gradient2AstUnSimp) =         revInterpretArtifact7 artifactsGradAstUnSimp-      !(!gradient2AstSUnSimp, !value2AstSUnSimp) =-        revInterpretArtifact7 (simplifyArtifact artifactsGradAstUnSimp)+      !(!value2AstSUnSimp, !gradient2AstSUnSimp) =+        revInterpretArtifact7 (simplifyArtifactRev artifactsGradAstUnSimp)       artifactsSimpleAstUnSimp =         fst $ revProduceArtifactWithoutInterpretation                 IgnoreIncomingCotangent (hAst unAstNoSimplify AstNoSimplify simplifyInlineContract)                 ftk-      !(!gradient3AstUnSimp, !value3AstUnSimp) =+      !(!value3AstUnSimp, !gradient3AstUnSimp) =         revInterpretArtifact7 artifactsSimpleAstUnSimp-      !(!gradient3AstSUnSimp, !value3AstSUnSimp) =-        revInterpretArtifact7 (simplifyArtifact artifactsSimpleAstUnSimp)+      !(!value3AstSUnSimp, !gradient3AstSUnSimp) =+        revInterpretArtifact7 (simplifyArtifactRev artifactsSimpleAstUnSimp)       artifactsPrimalAst =         fst $ revProduceArtifactWithoutInterpretation                 IgnoreIncomingCotangent (hAst unAstNoVectorize AstNoVectorize id) ftk-      !(!gradient4Ast, !value4Ast) =+      !(!value4Ast, !gradient4Ast) =         revInterpretArtifact7 artifactsPrimalAst-      !(!gradient4AstS, !value4AstS) =-        revInterpretArtifact7 (simplifyArtifact artifactsPrimalAst)+      !(!value4AstS, !gradient4AstS) =+        revInterpretArtifact7 (simplifyArtifactRev artifactsPrimalAst)       artifactsPSimpleAst =         fst $ revProduceArtifactWithoutInterpretation                 IgnoreIncomingCotangent (hAst unAstNoVectorize AstNoVectorize simplifyInlineContract)                 ftk-      !(!gradient5Ast, !value5Ast) =+      !(!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.-      (gradient5AstS, value5AstS) =-        revInterpretArtifact7 (simplifyArtifact artifactsPSimpleAst)+      (value5AstS, gradient5AstS) =+        revInterpretArtifact7 (simplifyArtifactRev artifactsPSimpleAst)       !cderivative = cjvp f vals vals       !derivative = jvp f vals vals       !derivativeRfwd1 = rfwd1ds @Concrete @r @n @m @r f vals@@ -294,22 +293,23 @@        , a ~ Concrete (ADTensorKind (TKR n r))        , AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v        , AssertEqualUpToEpsilon (ADTensorScalar r)-       , GoodScalar r, GoodScalar (ADTensorScalar r), HasCallStack)+       , 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 PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+         , 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 PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+         , 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@@ -321,11 +321,11 @@        , a ~ Concrete (ADTensorKind (TKR n r))        , AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v        , AssertEqualUpToEpsilon (ADTensorScalar r)-       , GoodScalar r, GoodScalar (ADTensorScalar r), HasCallStack)+       , 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 PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+       , 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@@ -440,7 +440,7 @@   -- 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 (rsum0 derivative) (rdot0 expected (toADTensorKindShared ftk vals))+                                 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.@@ -473,14 +473,16 @@ t128c = rreshape (2 :$: 2 :$: 8 :$: 4 :$: ZSR) t128  rrev1 :: forall g r n m r3.-         (ADReady g, GoodScalar r, KnownNat n, GoodScalar r3, KnownNat m)+         (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))-rrev1 f u = kgrad (kfromR . rsum0 . f) (tftk knownSTK u) u+{-# INLINE rrev1 #-}+rrev1 f u = kgrad (rsum0 . f) (tftk knownSTK u) u -rrevFTK :: forall g x z. (ADReady g, KnownSTK x)+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.@@ -488,22 +490,26 @@         => (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, GoodScalar (ADTensorScalar r), KnownNat n)+         (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, GoodScalar r3, KnownShS sh2)+         (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))-srev1 f u = kgrad (kfromS. ssum0 . f) (tftk knownSTK u) u+{-# INLINE srev1 #-}+srev1 f u = kgrad (ssum0 . f) (tftk knownSTK u) u  sfwd1 :: forall g r sh sh2 r3.-         (ADReady g, GoodScalar r, GoodScalar (ADTensorScalar r), KnownShS sh)+         (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)
test/tool/EqEpsilon.hs view
@@ -10,7 +10,7 @@  import Prelude -import Data.Int (Int64)+import Data.Int (Int16, Int32, Int64, Int8) import Data.IORef import Foreign.C (CInt) import System.IO.Unsafe (unsafePerformIO)@@ -160,6 +160,18 @@   assertEqualUpToEpsilonWithMsg = assert_close_eps ""  instance AssertEqualUpToEpsilon Int64 where+  assertEqualUpToEpsilonWithMsg = assert_close_eps ""++instance AssertEqualUpToEpsilon Int32 where+  assertEqualUpToEpsilonWithMsg = assert_close_eps ""++instance AssertEqualUpToEpsilon Int16 where+  assertEqualUpToEpsilonWithMsg = assert_close_eps ""++instance AssertEqualUpToEpsilon Int8 where+  assertEqualUpToEpsilonWithMsg = assert_close_eps ""++instance AssertEqualUpToEpsilon Int where   assertEqualUpToEpsilonWithMsg = assert_close_eps ""  instance AssertEqualUpToEpsilon CInt where
test/tool/Shared.hs view
@@ -6,16 +6,19 @@  import Prelude +import Data.Array.Internal.RankedS qualified as RS import Data.Char qualified-import Data.Foldable qualified-import Data.Int (Int64)+import Data.Foldable qualified as Foldable+import Data.Int (Int16, Int32, Int64, Int8) import Data.Vector.Storable qualified as VS import Foreign.C (CInt)-import GHC.Exts (IsList (..))-import GHC.TypeLits (KnownNat)  import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed qualified as Nested.Mixed+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped (stoPrimitive) import Data.Array.Nested.Shaped.Shape+import Data.Array.XArray qualified as X  import HordeAd.Core.CarriersConcrete import HordeAd.Core.Types@@ -27,11 +30,11 @@ class HasShape a where   shapeL :: a -> [Int] -instance (KnownNat n, Nested.PrimElt a) => HasShape (Nested.Ranked n a) where-  shapeL = toList . Nested.rshape+instance Nested.Elt a => HasShape (Nested.Ranked n a) where+  shapeL = shrToList . Nested.rshape -instance KnownShS sh => HasShape (Nested.Shaped sh a) where-  shapeL _ = toList $ knownShS @sh+instance Nested.Elt a => HasShape (Nested.Shaped sh a) where+  shapeL = shsToList . Nested.sshape  instance HasShape (RepConcrete y) => HasShape (Concrete y) where   shapeL = shapeL . unConcrete@@ -45,6 +48,18 @@ instance HasShape Int64 where   shapeL _ = [] +instance HasShape Int32 where+  shapeL _ = []++instance HasShape Int16 where+  shapeL _ = []++instance HasShape Int8 where+  shapeL _ = []++instance HasShape Int where+  shapeL _ = []+ instance HasShape CInt where   shapeL _ = [] @@ -60,11 +75,14 @@  instance (VS.Storable a, Nested.PrimElt a)          => Linearizable (Nested.Ranked n a) a where-  linearize = VS.toList . Nested.rtoVector+  linearize = RS.toList . Nested.rtoOrthotope  instance (VS.Storable a, Nested.PrimElt a)          => Linearizable (Nested.Shaped sh a) a where-  linearize = VS.toList . Nested.stoVector+  linearize a =+    let Nested.Shaped (Nested.Mixed.M_Primitive _ (X.XArray arr)) =+          stoPrimitive a+    in RS.toList arr  instance Linearizable (RepConcrete y) a          => Linearizable (Concrete y) a where@@ -79,6 +97,18 @@ instance Linearizable Int64 Int64 where   linearize x = [x] +instance Linearizable Int32 Int32 where+  linearize x = [x]++instance Linearizable Int16 Int16 where+  linearize x = [x]++instance Linearizable Int8 Int8 where+  linearize x = [x]++instance Linearizable Int Int where+  linearize x = [x]+ instance Linearizable CInt CInt where   linearize x = [x] @@ -86,4 +116,4 @@   linearize _ = []  instance {-# OVERLAPPABLE #-} (Foldable t) => Linearizable (t a) a where-  linearize = Data.Foldable.toList+  linearize = Foldable.toList