{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | Tests of the gather and scatter operations and of operations that expand
-- to gather and of fusion of all of them.
module TestGatherSimplified (testTrees) where

import Prelude

import Data.Int (Int64)
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.Ranked.Shape
import Data.Array.Nested.Shaped.Shape

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

testTrees :: [TestTree]
testTrees :: [TestTree]
testTrees =
  [ String -> Assertion -> TestTree
testCase String
"gatherNested1" Assertion
testGatherNested1
  , String -> Assertion -> TestTree
testCase String
"gatherNestedBuild1" Assertion
testGatherNestedBuild1
  , String -> Assertion -> TestTree
testCase String
"gather1" Assertion
testGather1
  , String -> Assertion -> TestTree
testCase String
"gatherBuild1" Assertion
testGatherBuild1
  , String -> Assertion -> TestTree
testCase String
"gatherSimpPP1" Assertion
testGatherSimpPP1
  , String -> Assertion -> TestTree
testCase String
"gatherSimp1" Assertion
testGatherSimp1
  , String -> Assertion -> TestTree
testCase String
"gatherNested02" Assertion
testGatherNested02
  , String -> Assertion -> TestTree
testCase String
"gatherNested2" Assertion
testGatherNested2
  , String -> Assertion -> TestTree
testCase String
"gatherNestedBuild2" Assertion
testGatherNestedBuild2
  , String -> Assertion -> TestTree
testCase String
"gather2" Assertion
testGather2
  , String -> Assertion -> TestTree
testCase String
"gatherBuild2" Assertion
testGatherBuild2
  , String -> Assertion -> TestTree
testCase String
"gatherSimpPP2" Assertion
testGatherSimpPP2
  , String -> Assertion -> TestTree
testCase String
"gatherSimp2" Assertion
testGatherSimp2
  , String -> Assertion -> TestTree
testCase String
"gatherNested12" Assertion
testGatherNested12
  , String -> Assertion -> TestTree
testCase String
"gatherNestedBuild12" Assertion
testGatherNestedBuild12
  , String -> Assertion -> TestTree
testCase String
"gather12" Assertion
testGather12
  , String -> Assertion -> TestTree
testCase String
"gatherBuild12" Assertion
testGatherBuild12
  , String -> Assertion -> TestTree
testCase String
"gatherSimpPP12" Assertion
testGatherSimpPP12
  , String -> Assertion -> TestTree
testCase String
"gatherSimp12" Assertion
testGatherSimp12
  , String -> Assertion -> TestTree
testCase String
"gatherReshape22" Assertion
testGatherReshape22
  , String -> Assertion -> TestTree
testCase String
"gatherReshapeBuild22" Assertion
testGatherReshapeBuild22
  , String -> Assertion -> TestTree
testCase String
"gatherSimpPP22" Assertion
testGatherSimpPP22
  , String -> Assertion -> TestTree
testCase String
"gatherSimp22" Assertion
testGatherSimp22
  , String -> Assertion -> TestTree
testCase String
"gatherSimpPP23" Assertion
testGatherSimpPP23
  , String -> Assertion -> TestTree
testCase String
"gatherSimp23" Assertion
testGatherSimp23
  , String -> Assertion -> TestTree
testCase String
"gatherTranspose33" Assertion
testGatherTranspose33
  , String -> Assertion -> TestTree
testCase String
"gatherTransposeBuild33" Assertion
testGatherTransposeBuild33
  , String -> Assertion -> TestTree
testCase String
"gatherTransposeBuild331" Assertion
testGatherTransposeBuild331
  , String -> Assertion -> TestTree
testCase String
"gatherTransposeBuild332" Assertion
testGatherTransposeBuild332
  , String -> Assertion -> TestTree
testCase String
"gatherTransposeBuild333" Assertion
testGatherTransposeBuild333
  , String -> Assertion -> TestTree
testCase String
"gatherTransposeBuild334" Assertion
testGatherTransposeBuild334
  , String -> Assertion -> TestTree
testCase String
"gatherTransposeBuild335" Assertion
testGatherTransposeBuild335
  , String -> Assertion -> TestTree
testCase String
"gatherTransposeBuild336" Assertion
testGatherTransposeBuild336
  , String -> Assertion -> TestTree
testCase String
"gatherSimpPP33" Assertion
testGatherSimpPP33
  , String -> Assertion -> TestTree
testCase String
"gatherSimpPP34" Assertion
testGatherSimpPP34
{- TODO: re-enable the tests once we drop GHC 9.10
   (they don't type-check with 9.10)
  , testCase "gatherCond" testGatherCond
  , testCase "gatherCondBuild" testGatherCondBuild
  , testCase "gatherCond2" testGatherCond2
  , testCase "gatherCondBuild2" testGatherCondBuild2
  , testCase "gatherSimpCond" testGatherSimpCond
  , testCase "gatherCond3" testGatherCond3
  , testCase "gatherCondBuild3" testGatherCondBuild3
  , testCase "gatherCond4" testGatherCond4
  , testCase "gatherCondBuild4" testGatherCondBuild4
  , testCase "gatherSimpCond3" testGatherSimpCond3
  , testCase "gatherCond5" testGatherCond5
  , testCase "gatherCondBuild5" testGatherCondBuild5
  , testCase "gatherCond6" testGatherCond6
  , testCase "gatherCondBuild6" testGatherCondBuild6
  , testCase "gatherSimpCond5" testGatherSimpCond5
-}

  , String -> Assertion -> TestTree
testCase String
"scatterNested1" Assertion
testScatterNested1
  , String -> Assertion -> TestTree
testCase String
"scatterNestedBuild1" Assertion
testScatterNestedBuild1
  , String -> Assertion -> TestTree
testCase String
"scatter1" Assertion
testScatter1
  , String -> Assertion -> TestTree
testCase String
"scatterBuild1" Assertion
testScatterBuild1
  , String -> Assertion -> TestTree
testCase String
"scatterSimpPP1" Assertion
testScatterSimpPP1
  , String -> Assertion -> TestTree
testCase String
"scatterSimp1" Assertion
testScatterSimp1
  , String -> Assertion -> TestTree
testCase String
"scatterNested2" Assertion
testScatterNested2
  , String -> Assertion -> TestTree
testCase String
"scatterNestedBuild2" Assertion
testScatterNestedBuild2
  , String -> Assertion -> TestTree
testCase String
"scatter2" Assertion
testScatter2
  , String -> Assertion -> TestTree
testCase String
"scatterBuild2" Assertion
testScatterBuild2
  , String -> Assertion -> TestTree
testCase String
"scatterSimpPP2" Assertion
testScatterSimpPP2
  , String -> Assertion -> TestTree
testCase String
"scatterSimp2" Assertion
testScatterSimp2
  , String -> Assertion -> TestTree
testCase String
"scatterNested12" Assertion
testScatterNested12
  , String -> Assertion -> TestTree
testCase String
"scatterNestedBuild12" Assertion
testScatterNestedBuild12
  , String -> Assertion -> TestTree
testCase String
"scatter12" Assertion
testScatter12
  , String -> Assertion -> TestTree
testCase String
"scatterBuild12" Assertion
testScatterBuild12
  , String -> Assertion -> TestTree
testCase String
"scatterSimpPP12" Assertion
testScatterSimpPP12
  , String -> Assertion -> TestTree
testCase String
"scatterSimp12" Assertion
testScatterSimp12

  , String -> Assertion -> TestTree
testCase String
"shmatterBarReluADVal320" Assertion
testBarReluADVal320
  , String -> Assertion -> TestTree
testCase String
"shmatterReluSimpPP" Assertion
testReluSimpPP

  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP2" Assertion
testCNNOPP2
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP2b" Assertion
testCNNOPP2b
--  , testCase "sminimizedCNNOPP3" testCNNOPP3
--  , testCase "sminimizedCNNOPP3b" testCNNOPP3b
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP4" Assertion
testCNNOPP4
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP4b" Assertion
testCNNOPP4b
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP5" Assertion
testCNNOPP5
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP5b" Assertion
testCNNOPP5b
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP6" Assertion
testCNNOPP6
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP6b" Assertion
testCNNOPP6b
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP7" Assertion
testCNNOPP7
  , String -> Assertion -> TestTree
testCase String
"sminimizedCNNOPP7b" Assertion
testCNNOPP7b
  , String -> Assertion -> TestTree
testCase String
"minimizedCNNOPP4bU" Assertion
testCNNOPP4bU
  ]


-- * Gathers

gatherNested1 :: forall target r. (ADReady target, GoodScalar r)
              => target (TKR 2 r) -> target (TKR 1 r)
gatherNested1 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gatherNested1 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @1
          (Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          (forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @1
                   (Int
4 Int -> ShR (0 + 1) Int -> ShR ((0 + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) target (TKR 2 r)
target (TKR2 (1 + 1) (TKScalar r))
t
                   (\(IntOf target
k3 :.: IxR n (IntOf target)
ZIR) -> IntOf target
k3 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR))
          (\(IntOf target
i2 :.: IxR n (IntOf target)
ZIR) -> IntOf target
i2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i2 IntOf target -> IxROf target 1 -> IxROf target (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
i2 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testGatherNested1 :: Assertion
testGatherNested1 :: Assertion
testGatherNested1 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @1 f (TKR 2 Double) -> f (TKR2 1 (TKScalar Double))
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gatherNested1 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherNestedBuild1 :: Assertion
testGatherNestedBuild1 :: Assertion
testGatherNestedBuild1 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
3.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
3.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (PrimalOf f (TKScalar Int64) -> f (TKR2 1 (TKScalar Double)))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
5 (\PrimalOf f (TKScalar Int64)
i ->
             BoolOf f
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double))
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (PrimalOf f (TKScalar Int64)
i PrimalOf f (TKScalar Int64)
-> PrimalOf f (TKScalar Int64) -> BoolOf (PrimalOf f)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
>. PrimalOf f (TKScalar Int64)
2) (f (TKR 2 Double) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gatherNested1 f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t) (f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR2 (1 + 1) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64))
-> f (TKR2 1 (TKScalar Double))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
i])))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

gather1 :: forall target r. (ADReady target, GoodScalar r)
        => target (TKR 2 r) -> target (TKR 1 r)
gather1 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gather1 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @1
          (Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          (Int
-> Int
-> target (TKR2 (1 + 1) (TKScalar r))
-> target (TKR2 (1 + 1) (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
Int -> Int -> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
rslice Int
0 Int
4 target (TKR 2 r)
target (TKR2 (1 + 1) (TKScalar r))
t)
          (\(IntOf target
i2 :.: IxR n (IntOf target)
ZIR) -> IntOf target
i2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i2 IntOf target -> IxROf target 1 -> IxROf target (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
i2 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testGather1 :: Assertion
testGather1 :: Assertion
testGather1 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @1 f (TKR 2 Double) -> f (TKR2 1 (TKScalar Double))
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gather1
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherBuild1 :: Assertion
testGatherBuild1 :: Assertion
testGatherBuild1 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
3.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
3.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (PrimalOf f (TKScalar Int64) -> f (TKR2 1 (TKScalar Double)))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
5 (\PrimalOf f (TKScalar Int64)
i ->
             BoolOf f
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double))
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (PrimalOf f (TKScalar Int64)
i PrimalOf f (TKScalar Int64)
-> PrimalOf f (TKScalar Int64) -> BoolOf (PrimalOf f)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
>. PrimalOf f (TKScalar Int64)
2) (f (TKR 2 Double) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gather1 f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t) (f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR2 (1 + 1) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64))
-> f (TKR2 1 (TKScalar Double))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
i])))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherSimpPP1 :: Assertion
testGatherSimpPP1 :: Assertion
testGatherSimpPP1 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gatherNested1 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 1 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 1 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
315
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gather1 (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 1 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 1 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
315
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 1 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1))
    Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 1 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2))

testGatherSimp1 :: Assertion
testGatherSimp1 :: Assertion
testGatherSimp1 = do
  let varName :: AstVarName s (TKR 2 Float)
varName = FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName s (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName s (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName s (TKR 2 Float))
-> Int -> AstVarName s (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000
      var :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var = AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar AstVarName PrimalSpan (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName
      vals :: [Float]
vals = [-Item [Float]
1, Item [Float]
0, Item [Float]
2.0,Item [Float]
5.0,Item [Float]
11.0,-Item [Float]
17.0,Item [Float]
23.0,Item [Float]
29.0,-Item [Float]
35.0,Item [Float]
41.0,Item [Float]
47.0,Item [Float]
33.0, Item [Float]
0.1, Item [Float]
0.007]
      env :: AstEnv Concrete
env = AstVarName (ZonkAny @AstSpanType 7) (TKR 2 Float)
-> Concrete (TKR 2 Float) -> AstEnv Concrete -> AstEnv Concrete
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName (ZonkAny @AstSpanType 7) (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName (IShR 2 -> [Float] -> Concrete (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals) AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gatherNested1 @(AstTensor AstMethodLet PrimalSpan) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gather1 (IShR 2
-> [Float] -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals)
  let !t1n :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n = AstNoSimplify PrimalSpan (TKR 1 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 1 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 1 Float))
-> AstNoSimplify PrimalSpan (TKR 1 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 1 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gatherNested1 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 1 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2n :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2n = AstNoSimplify PrimalSpan (TKR 1 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 1 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 1 Float))
-> AstNoSimplify PrimalSpan (TKR 1 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 1 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
gather1 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 1 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1)
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n)
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2)
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2n)
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2n

gatherNested02 :: forall target r. (ADReady target, GoodScalar r)
               => target (TKR 1 r) -> target (TKR 1 r)
gatherNested02 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 1 r) -> target (TKR 1 r)
gatherNested02 target (TKR 1 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @1
          (Int
1 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          (forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @1
                   (Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) target (TKR 1 r)
target (TKR2 (1 + 0) (TKScalar r))
t
                   (\(IntOf target
k3 :.: IxR n (IntOf target)
ZIR) -> IntOf target
k3 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
k3 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR))
          (\(IntOf target
i1 :.: IxR n (IntOf target)
ZIR) -> IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IxR 0 (IntOf target) -> IxROf target (1 + 0)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testGatherNested02 :: Assertion
testGatherNested02 :: Assertion
testGatherNested02 =
  Rational
-> Concrete (TKR2 1 (TKScalar Double))
-> ((Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double))),
    (Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
4] [Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @1 f (TKR2 1 (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (f :: Target).
ADReady f =>
f (TKR2 1 (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 1 r) -> target (TKR 1 r)
gatherNested02 (Int
-> Concrete (TKR 0 Double)
-> Concrete (TKR2 (1 + 0) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
4 (Double -> Concrete (TKR 0 Double)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKR 0 r)
rscalar Double
0.1)))

gatherNested2 :: forall target r. (ADReady target, GoodScalar r)
              => target (TKR 2 r) -> target (TKR 2 r)
gatherNested2 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested2 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @2
          (Int
2 Int -> IShR 1 -> IShR (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
3 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          (forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @3
                   (Int
2 Int
-> ShR (((0 + 1) + 1) + 1) Int -> ShR ((((0 + 1) + 1) + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
3 Int -> ShR ((0 + 1) + 1) Int -> ShR (((0 + 1) + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
4 Int -> ShR (0 + 1) Int -> ShR ((0 + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) target (TKR 2 r)
target (TKR2 (1 + 1) (TKScalar r))
t
                   (\(IntOf target
k1 :.: IntOf target
k2 :.: IntOf target
k3 :.: IxR n (IntOf target)
ZIR) -> IntOf target
k1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
k2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
k3 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR))
          (\(IntOf target
i1 :.: IntOf target
i2 :.: IxR n (IntOf target)
ZIR) -> IntOf target
i1 IntOf target -> IxROf target 3 -> IxROf target (3 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
i2 IntOf target -> IxROf target 2 -> IxROf target 3
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i2 IntOf target -> IxROf target 1 -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
i1 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testGatherNested2 :: Assertion
testGatherNested2 :: Assertion
testGatherNested2 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested2
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherNestedBuild2 :: Assertion
testGatherNestedBuild2 :: Assertion
testGatherNestedBuild2 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
6.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
6.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf f
i ->
             f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested2 (f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall a. Num a => a -> a -> a
* IShR 2 -> f (TKR 0 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

gather2 :: forall target r. (ADReady target, GoodScalar r)
        => target (TKR 2 r) -> target (TKR 2 r)
gather2 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather2 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @2
          (Int
2 Int -> IShR 1 -> IShR (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
3 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          target (TKR 2 r)
target (TKR2 (2 + 0) (TKScalar r))
t
          (\(IntOf target
i1 :.: IntOf target
i2 :.: IxR n (IntOf target)
ZIR) -> IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i2 IntOf target -> IxR 1 (IntOf target) -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
i1 IntOf target -> IxR 0 (IntOf target) -> IxR 1 (IntOf target)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testGather2 :: Assertion
testGather2 :: Assertion
testGather2 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather2 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherBuild2 :: Assertion
testGatherBuild2 :: Assertion
testGatherBuild2 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
6.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
6.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf f
i ->
             f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather2 (f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall a. Num a => a -> a -> a
* IShR 2 -> f (TKR 0 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherSimpPP2 :: Assertion
testGatherSimpPP2 :: Assertion
testGatherSimpPP2 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested2 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
582
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather2 (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
394
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
582
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
394

testGatherSimp2 :: Assertion
testGatherSimp2 :: Assertion
testGatherSimp2 = do
  let varName :: AstVarName s (TKR 2 Float)
varName = FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName s (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName s (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName s (TKR 2 Float))
-> Int -> AstVarName s (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000
      var :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var = AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar AstVarName PrimalSpan (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName
      vals :: [Float]
vals = [-Item [Float]
1, Item [Float]
0, Item [Float]
2.0,Item [Float]
5.0,Item [Float]
11.0,-Item [Float]
17.0,Item [Float]
23.0,Item [Float]
29.0,-Item [Float]
35.0,Item [Float]
41.0,Item [Float]
47.0,Item [Float]
33.0, Item [Float]
0.1, Item [Float]
0.007]
      env :: AstEnv Concrete
env = AstVarName (ZonkAny @AstSpanType 6) (TKR 2 Float)
-> Concrete (TKR 2 Float) -> AstEnv Concrete -> AstEnv Concrete
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName (ZonkAny @AstSpanType 6) (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName (IShR 2 -> [Float] -> Concrete (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals) AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested2 @(AstTensor AstMethodLet PrimalSpan) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather2 (IShR 2
-> [Float] -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals)
  let !t1n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested2 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather2 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n

gatherNested12 :: forall target r. (ADReady target, GoodScalar r)
               => target (TKR 2 r) -> target (TKR 2 r)
gatherNested12 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested12 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @1
          (Int
2 Int -> IShR 1 -> IShR (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
4 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          (forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @3
                   (Int
2 Int -> ShR ((0 + 1) + 1) Int -> ShR (((0 + 1) + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
3 Int -> ShR (0 + 1) Int -> ShR ((0 + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
4 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) target (TKR 2 r)
target (TKR2 (2 + 0) (TKScalar r))
t
                   (\(IntOf target
k1 :.: IntOf target
k2 :.: IntOf target
k3 :.: IxR n (IntOf target)
ZIR) -> IntOf target
k1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
k2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
k3 IntOf target -> IxROf target 1 -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
k1 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR))
          (\(IntOf target
i1 :.: IxR n (IntOf target)
ZIR) -> IntOf target
i1 IntOf target -> IxROf target 1 -> IxROf target ((0 + 1) + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testGatherNested12 :: Assertion
testGatherNested12 :: Assertion
testGatherNested12 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested12 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherNestedBuild12 :: Assertion
testGatherNestedBuild12 :: Assertion
testGatherNestedBuild12 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
4.0,Double
Item [Double]
4.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> f (TKR2 (1 + 2) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64)) -> f (TKR 2 Double)
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
rindex (Int
-> (PrimalOf f (TKScalar Int64) -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
5 (\PrimalOf f (TKScalar Int64)
i ->
             BoolOf f
-> f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (PrimalOf f (TKScalar Int64)
i PrimalOf f (TKScalar Int64)
-> PrimalOf f (TKScalar Int64) -> BoolOf (PrimalOf f)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
>. PrimalOf f (TKScalar Int64)
2) (f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested12 f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t)
                          (PermR -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0] (f (TKR 2 Double) -> f (TKR 2 Double))
-> f (TKR 2 Double) -> f (TKR 2 Double)
forall a b. (a -> b) -> a -> b
$ Int
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
4 (f (TKR2 1 (TKScalar Double))
 -> f (TKR2 (1 + 1) (TKScalar Double)))
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR2 (1 + 1) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64))
-> f (TKR2 1 (TKScalar Double))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
i]))) [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
1])
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

gather12 :: forall target r. (ADReady target, GoodScalar r)
         => target (TKR 2 r) -> target (TKR 2 r)
gather12 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather12 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
rgather @2
          (Int
2 Int -> IShR 1 -> IShR (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
4 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          target (TKR 2 r)
target (TKR2 (2 + 0) (TKScalar r))
t
          (\(IntOf target
i1 :.: IntOf target
k3 :.: IxR n (IntOf target)
ZIR) -> IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
k3 IntOf target -> IxR 1 (IntOf target) -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target
i1 IntOf target -> IxR 0 (IntOf target) -> IxR 1 (IntOf target)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testGather12 :: Assertion
testGather12 :: Assertion
testGather12 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0,Double
Item [Double]
0.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather12 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherBuild12 :: Assertion
testGatherBuild12 :: Assertion
testGatherBuild12 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
4.0,Double
Item [Double]
4.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> f (TKR2 (1 + 2) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64)) -> f (TKR 2 Double)
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
rindex (Int
-> (PrimalOf f (TKScalar Int64) -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
5 (\PrimalOf f (TKScalar Int64)
i ->
             BoolOf f
-> f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (PrimalOf f (TKScalar Int64)
i PrimalOf f (TKScalar Int64)
-> PrimalOf f (TKScalar Int64) -> BoolOf (PrimalOf f)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
>. PrimalOf f (TKScalar Int64)
2) (f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather12 f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t)
                          (PermR -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0] (f (TKR 2 Double) -> f (TKR 2 Double))
-> f (TKR 2 Double) -> f (TKR 2 Double)
forall a b. (a -> b) -> a -> b
$ Int
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
4 (f (TKR2 1 (TKScalar Double))
 -> f (TKR2 (1 + 1) (TKScalar Double)))
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR2 (1 + 1) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64))
-> f (TKR2 1 (TKScalar Double))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
i]))) [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
1])
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherSimpPP12 :: Assertion
testGatherSimpPP12 :: Assertion
testGatherSimpPP12 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested12 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
515
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather12 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
341
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
515
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
341

testGatherSimp12 :: Assertion
testGatherSimp12 :: Assertion
testGatherSimp12 = do
  let varName :: AstVarName s (TKR 2 Float)
varName = FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName s (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName s (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName s (TKR 2 Float))
-> Int -> AstVarName s (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000
      var :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var = AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar AstVarName PrimalSpan (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName
      vals :: [Float]
vals = [-Item [Float]
1, Item [Float]
0, Item [Float]
2.0,Item [Float]
5.0,Item [Float]
11.0,-Item [Float]
17.0,Item [Float]
23.0,Item [Float]
29.0,-Item [Float]
35.0,Item [Float]
41.0,Item [Float]
47.0,Item [Float]
33.0, Item [Float]
0.1, Item [Float]
0.007]
      env :: AstEnv Concrete
env = AstVarName (ZonkAny @AstSpanType 5) (TKR 2 Float)
-> Concrete (TKR 2 Float) -> AstEnv Concrete -> AstEnv Concrete
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName (ZonkAny @AstSpanType 5) (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName (IShR 2 -> [Float] -> Concrete (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals) AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested12 @(AstTensor AstMethodLet PrimalSpan) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather12 (IShR 2
-> [Float] -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals)
  let !t1n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherNested12 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gather12 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n

gatherReshape22 :: forall target r. (ADReady target, GoodScalar r)
                => target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22 target (TKR 2 r)
t =
  forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @6 [Int
Item (IShR 2)
2, Int
Item (IShR 2)
6]
  (target (TKR2 6 (TKScalar r)) -> target (TKR 2 r))
-> target (TKR2 6 (TKScalar r)) -> target (TKR 2 r)
forall a b. (a -> b) -> a -> b
$ IShR 6
-> target (TKR2 (((0 + 1) + 1) + 1) (TKScalar r))
-> target (TKR2 6 (TKScalar r))
forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape [Int
Item (IShR 6)
3, Int
Item (IShR 6)
1, Int
Item (IShR 6)
2, Int
Item (IShR 6)
1, Int
Item (IShR 6)
1, Int
Item (IShR 6)
2]
  (target (TKR2 (((0 + 1) + 1) + 1) (TKScalar r))
 -> target (TKR2 6 (TKScalar r)))
-> target (TKR2 (((0 + 1) + 1) + 1) (TKScalar r))
-> target (TKR2 6 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @4 (Int
1 Int -> ShR ((0 + 1) + 1) Int -> ShR (((0 + 1) + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
12 Int -> ShR (0 + 1) Int -> ShR ((0 + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
1 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
  (target (TKR2 4 (TKScalar r))
 -> target (TKR2 (((0 + 1) + 1) + 1) (TKScalar r)))
-> target (TKR2 4 (TKScalar r))
-> target (TKR2 (((0 + 1) + 1) + 1) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @3 [Int
Item (IShR 4)
3, Int
Item (IShR 4)
1, Int
Item (IShR 4)
1, Int
Item (IShR 4)
4]
  (target (TKR2 3 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 3 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ IShR 3 -> target (TKR 2 r) -> target (TKR2 3 (TKScalar r))
forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape [Int
Item (IShR 3)
2, Int
Item (IShR 3)
2, Int
Item (IShR 3)
3] target (TKR 2 r)
t

testGatherReshape22 :: Assertion
testGatherReshape22 :: Assertion
testGatherReshape22 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
6,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22
                               (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
6 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherReshapeBuild22 :: Assertion
testGatherReshapeBuild22 :: Assertion
testGatherReshapeBuild22 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
6,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf f
i ->
             f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22 (f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall a. Num a => a -> a -> a
* IShR 2 -> f (TKR 0 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
6 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testGatherSimpPP22 :: Assertion
testGatherSimpPP22 :: Assertion
testGatherSimpPP22 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
159
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
159
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @2 @2 [Int
Item (IShR 2)
2, Int
Item (IShR 2)
6]
            (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
159
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
159

testGatherSimp22 :: Assertion
testGatherSimp22 :: Assertion
testGatherSimp22 = do
  let varName :: AstVarName s (TKR 2 Float)
varName = FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName s (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName s (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName s (TKR 2 Float))
-> Int -> AstVarName s (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000
      var :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var = AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar AstVarName PrimalSpan (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName
      vals :: [Float]
vals = [-Item [Float]
1, Item [Float]
0, Item [Float]
2.0,Item [Float]
5.0,Item [Float]
11.0,-Item [Float]
17.0,Item [Float]
23.0,Item [Float]
29.0,-Item [Float]
35.0,Item [Float]
41.0,Item [Float]
47.0,Item [Float]
33.0]
      env :: AstEnv Concrete
env = AstVarName (ZonkAny @AstSpanType 3) (TKR 2 Float)
-> Concrete (TKR 2 Float) -> AstEnv Concrete -> AstEnv Concrete
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName (ZonkAny @AstSpanType 3) (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName (IShR 2 -> [Float] -> Concrete (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] [Float]
vals) AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22 @(AstTensor AstMethodLet PrimalSpan) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @2 @2 [Int
Item (IShR 2)
2, Int
Item (IShR 2)
6] (IShR 2
-> [Float] -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] [Float]
vals)
  let !t1n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @2 @2 [Int
Item (IShR 2)
2, Int
Item (IShR 2)
6] (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n

testGatherSimpPP23 :: Assertion
testGatherSimpPP23 :: Assertion
testGatherSimpPP23 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1 = (\AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t -> Int
-> (IntOf (AstTensor AstMethodLet PrimalSpan)
    -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor
     AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf (AstTensor AstMethodLet PrimalSpan)
i ->
              forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22 @(AstTensor AstMethodLet PrimalSpan)
                (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a. Num a => a -> a -> a
* IShR 2
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] (IntOf (AstTensor AstMethodLet PrimalSpan)
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf (AstTensor AstMethodLet PrimalSpan)
i))))
            (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float)))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
-> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
450
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
-> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 3 Float) AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
456
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2 = (\AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t -> Int
-> (IntOf (AstTensor AstMethodLet PrimalSpan)
    -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor
     AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf (AstTensor AstMethodLet PrimalSpan)
i ->
              forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @2 @2 [Int
Item (IShR 2)
2, Int
Item (IShR 2)
6]
                (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a. Num a => a -> a -> a
* IShR 2
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] (IntOf (AstTensor AstMethodLet PrimalSpan)
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf (AstTensor AstMethodLet PrimalSpan)
i))))
            (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float)))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
-> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
450
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
-> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 3 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
456

testGatherSimp23 :: Assertion
testGatherSimp23 :: Assertion
testGatherSimp23 = do
  let varName :: AstVarName s (TKR 2 Float)
varName = FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName s (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName s (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName s (TKR 2 Float))
-> Int -> AstVarName s (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000
      var :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var = AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar AstVarName PrimalSpan (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName
      vals :: [Float]
vals = [-Item [Float]
1, Item [Float]
0, Item [Float]
2.0,Item [Float]
5.0,Item [Float]
11.0,-Item [Float]
17.0,Item [Float]
23.0,Item [Float]
29.0,-Item [Float]
35.0,Item [Float]
41.0,Item [Float]
47.0,Item [Float]
33.0]
      env :: AstEnv Concrete
env = AstVarName (ZonkAny @AstSpanType 4) (TKR 2 Float)
-> Concrete (TKR 2 Float) -> AstEnv Concrete -> AstEnv Concrete
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName (ZonkAny @AstSpanType 4) (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName (IShR 2 -> [Float] -> Concrete (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] [Float]
vals) AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t1 = (\AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t -> Int
-> (IntOf (AstTensor AstMethodLet PrimalSpan)
    -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor
     AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf (AstTensor AstMethodLet PrimalSpan)
i ->
              forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22 @(AstTensor AstMethodLet PrimalSpan)
                (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a. Num a => a -> a -> a
* IShR 2
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] (IntOf (AstTensor AstMethodLet PrimalSpan)
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf (AstTensor AstMethodLet PrimalSpan)
i)))) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t2 = (\AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t -> Int
-> (IntOf (AstTensor AstMethodLet PrimalSpan)
    -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor
     AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf (AstTensor AstMethodLet PrimalSpan)
i ->
              forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @2 @2 [Int
Item (IShR 2)
2, Int
Item (IShR 2)
6]
                (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a. Num a => a -> a -> a
* IShR 2
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] (IntOf (AstTensor AstMethodLet PrimalSpan)
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf (AstTensor AstMethodLet PrimalSpan)
i)))) (IShR 2
-> [Float] -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] [Float]
vals)
  let !t1n :: AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1n = AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float))
 -> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float)))
-> AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ (\AstNoSimplify PrimalSpan (TKR 2 Float)
t -> Int
-> (IntOf (AstNoSimplify PrimalSpan)
    -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf (AstNoSimplify PrimalSpan)
i ->
              AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
gatherReshape22
                (AstNoSimplify PrimalSpan (TKR 2 Float)
t AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a. Num a => a -> a -> a
* IShR 2
-> AstNoSimplify PrimalSpan (TKR2 0 (TKScalar Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] (IntOf (AstNoSimplify PrimalSpan)
-> AstNoSimplify PrimalSpan (TKR2 0 (TKScalar Float))
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf (AstNoSimplify PrimalSpan)
i)))) (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float)))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2n :: AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2n = AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float))
 -> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float)))
-> AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ (\AstNoSimplify PrimalSpan (TKR 2 Float)
t -> Int
-> (IntOf (AstNoSimplify PrimalSpan)
    -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf (AstNoSimplify PrimalSpan)
i ->
              forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @2 @2 [Int
Item (IShR 2)
2, Int
Item (IShR 2)
6]
                (AstNoSimplify PrimalSpan (TKR 2 Float)
t AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a. Num a => a -> a -> a
* IShR 2
-> AstNoSimplify PrimalSpan (TKR2 0 (TKScalar Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
6, Int
Item (IShR 2)
2] (IntOf (AstNoSimplify PrimalSpan)
-> AstNoSimplify PrimalSpan (TKR2 0 (TKScalar Float))
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf (AstNoSimplify PrimalSpan)
i)))) (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float)))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR2 3 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t1
    Concrete (TKR2 3 (TKScalar Float))
-> Concrete (TKR2 3 (TKScalar Float)) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1n
    Concrete (TKR2 3 (TKScalar Float))
-> Concrete (TKR2 3 (TKScalar Float)) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t2
    Concrete (TKR2 3 (TKScalar Float))
-> Concrete (TKR2 3 (TKScalar Float)) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 3 Float) AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t1)
    Concrete (TKR2 3 (TKScalar Float))
-> Concrete (TKR2 3 (TKScalar Float)) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t1
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 3 Float) AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1n)
    Concrete (TKR2 3 (TKScalar Float))
-> Concrete (TKR2 3 (TKScalar Float)) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 3 Float) AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t2)
    Concrete (TKR2 3 (TKScalar Float))
-> Concrete (TKR2 3 (TKScalar Float)) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
AstTensor AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 3 Float) AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2n)
    Concrete (TKR2 3 (TKScalar Float))
-> Concrete (TKR2 3 (TKScalar Float)) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2n

-- 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)
                  => target (TKR 10 r) -> target (TKR 2 r)
gatherTranspose33 :: forall (target :: Target) r.
(ADReady target, GoodScalar r, RealFloat r) =>
target (TKR 10 r) -> target (TKR 2 r)
gatherTranspose33 target (TKR 10 r)
t =
  target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)
rmatmul2 (IShR 2 -> target (TKR2 7 (TKScalar r)) -> target (TKR 2 r)
forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape [Int
Item (IShR 2)
6, Int
Item (IShR 2)
8] (Ranked 7 r -> target (TKR2 7 (TKScalar r))
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 7 r -> target (TKR2 7 (TKScalar r)))
-> Ranked 7 r -> target (TKR2 7 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 7 (TKScalar r)) -> RepConcrete (TKR2 7 (TKScalar r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 7 (TKScalar r))
forall r. (GoodScalar r, Fractional r) => Concrete (TKR 7 r)
t48))
    (target (TKR2 (2 + 0) (TKScalar r))
-> target (TKR2 (2 + 0) (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
target (TKR2 (2 + n) x) -> target (TKR2 (2 + n) x)
rtr
     (target (TKR2 (2 + 0) (TKScalar r))
 -> target (TKR2 (2 + 0) (TKScalar r)))
-> target (TKR2 (2 + 0) (TKScalar r))
-> target (TKR2 (2 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @4 [Int
Item (IShR 2)
16, Int
Item (IShR 2)
8]
     (target (TKR2 4 (TKScalar r))
 -> target (TKR2 (2 + 0) (TKScalar r)))
-> target (TKR2 4 (TKScalar r))
-> target (TKR2 (2 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
2]
     (target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
2, Int
Item PermR
0, Int
Item PermR
1]
     (target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
0]
     (target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0, Int
Item PermR
2]
     (target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0]
     (target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
3]
     (target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
3, Int
Item PermR
0]
     (target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
3, Int
Item PermR
0, Int
Item PermR
2, Int
Item PermR
1]
     (target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r)))
-> target (TKR2 4 (TKScalar r)) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ IShR 4 -> target (TKR 10 r) -> target (TKR2 4 (TKScalar r))
forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
8, Int
Item (IShR 4)
4]
     (target (TKR 10 r) -> target (TKR2 4 (TKScalar r)))
-> target (TKR 10 r) -> target (TKR2 4 (TKScalar r))
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
3]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
3, Int
Item PermR
0]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0, Int
Item PermR
3, Int
Item PermR
2]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
3, Int
Item PermR
4, Int
Item PermR
5, Int
Item PermR
6, Int
Item PermR
7, Int
Item PermR
9, Int
Item PermR
8]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
3, Int
Item PermR
7, Int
Item PermR
5, Int
Item PermR
6, Int
Item PermR
4]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
3, Int
Item PermR
4, Int
Item PermR
5, Int
Item PermR
6]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
5, Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
3, Int
Item PermR
4]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
2, Int
Item PermR
4, Int
Item PermR
3, Int
Item PermR
5, Int
Item PermR
6, Int
Item PermR
7, Int
Item PermR
9, Int
Item PermR
8]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose []
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0]
     (target (TKR 10 r) -> target (TKR 10 r))
-> target (TKR 10 r) -> target (TKR 10 r)
forall a b. (a -> b) -> a -> b
$ PermR -> target (TKR 10 r) -> target (TKR 10 r)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
0, Int
Item PermR
1, Int
Item PermR
7, Int
Item PermR
4, Int
Item PermR
5, Int
Item PermR
3, Int
Item PermR
6, Int
Item PermR
2, Int
Item PermR
8]
       target (TKR 10 r)
t)

testGatherTranspose33 :: Assertion
testGatherTranspose33 :: Assertion
testGatherTranspose33 =
  Rational
-> Concrete (TKR 10 Double)
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double)),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double)))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 10 -> [Double] -> Concrete (TKR 10 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1] [Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
81.3003,Double
Item [Double]
71.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
80.0,Double
Item [Double]
79.0,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
166.8003,Double
Item [Double]
137.70326,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002,Double
Item [Double]
186.1003,Double
Item [Double]
162.3889400002])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 10 Double) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR 10 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r, RealFloat r) =>
target (TKR 10 r) -> target (TKR 2 r)
gatherTranspose33 Concrete (TKR 10 Double)
forall r. (GoodScalar r, Fractional r) => Concrete (TKR 10 r)
t128)

testGatherTransposeBuild33 :: Assertion
testGatherTransposeBuild33 :: Assertion
testGatherTransposeBuild33 =
  Rational
-> Concrete (TKR 10 Double)
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double)),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double), Concrete (TKR 10 Double),
     Concrete (TKR 10 Double)))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 10 -> [Double] -> Concrete (TKR 10 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1] [Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
487.80179999999996,Double
Item [Double]
426.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
480.0,Double
Item [Double]
474.0,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1000.8018,Double
Item [Double]
826.21956,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012,Double
Item [Double]
1116.6018,Double
Item [Double]
974.3336400012])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR 10 Double)
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf f
i ->
             f (TKR 10 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r, RealFloat r) =>
target (TKR 10 r) -> target (TKR 2 r)
gatherTranspose33 (f (TKR 10 Double)
t f (TKR 10 Double) -> f (TKR 10 Double) -> f (TKR 10 Double)
forall a. Num a => a -> a -> a
* IShR 10 -> f (TKR 0 Double) -> f (TKR 10 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
          Concrete (TKR 10 Double)
forall r. (GoodScalar r, Fractional r) => Concrete (TKR 10 r)
t128)

testGatherTransposeBuild331 :: Assertion
testGatherTransposeBuild331 :: Assertion
testGatherTransposeBuild331 =
  Rational
-> Concrete (TKR 2 Double)
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
3] [Double
Item [Double]
1,Double
Item [Double]
1,Double
Item [Double]
1,Double
Item [Double]
1,Double
Item [Double]
1,Double
Item [Double]
1])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR 2 Double)
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
2 (\IntOf f
i ->
             PermR -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0] (f (TKR 2 Double)
t f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall a. Num a => a -> a -> a
* IShR 2 -> f (TKR 0 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
2, Int
Item (IShR 2)
3] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
          (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
3] [Double
Item [Double]
1,Double
Item [Double]
2,Double
Item [Double]
3,Double
Item [Double]
4,Double
Item [Double]
5,Double
Item [Double]
6]))

testGatherTransposeBuild332 :: Assertion
testGatherTransposeBuild332 :: Assertion
testGatherTransposeBuild332 =
  Rational
-> Concrete (TKR 2 Double)
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
3] [Double
Item [Double]
1,Double
Item [Double]
1,Double
Item [Double]
1,Double
Item [Double]
1,Double
Item [Double]
1,Double
Item [Double]
1])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR 2 Double)
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
2 (\IntOf f
i ->
             PermR -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0] (f (TKR 2 Double)
t f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall a. Num a => a -> a -> a
* IShR 2 -> f (TKR 0 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
2, Int
Item (IShR 2)
3] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
          (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
3] [Double
Item [Double]
1,Double
Item [Double]
2,Double
Item [Double]
3,Double
Item [Double]
4,Double
Item [Double]
5,Double
Item [Double]
6]))

testGatherTransposeBuild333 :: Assertion
testGatherTransposeBuild333 :: Assertion
testGatherTransposeBuild333 =
  Rational
-> Concrete (TKR2 1 (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
1,Double
Item [Double]
1])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 1 (TKScalar Double))
t -> Int
-> (IntOf f -> f (TKR2 1 (TKScalar Double)))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
2 (\IntOf f
i ->
             f (TKR2 1 (TKScalar Double))
t f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall a. Num a => a -> a -> a
* IShR 1 -> f (TKR 0 Double) -> f (TKR2 1 (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 1)
2] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i)))
          (IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0,Double
Item [Double]
0]))

testGatherTransposeBuild334 :: Assertion
testGatherTransposeBuild334 :: Assertion
testGatherTransposeBuild334 =
  Rational
-> Concrete (TKR 2 Double)
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
1] [Double
Item [Double]
1,Double
Item [Double]
1])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR 2 Double)
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
2 (\IntOf f
i ->
             f (TKR 2 Double)
t f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall a. Num a => a -> a -> a
* Int
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
2 (Int -> f (TKR 0 Double) -> f (TKR2 (1 + 0) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
1 (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
         (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
1] [Double
Item [Double]
1,Double
Item [Double]
2]))

testGatherTransposeBuild335 :: Assertion
testGatherTransposeBuild335 :: Assertion
testGatherTransposeBuild335 =
  Rational
-> Concrete (TKR 2 Double)
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
1] [Double
Item [Double]
1,Double
Item [Double]
1])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR 2 Double)
t ->
             Int -> f (TKR 2 Double) -> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
2 f (TKR 2 Double)
t f (TKR 3 Double) -> f (TKR 3 Double) -> f (TKR 3 Double)
forall a. Num a => a -> a -> a
* PermR -> f (TKR 3 Double) -> f (TKR 3 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
2,Int
Item PermR
0,Int
Item PermR
1] (Int -> f (TKR 2 Double) -> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
2 (Int
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
1 (forall r1 r2 (target :: Target) (n :: Nat).
(GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target) =>
target (TKR n r1) -> target (TKR n r2)
rfromIntegral @Int64 (Ranked 1 Int64 -> f (TKR 1 Int64)
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 1 Int64 -> f (TKR 1 Int64))
-> Ranked 1 Int64 -> f (TKR 1 Int64)
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Int64] -> Ranked 1 Int64
forall a (n :: Nat). PrimElt a => IShR n -> [a] -> Ranked n a
Nested.rfromListPrimLinear ([Item (IShR 1)] -> IShR 1
forall l. IsList l => [Item l] -> l
fromList [Int
Item PermR
2]) [Int64
Item [Int64]
0, Int64
Item [Int64]
1])))))
         (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
1] [Double
Item [Double]
1,Double
Item [Double]
2]))

testGatherTransposeBuild336 :: Assertion
testGatherTransposeBuild336 :: Assertion
testGatherTransposeBuild336 =
  Rational
-> Concrete (TKR 2 Double)
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double)))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
1] [Double
Item [Double]
1,Double
Item [Double]
1])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR 2 Double)
t ->
             Int -> f (TKR 2 Double) -> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
2 f (TKR 2 Double)
t f (TKR 3 Double) -> f (TKR 3 Double) -> f (TKR 3 Double)
forall a. Num a => a -> a -> a
* PermR -> f (TKR 3 Double) -> f (TKR 3 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
2,Int
Item PermR
0,Int
Item PermR
1] (Int -> f (TKR 2 Double) -> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
2 (Int
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
1 (NonEmpty (f (TKR 0 Double)) -> f (TKR2 (1 + 0) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
NonEmpty (target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rfromList [Double -> f (TKR 0 Double)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKR 0 r)
rscalar Double
0, Double -> f (TKR 0 Double)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKR 0 r)
rscalar Double
1]))))
         (IShR 2 -> [Double] -> Concrete (TKR 2 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
2, Int
Item (IShR 2)
1] [Double
Item [Double]
1,Double
Item [Double]
2]))

-- These are different terms, but they should have similar lengths,
-- because they differ only by single transpose and reshape, most probably,
-- and all the rest of the element reordering should cancel out.
-- Still, probably impossible to lower the gap to zero.
testGatherSimpPP33 :: Assertion
testGatherSimpPP33 :: Assertion
testGatherSimpPP33 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r, RealFloat r) =>
target (TKR 10 r) -> target (TKR 2 r)
gatherTranspose33 @(AstTensor AstMethodLet PrimalSpan)
            (AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 10 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 10 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 10
-> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 10 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 10 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 10 Float))
-> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
1117
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
847
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = (\AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)
rmatmul2 (IShR 2
-> AstTensor AstMethodLet PrimalSpan (TKR2 7 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape [Int
Item (IShR 2)
6, Int
Item (IShR 2)
8] (Ranked 7 Float
-> AstTensor AstMethodLet PrimalSpan (TKR2 7 (TKScalar Float))
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 7 Float
 -> AstTensor AstMethodLet PrimalSpan (TKR2 7 (TKScalar Float)))
-> Ranked 7 Float
-> AstTensor AstMethodLet PrimalSpan (TKR2 7 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 7 (TKScalar Float))
-> RepConcrete (TKR2 7 (TKScalar Float))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 7 (TKScalar Float))
forall r. (GoodScalar r, Fractional r) => Concrete (TKR 7 r)
t48))
                            (forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @10 [Int
Item (IShR 2)
8, Int
Item (IShR 2)
16] AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t))
            (AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 10 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 10 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 10
-> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 10 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 10 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 10 Float))
-> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
796
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
526

testGatherSimpPP34 :: Assertion
testGatherSimpPP34 :: Assertion
testGatherSimpPP34 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1 = (\AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t -> Int
-> (IntOf (AstTensor AstMethodLet PrimalSpan)
    -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor
     AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf (AstTensor AstMethodLet PrimalSpan)
i ->
             forall (target :: Target) r.
(ADReady target, GoodScalar r, RealFloat r) =>
target (TKR 10 r) -> target (TKR 2 r)
gatherTranspose33 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall a. Num a => a -> a -> a
* IShR 10
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1] (IntOf (AstTensor AstMethodLet PrimalSpan)
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf (AstTensor AstMethodLet PrimalSpan)
i))))
            (AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float)))
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 10 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 10 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 10
-> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 10 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 10 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 10 Float))
-> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
-> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
2504
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
-> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 3 Float) AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
19838
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2 = (\AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t -> Int
-> (IntOf (AstTensor AstMethodLet PrimalSpan)
    -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor
     AstMethodLet PrimalSpan (TKR2 (1 + 2) (TKScalar Float))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf (AstTensor AstMethodLet PrimalSpan)
i ->
              (\AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t' -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)
rmatmul2 (IShR 2
-> AstTensor AstMethodLet PrimalSpan (TKR2 7 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape [Int
Item (IShR 2)
6, Int
Item (IShR 2)
8] (Ranked 7 Float
-> AstTensor AstMethodLet PrimalSpan (TKR2 7 (TKScalar Float))
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 7 Float
 -> AstTensor AstMethodLet PrimalSpan (TKR2 7 (TKScalar Float)))
-> Ranked 7 Float
-> AstTensor AstMethodLet PrimalSpan (TKR2 7 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 7 (TKScalar Float))
-> RepConcrete (TKR2 7 (TKScalar Float))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 7 (TKScalar Float))
forall r. (GoodScalar r, Fractional r) => Concrete (TKR 7 r)
t48))
                               (forall (n :: Nat) (m :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
rreshape @10 [Int
Item (IShR 2)
8, Int
Item (IShR 2)
16] AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t'))
                (AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall a. Num a => a -> a -> a
* IShR 10
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1] (IntOf (AstTensor AstMethodLet PrimalSpan)
-> AstTensor AstMethodLet PrimalSpan (TKR2 0 (TKScalar Float))
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf (AstTensor AstMethodLet PrimalSpan)
i))))
            (AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float)))
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 10 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 10 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 10
-> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 10 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
2, Int
Item (IShR 10)
1] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 10 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 10 Float))
-> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
-> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
2145
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
-> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 3 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR2 3 (TKScalar Float))
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
19479

{- TODO: re-enable the tests once we drop GHC 9.10
   (they don't type-check with 9.10)
gatherCond :: forall target r. (ADReady target, GoodScalar r)
           => target (TKR 2 r) -> target (TKR 2 r)
gatherCond u =
  let v = rtranspose [2, 0, 1] $ rreplicate (2 * rwidth u) u
  in rgather [rwidth u, 2] v (\(i :.: j :.: ZIR) ->
                                ifH (i ==. 3) 0 j :.: 2 * i :.: i :.: ZIR)

testGatherCond :: Assertion
testGatherCond =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [7,2]
                 [1.0,1.0,1.0,1.0,1.0,1.0,2.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0])
    (rev' @Double @2 gatherCond (rreplicate 7 $ ringestData [2] [0, 1]))

testGatherCondBuild :: Assertion
testGatherCondBuild =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [7,2]
                 [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))))
          (rreplicate 7 $ ringestData [2] [0, 1]))

gatherCond2 :: forall target r. (ADReady target, GoodScalar r)
            => target (TKR 2 r) -> target (TKR 2 r)
gatherCond2 u =
  let v = rreplicate (2 * rwidth u) u
  in rtr $ rgather [2, rwidth u] v (\(j :.: i :.: ZIR) ->
                                      2 * i :.: i :.: ifH (i ==. 3) 0 j :.: ZIR)

testGatherCond2 :: Assertion
testGatherCond2 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [7,2]
                 [1.0,1.0,1.0,1.0,1.0,1.0,2.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0])
    (rev' @Double @2 gatherCond2 (rreplicate 7 $ ringestData [2] [0, 1]))

testGatherCondBuild2 :: Assertion
testGatherCondBuild2 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [7,2]
                 [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))))
          (rreplicate 7 $ ringestData [2] [0, 1]))

testGatherSimpCond :: Assertion
testGatherSimpCond = do
  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . 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
  let !t1 = gatherCond @(AstTensor AstMethodLet PrimalSpan) var
  let !t2 = gatherCond2 (ringestData [7, 2] vals)
  let !t1n = unAstNoSimplify $ gatherCond $ AstNoSimplify var
  let !t2n = unAstNoSimplify $ gatherCond2 $ AstNoSimplify var
  interpretAstPrimal @Concrete env t1
    @?= interpretAstPrimal @Concrete env t1n
  interpretAstPrimal @Concrete env t1n
    @?= interpretAstPrimal @Concrete emptyEnv t2
  interpretAstPrimal @Concrete emptyEnv t2
    @?= interpretAstPrimal @Concrete env t2n
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t1)
    @?= interpretAstPrimal @Concrete env t1
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t1n)
    @?= interpretAstPrimal @Concrete env t1n
  interpretAstPrimal @Concrete emptyEnv
                     (simplifyInlineContract @(TKR 2 Float) t2)
    @?= interpretAstPrimal @Concrete emptyEnv t2
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t2n)
    @?= interpretAstPrimal @Concrete env t2n

gatherCond3 :: forall target r. (ADReady target, GoodScalar r)
            => target (TKR 2 r) -> target (TKR 2 r)
gatherCond3 u =
  let v = rtranspose [2, 0, 1] $ rreplicate (2 * rwidth u) u
  in rgather [rwidth u, 2] v (\(i :.: j :.: ZIR) ->
                                2 * i :.: i :.: ifH (i ==. 3) 0 j :.: ZIR)

testGatherCond3 :: Assertion
testGatherCond3 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [7,2]
                 [1.0,0.0,1.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 @2 gatherCond3 (rreplicate 7 $ ringestData [2] [0, 1]))

testGatherCondBuild3 :: Assertion
testGatherCondBuild3 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [7,2]
                 [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))))
          (rreplicate 7 $ ringestData [2] [0, 1]))

gatherCond4 :: forall target r. (ADReady target, GoodScalar r)
            => target (TKR 2 r) -> target (TKR 2 r)
gatherCond4 u =
  let v = rreplicate (2 * rwidth u) u
  in rtr $ rgather [2, rwidth u] v (\(j :.: i :.: ZIR) ->
                                      i :.: ifH (i ==. 3) 0 j :.: 2 * i :.: ZIR)

testGatherCond4 :: Assertion
testGatherCond4 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [7,2]
                 [1.0,0.0,1.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 @2 gatherCond4 (rreplicate 7 $ ringestData [2] [0, 1]))

testGatherCondBuild4 :: Assertion
testGatherCondBuild4 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [7,2]
                 [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))))
          (rreplicate 7 $ ringestData [2] [0, 1]))

testGatherSimpCond3 :: Assertion
testGatherSimpCond3 = do
  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . 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
  let !t1 = gatherCond3 @(AstTensor AstMethodLet PrimalSpan) var
  let !t2 = gatherCond4 (ringestData [7, 2] vals)
  let !t1n = unAstNoSimplify $ gatherCond3 $ AstNoSimplify var
  let !t2n = unAstNoSimplify $ gatherCond4 $ AstNoSimplify var
  interpretAstPrimal @Concrete env t1
    @?= interpretAstPrimal @Concrete env t1n
  interpretAstPrimal @Concrete env t1n
    @?= interpretAstPrimal @Concrete emptyEnv t2
  interpretAstPrimal @Concrete emptyEnv t2
    @?= interpretAstPrimal @Concrete env t2n
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t1)
    @?= interpretAstPrimal @Concrete env t1
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t1n)
    @?= interpretAstPrimal @Concrete env t1n
  interpretAstPrimal @Concrete emptyEnv
                     (simplifyInlineContract @(TKR 2 Float) t2)
    @?= interpretAstPrimal @Concrete emptyEnv t2
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t2n)
    @?= interpretAstPrimal @Concrete env t2n

gatherCond5 :: forall target r. (ADReady target, GoodScalar r)
            => target (TKR 3 r) -> target (TKR 2 r)
gatherCond5 v =
  rgather [rwidth v, 2] v (\(i :.: j :.: ZIR) ->
                             ifH (i ==. 1) 0 j :.: 2 * i :.: i :.: ZIR)

testGatherCond5 :: Assertion
testGatherCond5 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [2,4,2]
                 [1.0,0.0,0.0,0.0,0.0,2.0,0.0,0.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])
    (rev' @Double @2 gatherCond5 (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))

testGatherCondBuild5 :: Assertion
testGatherCondBuild5 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [2,4,2]
                 [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))))
          (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))

gatherCond6 :: forall target r. (ADReady target, GoodScalar r)
            => target (TKR 3 r) -> target (TKR 2 r)
gatherCond6 u =
  let v = rtranspose [2, 0, 1] u
  in rtr $ rgather [2, rwidth v] v (\(j :.: i :.: ZIR) ->
                                      i :.: ifH (i ==. 1) 0 j :.: 2 * i :.: ZIR)

testGatherCond6 :: Assertion
testGatherCond6 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [2,4,2]
                 [1.0,0.0,0.0,0.0,0.0,2.0,0.0,0.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])
    (rev' @Double @2 gatherCond6 (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))

testGatherCondBuild6 :: Assertion
testGatherCondBuild6 =
  assertEqualUpToEpsilon' 1e-10
    (ringestData [2,4,2]
                 [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))))
          (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))

testGatherSimpCond5 :: Assertion
testGatherSimpCond5 = do
  let varName = mkAstVarName (FTKR [2,4,2] FTKScalar) Nothing . 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
  let !t1 = gatherCond5 @(AstTensor AstMethodLet PrimalSpan) var
  let !t2 = gatherCond6 (ringestData [2,4,2] vals)
  let !t1n = unAstNoSimplify $ gatherCond5 $ AstNoSimplify var
  let !t2n = unAstNoSimplify $ gatherCond6 $ AstNoSimplify var
  interpretAstPrimal @Concrete env t1
    @?= interpretAstPrimal @Concrete env t1n
  interpretAstPrimal @Concrete env t1n
    @?= interpretAstPrimal @Concrete emptyEnv t2
  interpretAstPrimal @Concrete emptyEnv t2
    @?= interpretAstPrimal @Concrete env t2n
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t1)
    @?= interpretAstPrimal @Concrete env t1
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t1n)
    @?= interpretAstPrimal @Concrete env t1n
  interpretAstPrimal @Concrete emptyEnv
                     (simplifyInlineContract @(TKR 2 Float) t2)
    @?= interpretAstPrimal @Concrete emptyEnv t2
  interpretAstPrimal @Concrete env
                     (simplifyInlineContract @(TKR 2 Float) t2n)
    @?= interpretAstPrimal @Concrete env t2n
-}


-- * Scatters instead of gathers

scatterNested1 :: forall target r. (ADReady target, GoodScalar r)
               => target (TKR 2 r) -> target (TKR 1 r)
scatterNested1 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatterNested1 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @2
          (Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          (forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @1
                   (Int
7 Int -> ShR (0 + 1) Int -> ShR ((0 + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) target (TKR 2 r)
target (TKR2 (1 + 1) (TKScalar r))
t
                   (\(IntOf target
k3 :.: IxR n (IntOf target)
ZIR) -> IntOf target
k3 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR))
          (\(IntOf target
i1 :.: IntOf target
i2 :.: IxR n (IntOf target)
ZIR) -> IntOf target
i2 IntOf target -> IntOf target -> IntOf target
forall a. IntegralH a => a -> a -> a
`quotH` (IntOf target
1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1) IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testScatterNested1 :: Assertion
testScatterNested1 :: Assertion
testScatterNested1 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @1 f (TKR 2 Double) -> f (TKR2 1 (TKScalar Double))
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatterNested1 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterNestedBuild1 :: Assertion
testScatterNestedBuild1 :: Assertion
testScatterNestedBuild1 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (PrimalOf f (TKScalar Int64) -> f (TKR2 1 (TKScalar Double)))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
5 (\PrimalOf f (TKScalar Int64)
i ->
             BoolOf f
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double))
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (PrimalOf f (TKScalar Int64)
i PrimalOf f (TKScalar Int64)
-> PrimalOf f (TKScalar Int64) -> BoolOf (PrimalOf f)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
>. PrimalOf f (TKScalar Int64)
2) (f (TKR 2 Double) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatterNested1 f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t) (f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR2 (1 + 1) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64))
-> f (TKR2 1 (TKScalar Double))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
i])))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

scatter1 :: forall target r. (ADReady target, GoodScalar r)
         => target (TKR 2 r) -> target (TKR 1 r)
scatter1 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatter1 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @2
          (Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          target (TKR 2 r)
target (TKR2 (2 + 0) (TKScalar r))
t
          (\(IntOf target
i1 :.: IntOf target
i2 :.: IxR n (IntOf target)
ZIR) -> IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH (IntOf target
i2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
* IntOf target
i1) IntOf target
1 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testScatter1 :: Assertion
testScatter1 :: Assertion
testScatter1 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 1 (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @1 f (TKR 2 Double) -> f (TKR2 1 (TKScalar Double))
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatter1 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterBuild1 :: Assertion
testScatterBuild1 :: Assertion
testScatterBuild1 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
3.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0,Double
Item [Double]
2.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (PrimalOf f (TKScalar Int64) -> f (TKR2 1 (TKScalar Double)))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
5 (\PrimalOf f (TKScalar Int64)
i ->
             BoolOf f
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 1 (TKScalar Double))
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (PrimalOf f (TKScalar Int64)
i PrimalOf f (TKScalar Int64)
-> PrimalOf f (TKScalar Int64) -> BoolOf (PrimalOf f)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
>. PrimalOf f (TKScalar Int64)
2) (f (TKR 2 Double) -> f (TKR2 1 (TKScalar Double))
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatter1 f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t) (f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR2 (1 + 1) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64))
-> f (TKR2 1 (TKScalar Double))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
i])))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterSimpPP1 :: Assertion
testScatterSimpPP1 :: Assertion
testScatterSimpPP1 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatterNested1 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 1 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 1 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
397
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatter1 (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 1 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 1 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
492
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 1 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
397
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 1 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) @PrimalSpan AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
492

testScatterSimp1 :: Assertion
testScatterSimp1 :: Assertion
testScatterSimp1 = do
  let varName :: AstVarName s (TKR 2 Float)
varName = FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName s (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName s (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName s (TKR 2 Float))
-> Int -> AstVarName s (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000
      var :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var = AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar AstVarName PrimalSpan (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName
      vals :: [Float]
vals = [-Item [Float]
1, Item [Float]
0, Item [Float]
2.0,Item [Float]
5.0,Item [Float]
11.0,-Item [Float]
17.0,Item [Float]
23.0,Item [Float]
29.0,-Item [Float]
35.0,Item [Float]
41.0,Item [Float]
47.0,Item [Float]
33.0, Item [Float]
0.1, Item [Float]
0.007]
      env :: AstEnv Concrete
env = AstVarName (ZonkAny @AstSpanType 2) (TKR 2 Float)
-> Concrete (TKR 2 Float) -> AstEnv Concrete -> AstEnv Concrete
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName (ZonkAny @AstSpanType 2) (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName (IShR 2 -> [Float] -> Concrete (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals) AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatterNested1 @(AstTensor AstMethodLet PrimalSpan) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatter1 (IShR 2
-> [Float] -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals)
  let !t1n :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n = AstNoSimplify PrimalSpan (TKR 1 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 1 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 1 Float))
-> AstNoSimplify PrimalSpan (TKR 1 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 1 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatterNested1 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 1 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2n :: AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2n = AstNoSimplify PrimalSpan (TKR 1 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 1 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 1 Float))
-> AstNoSimplify PrimalSpan (TKR 1 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 1 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 1 r)
scatter1 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 1 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 1 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n
  -- TODO: scatter fusion isn't sound? or just incorrectly manually done here?
  -- interpretAstPrimal @Concrete env t1n
  --   @?= interpretAstPrimal @Concrete emptyEnv t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1)
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n)
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2)
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 1 Float) AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2n)
    Concrete (TKR 1 Float) -> Concrete (TKR 1 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 1 Float)
t2n

scatterNested2 :: forall target r. (ADReady target, GoodScalar r)
              => target (TKR 2 r) -> target (TKR 2 r)
scatterNested2 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested2 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @4
          (Int
2 Int -> IShR 1 -> IShR (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
3 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          (forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @1
                   (Int
2 Int
-> ShR (((0 + 1) + 1) + 1) Int -> ShR ((((0 + 1) + 1) + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
3 Int -> ShR ((0 + 1) + 1) Int -> ShR (((0 + 1) + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
4 Int -> ShR (0 + 1) Int -> ShR ((0 + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
2 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) target (TKR 2 r)
target (TKR2 (1 + 1) (TKScalar r))
t
                   (\(IntOf target
k1 :.: IxR n (IntOf target)
ZIR) -> IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH IntOf target
k1 IntOf target
1 IntOf target -> IxROf target 2 -> IxROf target 3
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH IntOf target
k1 IntOf target
2  IntOf target -> IxROf target 1 -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH IntOf target
k1 IntOf target
3 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR))
          (\(IntOf target
i1 :.: IntOf target
i2 :.: IntOf target
_i3 :.: IntOf target
i4 :.: IxR n (IntOf target)
ZIR) ->
            IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH (IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i2) IntOf target
1 IntOf target -> IxROf target 1 -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH (IntOf target
i4 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1) IntOf target
2 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testScatterNested2 :: Assertion
testScatterNested2 :: Assertion
testScatterNested2 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested2 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterNestedBuild2 :: Assertion
testScatterNestedBuild2 :: Assertion
testScatterNestedBuild2 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf f
i ->
             f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested2 (f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall a. Num a => a -> a -> a
* IShR 2 -> f (TKR 0 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

scatter2 :: forall target r. (ADReady target, GoodScalar r)
        => target (TKR 2 r) -> target (TKR 2 r)
scatter2 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter2 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @2
          (Int
2 Int -> IShR 1 -> IShR (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
3 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          target (TKR 2 r)
target (TKR2 (2 + 0) (TKScalar r))
t
          (\(IntOf target
i1 :.: IntOf target
i2 :.: IxR n (IntOf target)
ZIR) -> IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH (IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i2) IntOf target
1 IntOf target -> IxR 1 (IntOf target) -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH IntOf target
i1 IntOf target
2 IntOf target -> IxR 0 (IntOf target) -> IxR 1 (IntOf target)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testScatter2 :: Assertion
testScatter2 :: Assertion
testScatter2 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter2 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterBuild2 :: Assertion
testScatterBuild2 :: Assertion
testScatterBuild2 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 3 Double),
     Concrete (TKR 3 Double), Concrete (TKR 3 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0,Double
Item [Double]
6.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @3
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> Int
-> (IntOf f -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
4 (\IntOf f
i ->
             f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter2 (f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall a. Num a => a -> a -> a
* IShR 2 -> f (TKR 0 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] (IntOf f -> f (TKR 0 Double)
forall r (target :: Target).
(BaseTensor target, ConvertTensor target, GoodScalar r) =>
IntOf target -> target (TKR 0 r)
rfromIndex0 IntOf f
i))))
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterSimpPP2 :: Assertion
testScatterSimpPP2 :: Assertion
testScatterSimpPP2 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested2 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
1022
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter2 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
782
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
1022
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
782

testScatterSimp2 :: Assertion
testScatterSimp2 :: Assertion
testScatterSimp2 = do
  let varName :: AstVarName s (TKR 2 Float)
varName = FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName s (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName s (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName s (TKR 2 Float))
-> Int -> AstVarName s (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000
      var :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var = AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar AstVarName PrimalSpan (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName
      vals :: [Float]
vals = [-Item [Float]
1, Item [Float]
0, Item [Float]
2.0,Item [Float]
5.0,Item [Float]
11.0,-Item [Float]
17.0,Item [Float]
23.0,Item [Float]
29.0,-Item [Float]
35.0,Item [Float]
41.0,Item [Float]
47.0,Item [Float]
33.0, Item [Float]
0.1, Item [Float]
0.007]
      env :: AstEnv Concrete
env = AstVarName (ZonkAny @AstSpanType 1) (TKR 2 Float)
-> Concrete (TKR 2 Float) -> AstEnv Concrete -> AstEnv Concrete
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName (ZonkAny @AstSpanType 1) (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName (IShR 2 -> [Float] -> Concrete (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals) AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested2 @(AstTensor AstMethodLet PrimalSpan) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter2 (IShR 2
-> [Float] -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals)
  let !t1n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested2 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter2 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  -- TODO: scatter fusion isn't sound? or just incorrectly manually done here?
  -- interpretAstPrimal @Concrete env t1n
  --  @?= interpretAstPrimal @Concrete emptyEnv t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n

scatterNested12 :: forall target r. (ADReady target, GoodScalar r)
               => target (TKR 2 r) -> target (TKR 2 r)
scatterNested12 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested12 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @2
          (Int
2 Int -> IShR 1 -> IShR (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
4 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          (forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @2
                   (Int
2 Int -> ShR ((0 + 1) + 1) Int -> ShR (((0 + 1) + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
3 Int -> ShR (0 + 1) Int -> ShR ((0 + 1) + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
4 Int -> ShR 0 Int -> ShR (0 + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) target (TKR 2 r)
target (TKR2 (2 + 0) (TKScalar r))
t
                   (\(IntOf target
k1 :.: IntOf target
k2 :.: IxR n (IntOf target)
ZIR) ->
                     IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH IntOf target
k1 IntOf target
1 IntOf target -> IxROf target 2 -> IxROf target 3
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH (IntOf target
k2 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
k1) IntOf target
2 IntOf target -> IxROf target 1 -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH IntOf target
k1 IntOf target
3 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR))
          (\(IntOf target
i1 :.: IntOf target
_i2 :.: IxR n (IntOf target)
ZIR) -> IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH (IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1) IntOf target
1 IntOf target -> IxR 0 (IntOf target) -> IxROf target 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testScatterNested12 :: Assertion
testScatterNested12 :: Assertion
testScatterNested12 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested12 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterNestedBuild12 :: Assertion
testScatterNestedBuild12 :: Assertion
testScatterNestedBuild12 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
4.0,Double
Item [Double]
4.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> f (TKR2 (1 + 2) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64)) -> f (TKR 2 Double)
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
rindex (Int
-> (PrimalOf f (TKScalar Int64) -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
5 (\PrimalOf f (TKScalar Int64)
i ->
             BoolOf f
-> f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (PrimalOf f (TKScalar Int64)
i PrimalOf f (TKScalar Int64)
-> PrimalOf f (TKScalar Int64) -> BoolOf (PrimalOf f)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
>. PrimalOf f (TKScalar Int64)
2) (f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested12 f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t)
                          (PermR -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0] (f (TKR 2 Double) -> f (TKR 2 Double))
-> f (TKR 2 Double) -> f (TKR 2 Double)
forall a b. (a -> b) -> a -> b
$ Int
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
4 (f (TKR2 1 (TKScalar Double))
 -> f (TKR2 (1 + 1) (TKScalar Double)))
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR2 (1 + 1) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64))
-> f (TKR2 1 (TKScalar Double))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
i]))) [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
1])
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

scatter12 :: forall target r. (ADReady target, GoodScalar r)
         => target (TKR 2 r) -> target (TKR 2 r)
scatter12 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter12 target (TKR 2 r)
t =
  forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK)
       (target :: Target).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x,
 BaseTensor target) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
rscatter @2
          (Int
2 Int -> IShR 1 -> IShR (1 + 1)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
4 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR)
          target (TKR 2 r)
target (TKR2 (2 + 0) (TKScalar r))
t
          (\(IntOf target
i1 :.: IntOf target
k3 :.: IxR n (IntOf target)
ZIR) -> IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH (IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
i1 IntOf target -> IntOf target -> IntOf target
forall a. Num a => a -> a -> a
+ IntOf target
k3) IntOf target
1 IntOf target -> IxR 1 (IntOf target) -> IxROf target 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IntOf target -> IntOf target -> IntOf target
forall (y :: TK) (target :: Target).
(KnownSTK y, OrdH target y, BaseTensor target) =>
target y -> target y -> target y
minH IntOf target
i1 IntOf target
3 IntOf target -> IxR 0 (IntOf target) -> IxR 1 (IntOf target)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (IntOf target)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)

testScatter12 :: Assertion
testScatter12 :: Assertion
testScatter12 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0,Double
Item [Double]
1.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2 f (TKR 2 Double) -> f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (f :: Target).
ADReady f =>
f (TKR2 (1 + 1) (TKScalar Double)) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter12 (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterBuild12 :: Assertion
testScatterBuild12 :: Assertion
testScatterBuild12 =
  Rational
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
-> ((Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))),
    (Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     AstTensor AstMethodLet PrimalSpan (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)), Concrete (TKR 2 Double),
     Concrete (TKR 2 Double), Concrete (TKR 2 Double),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double)),
     Concrete (TKR2 (1 + 1) (TKScalar Double))))
-> Assertion
forall (n :: Nat) (m :: Nat) v r w a.
(KnownNat n, KnownNat m,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type),
 AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v,
 AssertEqualUpToEpsilon (ADTensorScalar r), GoodScalar r,
 GoodScalar (ADTensorScalar r), HasCallStack) =>
Rational
-> 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), 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), 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
assertEqualUpToEpsilon' Rational
1e-10
    (IShR (1 + 1)
-> [Double] -> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR (1 + 1))
7,Int
Item (IShR (1 + 1))
2]
                 [Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
4.0,Double
Item [Double]
4.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0])
    (forall r (m :: Nat) (n :: Nat) v a w.
(KnownNat m, KnownNat n, GoodScalar r,
 (v :: Type) ~ (Concrete (TKR m r) :: Type),
 (w :: Type) ~ (Concrete (ADTensorKind (TKR m r)) :: Type),
 (a :: Type) ~ (Concrete (ADTensorKind (TKR n r)) :: Type)) =>
(forall (f :: Target). 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), 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), 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))
rev' @Double @2
          (\f (TKR2 (1 + 1) (TKScalar Double))
t -> f (TKR2 (1 + 2) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64)) -> f (TKR 2 Double)
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
rindex (Int
-> (PrimalOf f (TKScalar Int64) -> f (TKR 2 Double))
-> f (TKR2 (1 + 2) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
rbuild1 Int
5 (\PrimalOf f (TKScalar Int64)
i ->
             BoolOf f
-> f (TKR 2 Double) -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (PrimalOf f (TKScalar Int64)
i PrimalOf f (TKScalar Int64)
-> PrimalOf f (TKScalar Int64) -> BoolOf (PrimalOf f)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
>. PrimalOf f (TKScalar Int64)
2) (f (TKR 2 Double) -> f (TKR 2 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter12 f (TKR 2 Double)
f (TKR2 (1 + 1) (TKScalar Double))
t)
                          (PermR -> f (TKR 2 Double) -> f (TKR 2 Double)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
rtranspose [Int
Item PermR
1, Int
Item PermR
0] (f (TKR 2 Double) -> f (TKR 2 Double))
-> f (TKR 2 Double) -> f (TKR 2 Double)
forall a b. (a -> b) -> a -> b
$ Int
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
4 (f (TKR2 1 (TKScalar Double))
 -> f (TKR2 (1 + 1) (TKScalar Double)))
-> f (TKR2 1 (TKScalar Double))
-> f (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ f (TKR2 (1 + 1) (TKScalar Double))
t f (TKR2 (1 + 1) (TKScalar Double))
-> IxR 1 (PrimalOf f (TKScalar Int64))
-> f (TKR2 1 (TKScalar Double))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
i]))) [Item (IxR 1 (PrimalOf f (TKScalar Int64)))
PrimalOf f (TKScalar Int64)
1])
          (Int
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
rreplicate Int
7 (Concrete (TKR2 1 (TKScalar Double))
 -> Concrete (TKR2 (1 + 1) (TKScalar Double)))
-> Concrete (TKR2 1 (TKScalar Double))
-> Concrete (TKR2 (1 + 1) (TKScalar Double))
forall a b. (a -> b) -> a -> b
$ IShR 1 -> [Double] -> Concrete (TKR2 1 (TKScalar Double))
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 1)
2] [Double
Item [Double]
0, Double
Item [Double]
1]))

testScatterSimpPP12 :: Assertion
testScatterSimpPP12 :: Assertion
testScatterSimpPP12 = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested12 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
952
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter12 @(AstTensor AstMethodLet PrimalSpan) (AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 2 Float))
-> Int -> AstVarName PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
625
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
952
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 2 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
625

testScatterSimp12 :: Assertion
testScatterSimp12 :: Assertion
testScatterSimp12 = do
  let varName :: AstVarName s (TKR 2 Float)
varName = FullShapeTK (TKR 2 Float)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 2 -> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 2 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName s (TKR 2 Float))
-> (Int -> AstVarId) -> Int -> AstVarName s (TKR 2 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName s (TKR 2 Float))
-> Int -> AstVarName s (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000
      var :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var = AstVarName PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar AstVarName PrimalSpan (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName
      vals :: [Float]
vals = [-Item [Float]
1, Item [Float]
0, Item [Float]
2.0,Item [Float]
5.0,Item [Float]
11.0,-Item [Float]
17.0,Item [Float]
23.0,Item [Float]
29.0,-Item [Float]
35.0,Item [Float]
41.0,Item [Float]
47.0,Item [Float]
33.0, Item [Float]
0.1, Item [Float]
0.007]
      env :: AstEnv Concrete
env = AstVarName (ZonkAny @AstSpanType 0) (TKR 2 Float)
-> Concrete (TKR 2 Float) -> AstEnv Concrete -> AstEnv Concrete
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName (ZonkAny @AstSpanType 0) (TKR 2 Float)
forall {s :: AstSpanType}. AstVarName s (TKR 2 Float)
varName (IShR 2 -> [Float] -> Concrete (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals) AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1 = forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested12 @(AstTensor AstMethodLet PrimalSpan) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2 = AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter12 (IShR 2
-> [Float] -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 2)
7, Int
Item (IShR 2)
2] [Float]
vals)
  let !t1n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatterNested12 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  let !t2n :: AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n = AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 2 r) -> target (TKR 2 r)
scatter12 (AstNoSimplify PrimalSpan (TKR 2 Float)
 -> AstNoSimplify PrimalSpan (TKR 2 Float))
-> AstNoSimplify PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
-> AstNoSimplify PrimalSpan (TKR 2 Float)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
var
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t1n
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
forall (target :: Target). AstEnv target
emptyEnv AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2
  forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env
                     (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 2 Float) AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n)
    Concrete (TKR 2 Float) -> Concrete (TKR 2 Float) -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target
-> AstTensor AstMethodLet PrimalSpan y -> PrimalOf target y
interpretAstPrimal @Concrete AstEnv Concrete
env AstTensor AstMethodLet PrimalSpan (TKR 2 Float)
t2n

foo :: RealFloatH a => (a,a,a) -> a
foo :: forall a. RealFloatH a => (a, a, a) -> a
foo (a
x,a
y,a
z) =
  let w :: a
w = a
x a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sin a
y
  in a -> a -> a
forall a. RealFloatH a => a -> a -> a
atan2H a
z a
w a -> a -> a
forall a. Num a => a -> a -> a
+ a
z a -> a -> a
forall a. Num a => a -> a -> a
* a
w

bar :: forall a. RealFloatH a => (a, a) -> a
bar :: forall a. RealFloatH a => (a, a) -> a
bar (a
x, a
y) =
  let w :: a
w = (a, a, a) -> a
forall a. RealFloatH a => (a, a, a) -> a
foo (a
x, a
y, a
x) a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sin a
y
  in a -> a -> a
forall a. RealFloatH a => a -> a -> a
atan2H a
x a
w a -> a -> a
forall a. Num a => a -> a -> a
+ a
y a -> a -> a
forall a. Num a => a -> a -> a
* a
w

barRelu
  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )
  => target (TKR n r) -> target (TKR n r)
barRelu :: forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
barRelu target (TKR2 n (TKScalar r))
x = let t :: target (TKR2 n (TKScalar r))
t = IShR n
-> target (TKR2 0 (TKScalar r)) -> target (TKR2 n (TKScalar r))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
IShR n -> target (TKR2 0 x) -> target (TKR2 n x)
rreplicate0N (target (TKR2 n (TKScalar r)) -> IShR n
forall (n :: Nat) (x :: TK).
KnownSTK x =>
target (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape target (TKR2 n (TKScalar r))
x) (r -> target (TKR2 0 (TKScalar r))
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKR 0 r)
rscalar r
0.001) target (TKR2 n (TKScalar r))
-> target (TKR2 n (TKScalar r)) -> target (TKR2 n (TKScalar r))
forall a. Num a => a -> a -> a
* target (TKR2 n (TKScalar r))
x
            in target (TKR2 n (TKScalar r)) -> target (TKR2 n (TKScalar r))
forall (target :: Target) (n :: Nat) r.
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
relu (target (TKR2 n (TKScalar r)) -> target (TKR2 n (TKScalar r)))
-> target (TKR2 n (TKScalar r)) -> target (TKR2 n (TKScalar r))
forall a b. (a -> b) -> a -> b
$ (target (TKR2 n (TKScalar r)), target (TKR2 n (TKScalar r)))
-> target (TKR2 n (TKScalar r))
forall a. RealFloatH a => (a, a) -> a
bar (target (TKR2 n (TKScalar r))
t, target (TKR2 n (TKScalar r)) -> target (TKR2 n (TKScalar r))
forall (target :: Target) (n :: Nat) r.
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
relu target (TKR2 n (TKScalar r))
t)

barRelu10xSlower
  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )
  => target (TKR n r) -> target (TKR n r)
barRelu10xSlower :: forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
barRelu10xSlower target (TKR n r)
x = let t :: target (TKR n r)
t = (target (TKR2 0 (TKScalar r)) -> target (TKR2 0 (TKScalar r)))
-> target (TKR n r) -> target (TKR n r)
forall (n :: Nat) (x :: TK) (x1 :: TK) (target :: Target).
(KnownNat n, KnownSTK x, KnownSTK x1, BaseTensor target) =>
(target (TKR2 0 x1) -> target (TKR2 0 x))
-> target (TKR2 n x1) -> target (TKR2 n x)
rmap0N (target (TKR2 0 (TKScalar r))
-> target (TKR2 0 (TKScalar r)) -> target (TKR2 0 (TKScalar r))
forall a. Num a => a -> a -> a
* r -> target (TKR2 0 (TKScalar r))
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKR 0 r)
rscalar r
0.001) target (TKR n r)
x
                     in target (TKR n r) -> target (TKR n r)
forall (target :: Target) (n :: Nat) r.
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
relu (target (TKR n r) -> target (TKR n r))
-> target (TKR n r) -> target (TKR n r)
forall a b. (a -> b) -> a -> b
$ (target (TKR n r), target (TKR n r)) -> target (TKR n r)
forall a. RealFloatH a => (a, a) -> a
bar (target (TKR n r)
t, target (TKR n r) -> target (TKR n r)
forall (target :: Target) (n :: Nat) r.
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
relu target (TKR n r)
t)

testBarReluADVal320 :: Assertion
testBarReluADVal320 :: Assertion
testBarReluADVal320 =
  Rational
-> Concrete (TKR 10 Double)
-> Concrete (TKR 10 Double)
-> Assertion
forall a.
(AssertEqualUpToEpsilon a, HasCallStack) =>
Rational -> a -> a -> Assertion
assertEqualUpToEpsilon Rational
1e-10
    (IShR 10 -> [Double] -> Concrete (TKR 10 Double)
forall (n :: Nat) r (target :: Target).
(GoodScalar r, BaseTensor target) =>
IShR n -> [r] -> target (TKR n r)
ringestData [Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1] [Double
Item [Double]
2.885038541771792e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.885034988157713e-4,Double
Item [Double]
2.885923176600045e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.886097295122454e-4,Double
Item [Double]
2.8846476339094805e-4,Double
Item [Double]
2.885038541771792e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.8851415976532735e-4,Double
Item [Double]
2.885923176600045e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.8849246223035154e-4,Double
Item [Double]
2.884182085399516e-4,Double
Item [Double]
2.884075468755327e-4,Double
Item [Double]
2.8842176240868867e-4,Double
Item [Double]
2.8840399312321096e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.886097295122454e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.88599069218435e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.886097295122454e-4,Double
Item [Double]
2.8846476339094805e-4,Double
Item [Double]
2.885038541771792e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.8858878438222746e-4,Double
Item [Double]
2.885923176600045e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.884007943794131e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.884469945274759e-4,Double
Item [Double]
2.8843242392031246e-4,Double
Item [Double]
2.884288700806792e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.885034988157713e-4,Double
Item [Double]
2.884110805753153e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.8849283778617973e-4,Double
Item [Double]
2.884075468755327e-4,Double
Item [Double]
2.884075468755327e-4,Double
Item [Double]
2.884075468755327e-4,Double
Item [Double]
2.884075468755327e-4,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
2.884892851579934e-4,Double
Item [Double]
2.884892851579934e-4,Double
Item [Double]
2.884892851579934e-4,Double
Item [Double]
2.884892851579934e-4,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
0.0,Double
Item [Double]
2.884892851579934e-4,Double
Item [Double]
2.884892851579934e-4,Double
Item [Double]
2.884892851579934e-4,Double
Item [Double]
2.884892851579934e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.884288700806792e-4,Double
Item [Double]
2.884395315486472e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.8849246223035154e-4,Double
Item [Double]
2.8850276789489724e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.8849212704517413e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.884288700806792e-4,Double
Item [Double]
2.884395315486472e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.8849246223035154e-4,Double
Item [Double]
2.8850276789489724e-4,Double
Item [Double]
0.0,Double
Item [Double]
2.8849212704517413e-4,Double
Item [Double]
2.8842922547482884e-4,Double
Item [Double]
2.885038541771792e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.894378297730782e-4,Double
Item [Double]
2.885923176600045e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.88599069218435e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.887056688523444e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.887056688523444e-4,Double
Item [Double]
2.8846476339094805e-4,Double
Item [Double]
2.885038541771792e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.884786229769816e-4,Double
Item [Double]
2.885923176600045e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.886950092188272e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.884818011261814e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.886097295122454e-4,Double
Item [Double]
2.8846476339094805e-4,Double
Item [Double]
2.885038541771792e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.885145151321922e-4,Double
Item [Double]
2.8854294397024206e-4,Double
Item [Double]
2.887167039107226e-4,Double
Item [Double]
2.885923176600045e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.8860262265516213e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.885884088500461e-4,Double
Item [Double]
2.887454843457817e-4,Double
Item [Double]
2.88599069218435e-4])
    ((AstTensor AstMethodLet FullSpan (TKR 10 Double)
 -> AstTensor AstMethodLet FullSpan (TKScalar Double))
-> Value (AstTensor AstMethodLet FullSpan (TKR 10 Double))
-> Value (AstTensor AstMethodLet FullSpan (TKR 10 Double))
forall src r tgt.
((X src :: TK) ~ (X (Value src) :: TK), KnownSTK (X src),
 AdaptableTarget (AstTensor AstMethodLet FullSpan) src,
 AdaptableTarget Concrete (Value src),
 (tgt :: Type)
 ~ (AstTensor AstMethodLet FullSpan (TKScalar r) :: Type)) =>
(src -> tgt) -> Value src -> Value src
grad (AstTensor AstMethodLet FullSpan (TKR 0 Double)
-> AstTensor AstMethodLet FullSpan (TKScalar Double)
forall r.
GoodScalar r =>
AstTensor AstMethodLet FullSpan (TKR 0 r)
-> AstTensor AstMethodLet FullSpan (TKScalar r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKR 0 r) -> target (TKScalar r)
kfromR (AstTensor AstMethodLet FullSpan (TKR 0 Double)
 -> AstTensor AstMethodLet FullSpan (TKScalar Double))
-> (AstTensor AstMethodLet FullSpan (TKR 10 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 0 Double))
-> AstTensor AstMethodLet FullSpan (TKR 10 Double)
-> AstTensor AstMethodLet FullSpan (TKScalar Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) (x :: TK) (target :: Target).
(KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 n x) -> target (TKR2 0 x)
rsum0 @10 @(TKScalar Double) (AstTensor AstMethodLet FullSpan (TKR 10 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 0 Double))
-> (AstTensor AstMethodLet FullSpan (TKR 10 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 10 Double))
-> AstTensor AstMethodLet FullSpan (TKR 10 Double)
-> AstTensor AstMethodLet FullSpan (TKR 0 Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet FullSpan (TKR 10 Double)
-> AstTensor AstMethodLet FullSpan (TKR 10 Double)
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
barRelu10xSlower)
          ((Concrete (TKR 0 Double) -> Concrete (TKR 0 Double))
-> Concrete (TKR 10 Double) -> Concrete (TKR 10 Double)
forall (n :: Nat) (x :: TK) (x1 :: TK) (target :: Target).
(KnownNat n, KnownSTK x, KnownSTK x1, BaseTensor target) =>
(target (TKR2 0 x1) -> target (TKR2 0 x))
-> target (TKR2 n x1) -> target (TKR2 n x)
rmap0N (Concrete (TKR 0 Double)
-> Concrete (TKR 0 Double) -> Concrete (TKR 0 Double)
forall a. Num a => a -> a -> a
* Double -> Concrete (TKR 0 Double)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKR 0 r)
rscalar Double
0.001) Concrete (TKR 10 Double)
forall r. (GoodScalar r, Fractional r) => Concrete (TKR 10 r)
t128))

testReluSimpPP :: Assertion
testReluSimpPP :: Assertion
testReluSimpPP = do
  Assertion
resetVarCounter
  let !t1 :: AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t1 = forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
barRelu10xSlower @(AstTensor AstMethodLet PrimalSpan)
            (AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 10 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 10 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 10 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 10
-> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 10 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 10 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 10 Float))
-> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 10 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t1) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
22675
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 10 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 10 Float) AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t1)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
22675
  Assertion
resetVarCounter
  let !t2 :: AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t2 = forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n, Differentiable r) =>
target (TKR n r) -> target (TKR n r)
barRelu @(AstTensor AstMethodLet PrimalSpan)
            (AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
 -> AstTensor AstMethodLet PrimalSpan (TKR 10 Float))
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKR 10 Float)
-> AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
AstVar (FullShapeTK (TKR 10 Float)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKR 10 Float)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IShR 10
-> FullShapeTK (TKScalar Float) -> FullShapeTK (TKR 10 Float)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
2,Int
Item (IShR 10)
1] FullShapeTK (TKScalar Float)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarId -> AstVarName PrimalSpan (TKR 10 Float))
-> (Int -> AstVarId) -> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstVarId
intToAstVarId (Int -> AstVarName PrimalSpan (TKR 10 Float))
-> Int -> AstVarName PrimalSpan (TKR 10 Float)
forall a b. (a -> b) -> a -> b
$ Int
100000000)
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 10 Float) -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t2) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
12258
  String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length (AstTensor AstMethodLet PrimalSpan (TKR 10 Float) -> String
forall a. Show a => a -> String
show (forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract @(TKR 10 Float) AstTensor AstMethodLet PrimalSpan (TKR 10 Float)
t2)) Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
12258

testCNNOPP2 :: Assertion
testCNNOPP2 :: Assertion
testCNNOPP2 = do
  Assertion
resetVarCounter
  let t :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
t = AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type)) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded2
            (Ranked 4 Double -> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 4 Double -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> Ranked 4 Double
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall a b. (a -> b) -> a -> b
$ IShR 4 -> Double -> Ranked 4 Double
forall (n :: Nat) a. PrimElt a => IShR n -> a -> Ranked n a
Nested.rreplicateScal (Int
1 Int -> IShR 3 -> IShR 4
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
1 Int -> IShR 2 -> IShR 3
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
2 Int -> IShR 1 -> IShR 2
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
2 Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) Double
1)
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty (AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract AstTensor AstMethodLet FullSpan (TKR 4 Double)
t)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"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]))"
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty AstTensor AstMethodLet FullSpan (TKR 4 Double)
t
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"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])))"

testCNNOPP2b :: Assertion
testCNNOPP2b :: Assertion
testCNNOPP2b = do
  Assertion
resetVarCounter
  let artifactRev :: AstArtifactRev
  (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
artifactRev = IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> FullShapeTK (X (AstTensor AstMethodLet FullSpan (TKR 4 Double)))
-> AstArtifactRev
     (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
forall src (ztgt :: TK) tgt.
(AdaptableTarget (AstTensor AstMethodLet FullSpan) src,
 (tgt :: Type) ~ (AstTensor AstMethodLet FullSpan ztgt :: Type)) =>
IncomingCotangentHandling
-> (src -> tgt)
-> FullShapeTK (X src)
-> AstArtifactRev (X src) ztgt
revArtifactAdapt IncomingCotangentHandling
UseIncomingCotangent AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type)) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded2 (IShR 4
-> FullShapeTK (TKScalar Double) -> FullShapeTK (TKR 4 Double)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 4)
1, Int
Item (IShR 4)
1, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] (forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar @Double))
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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]))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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]))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\dret u1 -> rfromS (sconcrete (sreplicate [1,1,2,2] 0.0))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\dret u1 -> rfromS (sconcrete (sreplicate [1,1,2,2] 0.0))"

maxPool2dUnpadded2
  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)
  => target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded2 :: forall (target :: Target) r.
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type)) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded2 target (TKR 4 r)
a =
  IShR (4 + 0)
-> (IxR 4 (IntOf target) -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] ((IxR 4 (IntOf target) -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (IntOf target) -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
_, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
_, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBh, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBw] ->
      let arrt :: target (TKR 4 r)
arrt = target (TKR 4 r) -> IxR 4 (IntOf target) -> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type), (n :: Nat) ~ (4 :: Nat)) =>
target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez2 (target (TKR 4 r) -> target (TKR 4 r)
forall (target :: Target) r.
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type)) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded2 target (TKR 4 r)
a) [Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBw, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
1, AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
2 AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
forall a. Num a => a -> a -> a
* Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
iBh, AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
2 AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
forall a. Num a => a -> a -> a
* Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
iBw]
      in target (TKR 4 r) -> target (TKR2 0 (TKScalar r))
forall (target :: Target) r.
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type)) =>
target (TKR 4 r) -> target (TKR 0 r)
rmaximum2 target (TKR 4 r)
arrt
    IxR 4 (IntOf target)
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"maxPool2dUnpadded2: impossible pattern needlessly required"

conv2dUnpadded2
  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)
  => target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded2 :: forall (target :: Target) r.
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type)) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded2 target (TKR 4 r)
a =
  IShR (4 + 0)
-> (IxR 4 (IntOf target) -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild [Int
Item (IShR 4)
3, Int
Item (IShR 4)
3, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] ((IxR 4 (IntOf target) -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (IntOf target) -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iImg, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
_, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBh, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBw] ->
      let arrAt :: target (TKR 4 r)
arrAt = target (TKR 4 r) -> IxR 4 (IntOf target) -> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type), (n :: Nat) ~ (4 :: Nat)) =>
target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez2 target (TKR 4 r)
a [Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iImg, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
0, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBh, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBw]
      in target (TKR 4 r)
-> IxR 4 (IntOf target) -> target (TKR2 0 (TKScalar r))
forall (m :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownSTK x, BaseTensor target) =>
target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)
rindex0 target (TKR 4 r)
arrAt [Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
0, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBw, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
iBw, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
0]
    IxR 4 (IntOf target)
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"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 :: forall (target :: Target) r (n :: Nat).
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type), (n :: Nat) ~ (4 :: Nat)) =>
target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez2 target (TKR n r)
d IxROf target n
ixBase =
  IShR (n + 0)
-> (IxROf target n -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (n + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild [Int
Item (IShR 4)
1, Int
Item (IShR 4)
1, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] ((IxROf target n -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (n + 0) (TKScalar r)))
-> (IxROf target n -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (n + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \IxROf target n
ixResult -> target (TKR n r) -> IxROf target n -> target (TKR2 0 (TKScalar r))
forall (target :: Target) r (n :: Nat).
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type), (n :: Nat) ~ (4 :: Nat)) =>
target (TKR n r) -> IxROf target n -> target (TKR 0 r)
indexz02 target (TKR n r)
d ((AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
-> IxR n (AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
-> IxR n (AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
-> IxR n (AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall i j k (n :: Nat).
(i -> j -> k) -> IxR n i -> IxR n j -> IxR n k
ixrZipWith AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
forall a. Num a => a -> a -> a
(+) IxR n (AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
IxROf target n
ixBase IxR n (AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
IxROf target n
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 :: forall (target :: Target) r (n :: Nat).
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type), (n :: Nat) ~ (4 :: Nat)) =>
target (TKR n r) -> IxROf target n -> target (TKR 0 r)
indexz02 target (TKR n r)
d IxROf target n
ix = BoolOf target
-> target (TKR 0 r) -> target (TKR 0 r) -> target (TKR 0 r)
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
1 AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). EqH f y => f y -> f y -> BoolOf f
==. (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
-> [Item
      (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxROf target n
IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
ix [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
-> Int -> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
forall a. HasCallStack => [a] -> Int -> a
!! Int
0)) (target (TKR n r)
target (TKR2 (n + 0) (TKScalar r))
d target (TKR2 (n + 0) (TKScalar r))
-> IxROf target n -> target (TKR 0 r)
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! IxROf target n
ix) (r -> target (TKR 0 r)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKR 0 r)
rscalar r
0)

rmaximum2 :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)
         => target (TKR 4 r) -> target (TKR 0 r)
rmaximum2 :: forall (target :: Target) r.
((target :: Target) ~ (AstTensor AstMethodLet FullSpan :: Target),
 (r :: Type) ~ (Double :: Type)) =>
target (TKR 4 r) -> target (TKR 0 r)
rmaximum2 target (TKR 4 r)
t0 = target (TKR 4 r)
-> (target (TKR 4 r) -> target (TKR 0 r)) -> target (TKR 0 r)
forall (x :: TK) (z :: TK) (target :: Target).
LetTensor target =>
target x -> (target x -> target z) -> target z
tlet target (TKR 4 r)
t0 ((target (TKR 4 r) -> target (TKR 0 r)) -> target (TKR 0 r))
-> (target (TKR 4 r) -> target (TKR 0 r)) -> target (TKR 0 r)
forall a b. (a -> b) -> a -> b
$ \target (TKR 4 r)
t -> target (TKR 4 r) -> IxROf target 4 -> target (TKR 0 r)
forall (m :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownSTK x, BaseTensor target) =>
target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)
rindex0 target (TKR 4 r)
t [Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
0, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
0, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
0, Item (IxR 4 (AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
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
                       (rconcrete $ Nested.rscalar 7
                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      afcnn2T = maxPool2dUnpadded3 $ conv2dUnpadded3 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])))"
-}

-- TODO: OOMs
_testCNNOPP3b :: Assertion
_testCNNOPP3b :: Assertion
_testCNNOPP3b = do
  Assertion
resetVarCounter
  let artifactRev :: AstArtifactRev
  (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
artifactRev = IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> FullShapeTK (X (AstTensor AstMethodLet FullSpan (TKR 4 Double)))
-> AstArtifactRev
     (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
forall src (ztgt :: TK) tgt.
(AdaptableTarget (AstTensor AstMethodLet FullSpan) src,
 (tgt :: Type) ~ (AstTensor AstMethodLet FullSpan ztgt :: Type)) =>
IncomingCotangentHandling
-> (src -> tgt)
-> FullShapeTK (X src)
-> AstArtifactRev (X src) ztgt
revArtifactAdapt IncomingCotangentHandling
UseIncomingCotangent (AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3 (AstTensor AstMethodLet FullSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3) (IShR 4
-> FullShapeTK (TKScalar Double) -> FullShapeTK (TKR 4 Double)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 4)
3, Int
Item (IShR 4)
3, Int
Item (IShR 4)
3, Int
Item (IShR 4)
3] (forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar @Double))
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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)))))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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)))))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"

maxPool2dUnpadded3
  :: (ADReady target, GoodScalar r)
  => target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3 target (TKR 4 r)
arr =
  IShR (4 + 0)
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] ((IxR 4 (PrimalOf target (TKScalar Int64))
  -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
aa, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
bb, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw] ->
      let arrt :: target (TKR 4 r)
arrt = IShR 4
-> target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez3 [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] target (TKR 4 r)
arr [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
iBh PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. IntegralH a => a -> a -> a
`quotH` PrimalOf target (TKScalar Int64)
2, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
aa, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
bb, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw]
      in target (TKR 4 r) -> target (TKR2 0 (TKScalar r))
forall (target :: Target) (n :: Nat) r.
(BaseTensor target, LetTensor target, KnownNat n, GoodScalar r) =>
target (TKR n r) -> target (TKR 0 r)
rmaximum3 target (TKR 4 r)
arrt
    IxR 4 (PrimalOf target (TKScalar Int64))
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"maxPool2dUnpadded3: impossible pattern needlessly required"

conv2dUnpadded3
  :: (ADReady target, GoodScalar r)
  => target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3 target (TKR 4 r)
arrA =
  let shB :: IShR 4
shB = [Item (IShR 4)
2, Item (IShR 4)
2, Item (IShR 4)
2, Item (IShR 4)
2]
  in IShR (4 + 0)
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild IShR 4
IShR (4 + 0)
shB ((IxR 4 (PrimalOf target (TKScalar Int64))
  -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
_, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw] ->
      let arrAt :: target (TKR 4 r)
arrAt = IShR 4
-> target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez4 IShR 4
shB target (TKR 4 r)
arrA [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
1]
      in target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR2 0 (TKScalar r))
forall (m :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownSTK x, BaseTensor target) =>
target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)
rindex0 target (TKR 4 r)
arrAt [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh] target (TKR2 0 (TKScalar r))
-> target (TKR2 0 (TKScalar r)) -> target (TKR2 0 (TKScalar r))
forall a. Num a => a -> a -> a
+ target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR2 0 (TKScalar r))
forall (m :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownSTK x, BaseTensor target) =>
target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)
rindex0 target (TKR 4 r)
arrAt [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
1, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
iBw PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
+ PrimalOf target (TKScalar Int64)
1, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh]
    IxR 4 (PrimalOf target (TKScalar Int64))
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"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 :: forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez3 IShR n
shOut target (TKR n r)
d IxROf target n
ixBase =
  IShR (n + 0)
-> (IxROf target n -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (n + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild IShR n
IShR (n + 0)
shOut ((IxROf target n -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (n + 0) (TKScalar r)))
-> (IxROf target n -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (n + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \IxROf target n
_ -> target (TKR n r) -> IxROf target n -> target (TKR2 0 (TKScalar r))
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
target (TKR n r) -> IxROf target n -> target (TKR 0 r)
indexz03 target (TKR n r)
d ((PrimalOf target (TKScalar Int64)
 -> PrimalOf target (TKScalar Int64)
 -> PrimalOf target (TKScalar Int64))
-> IxROf target n -> IxROf target n -> IxROf target n
forall i j k (n :: Nat).
(i -> j -> k) -> IxR n i -> IxR n j -> IxR n k
ixrZipWith PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
(+) IxROf target n
ixBase IxROf target n
ixBase)

indexz03
  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)
  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)
indexz03 :: forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
target (TKR n r) -> IxROf target n -> target (TKR 0 r)
indexz03 target (TKR n r)
d IxROf target n
ix = BoolOf target
-> target (TKR 0 r) -> target (TKR 0 r) -> target (TKR 0 r)
forall (y :: TK) (target :: Target).
(KnownSTK y, Boolean (BoolOf target), BaseTensor target) =>
BoolOf target -> target y -> target y -> target y
ifH (forall (target :: Target) (n :: Nat).
(ADReady target, KnownNat n) =>
IShR n -> IxROf target n -> BoolOf target
within3 @target (forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape @target target (TKR n r)
d) IxROf target n
ix) (target (TKR n r)
target (TKR2 (n + 0) (TKScalar r))
d target (TKR2 (n + 0) (TKScalar r))
-> IxROf target n -> target (TKR 0 r)
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
! IxROf target n
ix) (r -> target (TKR 0 r)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKR 0 r)
rscalar r
0)

within3
  :: forall target n. (ADReady target, KnownNat n)
  => IShR n -> IxROf target n -> BoolOf target
within3 :: forall (target :: Target) (n :: Nat).
(ADReady target, KnownNat n) =>
IShR n -> IxROf target n -> BoolOf target
within3 IShR n
sh IxR n (PrimalOf target (TKScalar Int64))
ix =
  let within :: IntOf target -> IntOf target -> BoolOf target
      within :: PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64) -> BoolOf target
within PrimalOf target (TKScalar Int64)
i PrimalOf target (TKScalar Int64)
dim = PrimalOf target (TKScalar Int64)
0 PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64) -> BoolOf (PrimalOf target)
forall (f :: Target) (y :: TK). EqH f y => f y -> f y -> BoolOf f
==. PrimalOf target (TKScalar Int64)
i BoolOf target -> BoolOf target -> BoolOf target
forall b. Boolean b => b -> b -> b
||* PrimalOf target (TKScalar Int64)
dim PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
- PrimalOf target (TKScalar Int64)
2 PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64) -> BoolOf (PrimalOf target)
forall (f :: Target) (y :: TK). EqH f y => f y -> f y -> BoolOf f
==. PrimalOf target (TKScalar Int64)
i
  in (BoolOf target -> BoolOf target -> BoolOf target)
-> BoolOf target -> [BoolOf target] -> BoolOf target
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr BoolOf target -> BoolOf target -> BoolOf target
forall b. Boolean b => b -> b -> b
(&&*) BoolOf target
forall b. Boolean b => b
true
     ([BoolOf target] -> BoolOf target)
-> [BoolOf target] -> BoolOf target
forall a b. (a -> b) -> a -> b
$ (PrimalOf target (TKScalar Int64)
 -> PrimalOf target (TKScalar Int64) -> BoolOf target)
-> [PrimalOf target (TKScalar Int64)]
-> [PrimalOf target (TKScalar Int64)]
-> [BoolOf target]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64) -> BoolOf target
within (IxR n (PrimalOf target (TKScalar Int64))
-> [Item (IxR n (PrimalOf target (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxR n (PrimalOf target (TKScalar Int64))
ix) ((Int -> PrimalOf target (TKScalar Int64))
-> PermR -> [PrimalOf target (TKScalar Int64)]
forall a b. (a -> b) -> [a] -> [b]
map Int -> PrimalOf target (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PermR -> [PrimalOf target (TKScalar Int64)])
-> PermR -> [PrimalOf target (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IShR n -> [Item (IShR n)]
forall l. IsList l => l -> [Item l]
toList IShR n
sh)

rmaximum3 :: (BaseTensor target, LetTensor target, KnownNat n, GoodScalar r)
         => target (TKR n r) -> target (TKR 0 r)
rmaximum3 :: forall (target :: Target) (n :: Nat) r.
(BaseTensor target, LetTensor target, KnownNat n, GoodScalar r) =>
target (TKR n r) -> target (TKR 0 r)
rmaximum3 target (TKR n r)
t0 = target (TKR n r)
-> (target (TKR n r) -> target (TKR 0 r)) -> target (TKR 0 r)
forall (x :: TK) (z :: TK) (target :: Target).
LetTensor target =>
target x -> (target x -> target z) -> target z
tlet target (TKR n r)
t0 ((target (TKR n r) -> target (TKR 0 r)) -> target (TKR 0 r))
-> (target (TKR n r) -> target (TKR 0 r)) -> target (TKR 0 r)
forall a b. (a -> b) -> a -> b
$ \target (TKR n r)
t -> target (TKR n r)
-> IxR n (PrimalOf target (TKScalar Int64)) -> target (TKR 0 r)
forall (m :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownSTK x, BaseTensor target) =>
target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)
rindex0 target (TKR n r)
t [Item (IxR n (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR n (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR n (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR n (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0]

testCNNOPP4 :: Assertion
testCNNOPP4 :: Assertion
testCNNOPP4 = do
  Assertion
resetVarCounter
  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
blackGlyph = AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall a b. (a -> b) -> a -> b
$ SNat 3
-> SingletonTK
     (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind
        3
        (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @3) SingletonTK
  (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor
   AstMethodLet
   PrimalSpan
   (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind
         3
         (BuildTensorKind
            3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind
        3
        (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))))
forall a b. (a -> b) -> a -> b
$ SNat 3
-> SingletonTK (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @3) SingletonTK (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor
   AstMethodLet
   PrimalSpan
   (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 3 (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))
forall a b. (a -> b) -> a -> b
$ SNat 3
-> SingletonTK (TKR2 1 (TKScalar Double))
-> AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @3) SingletonTK (TKR2 1 (TKScalar Double))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind 3 (TKR2 1 (TKScalar Double))))
-> AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 3 (TKR2 1 (TKScalar Double)))
forall a b. (a -> b) -> a -> b
$ SNat 3
-> SingletonTK (TKR 0 Double)
-> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind 3 (TKR 0 Double))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @3) SingletonTK (TKR 0 Double)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                       (Ranked 0 Double -> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 0 Double
 -> AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
-> Ranked 0 Double
-> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
forall a b. (a -> b) -> a -> b
$ Double -> Ranked 0 Double
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar Double
7
                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T = AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded4 AstTensor AstMethodLet FullSpan (TKR 4 Double)
blackGlyph
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty (AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"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]))))"
      -- 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]))"
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"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])"

-- In this test primal is trivial but gradient is not, so we know how much
-- scatters should be able to simplify in the future.
testCNNOPP4b :: Assertion
testCNNOPP4b :: Assertion
testCNNOPP4b = do
  Assertion
resetVarCounter
  let artifactRev :: AstArtifactRev
  (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
artifactRev = IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> FullShapeTK (X (AstTensor AstMethodLet FullSpan (TKR 4 Double)))
-> AstArtifactRev
     (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
forall src (ztgt :: TK) tgt.
(AdaptableTarget (AstTensor AstMethodLet FullSpan) src,
 (tgt :: Type) ~ (AstTensor AstMethodLet FullSpan ztgt :: Type)) =>
IncomingCotangentHandling
-> (src -> tgt)
-> FullShapeTK (X src)
-> AstArtifactRev (X src) ztgt
revArtifactAdapt IncomingCotangentHandling
UseIncomingCotangent AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded4 (IShR 4
-> FullShapeTK (TKScalar Double) -> FullShapeTK (TKR 4 Double)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 4)
3, Int
Item (IShR 4)
3, Int
Item (IShR 4)
3, Int
Item (IShR 4)
3] (forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar @Double))
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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]))))"
      -- 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]))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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]))))"
      -- 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]))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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] [])))))"
      -- TODO: was once "\\dret u1 -> rfromS (sconcrete (sreplicate [3,3,3,3] 0.0))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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]))))"
      -- TODO: was once "\\dret u1 -> rfromS (sconcrete (sreplicate [3,3,3,3] 0.0))"

testCNNOPP5 :: Assertion
testCNNOPP5 :: Assertion
testCNNOPP5 = do
  Assertion
resetVarCounter
  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
blackGlyph = AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall a b. (a -> b) -> a -> b
$ SNat 6
-> SingletonTK
     (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind
        6
        (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @6) SingletonTK
  (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor
   AstMethodLet
   PrimalSpan
   (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind
         6
         (BuildTensorKind
            6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind
        6
        (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))))
forall a b. (a -> b) -> a -> b
$ SNat 6
-> SingletonTK (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @6) SingletonTK (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor
   AstMethodLet
   PrimalSpan
   (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 6 (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))
forall a b. (a -> b) -> a -> b
$ SNat 6
-> SingletonTK (TKR2 1 (TKScalar Double))
-> AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @6) SingletonTK (TKR2 1 (TKScalar Double))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind 6 (TKR2 1 (TKScalar Double))))
-> AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 6 (TKR2 1 (TKScalar Double)))
forall a b. (a -> b) -> a -> b
$ SNat 6
-> SingletonTK (TKR 0 Double)
-> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind 6 (TKR 0 Double))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @6) SingletonTK (TKR 0 Double)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                       (Ranked 0 Double -> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 0 Double
 -> AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
-> Ranked 0 Double
-> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
forall a b. (a -> b) -> a -> b
$ Double -> Ranked 0 Double
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar Double
7
                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T = AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded4 AstTensor AstMethodLet FullSpan (TKR 4 Double)
blackGlyph
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty (AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"rfromS (sconcrete (sfromListLinear [1,1,2,2] [7.0,0.0,0.0,0.0]))"
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"rfromS (sconcrete (sfromListLinear [1,1,2,2] [7.0,0.0,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.
testCNNOPP5b :: Assertion
testCNNOPP5b :: Assertion
testCNNOPP5b = do
  Assertion
resetVarCounter
  let artifactRev :: AstArtifactRev
  (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
artifactRev = IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> FullShapeTK (X (AstTensor AstMethodLet FullSpan (TKR 4 Double)))
-> AstArtifactRev
     (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
forall src (ztgt :: TK) tgt.
(AdaptableTarget (AstTensor AstMethodLet FullSpan) src,
 (tgt :: Type) ~ (AstTensor AstMethodLet FullSpan ztgt :: Type)) =>
IncomingCotangentHandling
-> (src -> tgt)
-> FullShapeTK (X src)
-> AstArtifactRev (X src) ztgt
revArtifactAdapt IncomingCotangentHandling
UseIncomingCotangent AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded4 (IShR 4
-> FullShapeTK (TKScalar Double) -> FullShapeTK (TKR 4 Double)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 4)
5, Int
Item (IShR 4)
5, Int
Item (IShR 4)
5, Int
Item (IShR 4)
5] (forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar @Double))
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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))))))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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))))))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\dret u1 -> rfromS (soneHot (sfromR dret !$ [0, 0, 0, 0]) [0, 0, 0, 0])"

maxPool2dUnpadded4
  :: (ADReady target, GoodScalar r)
  => target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded4 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded4 target (TKR 4 r)
arr =
  IShR (4 + 0)
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] ((IxR 4 (PrimalOf target (TKScalar Int64))
  -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
aa, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
bb, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw] ->
      let arrt :: target (TKR 4 r)
arrt = IShR 4
-> target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez4 [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] target (TKR 4 r)
arr [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
bb PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
+ PrimalOf target (TKScalar Int64)
1, PrimalOf target (TKScalar Int64)
2 PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
- Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
bb, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
aa PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
* Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
iBh, PrimalOf target (TKScalar Int64)
2 PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
* Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
iBw]
      in target (TKR 4 r) -> target (TKR2 0 (TKScalar r))
forall (target :: Target) (n :: Nat) r.
(BaseTensor target, LetTensor target, KnownNat n, GoodScalar r) =>
target (TKR n r) -> target (TKR 0 r)
rmaximum3 target (TKR 4 r)
arrt
    IxR 4 (PrimalOf target (TKScalar Int64))
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"maxPool2dUnpadded4: impossible pattern needlessly required"

conv2dUnpadded4
  :: (ADReady target, GoodScalar r)
  => target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded4 :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded4 target (TKR 4 r)
arrA =
  let shB :: IShR 4
shB = [Item (IShR 4)
1, Item (IShR 4)
1, Item (IShR 4)
2, Item (IShR 4)
2]
  in IShR (4 + 0)
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild IShR 4
IShR (4 + 0)
shB ((IxR 4 (PrimalOf target (TKScalar Int64))
  -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
_, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw] ->
      let arrAt :: target (TKR 4 r)
arrAt = IShR 4
-> target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez4 IShR 4
shB target (TKR 4 r)
arrA [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw]
      in target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR2 0 (TKScalar r))
forall (m :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownSTK x, BaseTensor target) =>
target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)
rindex0 target (TKR 4 r)
arrAt [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
0]
    IxR 4 (PrimalOf target (TKScalar Int64))
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"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 :: forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez4 IShR n
shOut target (TKR n r)
d IxROf target n
ixBase =
  IShR (n + 0)
-> (IxROf target n -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (n + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild IShR n
IShR (n + 0)
shOut ((IxROf target n -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (n + 0) (TKScalar r)))
-> (IxROf target n -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (n + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \IxROf target n
ixResult -> target (TKR n r) -> IxROf target n -> target (TKR2 0 (TKScalar r))
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
target (TKR n r) -> IxROf target n -> target (TKR 0 r)
indexz03 target (TKR n r)
d ((PrimalOf target (TKScalar Int64)
 -> PrimalOf target (TKScalar Int64)
 -> PrimalOf target (TKScalar Int64))
-> IxROf target n -> IxROf target n -> IxROf target n
forall i j k (n :: Nat).
(i -> j -> k) -> IxR n i -> IxR n j -> IxR n k
ixrZipWith PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
(+) IxROf target n
ixBase IxROf target n
ixResult)

testCNNOPP6 :: Assertion
testCNNOPP6 :: Assertion
testCNNOPP6 = do
  Assertion
resetVarCounter
  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
blackGlyph = AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind
        2
        (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @2) SingletonTK
  (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor
   AstMethodLet
   PrimalSpan
   (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind
         2
         (BuildTensorKind
            2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind
        2
        (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))))
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @2) SingletonTK (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor
   AstMethodLet
   PrimalSpan
   (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK (TKR2 1 (TKScalar Double))
-> AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @2) SingletonTK (TKR2 1 (TKScalar Double))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
-> AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK (TKR 0 Double)
-> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind 2 (TKR 0 Double))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @2) SingletonTK (TKR 0 Double)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                       (Ranked 0 Double -> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 0 Double
 -> AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
-> Ranked 0 Double
-> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
forall a b. (a -> b) -> a -> b
$ Double -> Ranked 0 Double
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar Double
7
                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T = AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3 (AstTensor AstMethodLet FullSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3z AstTensor AstMethodLet FullSpan (TKR 4 Double)
blackGlyph
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty (AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"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]))"
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"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]))))"

testCNNOPP6b :: Assertion
testCNNOPP6b :: Assertion
testCNNOPP6b = do
  Assertion
resetVarCounter
  let artifactRev :: AstArtifactRev
  (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
artifactRev = IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> FullShapeTK (X (AstTensor AstMethodLet FullSpan (TKR 4 Double)))
-> AstArtifactRev
     (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
forall src (ztgt :: TK) tgt.
(AdaptableTarget (AstTensor AstMethodLet FullSpan) src,
 (tgt :: Type) ~ (AstTensor AstMethodLet FullSpan ztgt :: Type)) =>
IncomingCotangentHandling
-> (src -> tgt)
-> FullShapeTK (X src)
-> AstArtifactRev (X src) ztgt
revArtifactAdapt IncomingCotangentHandling
UseIncomingCotangent (AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3 (AstTensor AstMethodLet FullSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3z) (IShR 4
-> FullShapeTK (TKScalar Double) -> FullShapeTK (TKR 4 Double)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] (forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar @Double))
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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)))))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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)))))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"

  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"

conv2dUnpadded3z
  :: (ADReady target, GoodScalar r)
  => target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3z :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3z target (TKR 4 r)
arrA =
  let shB :: IShR 4
shB = [Item (IShR 4)
2, Item (IShR 4)
2, Item (IShR 4)
2, Item (IShR 4)
2]
  in IShR (4 + 0)
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild IShR 4
IShR (4 + 0)
shB ((IxR 4 (PrimalOf target (TKScalar Int64))
  -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
_, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw] ->
      let arrAt :: target (TKR 4 r)
arrAt = IShR 4
-> target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez3 IShR 4
shB target (TKR 4 r)
arrA [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw]
      in target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR2 0 (TKScalar r))
forall (m :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownSTK x, BaseTensor target) =>
target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)
rindex0 target (TKR 4 r)
arrAt [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh]
    IxR 4 (PrimalOf target (TKScalar Int64))
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"conv2dUnpadded3z: impossible pattern needlessly required"

testCNNOPP7 :: Assertion
testCNNOPP7 :: Assertion
testCNNOPP7 = do
  Assertion
resetVarCounter
  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
blackGlyph = AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet PrimalSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind
        2
        (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @2) SingletonTK
  (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor
   AstMethodLet
   PrimalSpan
   (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind
         2
         (BuildTensorKind
            2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind
        2
        (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))))
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @2) SingletonTK (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor
   AstMethodLet
   PrimalSpan
   (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK (TKR2 1 (TKScalar Double))
-> AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @2) SingletonTK (TKR2 1 (TKScalar Double))
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                   (AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (BuildTensorKind 2 (TKR2 1 (TKScalar Double))))
-> AstTensor AstMethodLet PrimalSpan (TKR2 1 (TKScalar Double))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind 2 (TKR2 1 (TKScalar Double)))
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK (TKR 0 Double)
-> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind 2 (TKR 0 Double))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @2) SingletonTK (TKR 0 Double)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                       (Ranked 0 Double -> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
forall r (target :: Target) (n :: Nat).
(GoodScalar r, BaseTensor target) =>
Ranked n r -> target (TKR n r)
rconcrete (Ranked 0 Double
 -> AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
-> Ranked 0 Double
-> AstTensor AstMethodLet PrimalSpan (TKR 0 Double)
forall a b. (a -> b) -> a -> b
$ Double -> Ranked 0 Double
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar Double
7
                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T = AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3y (AstTensor AstMethodLet FullSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3y AstTensor AstMethodLet FullSpan (TKR 4 Double)
blackGlyph
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty (AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInlineContract AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"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]))"
  AstTensor AstMethodLet FullSpan (TKR 4 Double) -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstPretty AstTensor AstMethodLet FullSpan (TKR 4 Double)
afcnn2T
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"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]))"

testCNNOPP7b :: Assertion
testCNNOPP7b :: Assertion
testCNNOPP7b = do
  Assertion
resetVarCounter
  let artifactRev :: AstArtifactRev
  (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
artifactRev = IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> FullShapeTK (X (AstTensor AstMethodLet FullSpan (TKR 4 Double)))
-> AstArtifactRev
     (X (AstTensor AstMethodLet FullSpan (TKR 4 Double))) (TKR 4 Double)
forall src (ztgt :: TK) tgt.
(AdaptableTarget (AstTensor AstMethodLet FullSpan) src,
 (tgt :: Type) ~ (AstTensor AstMethodLet FullSpan ztgt :: Type)) =>
IncomingCotangentHandling
-> (src -> tgt)
-> FullShapeTK (X src)
-> AstArtifactRev (X src) ztgt
revArtifactAdapt IncomingCotangentHandling
UseIncomingCotangent (AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3y (AstTensor AstMethodLet FullSpan (TKR 4 Double)
 -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> (AstTensor AstMethodLet FullSpan (TKR 4 Double)
    -> AstTensor AstMethodLet FullSpan (TKR 4 Double))
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet FullSpan (TKR 4 Double)
-> AstTensor AstMethodLet FullSpan (TKR 4 Double)
forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3y) (IShR 4
-> FullShapeTK (TKScalar Double) -> FullShapeTK (TKR 4 Double)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] (forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar @Double))
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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))))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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))))"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"
  AstArtifactRev (TKR 4 Double) (TKR 4 Double) -> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty (AstArtifactRev (TKR 4 Double) (TKR 4 Double)
-> AstArtifactRev (TKR 4 Double) (TKR 4 Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev (TKR 4 Double) (TKR 4 Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"

maxPool2dUnpadded3y
  :: (ADReady target, GoodScalar r)
  => target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3y :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
maxPool2dUnpadded3y target (TKR 4 r)
arr =
  IShR (4 + 0)
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] ((IxR 4 (PrimalOf target (TKScalar Int64))
  -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
aa, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
bb, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw] ->
      let arrt :: target (TKR 4 r)
arrt = IShR 4
-> target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez3 [Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2, Int
Item (IShR 4)
2] target (TKR 4 r)
arr [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
aa, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
bb, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw]
      in target (TKR 4 r) -> target (TKR2 0 (TKScalar r))
forall (target :: Target) (n :: Nat) r.
(BaseTensor target, LetTensor target, KnownNat n, GoodScalar r) =>
target (TKR n r) -> target (TKR 0 r)
rmaximum3 target (TKR 4 r)
arrt
    IxR 4 (PrimalOf target (TKScalar Int64))
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"maxPool2dUnpadded3y: impossible pattern needlessly required"

conv2dUnpadded3y
  :: (ADReady target, GoodScalar r)
  => target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3y :: forall (target :: Target) r.
(ADReady target, GoodScalar r) =>
target (TKR 4 r) -> target (TKR 4 r)
conv2dUnpadded3y target (TKR 4 r)
arrA =
  let shB :: IShR 4
shB = [Item (IShR 4)
2, Item (IShR 4)
2, Item (IShR 4)
2, Item (IShR 4)
2]
  in IShR (4 + 0)
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall (m :: Nat) (n :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownNat n, KnownSTK x, BaseTensor target) =>
IShR (m + n)
-> (IxROf target m -> target (TKR2 n x)) -> target (TKR2 (m + n) x)
rbuild IShR 4
IShR (4 + 0)
shB ((IxR 4 (PrimalOf target (TKScalar Int64))
  -> target (TKR2 0 (TKScalar r)))
 -> target (TKR2 (4 + 0) (TKScalar r)))
-> (IxR 4 (PrimalOf target (TKScalar Int64))
    -> target (TKR2 0 (TKScalar r)))
-> target (TKR2 (4 + 0) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
_, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw] ->
      let arrAt :: target (TKR 4 r)
arrAt = IShR 4
-> target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR 4 r)
forall (target :: Target) r (n :: Nat).
(ADReady target, GoodScalar r, KnownNat n) =>
IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)
slicez3 IShR 4
shB target (TKR 4 r)
arrA [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh]
      in target (TKR 4 r)
-> IxR 4 (PrimalOf target (TKScalar Int64))
-> target (TKR2 0 (TKScalar r))
forall (m :: Nat) (x :: TK) (target :: Target).
(KnownNat m, KnownSTK x, BaseTensor target) =>
target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)
rindex0 target (TKR 4 r)
arrAt [Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBw, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iImg, Item (IxR 4 (PrimalOf target (TKScalar Int64)))
iBh]
    IxR 4 (PrimalOf target (TKScalar Int64))
_ -> String -> target (TKR2 0 (TKScalar r))
forall a. HasCallStack => String -> a
error String
"conv2dUnpadded3y: impossible pattern needlessly required"

-- This test uses a disastrous version of smaximum, but shows how
-- smaxIndex gets non-trivially vectorized, preserving sharing, too.
testCNNOPP4bU :: Assertion
testCNNOPP4bU :: Assertion
testCNNOPP4bU = do
  Assertion
resetVarCounter
  let artifactRev :: AstArtifactRev
  (X (AstTensor
        AstMethodLet
        FullSpan
        (TKS
           ((':)
              @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
           Double)))
  (TKS
     ((':)
        @Nat
        31
        ((':)
           @Nat 31 ((':) @Nat (Div 31 2) ((':) @Nat (Div 31 2) ('[] @Nat)))))
     Double)
artifactRev = IncomingCotangentHandling
-> (AstTensor
      AstMethodLet
      FullSpan
      (TKS
         ((':)
            @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
         Double)
    -> AstTensor
         AstMethodLet
         FullSpan
         (TKS
            ((':)
               @Nat
               31
               ((':)
                  @Nat 31 ((':) @Nat (Div 31 2) ((':) @Nat (Div 31 2) ('[] @Nat)))))
            Double))
-> FullShapeTK
     (X (AstTensor
           AstMethodLet
           FullSpan
           (TKS
              ((':)
                 @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
              Double)))
-> AstArtifactRev
     (X (AstTensor
           AstMethodLet
           FullSpan
           (TKS
              ((':)
                 @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
              Double)))
     (TKS
        ((':)
           @Nat
           31
           ((':)
              @Nat 31 ((':) @Nat (Div 31 2) ((':) @Nat (Div 31 2) ('[] @Nat)))))
        Double)
forall src (ztgt :: TK) tgt.
(AdaptableTarget (AstTensor AstMethodLet FullSpan) src,
 (tgt :: Type) ~ (AstTensor AstMethodLet FullSpan ztgt :: Type)) =>
IncomingCotangentHandling
-> (src -> tgt)
-> FullShapeTK (X src)
-> AstArtifactRev (X src) ztgt
revArtifactAdapt IncomingCotangentHandling
UseIncomingCotangent (forall (ksize :: Nat) (stride :: Nat) (batch_size :: Nat)
       (channels :: Nat) (h :: Nat) (w :: Nat) (target :: Target) r
       (shOut :: [Nat]) (shK1 :: [Nat]).
(KnownNat ksize, KnownNat stride, KnownNat batch_size,
 KnownNat channels, KnownNat h, KnownNat w, (<=) @Nat 1 stride,
 ADReady target, GoodScalar r,
 (shOut :: [Nat])
 ~ ((':)
      @Nat
      batch_size
      ((':)
         @Nat
         channels
         ((':)
            @Nat
            (Div h stride)
            ((':) @Nat (Div w stride) ('[] @Nat)))) :: [Nat]),
 (shK1 :: [Nat])
 ~ ((':)
      @Nat
      1
      ((':)
         @Nat
         1
         ((':) @Nat ksize ((':) @Nat ksize ('[] @Nat)))) :: [Nat])) =>
target
  (TKS
     ((':)
        @Nat
        batch_size
        ((':) @Nat channels ((':) @Nat h ((':) @Nat w ('[] @Nat)))))
     r)
-> target (TKS shOut r)
maxPool2dUnpaddedS4 @4 @2) (ShS
  ((':)
     @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
-> FullShapeTK (TKScalar Double)
-> FullShapeTK
     (TKS
        ((':)
           @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
        Double)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (forall (n :: Nat). KnownNat n => SNat n
SNat @31 SNat 31
-> ShS ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat))))
-> ShS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ forall (n :: Nat). KnownNat n => SNat n
SNat @31 SNat 31
-> ShS ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))
-> ShS ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat))))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ forall (n :: Nat). KnownNat n => SNat n
SNat @31 SNat 31
-> ShS ((':) @Nat 31 ('[] @Nat))
-> ShS ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ forall (n :: Nat). KnownNat n => SNat n
SNat @31 SNat 31 -> ShS ('[] @Nat) -> ShS ((':) @Nat 31 ('[] @Nat))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS) (forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar @Double))
  AstArtifactRev
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
     Double)
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 15 ((':) @Nat 15 ('[] @Nat)))))
     Double)
-> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty (AstArtifactRev
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
     Double)
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 15 ((':) @Nat 15 ('[] @Nat)))))
     Double)
-> AstArtifactRev
     (TKS
        ((':)
           @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
        Double)
     (TKS
        ((':)
           @Nat 31 ((':) @Nat 31 ((':) @Nat 15 ((':) @Nat 15 ('[] @Nat)))))
        Double)
forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstArtifactRev x z
simplifyArtifact AstArtifactRev
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
     Double)
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 15 ((':) @Nat 15 ('[] @Nat)))))
     Double)
artifactRev)
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"
  AstArtifactRev
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
     Double)
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 15 ((':) @Nat 15 ('[] @Nat)))))
     Double)
-> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPrimalPretty AstArtifactRev
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
     Double)
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 15 ((':) @Nat 15 ('[] @Nat)))))
     Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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])"
  AstArtifactRev
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
     Double)
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 15 ((':) @Nat 15 ('[] @Nat)))))
     Double)
-> String
forall (x :: TK) (z :: TK). AstArtifactRev x z -> String
printArtifactPretty AstArtifactRev
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ((':) @Nat 31 ('[] @Nat)))))
     Double)
  (TKS
     ((':)
        @Nat 31 ((':) @Nat 31 ((':) @Nat 15 ((':) @Nat 15 ('[] @Nat)))))
     Double)
artifactRev
    String -> String -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= String
"\\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]))"

smaximum4 :: forall r sh target. (ADReady target, GoodScalar r, KnownShS sh)
          => target (TKS sh r) -> target (TKS '[] r)
smaximum4 :: forall r (sh :: [Nat]) (target :: Target).
(ADReady target, GoodScalar r, KnownShS sh) =>
target (TKS sh r) -> target (TKS ('[] @Nat) r)
smaximum4 target (TKS sh r)
t0 =
  target (TKS sh r)
-> (target (TKS sh r) -> target (TKS ('[] @Nat) r))
-> target (TKS ('[] @Nat) r)
forall (x :: TK) (z :: TK) (target :: Target).
LetTensor target =>
target x -> (target x -> target z) -> target z
tlet target (TKS sh r)
t0 ((target (TKS sh r) -> target (TKS ('[] @Nat) r))
 -> target (TKS ('[] @Nat) r))
-> (target (TKS sh r) -> target (TKS ('[] @Nat) r))
-> target (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ \target (TKS sh r)
t ->
  PrimalOf target (TKScalar Int64)
-> (PrimalOf target (TKScalar Int64) -> target (TKS ('[] @Nat) r))
-> target (TKS ('[] @Nat) r)
forall (x :: TK) (z :: TK).
PrimalOf target x -> (PrimalOf target x -> target z) -> target z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
PrimalOf target x -> (PrimalOf target x -> target z) -> target z
ttletPrimal (target (TKScalar Int64) -> PrimalOf target (TKScalar Int64)
forall (y :: TK). target y -> PrimalOf target y
forall (target :: Target) (y :: TK).
BaseTensor target =>
target y -> PrimalOf target y
tprimalPart (target (TKScalar Int64) -> PrimalOf target (TKScalar Int64))
-> target (TKScalar Int64) -> PrimalOf target (TKScalar Int64)
forall a b. (a -> b) -> a -> b
$ target (TKS ('[] @Nat) Int64) -> target (TKScalar Int64)
forall r.
GoodScalar r =>
target (TKS ('[] @Nat) r) -> target (TKScalar r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKS ('[] @Nat) r) -> target (TKScalar r)
kfromS (target (TKS ('[] @Nat) Int64) -> target (TKScalar Int64))
-> target (TKS ('[] @Nat) Int64) -> target (TKScalar Int64)
forall a b. (a -> b) -> a -> b
$ target (TKS ((':) @Nat (Product sh) ('[] @Nat)) r)
-> target
     (TKS (Init @Nat ((':) @Nat (Product sh) ('[] @Nat))) Int64)
forall (n :: Nat) (sh :: [Nat]) r r2 (target :: Target).
(GoodScalar r, GoodScalar r2, BaseTensor target) =>
target (TKS ((':) @Nat n sh) r)
-> target (TKS (Init @Nat ((':) @Nat n sh)) r2)
smaxIndex (target (TKS sh r)
-> target (TKS ((':) @Nat (Product sh) ('[] @Nat)) r)
forall (sh :: [Nat]) (x :: TK) (target :: Target).
(KnownShS sh, KnownSTK x, BaseTensor target) =>
target (TKS2 sh x)
-> target (TKS2 ((':) @Nat (Product sh) ('[] @Nat)) x)
sflatten target (TKS sh r)
t)) ((PrimalOf target (TKScalar Int64) -> target (TKS ('[] @Nat) r))
 -> target (TKS ('[] @Nat) r))
-> (PrimalOf target (TKScalar Int64) -> target (TKS ('[] @Nat) r))
-> target (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ \PrimalOf target (TKScalar Int64)
maxIndex ->
    target (TKS sh r)
-> IxS sh (PrimalOf target (TKScalar Int64))
-> target (TKS ('[] @Nat) r)
forall (sh1 :: [Nat]) (x :: TK) (target :: Target).
(KnownShS sh1, KnownSTK x, BaseTensor target) =>
target (TKS2 sh1 x)
-> IxSOf target sh1 -> target (TKS2 ('[] @Nat) x)
sindex0 target (TKS sh r)
t
    (IxS sh (PrimalOf target (TKScalar Int64))
 -> target (TKS ('[] @Nat) r))
-> IxS sh (PrimalOf target (TKScalar Int64))
-> target (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ (Int -> PrimalOf target (TKScalar Int64))
-> ShS sh
-> PrimalOf target (TKScalar Int64)
-> IxS sh (PrimalOf target (TKScalar Int64))
forall (sh :: [Nat]) j.
IntegralH j =>
(Int -> j) -> ShS sh -> j -> IxS sh j
fromLinearIdxS (forall (target :: Target) (y :: TK).
BaseTensor target =>
target y -> PrimalOf target y
tprimalPart @target (target (TKScalar Int64) -> PrimalOf target (TKScalar Int64))
-> (Int -> target (TKScalar Int64))
-> Int
-> PrimalOf target (TKScalar Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> target (TKScalar Int64)
forall r (target :: Target).
(GoodScalar r, BaseTensor target) =>
r -> target (TKScalar r)
kconcrete (Int64 -> target (TKScalar Int64))
-> (Int -> Int64) -> Int -> target (TKScalar Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral)
                     (target (TKS sh r) -> ShS sh
forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
target (TKS2 sh x) -> ShS sh
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 sh x) -> ShS sh
sshape target (TKS sh r)
t)
                     PrimalOf target (TKScalar Int64)
maxIndex

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
     , shOut ~ '[batch_size, channels, h `Div` stride, w `Div` stride]
     , shK1 ~ '[1, 1, ksize, ksize]
     )
  => target (TKS '[batch_size, channels, h, w] r)
  -> target (TKS shOut r)
maxPool2dUnpaddedS4 :: forall (ksize :: Nat) (stride :: Nat) (batch_size :: Nat)
       (channels :: Nat) (h :: Nat) (w :: Nat) (target :: Target) r
       (shOut :: [Nat]) (shK1 :: [Nat]).
(KnownNat ksize, KnownNat stride, KnownNat batch_size,
 KnownNat channels, KnownNat h, KnownNat w, (<=) @Nat 1 stride,
 ADReady target, GoodScalar r,
 (shOut :: [Nat])
 ~ ((':)
      @Nat
      batch_size
      ((':)
         @Nat
         channels
         ((':)
            @Nat
            (Div h stride)
            ((':) @Nat (Div w stride) ('[] @Nat)))) :: [Nat]),
 (shK1 :: [Nat])
 ~ ((':)
      @Nat
      1
      ((':)
         @Nat
         1
         ((':) @Nat ksize ((':) @Nat ksize ('[] @Nat)))) :: [Nat])) =>
target
  (TKS
     ((':)
        @Nat
        batch_size
        ((':) @Nat channels ((':) @Nat h ((':) @Nat w ('[] @Nat)))))
     r)
-> target (TKS shOut r)
maxPool2dUnpaddedS4 target
  (TKS
     ((':)
        @Nat
        batch_size
        ((':) @Nat channels ((':) @Nat h ((':) @Nat w ('[] @Nat)))))
     r)
arr =
  let stride :: Int
stride = forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @stride :: Int
  in forall (m :: Nat) (sh :: [Nat]) (x :: TK) (target :: Target).
(KnownShS (Take @Nat m sh), KnownShS sh, KnownSTK x,
 BaseTensor target) =>
(IxSOf target (Take @Nat m sh) -> target (TKS2 (Drop @Nat m sh) x))
-> target (TKS2 sh x)
sbuild @(Rank shOut) ((IxSOf target (Take @Nat (Rank @Nat shOut) shOut)
  -> target (TKS2 (Drop @Nat (Rank @Nat shOut) shOut) (TKScalar r)))
 -> target (TKS2 shOut (TKScalar r)))
-> (IxSOf target (Take @Nat (Rank @Nat shOut) shOut)
    -> target (TKS2 (Drop @Nat (Rank @Nat shOut) shOut) (TKScalar r)))
-> target (TKS2 shOut (TKScalar r))
forall a b. (a -> b) -> a -> b
$ \case
    [Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':)
           @Nat
           channels
           ((':) @Nat (Div h stride) ((':) @Nat (Div w stride) ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
iImg, Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':)
           @Nat
           channels
           ((':) @Nat (Div h stride) ((':) @Nat (Div w stride) ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
iChan, Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':)
           @Nat
           channels
           ((':) @Nat (Div h stride) ((':) @Nat (Div w stride) ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
iBh, Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':)
           @Nat
           channels
           ((':) @Nat (Div h stride) ((':) @Nat (Div w stride) ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
iBw] ->
      target (TKS shK1 r) -> target (TKS ('[] @Nat) r)
forall r (sh :: [Nat]) (target :: Target).
(ADReady target, GoodScalar r, KnownShS sh) =>
target (TKS sh r) -> target (TKS ('[] @Nat) r)
smaximum4 (target (TKS shK1 r) -> target (TKS ('[] @Nat) r))
-> target (TKS shK1 r) -> target (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ forall (shOut :: [Nat]) (sh :: [Nat]) (target :: Target) r.
(KnownShS sh, KnownShS shOut,
 KnownShS (Take @Nat (Rank @Nat sh) shOut), KnownNat (Rank @Nat sh),
 (Rank @Nat shOut :: Nat) ~ (Rank @Nat sh :: Nat), ADReady target,
 GoodScalar r) =>
target (TKS sh r) -> IxSOf target sh -> target (TKS shOut r)
slicezS @shK1 target
  (TKS
     ((':)
        @Nat
        batch_size
        ((':) @Nat channels ((':) @Nat h ((':) @Nat w ('[] @Nat)))))
     r)
arr [ Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':) @Nat channels ((':) @Nat h ((':) @Nat w ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':)
           @Nat
           channels
           ((':) @Nat (Div h stride) ((':) @Nat (Div w stride) ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
iImg, Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':) @Nat channels ((':) @Nat h ((':) @Nat w ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':)
           @Nat
           channels
           ((':) @Nat (Div h stride) ((':) @Nat (Div w stride) ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
iChan
                                    , Int -> PrimalOf target (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
stride PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
* Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':)
           @Nat
           channels
           ((':) @Nat (Div h stride) ((':) @Nat (Div w stride) ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
iBh
                                    , Int -> PrimalOf target (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
stride PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
-> PrimalOf target (TKScalar Int64)
forall a. Num a => a -> a -> a
* Item
  (IxS
     ((':)
        @Nat
        batch_size
        ((':)
           @Nat
           channels
           ((':) @Nat (Div h stride) ((':) @Nat (Div w stride) ('[] @Nat)))))
     (PrimalOf target (TKScalar Int64)))
PrimalOf target (TKScalar Int64)
iBw ]
    IxSOf target (Take @Nat (Rank @Nat shOut) shOut)
_ -> String -> target (TKS ('[] @Nat) r)
forall a. HasCallStack => String -> a
error String
"maxPool2dUnpaddedS4: impossible pattern needlessly required"